Upload 2 files

app/api/image/route.ts

@@ -1,8 +1,14 @@
 import { NextRequest, NextResponse } from "next/server";
 import { createGoogleGenerativeAI } from "@ai-sdk/google";
-import { generateText, CoreMessage, MessagePart } from "ai";
+import { generateText } from "ai";
 
 const MODEL_ID = "gemini-2.0-flash-exp";
+const MODEL_IDS = [
+  MODEL_ID,
+  "gemini-2.5-flash-image",
+  "gemini-3.1-flash-image-preview",
+  "gemini-3-pro-image-preview",
+];
 
 interface ImageHistoryPart {
   text?: string;
@@ -10,7 +16,7 @@ interface ImageHistoryPart {
 }
 
 interface HistoryMessage {
-  role: "user" | "assistant";
+  role: "user" | "model" | "assistant";
   parts: ImageHistoryPart[];
 }
 
@@ -18,12 +24,119 @@ interface RequestBody {
   prompt: string;
   image?: string;
   history?: HistoryMessage[];
+  model?: string;
+}
+
+function getModelCandidates(model?: string) {
+  const selectedModel = model && MODEL_IDS.includes(model) ? model : MODEL_ID;
+  return [selectedModel, ...MODEL_IDS.filter((id) => id !== selectedModel)];
+}
+
+function getMimeType(dataUrl: string) {
+  return dataUrl.match(/^data:([^;]+);base64,/)?.[1] ?? "image/png";
+}
+
+function getBase64Data(dataUrl: string) {
+  return dataUrl.includes(",") ? dataUrl.split(",")[1] : dataUrl;
+}
+
+function toDataUrl(data: unknown, mediaType = "image/png") {
+  if (typeof data === "string") {
+    if (data.startsWith("data:")) {
+      return data;
+    }
+    if (data.startsWith("http://") || data.startsWith("https://")) {
+      return data;
+    }
+    return `data:${mediaType};base64,${data}`;
+  }
+
+  if (data instanceof Uint8Array) {
+    return `data:${mediaType};base64,${Buffer.from(data).toString("base64")}`;
+  }
+
+  if (data instanceof ArrayBuffer) {
+    return `data:${mediaType};base64,${Buffer.from(data).toString("base64")}`;
+  }
+
+  return null;
+}
+
+function extractImage(result: any) {
+  for (const file of result.files ?? []) {
+    const mediaType = file.mediaType ?? file.mimeType ?? "image/png";
+
+    if (!mediaType.startsWith("image/")) {
+      continue;
+    }
+
+    const image =
+      toDataUrl(file.base64, mediaType) ??
+      toDataUrl(file.data, mediaType) ??
+      toDataUrl(file.uint8Array, mediaType) ??
+      toDataUrl(file.url, mediaType);
+
+    if (image) {
+      return image;
+    }
+  }
+
+  for (const part of result.parts ?? []) {
+    const mediaType = part.image?.mediaType ?? part.image?.mimeType ?? "image/png";
+    const image =
+      toDataUrl(part.image?.data, mediaType) ??
+      toDataUrl(part.image?.base64, mediaType) ??
+      toDataUrl(part.image?.url, mediaType);
+
+    if (part.type === "output_image" && image) {
+      return image;
+    }
+
+    if ((part.type === "image" || part.type === "file") && mediaType.startsWith("image/") && image) {
+      return image;
+    }
+  }
+
+  return null;
+}
+
+function formatHistory(history?: HistoryMessage[]) {
+  return (
+    history
+      ?.map((msg) => {
+        const role = msg.role === "user" ? "user" : "assistant";
+        const parts = msg.parts
+          .map((p) => {
+            if (p.text) {
+              return { type: "text", text: p.text };
+            }
+
+            if (p.image && role === "user") {
+              return {
+                type: "image",
+                image: getBase64Data(p.image),
+                mediaType: getMimeType(p.image),
+              };
+            }
+
+            return null;
+          })
+          .filter(Boolean);
+
+        if (parts.length === 0) {
+          return null;
+        }
+
+        return { role, content: parts };
+      })
+      .filter(Boolean) ?? []
+  );
 }
 
 export async function POST(req: NextRequest) {
   try {
     const body: RequestBody = await req.json();
-    const { prompt, image: inputImage, history } = body;
+    const { prompt, image: inputImage, history, model } = body;
 
     if (!prompt) {
       return NextResponse.json(
@@ -32,7 +145,10 @@ export async function POST(req: NextRequest) {
       );
     }
 
-    const parts: MessagePart[] = [{ type: "text", text: prompt }];
+    const google = createGoogleGenerativeAI({
+      apiKey: process.env.GEMINI_API_KEY ?? process.env.GOOGLE_GENERATIVE_AI_API_KEY,
+    });
+    const parts: any[] = [{ type: "text", text: prompt }];
 
     if (inputImage) {
       if (!inputImage.startsWith("data:")) {
@@ -42,68 +158,67 @@ export async function POST(req: NextRequest) {
         );
       }
 
-      const [meta, base64] = inputImage.split(",");
-      const mimeType = meta.includes("png") ? "image/png" : "image/jpeg";
-
       parts.push({
-        type: "input_image",
-        image: {
-          data: base64,
-          mimeType,
-        },
+        type: "image",
+        image: getBase64Data(inputImage),
+        mediaType: getMimeType(inputImage),
       });
     }
 
-    const formattedHistory: CoreMessage[] =
-      history?.map((msg) => ({
-        role: msg.role,
-        parts: msg.parts.map<MessagePart>((p) => {
-          if (p.text) {
-            return { type: "text", text: p.text };
-          }
-          if (p.image) {
-            const [meta, base64] = p.image.split(",");
-            const mimeType = meta.includes("png")
-              ? "image/png"
-              : "image/jpeg";
-
-            return {
-              type: "input_image",
-              image: {
-                data: base64,
-                mimeType,
-              },
-            };
-          }
-          return { type: "text", text: "" };
-        }),
-      })) ?? [];
-
-    const result = await generateText({
-      model: GoogleGenerativeAI(MODEL_ID, {
-        responseModalities: ["text", "image"],
-      }),
-      messages: [...formattedHistory, { role: "user", parts }],
-    });
-
-    const textResponse = result.text;
-
-    let imageBase64: string | null = null;
-    let mimeType = "image/png";
-
-    for (const part of result.parts ?? []) {
-      if (part.type === "output_image" && part.image) {
-        imageBase64 = part.image.data;
-        mimeType = part.image.mimeType ?? "image/png";
+    const formattedHistory = formatHistory(history);
+    const modelCandidates = getModelCandidates(model);
+    const attemptedModels: { model: string; error?: string }[] = [];
+
+    for (const modelId of modelCandidates) {
+      try {
+        const result = await generateText({
+          model: google(modelId),
+          messages: [...formattedHistory, { role: "user", content: parts }] as any,
+          providerOptions: {
+            google: {
+              responseModalities: ["TEXT", "IMAGE"],
+            },
+          },
+        });
+
+        const image = extractImage(result);
+
+        if (!image) {
+          attemptedModels.push({ model: modelId, error: "No image returned" });
+          continue;
+        }
+
+        const fallback = modelId !== modelCandidates[0];
+
+        return NextResponse.json({
+          description: result.text ?? null,
+          image,
+          model: modelId,
+          fallback,
+          attemptedModels,
+          notice: fallback
+            ? {
+                type: "info",
+                message: `${modelCandidates[0]} did not return a usable image. The result was generated with ${modelId} instead.`,
+              }
+            : null,
+        });
+      } catch (err) {
+        attemptedModels.push({
+          model: modelId,
+          error: err instanceof Error ? err.message : "Unknown error",
+        });
       }
     }
 
-    return NextResponse.json({
-      description: textResponse ?? null,
-      image: imageBase64
-        ? `data:${mimeType};base64,${imageBase64}`
-        : null,
-    });
+    return NextResponse.json(
+      {
+        error: "Failed to generate image",
+        details: attemptedModels.map((item) => `${item.model}: ${item.error}`).join("\n"),
+        attemptedModels,
+      },
+      { status: 500 }
+    );
   } catch (err) {
     const message =
       err instanceof Error ? err.message : "Unknown error";

app/page.tsx

@@ -7,13 +7,22 @@ import { ImageIcon, Wand2, ExternalLink } from "lucide-react";
 import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card";
 import { HistoryItem } from "@/lib/types";
 
+const MODEL_OPTIONS = [
+  "gemini-2.0-flash-exp",
+  "gemini-2.5-flash-image",
+  "gemini-3.1-flash-image-preview",
+  "gemini-3-pro-image-preview",
+];
+
 export default function Home() {
   const [image, setImage] = useState<string | null>(null);
   const [generatedImage, setGeneratedImage] = useState<string | null>(null);
   const [description, setDescription] = useState<string | null>(null);
   const [loading, setLoading] = useState(false);
   const [error, setError] = useState<string | null>(null);
+  const [notice, setNotice] = useState<string | null>(null);
   const [history, setHistory] = useState<HistoryItem[]>([]);
+  const [selectedModel, setSelectedModel] = useState(MODEL_OPTIONS[0]);
 
   const handleImageSelect = (imageData: string) => {
     setImage(imageData || null);
@@ -23,6 +32,7 @@ export default function Home() {
     try {
       setLoading(true);
       setError(null);
+      setNotice(null);
 
       // If we have a generated image, use that for editing, otherwise use the uploaded image
       const imageToEdit = generatedImage || image;
@@ -32,6 +42,7 @@ export default function Home() {
         prompt,
         image: imageToEdit,
         history: history.length > 0 ? history : undefined,
+        model: selectedModel,
       };
 
       const response = await fetch("/api/image", {
@@ -48,6 +59,7 @@ export default function Home() {
       }
 
       const data = await response.json();
+      setNotice(data.notice?.message || null);
 
       if (data.image) {
         // Update the generated image and description
@@ -91,6 +103,7 @@ export default function Home() {
     setDescription(null);
     setLoading(false);
     setError(null);
+    setNotice(null);
     setHistory([]);
   };
 
@@ -128,6 +141,31 @@ export default function Home() {
             </div>
           )}
 
+          {notice && (
+            <div className="p-4 mb-4 text-sm text-blue-700 bg-blue-100 rounded-lg">
+              {notice}
+            </div>
+          )}
+
+          <div className="space-y-2">
+            <label className="text-sm font-medium text-foreground" htmlFor="model-select">
+              Model
+            </label>
+            <select
+              id="model-select"
+              value={selectedModel}
+              onChange={(event) => setSelectedModel(event.target.value)}
+              disabled={loading}
+              className="w-full rounded-md border border-input bg-background px-3 py-2 text-sm"
+            >
+              {MODEL_OPTIONS.map((model) => (
+                <option key={model} value={model}>
+                  {model}
+                </option>
+              ))}
+            </select>
+          </div>
+
           {!displayImage && !loading ? (
             <>
               <ImageUpload