1import path from "node:path";
2import { readFile } from "node:fs/promises";
3import type { CliArgs } from "../types";
4 
5const DEFAULT_MODEL = "google/nano-banana-2";
6const SYNC_WAIT_SECONDS = 60;
7const POLL_INTERVAL_MS = 2000;
8const MAX_POLL_MS = 300_000;
9const DOCUMENTED_REPLICATE_ASPECT_RATIOS = new Set([
10  "1:1",
11  "2:3",
12  "3:2",
13  "3:4",
14  "4:3",
15  "5:4",
16  "4:5",
17  "9:16",
18  "16:9",
19  "21:9",
20]);
21 
22export type ReplicateModelFamily =
23  | "nano-banana"
24  | "seedream45"
25  | "seedream5lite"
26  | "wan27image"
27  | "wan27imagepro"
28  | "unknown";
29 
30type PixelSize = {
31  width: number;
32  height: number;
33};
34 
35type Seedream45Size = "2K" | "4K" | { width: number; height: number };
36 
37export function getDefaultModel(): string {
38  return process.env.REPLICATE_IMAGE_MODEL || DEFAULT_MODEL;
39}
40 
41function getApiToken(): string | null {
42  return process.env.REPLICATE_API_TOKEN || null;
43}
44 
45function getBaseUrl(): string {
46  const base = process.env.REPLICATE_BASE_URL || "https://api.replicate.com";
47  return base.replace(/\/+$/g, "");
48}
49 
50function normalizeModelId(model: string): string {
51  return model.trim().toLowerCase().split(":")[0]!;
52}
53 
54export function getModelFamily(model: string): ReplicateModelFamily {
55  const normalized = normalizeModelId(model);
56 
57  if (
58    normalized === "google/nano-banana" ||
59    normalized === "google/nano-banana-pro" ||
60    normalized === "google/nano-banana-2"
61  ) {
62    return "nano-banana";
63  }
64 
65  if (normalized === "bytedance/seedream-4.5") {
66    return "seedream45";
67  }
68 
69  if (normalized === "bytedance/seedream-5-lite") {
70    return "seedream5lite";
71  }
72 
73  if (normalized === "wan-video/wan-2.7-image") {
74    return "wan27image";
75  }
76 
77  if (normalized === "wan-video/wan-2.7-image-pro") {
78    return "wan27imagepro";
79  }
80 
81  return "unknown";
82}
83 
84export function parseModelId(model: string): { owner: string; name: string; version: string | null } {
85  const [ownerName, version] = model.split(":");
86  const parts = ownerName!.split("/");
87  if (parts.length !== 2 || !parts[0] || !parts[1]) {
88    throw new Error(
89      `Invalid Replicate model format: "${model}". Expected "owner/name" or "owner/name:version".`
90    );
91  }
92  return { owner: parts[0], name: parts[1], version: version || null };
93}
94 
95function parsePixelSize(value: string): PixelSize | null {
96  const match = value.trim().match(/^(\d+)\s*[xX*]\s*(\d+)$/);
97  if (!match) return null;
98 
99  const width = parseInt(match[1]!, 10);
100  const height = parseInt(match[2]!, 10);
101  if (!Number.isFinite(width) || !Number.isFinite(height) || width <= 0 || height <= 0) {
102    return null;
103  }
104 
105  return { width, height };
106}
107 
108function parseAspectRatio(value: string): PixelSize | null {
109  const match = value.trim().match(/^(\d+)\s*:\s*(\d+)$/);
110  if (!match) return null;
111 
112  const width = parseInt(match[1]!, 10);
113  const height = parseInt(match[2]!, 10);
114  if (!Number.isFinite(width) || !Number.isFinite(height) || width <= 0 || height <= 0) {
115    return null;
116  }
117 
118  return { width, height };
119}
120 
121function gcd(a: number, b: number): number {
122  let x = Math.abs(a);
123  let y = Math.abs(b);
124 
125  while (y !== 0) {
126    const next = x % y;
127    x = y;
128    y = next;
129  }
130 
131  return x || 1;
132}
133 
134function inferAspectRatioFromSize(size: string): string | null {
135  const parsed = parsePixelSize(size);
136  if (!parsed) return null;
137 
138  const divisor = gcd(parsed.width, parsed.height);
139  const normalized = `${parsed.width / divisor}:${parsed.height / divisor}`;
140  if (!DOCUMENTED_REPLICATE_ASPECT_RATIOS.has(normalized)) {
141    return null;
142  }
143 
144  return normalized;
145}
146 
147function getQualityPreset(args: CliArgs): "normal" | "2k" {
148  return args.quality === "normal" ? "normal" : "2k";
149}
150 
151function validateDocumentedAspectRatio(model: string, aspectRatio: string): void {
152  if (aspectRatio === "match_input_image") {
153    return;
154  }
155 
156  if (DOCUMENTED_REPLICATE_ASPECT_RATIOS.has(aspectRatio)) {
157    return;
158  }
159 
160  throw new Error(
161    `Replicate model ${model} does not support aspect ratio ${aspectRatio}. Supported values: ${Array.from(DOCUMENTED_REPLICATE_ASPECT_RATIOS).join(", ")}`
162  );
163}
164 
165function getRequestedAspectRatio(model: string, args: CliArgs): string | null {
166  if (args.aspectRatio) {
167    validateDocumentedAspectRatio(model, args.aspectRatio);
168    return args.aspectRatio;
169  }
170 
171  if (!args.size) return null;
172 
173  const inferred = inferAspectRatioFromSize(args.size);
174  if (!inferred) {
175    throw new Error(
176      `Replicate model ${model} cannot derive a supported aspect ratio from --size ${args.size}. Use one of: ${Array.from(DOCUMENTED_REPLICATE_ASPECT_RATIOS).join(", ")}`
177    );
178  }
179 
180  return inferred;
181}
182 
183function getNanoBananaResolution(args: CliArgs): "1K" | "2K" {
184  if (args.size) {
185    const parsed = parsePixelSize(args.size);
186    if (!parsed) {
187      throw new Error("Replicate nano-banana --size must be in WxH format, for example 1536x1024.");
188    }
189 
190    const longestEdge = Math.max(parsed.width, parsed.height);
191    if (longestEdge <= 1024) return "1K";
192    if (longestEdge <= 2048) return "2K";
193    throw new Error("Replicate nano-banana only supports sizes that map to 1K or 2K output.");
194  }
195 
196  return getQualityPreset(args) === "normal" ? "1K" : "2K";
197}
198 
199function resolveSeedream45Size(args: CliArgs): Seedream45Size {
200  if (args.size) {
201    const upper = args.size.trim().toUpperCase();
202    if (upper === "2K" || upper === "4K") {
203      return upper;
204    }
205 
206    const parsed = parsePixelSize(args.size);
207    if (!parsed) {
208      throw new Error("Replicate Seedream 4.5 --size must be 2K, 4K, or an explicit WxH size.");
209    }
210    if (parsed.width < 1024 || parsed.width > 4096 || parsed.height < 1024 || parsed.height > 4096) {
211      throw new Error("Replicate Seedream 4.5 custom --size must keep width and height between 1024 and 4096.");
212    }
213    return parsed;
214  }
215 
216  return getQualityPreset(args) === "normal" ? "2K" : "4K";
217}
218 
219function resolveSeedream5LiteSize(args: CliArgs): "2K" | "3K" {
220  if (args.size) {
221    const upper = args.size.trim().toUpperCase();
222    if (upper === "2K" || upper === "3K") {
223      return upper;
224    }
225 
226    throw new Error("Replicate Seedream 5 Lite currently supports 2K or 3K output in this tool.");
227  }
228 
229  return getQualityPreset(args) === "normal" ? "2K" : "3K";
230}
231 
232function formatCustomWanSize(size: PixelSize): string {
233  return `${size.width}*${size.height}`;
234}
235 
236function resolveWanSizeFromAspectRatio(
237  aspectRatio: string,
238  maxDimension: number,
239): string {
240  const parsedRatio = parseAspectRatio(aspectRatio);
241  if (!parsedRatio) {
242    throw new Error(`Replicate Wan aspect ratio must be in W:H format, got ${aspectRatio}.`);
243  }
244 
245  const scale = Math.min(maxDimension / parsedRatio.width, maxDimension / parsedRatio.height);
246  const width = Math.max(1, Math.floor(parsedRatio.width * scale));
247  const height = Math.max(1, Math.floor(parsedRatio.height * scale));
248  return formatCustomWanSize({ width, height });
249}
250 
251function resolveWanSize(family: "wan27image" | "wan27imagepro", args: CliArgs): "1K" | "2K" | "4K" | string {
252  const referenceMode = args.referenceImages.length > 0;
253  const maxDimension = family === "wan27imagepro" && !referenceMode ? 4096 : 2048;
254 
255  if (args.size) {
256    const upper = args.size.trim().toUpperCase();
257    if (upper === "1K" || upper === "2K" || upper === "4K") {
258      if (upper === "4K" && family !== "wan27imagepro") {
259        throw new Error("Replicate Wan 2.7 Image only supports 1K, 2K, or custom sizes up to 2048px.");
260      }
261      if (upper === "4K" && referenceMode) {
262        throw new Error("Replicate Wan 2.7 Image Pro only supports 4K text-to-image. Remove --ref or lower the size.");
263      }
264      return upper;
265    }
266 
267    const parsed = parsePixelSize(args.size);
268    if (!parsed) {
269      throw new Error("Replicate Wan --size must be 1K, 2K, 4K, or an explicit WxH size.");
270    }
271    if (parsed.width > maxDimension || parsed.height > maxDimension) {
272      throw new Error(
273        `Replicate ${family === "wan27imagepro" ? "Wan 2.7 Image Pro" : "Wan 2.7 Image"} custom --size must keep width and height at or below ${maxDimension}px in the current mode.`
274      );
275    }
276    return formatCustomWanSize(parsed);
277  }
278 
279  if (args.aspectRatio) {
280    return resolveWanSizeFromAspectRatio(
281      args.aspectRatio,
282      getQualityPreset(args) === "normal" ? 1024 : 2048,
283    );
284  }
285 
286  return getQualityPreset(args) === "normal" ? "1K" : "2K";
287}
288 
289function buildNanoBananaInput(
290  prompt: string,
291  model: string,
292  args: CliArgs,
293  referenceImages: string[],
294): Record<string, unknown> {
295  const input: Record<string, unknown> = {
296    prompt,
297    resolution: getNanoBananaResolution(args),
298    output_format: "png",
299  };
300 
301  const aspectRatio = getRequestedAspectRatio(model, args);
302  if (aspectRatio) {
303    input.aspect_ratio = aspectRatio;
304  } else if (referenceImages.length > 0) {
305    input.aspect_ratio = "match_input_image";
306  }
307 
308  if (referenceImages.length > 0) {
309    input.image_input = referenceImages;
310  }
311 
312  return input;
313}
314 
315function buildSeedreamInput(
316  family: "seedream45" | "seedream5lite",
317  prompt: string,
318  model: string,
319  args: CliArgs,
320  referenceImages: string[],
321): Record<string, unknown> {
322  const size = family === "seedream45" ? resolveSeedream45Size(args) : resolveSeedream5LiteSize(args);
323  const input: Record<string, unknown> = {
324    prompt,
325  };
326 
327  if (family === "seedream45" && typeof size === "object") {
328    input.size = "custom";
329    input.width = size.width;
330    input.height = size.height;
331  } else {
332    input.size = size;
333  }
334 
335  if (referenceImages.length > 0) {
336    input.image_input = referenceImages;
337  }
338 
339  if (args.aspectRatio) {
340    validateDocumentedAspectRatio(model, args.aspectRatio);
341    input.aspect_ratio = args.aspectRatio;
342  } else if (referenceImages.length > 0 && family === "seedream45") {
343    input.aspect_ratio = "match_input_image";
344  }
345 
346  return input;
347}
348 
349function buildWanInput(
350  family: "wan27image" | "wan27imagepro",
351  prompt: string,
352  args: CliArgs,
353  referenceImages: string[],
354): Record<string, unknown> {
355  const input: Record<string, unknown> = {
356    prompt,
357    size: resolveWanSize(family, args),
358  };
359 
360  if (referenceImages.length > 0) {
361    input.images = referenceImages;
362  }
363 
364  return input;
365}
366 
367export function validateArgs(model: string, args: CliArgs): void {
368  parseModelId(model);
369 
370  if (args.n !== 1) {
371    throw new Error("Replicate integration currently supports exactly one output image per request. Remove --n or use --n 1.");
372  }
373 
374  if (args.imageSize && args.imageSizeSource !== "config") {
375    throw new Error("Replicate models in baoyu-image-gen do not use --imageSize. Use --quality, --ar, or --size instead.");
376  }
377 
378  const family = getModelFamily(model);
379 
380  if (family === "nano-banana") {
381    if (args.referenceImages.length > 14) {
382      throw new Error("Replicate nano-banana supports at most 14 reference images.");
383    }
384    if (args.aspectRatio) {
385      validateDocumentedAspectRatio(model, args.aspectRatio);
386    }
387    if (args.size) {
388      getRequestedAspectRatio(model, args);
389      getNanoBananaResolution(args);
390    }
391    return;
392  }
393 
394  if (family === "seedream45") {
395    if (args.referenceImages.length > 14) {
396      throw new Error("Replicate Seedream 4.5 supports at most 14 reference images.");
397    }
398    if (args.aspectRatio) {
399      validateDocumentedAspectRatio(model, args.aspectRatio);
400    }
401    resolveSeedream45Size(args);
402    return;
403  }
404 
405  if (family === "seedream5lite") {
406    if (args.referenceImages.length > 14) {
407      throw new Error("Replicate Seedream 5 Lite supports at most 14 reference images.");
408    }
409    if (args.aspectRatio) {
410      validateDocumentedAspectRatio(model, args.aspectRatio);
411    }
412    resolveSeedream5LiteSize(args);
413    return;
414  }
415 
416  if (family === "wan27image" || family === "wan27imagepro") {
417    if (args.referenceImages.length > 9) {
418      throw new Error("Replicate Wan 2.7 image models support at most 9 reference images.");
419    }
420    if (args.aspectRatio) {
421      const parsed = parseAspectRatio(args.aspectRatio);
422      if (!parsed) {
423        throw new Error(`Replicate Wan aspect ratio must be in W:H format, got ${args.aspectRatio}.`);
424      }
425    }
426    resolveWanSize(family, args);
427    return;
428  }
429 
430  const hasExplicitAspectRatio = !!args.aspectRatio && args.aspectRatioSource !== "config";
431 
432  if (args.referenceImages.length > 0 || hasExplicitAspectRatio || args.size) {
433    throw new Error(
434      `Replicate model ${model} is not in the baoyu-image-gen compatibility list. Supported families: google/nano-banana*, bytedance/seedream-4.5, bytedance/seedream-5-lite, wan-video/wan-2.7-image, wan-video/wan-2.7-image-pro.`
435    );
436  }
437}
438 
439export function getDefaultOutputExtension(model: string): ".png" {
440  const _family = getModelFamily(model);
441  return ".png";
442}
443 
444export function buildInput(
445  model: string,
446  prompt: string,
447  args: CliArgs,
448  referenceImages: string[],
449): Record<string, unknown> {
450  const family = getModelFamily(model);
451 
452  if (family === "nano-banana") {
453    return buildNanoBananaInput(prompt, model, args, referenceImages);
454  }
455 
456  if (family === "seedream45" || family === "seedream5lite") {
457    return buildSeedreamInput(family, prompt, model, args, referenceImages);
458  }
459 
460  if (family === "wan27image" || family === "wan27imagepro") {
461    return buildWanInput(family, prompt, args, referenceImages);
462  }
463 
464  return { prompt };
465}
466 
467async function readImageAsDataUrl(p: string): Promise<string> {
468  const buf = await readFile(p);
469  const ext = path.extname(p).toLowerCase();
470  let mimeType = "image/png";
471  if (ext === ".jpg" || ext === ".jpeg") mimeType = "image/jpeg";
472  else if (ext === ".gif") mimeType = "image/gif";
473  else if (ext === ".webp") mimeType = "image/webp";
474  return `data:${mimeType};base64,${buf.toString("base64")}`;
475}
476 
477type PredictionResponse = {
478  id: string;
479  status: string;
480  output: unknown;
481  error: string | null;
482  urls?: { get?: string };
483};
484 
485async function createPrediction(
486  apiToken: string,
487  model: { owner: string; name: string; version: string | null },
488  input: Record<string, unknown>,
489  sync: boolean
490): Promise<PredictionResponse> {
491  const baseUrl = getBaseUrl();
492 
493  let url: string;
494  const body: Record<string, unknown> = { input };
495 
496  if (model.version) {
497    url = `${baseUrl}/v1/predictions`;
498    body.version = model.version;
499  } else {
500    url = `${baseUrl}/v1/models/${model.owner}/${model.name}/predictions`;
501  }
502 
503  const headers: Record<string, string> = {
504    Authorization: `Bearer ${apiToken}`,
505    "Content-Type": "application/json",
506  };
507 
508  if (sync) {
509    headers["Prefer"] = `wait=${SYNC_WAIT_SECONDS}`;
510  }
511 
512  const res = await fetch(url, {
513    method: "POST",
514    headers,
515    body: JSON.stringify(body),
516  });
517 
518  if (!res.ok) {
519    const err = await res.text();
520    throw new Error(`Replicate API error (${res.status}): ${err}`);
521  }
522 
523  return (await res.json()) as PredictionResponse;
524}
525 
526async function pollPrediction(apiToken: string, getUrl: string): Promise<PredictionResponse> {
527  const start = Date.now();
528 
529  while (Date.now() - start < MAX_POLL_MS) {
530    const res = await fetch(getUrl, {
531      headers: { Authorization: `Bearer ${apiToken}` },
532    });
533 
534    if (!res.ok) {
535      const err = await res.text();
536      throw new Error(`Replicate poll error (${res.status}): ${err}`);
537    }
538 
539    const prediction = (await res.json()) as PredictionResponse;
540 
541    if (prediction.status === "succeeded") return prediction;
542    if (prediction.status === "failed" || prediction.status === "canceled") {
543      throw new Error(`Replicate prediction ${prediction.status}: ${prediction.error || "unknown error"}`);
544    }
545 
546    await new Promise((r) => setTimeout(r, POLL_INTERVAL_MS));
547  }
548 
549  throw new Error(`Replicate prediction timed out after ${MAX_POLL_MS / 1000}s`);
550}
551 
552export function extractOutputUrl(prediction: PredictionResponse): string {
553  const output = prediction.output;
554 
555  if (typeof output === "string") return output;
556 
557  if (Array.isArray(output)) {
558    if (output.length !== 1) {
559      throw new Error(
560        `Replicate returned ${output.length} outputs, but baoyu-image-gen currently supports saving exactly one image per request.`
561      );
562    }
563    const first = output[0];
564    if (typeof first === "string") return first;
565  }
566 
567  if (output && typeof output === "object" && "url" in output) {
568    const url = (output as Record<string, unknown>).url;
569    if (typeof url === "string") return url;
570  }
571 
572  throw new Error(`Unexpected Replicate output format: ${JSON.stringify(output)}`);
573}
574 
575async function downloadImage(url: string): Promise<Uint8Array> {
576  const res = await fetch(url);
577  if (!res.ok) throw new Error(`Failed to download image from Replicate: ${res.status}`);
578  const buf = await res.arrayBuffer();
579  return new Uint8Array(buf);
580}
581 
582export async function generateImage(
583  prompt: string,
584  model: string,
585  args: CliArgs
586): Promise<Uint8Array> {
587  const apiToken = getApiToken();
588  if (!apiToken) throw new Error("REPLICATE_API_TOKEN is required. Get one at https://replicate.com/account/api-tokens");
589 
590  const parsedModel = parseModelId(model);
591  validateArgs(model, args);
592 
593  const refDataUrls: string[] = [];
594  for (const refPath of args.referenceImages) {
595    refDataUrls.push(await readImageAsDataUrl(refPath));
596  }
597 
598  const input = buildInput(model, prompt, args, refDataUrls);
599 
600  console.log(`Generating image with Replicate (${model})...`);
601 
602  let prediction = await createPrediction(apiToken, parsedModel, input, true);
603 
604  if (prediction.status !== "succeeded") {
605    if (!prediction.urls?.get) {
606      throw new Error("Replicate prediction did not return a poll URL");
607    }
608    console.log("Waiting for prediction to complete...");
609    prediction = await pollPrediction(apiToken, prediction.urls.get);
610  }
611 
612  console.log("Generation completed.");
613 
614  const outputUrl = extractOutputUrl(prediction);
615  return downloadImage(outputUrl);
616}
617