retrotoon-studio/server/_core/imageGeneration.ts

import { storagePut } from "server/storage";

export type GenerateImageOptions = {
  prompt: string;
  originalImages?: Array<{
    url?: string;
    b64Json?: string;
    mimeType?: string;
  }>;
};

export type GenerateImageResponse = {
  url?: string;
  provider?: "gemini" | "openai";
};

/**
 * Try Gemini first, fallback to OpenAI gpt-image-1 if Gemini fails (quota, error, etc.)
 */
export async function generateImage(
  options: GenerateImageOptions
): Promise<GenerateImageResponse> {
  // Try Gemini first
  if (process.env.GEMINI_API_KEY) {
    try {
      const result = await generateWithGemini(options);
      if (result.url) return { ...result, provider: "gemini" };
    } catch (err: any) {
      console.warn(`[ImageGen] Gemini failed (${err.message?.slice(0, 100)}), falling back to OpenAI...`);
    }
  }

  // Fallback to OpenAI
  if (process.env.OPENAI_API_KEY) {
    const result = await generateWithOpenAI(options);
    return { ...result, provider: "openai" };
  }

  throw new Error("No image generation provider configured (need GEMINI_API_KEY or OPENAI_API_KEY)");
}

async function generateWithGemini(options: GenerateImageOptions): Promise<GenerateImageResponse> {
  const apiKey = process.env.GEMINI_API_KEY!;
  const url = "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-image:generateContent";

  const parts: any[] = [{ text: options.prompt }];

  if (options.originalImages) {
    for (const img of options.originalImages) {
      if (img.url && img.url.startsWith("http")) {
        try {
          const resp = await fetch(img.url);
          if (resp.ok) {
            const buffer = Buffer.from(await resp.arrayBuffer());
            const mime = img.mimeType || resp.headers.get("content-type") || "image/jpeg";
            parts.push({
              inlineData: { mimeType: mime, data: buffer.toString("base64") },
            });
          }
        } catch (e) {
          console.warn("[Gemini] Failed to download reference image:", e);
        }
      } else if (img.b64Json) {
        parts.push({
          inlineData: { mimeType: img.mimeType || "image/png", data: img.b64Json },
        });
      }
    }
  }

  const payload = {
    contents: [{ parts }],
    generationConfig: { responseModalities: ["TEXT", "IMAGE"] },
  };

  const response = await fetch(`${url}?key=${apiKey}`, {
    method: "POST",
    headers: { "Content-Type": "application/json" },
    body: JSON.stringify(payload),
  });

  if (!response.ok) {
    const detail = await response.text().catch(() => "");
    throw new Error(`Gemini ${response.status}: ${detail.slice(0, 150)}`);
  }

  const result = await response.json() as any;
  for (const candidate of result.candidates || []) {
    for (const part of candidate.content?.parts || []) {
      if (part.inlineData?.data) {
        const buffer = Buffer.from(part.inlineData.data, "base64");
        const ext = part.inlineData.mimeType?.includes("png") ? "png" : "jpg";
        const { url: storageUrl } = await storagePut(
          `generated/${Date.now()}.${ext}`,
          buffer,
          part.inlineData.mimeType || "image/png"
        );
        return { url: storageUrl };
      }
    }
  }

  throw new Error("Gemini: no image in response");
}

async function generateWithOpenAI(options: GenerateImageOptions): Promise<GenerateImageResponse> {
  const apiKey = process.env.OPENAI_API_KEY!;
  const hasReferenceImage = options.originalImages && options.originalImages.length > 0 && options.originalImages[0].url;

  let response: Response;

  if (hasReferenceImage) {
    // Use images.edits with reference image (multipart)
    const formData = new FormData();
    formData.append("model", "gpt-image-1");
    formData.append("prompt", options.prompt);
    formData.append("n", "1");
    formData.append("size", "1024x1024");

    // Download reference image(s) and attach
    for (let i = 0; i < options.originalImages!.length; i++) {
      const img = options.originalImages![i];
      if (img.url && img.url.startsWith("http")) {
        try {
          const resp = await fetch(img.url);
          if (!resp.ok) continue;
          const arrayBuffer = await resp.arrayBuffer();
          const mime = img.mimeType || resp.headers.get("content-type") || "image/png";
          const ext = mime.includes("png") ? "png" : mime.includes("jpeg") ? "jpg" : "png";
          const blob = new Blob([arrayBuffer], { type: mime });
          formData.append("image[]", blob, `reference_${i}.${ext}`);
        } catch (e) {
          console.warn("[OpenAI] Failed to download reference:", e);
        }
      }
    }

    response = await fetch("https://api.openai.com/v1/images/edits", {
      method: "POST",
      headers: { Authorization: `Bearer ${apiKey}` },
      body: formData,
    });
  } else {
    // Use images.generations (no reference)
    response = await fetch("https://api.openai.com/v1/images/generations", {
      method: "POST",
      headers: {
        Authorization: `Bearer ${apiKey}`,
        "Content-Type": "application/json",
      },
      body: JSON.stringify({
        model: "gpt-image-1",
        prompt: options.prompt,
        n: 1,
        size: "1024x1024",
      }),
    });
  }

  if (!response.ok) {
    const detail = await response.text().catch(() => "");
    throw new Error(`OpenAI ${response.status}: ${detail.slice(0, 200)}`);
  }

  const result = await response.json() as any;
  const img = result.data?.[0];
  if (!img) throw new Error("OpenAI: no image in response");

  let buffer: Buffer;
  if (img.b64_json) {
    buffer = Buffer.from(img.b64_json, "base64");
  } else if (img.url) {
    const downloadResp = await fetch(img.url);
    if (!downloadResp.ok) throw new Error("OpenAI: failed to download generated image");
    buffer = Buffer.from(await downloadResp.arrayBuffer());
  } else {
    throw new Error("OpenAI: no b64_json or url in response");
  }

  const { url: storageUrl } = await storagePut(
    `generated/${Date.now()}.png`,
    buffer,
    "image/png"
  );
  return { url: storageUrl };
}