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;
};

export async function generateImage(
  options: GenerateImageOptions
): Promise<GenerateImageResponse> {
  const apiKey = process.env.GEMINI_API_KEY;
  if (!apiKey) {
    throw new Error("GEMINI_API_KEY is not configured");
  }

  const url = "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash-exp: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("[ImageGen] 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(
      `Image generation failed (${response.status}): ${detail.slice(0, 200)}`
    );
  }

  const result = await response.json() as any;

  const candidates = result.candidates || [];
  for (const candidate of 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("No image generated in response");
}