update app

Home.tsx

@@ -2,10 +2,9 @@
  * @license
  * SPDX-License-Identifier: Apache-2.0
  */
-/* tlint:disable */
-import {Content, GoogleGenAI, Modality} from '@google/genai';
+/* tslint:disable */
 import {
-  Eraser,
+  ChevronDown,
   Library,
   LoaderCircle,
   Paintbrush,
@@ -19,78 +18,110 @@ import {
 } from 'lucide-react';
 import {useEffect, useRef, useState} from 'react';
 
+// This function remains useful for parsing potential error messages
 function parseError(error: string) {
-  const regex = /{"error":(.*)}/gm;
-  const m = regex.exec(error);
   try {
-    const e = m[1];
-    const err = JSON.parse(e);
-    return err.message || error;
+    // Attempt to parse the error as a JSON object which the proxy might send
+    const errObj = JSON.parse(error);
+    return errObj.message || error;
   } catch (e) {
-    return error;
+    // If it's not JSON, return the original error string
+    const regex = /{"error":(.*)}/gm;
+    const m = regex.exec(error);
+    try {
+      const e = m[1];
+      const err = JSON.parse(e);
+      return err.message || error;
+    } catch (e) {
+      return error;
+    }
   }
 }
 
 export default function Home() {
-  const canvasRef = useRef(null);
-  const fileInputRef = useRef(null);
-  const backgroundImageRef = useRef(null);
+  const canvasRef = useRef<HTMLCanvasElement>(null);
+  const fileInputRef = useRef<HTMLInputElement>(null);
+  const backgroundImageRef = useRef<HTMLImageElement | null>(null);
+  const dropdownRef = useRef<HTMLDivElement>(null);
   const [isDrawing, setIsDrawing] = useState(false);
   const [prompt, setPrompt] = useState('');
   const [generatedImage, setGeneratedImage] = useState<string | null>(null);
-  const [multiImages, setMultiImages] = useState<string[]>([]);
+  const [multiImages, setMultiImages] = useState<
+    {url: string; type: string}[]
+  >([]);
   const [isLoading, setIsLoading] = useState(false);
   const [showErrorModal, setShowErrorModal] = useState(false);
   const [errorMessage, setErrorMessage] = useState('');
   const [mode, setMode] = useState<
     'canvas' | 'editor' | 'imageGen' | 'multi-img-edit'
   >('editor');
-
+  const [isDropdownOpen, setIsDropdownOpen] = useState(false);
   const [apiKey, setApiKey] = useState('');
   const [showApiKeyModal, setShowApiKeyModal] = useState(false);
-  const [tempApiKey, setTempApiKey] = useState('');
-  const submissionRef = useRef<(() => Promise<void>) | null>(null);
 
   // State for canvas history
   const [history, setHistory] = useState<string[]>([]);
   const [historyIndex, setHistoryIndex] = useState(-1);
-  const [drawingTool, setDrawingTool] = useState<'pen' | 'eraser'>('pen');
 
   // When switching to canvas mode, initialize it and its history
   useEffect(() => {
     if (mode === 'canvas' && canvasRef.current) {
       const canvas = canvasRef.current;
-      // Draw current background if it exists
-      if (backgroundImageRef.current) {
-        drawImageToCanvas();
+      const ctx = canvas.getContext('2d');
+      ctx.fillStyle = '#FFFFFF';
+      ctx.fillRect(0, 0, canvas.width, canvas.height);
+
+      // If an image already exists from another mode, draw it.
+      if (generatedImage) {
+        const img = new window.Image();
+        img.onload = () => {
+          ctx.drawImage(img, 0, 0, canvas.width, canvas.height);
+          // Save this as the initial state for this session
+          const dataUrl = canvas.toDataURL();
+          setHistory([dataUrl]);
+          setHistoryIndex(0);
+        };
+        img.src = generatedImage;
       } else {
-        initializeCanvas();
+        // Otherwise, save the blank state as initial
+        const dataUrl = canvas.toDataURL();
+        setHistory([dataUrl]);
+        setHistoryIndex(0);
       }
-      // Set this as the initial state for this canvas session
-      const dataUrl = canvas.toDataURL();
-      setHistory([dataUrl]);
-      setHistoryIndex(0);
-      setDrawingTool('pen'); // Reset tool
     }
-  }, [mode]);
+  }, [mode, generatedImage]);
 
   // Load background image when generatedImage changes
   useEffect(() => {
-    if (generatedImage) {
+    if (generatedImage && canvasRef.current) {
       const img = new window.Image();
       img.onload = () => {
         backgroundImageRef.current = img;
-        if (mode === 'canvas' && canvasRef.current) {
-          drawImageToCanvas();
-          // This new image becomes a new state in history
-          saveCanvasState();
+        drawImageToCanvas();
+        if (mode === 'canvas') {
+          // A small timeout to let the draw happen before saving
+          setTimeout(saveCanvasState, 50);
         }
       };
       img.src = generatedImage;
-    } else {
-      backgroundImageRef.current = null;
     }
-  }, [generatedImage]);
+  }, [generatedImage, mode]);
+
+  // Handle clicks outside the dropdown to close it
+  useEffect(() => {
+    function handleClickOutside(event: MouseEvent) {
+      if (
+        dropdownRef.current &&
+        !dropdownRef.current.contains(event.target as Node)
+      ) {
+        setIsDropdownOpen(false);
+      }
+    }
+    document.addEventListener('mousedown', handleClickOutside);
+    return () => {
+      document.removeEventListener('mousedown', handleClickOutside);
+    };
+  }, [dropdownRef]);
 
   // Initialize canvas with white background
   const initializeCanvas = () => {
@@ -104,6 +135,7 @@ export default function Home() {
   // Draw the background image to the canvas
   const drawImageToCanvas = () => {
     if (!canvasRef.current || !backgroundImageRef.current) return;
+
     const canvas = canvasRef.current;
     const ctx = canvas.getContext('2d');
     ctx.fillStyle = '#FFFFFF';
@@ -119,7 +151,7 @@ export default function Home() {
 
   // Canvas history functions
   const saveCanvasState = () => {
-    if (!canvasRef.current || mode !== 'canvas') return;
+    if (!canvasRef.current) return;
     const canvas = canvasRef.current;
     const dataUrl = canvas.toDataURL();
     const newHistory = history.slice(0, historyIndex + 1);
@@ -158,8 +190,8 @@ export default function Home() {
   };
 
   // Get the correct coordinates based on canvas scaling
-  const getCoordinates = (e) => {
-    const canvas = canvasRef.current;
+  const getCoordinates = (e: any) => {
+    const canvas = canvasRef.current!;
     const rect = canvas.getBoundingClientRect();
     const scaleX = canvas.width / rect.width;
     const scaleY = canvas.height / rect.height;
@@ -173,9 +205,9 @@ export default function Home() {
     };
   };
 
-  const startDrawing = (e) => {
-    const canvas = canvasRef.current;
-    const ctx = canvas.getContext('2d');
+  const startDrawing = (e: any) => {
+    const canvas = canvasRef.current!;
+    const ctx = canvas.getContext('2d')!;
     const {x, y} = getCoordinates(e);
     if (e.type === 'touchstart') {
       e.preventDefault();
@@ -185,24 +217,17 @@ export default function Home() {
     setIsDrawing(true);
   };
 
-  const draw = (e) => {
+  const draw = (e: any) => {
     if (!isDrawing) return;
     if (e.type === 'touchmove') {
       e.preventDefault();
     }
-    const canvas = canvasRef.current;
-    const ctx = canvas.getContext('2d');
+    const canvas = canvasRef.current!;
+    const ctx = canvas.getContext('2d')!;
     const {x, y} = getCoordinates(e);
-
+    ctx.lineWidth = 5;
     ctx.lineCap = 'round';
-    if (drawingTool === 'pen') {
-      ctx.strokeStyle = '#000000';
-      ctx.lineWidth = 5;
-    } else {
-      // Eraser
-      ctx.strokeStyle = '#FFFFFF';
-      ctx.lineWidth = 20;
-    }
+    ctx.strokeStyle = '#000000';
     ctx.lineTo(x, y);
     ctx.stroke();
   };
@@ -223,6 +248,7 @@ export default function Home() {
     setGeneratedImage(null);
     setMultiImages([]);
     backgroundImageRef.current = null;
+    setPrompt('');
   };
 
   const processFiles = (files: FileList | null) => {
@@ -232,11 +258,16 @@ export default function Home() {
     );
     if (fileArray.length === 0) return;
 
+    if (!apiKey) {
+      setShowApiKeyModal(true);
+    }
+
     if (mode === 'multi-img-edit') {
       const readers = fileArray.map((file) => {
-        return new Promise<string>((resolve, reject) => {
+        return new Promise<{url: string; type: string}>((resolve, reject) => {
           const reader = new FileReader();
-          reader.onload = () => resolve(reader.result as string);
+          reader.onload = () =>
+            resolve({url: reader.result as string, type: file.type});
           reader.onerror = reject;
           reader.readAsDataURL(file);
         });
@@ -245,16 +276,18 @@ export default function Home() {
         setMultiImages((prev) => [...prev, ...newImages]);
       });
     } else {
+      const file = fileArray[0];
       const reader = new FileReader();
       reader.onload = () => {
         setGeneratedImage(reader.result as string);
       };
-      reader.readAsDataURL(fileArray[0]);
+      reader.readAsDataURL(file);
     }
   };
 
   const handleFileChange = (e: React.ChangeEvent<HTMLInputElement>) => {
     processFiles(e.target.files);
+    e.target.value = '';
   };
 
   const handleDrop = (e: React.DragEvent<HTMLDivElement>) => {
@@ -281,145 +314,128 @@ export default function Home() {
     );
   };
 
-  const handleApiKeySubmit = async (e: React.FormEvent) => {
-    e.preventDefault();
-    if (tempApiKey) {
-      setErrorMessage('');
-      setApiKey(tempApiKey);
-      setShowApiKeyModal(false);
-      setTempApiKey('');
-      if (submissionRef.current) {
-        await submissionRef.current();
-        submissionRef.current = null;
-      }
-    }
-  };
-
-  const handleSubmit = async (e) => {
+  // *** MODIFIED FUNCTION ***
+  const handleSubmit = async (e: React.FormEvent) => {
     e.preventDefault();
 
-    if (mode === 'editor' && !generatedImage) {
-      setErrorMessage('Please upload an image to edit.');
-      setShowErrorModal(true);
+    if (!apiKey) {
+      setShowApiKeyModal(true);
       return;
     }
 
-    if (mode === 'multi-img-edit' && multiImages.length === 0) {
-      setErrorMessage('Please upload at least one image to edit.');
-      setShowErrorModal(true);
-      return;
-    }
+    setIsLoading(true);
 
-    const submitAction = async () => {
-      setIsLoading(true);
-      try {
-        const ai = new GoogleGenAI(apiKey);
-
-        if (mode === 'imageGen') {
-          const response = await ai.models.generateImages({
-            model: 'imagen-4.0-generate-001',
-            prompt: prompt,
-            config: {
-              numberOfImages: 1,
-            },
-          });
+    try {
+      if (mode === 'editor' && !generatedImage) {
+        setErrorMessage('Please upload an image to edit.');
+        setShowErrorModal(true);
+        return;
+      }
 
-          const base64ImageBytes: string =
-            response.generatedImages[0].image.imageBytes;
-          const imageUrl = `data:image/png;base64,${base64ImageBytes}`;
-          setGeneratedImage(imageUrl);
-        } else {
-          const parts: any[] = [];
-          if (mode === 'canvas') {
-            if (!canvasRef.current) return;
-            const canvas = canvasRef.current;
-            const tempCanvas = document.createElement('canvas');
-            tempCanvas.width = canvas.width;
-            tempCanvas.height = canvas.height;
-            const tempCtx = tempCanvas.getContext('2d');
-            tempCtx.fillStyle = '#FFFFFF';
-            tempCtx.fillRect(0, 0, tempCanvas.width, tempCanvas.height);
-            tempCtx.drawImage(canvas, 0, 0);
-            const imageB64 = tempCanvas.toDataURL('image/png').split(',')[1];
-            parts.push({inlineData: {data: imageB64, mimeType: 'image/png'}});
-          } else if (mode === 'editor') {
-            const imageB64 = generatedImage.split(',')[1];
-            parts.push({inlineData: {data: imageB64, mimeType: 'image/png'}});
-          } else if (mode === 'multi-img-edit') {
-            multiImages.forEach((img) => {
-              parts.push({
-                inlineData: {data: img.split(',')[1], mimeType: 'image/png'},
-              });
-            });
-          }
+      if (mode === 'multi-img-edit' && multiImages.length === 0) {
+        setErrorMessage('Please upload at least one image to edit.');
+        setShowErrorModal(true);
+        return;
+      }
 
-          parts.push({text: prompt});
-          const contents: Content[] = [{role: 'USER', parts}];
-          const response = await ai.models.generateContent({
-            model: 'gemini-2.5-flash-image-preview',
-            contents,
-            config: {
-              responseModalities: [Modality.TEXT, Modality.IMAGE],
-            },
+      const parts: any[] = [];
+
+      // This logic for building the 'parts' array is correct.
+      if (mode === 'imageGen') {
+        const tempCanvas = document.createElement('canvas');
+        tempCanvas.width = 960;
+        tempCanvas.height = 540;
+        const tempCtx = tempCanvas.getContext('2d')!;
+        tempCtx.fillStyle = '#FFFFFF';
+        tempCtx.fillRect(0, 0, tempCanvas.width, tempCanvas.height);
+        tempCtx.fillStyle = '#FEFEFE';
+        tempCtx.fillRect(0, 0, 1, 1);
+        const imageB64 = tempCanvas.toDataURL('image/png').split(',')[1];
+        parts.push({inlineData: {data: imageB64, mimeType: 'image/png'}});
+      } else if (mode === 'canvas') {
+        if (!canvasRef.current) return;
+        const canvas = canvasRef.current;
+        const imageB64 = canvas.toDataURL('image/png').split(',')[1];
+        parts.push({inlineData: {data: imageB64, mimeType: 'image/png'}});
+      } else if (mode === 'editor' && generatedImage) {
+        const mimeType = generatedImage.substring(
+          generatedImage.indexOf(':') + 1,
+          generatedImage.indexOf(';'),
+        );
+        const imageB64 = generatedImage.split(',')[1];
+        parts.push({inlineData: {data: imageB64, mimeType}});
+      } else if (mode === 'multi-img-edit') {
+        multiImages.forEach((img) => {
+          parts.push({
+            inlineData: {data: img.url.split(',')[1], mimeType: img.type},
           });
+        });
+      }
 
-          const data = {
-            success: true,
-            message: '',
-            imageData: null,
-            error: undefined,
-          };
-          for (const part of response.candidates[0].content.parts) {
-            if (part.text) {
-              data.message = part.text;
-            } else if (part.inlineData) {
-              data.imageData = part.inlineData.data;
-            }
-          }
+      parts.push({text: prompt});
+
+      // Construct the request body for the Gemini REST API
+      const requestBody = {
+        contents: [{role: 'USER', parts}],
+      };
 
-          if (data.imageData) {
-            const imageUrl = `data:image/png;base64,${data.imageData}`;
-            if (mode === 'multi-img-edit') {
-              setGeneratedImage(imageUrl);
-              setMultiImages([]);
-              setMode('editor');
-            } else {
-              setGeneratedImage(imageUrl);
-            }
-          } else {
-            setErrorMessage(
-              data.message || 'Failed to generate image. Please try again.',
-            );
-            setShowErrorModal(true);
+      // Define the proxy endpoint
+      const proxyUrl = `/api-proxy/v1beta/models/gemini-2.5-flash-image-preview:generateContent?key=${apiKey}`;
+
+      // Use fetch to send the request to your proxy server
+      const response = await fetch(proxyUrl, {
+        method: 'POST',
+        headers: {
+          'Content-Type': 'application/json',
+        },
+        body: JSON.stringify(requestBody),
+      });
+
+      if (!response.ok) {
+        const errorData = await response.json();
+        throw new Error(
+          errorData.error?.message || `HTTP error! status: ${response.status}`,
+        );
+      }
+
+      const responseData = await response.json();
+
+      // Process the response
+      const result = {message: '', imageData: null};
+
+      if (responseData.candidates && responseData.candidates.length > 0) {
+        for (const part of responseData.candidates[0].content.parts) {
+          if (part.text) {
+            result.message = part.text;
+          } else if (part.inlineData) {
+            result.imageData = part.inlineData.data;
           }
         }
-      } catch (error) {
-        console.error('Error submitting:', error);
-        const parsedError = parseError(error.message);
-        if (
-          parsedError &&
-          (parsedError.includes('API_KEY_INVALID') ||
-            parsedError.includes('API key not valid'))
-        ) {
-          setErrorMessage(
-            'Your API key is not valid. Please enter a valid key to continue.',
-          );
-          setApiKey('');
-          setShowApiKeyModal(true);
+      } else {
+        throw new Error('Invalid response structure from API.');
+      }
+
+      if (result.imageData) {
+        const imageUrl = `data:image/png;base64,${result.imageData}`;
+        if (mode === 'multi-img-edit') {
+          setGeneratedImage(imageUrl);
+          setMultiImages([]);
+          setMode('editor');
         } else {
-          setErrorMessage(parsedError || 'An unexpected error occurred.');
-          setShowErrorModal(true);
+          setGeneratedImage(imageUrl);
         }
-      } finally {
-        setIsLoading(false);
+      } else {
+        setErrorMessage(
+          result.message || 'Failed to generate image. Please try again.',
+        );
+        setShowErrorModal(true);
       }
-    };
-
-    if (!apiKey) {
-      submissionRef.current = submitAction;
-      setShowApiKeyModal(true);
-    } else {
-      await submitAction();
+    } catch (error: any) {
+      console.error('Error submitting:', error);
+      setErrorMessage(error.message || 'An unexpected error occurred.');
+      setShowErrorModal(true);
+    } finally {
+      setIsLoading(false);
     }
   };
 
@@ -427,20 +443,32 @@ export default function Home() {
     setShowErrorModal(false);
   };
 
+  const handleApiKeySubmit = (e: React.FormEvent) => {
+    e.preventDefault();
+    const newApiKey = (e.target as any).apiKey.value;
+    if (newApiKey) {
+      setApiKey(newApiKey);
+      setShowApiKeyModal(false);
+    }
+  };
+
   useEffect(() => {
     const canvas = canvasRef.current;
     if (!canvas) return;
-    const preventTouchDefault = (e) => {
+
+    const preventTouchDefault = (e: TouchEvent) => {
       if (isDrawing) {
         e.preventDefault();
       }
     };
+
     canvas.addEventListener('touchstart', preventTouchDefault, {
       passive: false,
     });
     canvas.addEventListener('touchmove', preventTouchDefault, {
       passive: false,
     });
+
     return () => {
       canvas.removeEventListener('touchstart', preventTouchDefault);
       canvas.removeEventListener('touchmove', preventTouchDefault);
@@ -448,15 +476,15 @@ export default function Home() {
   }, [isDrawing]);
 
   const baseDisplayClass =
-    'w-full sm:h-[60vh] h-[30vh] min-h-[320px] bg-white/90 touch-none flex items-center justify-center p-4 transition-colors';
+    'w-full sm:h-[60vh] h-[40vh] min-h-[320px] bg-white/90 touch-none flex items-center justify-center p-4 transition-colors';
 
   return (
     <>
-      <div className="min-h-screen notebook-paper-bg text-gray-900 flex flex-col justify-start items-center">
+      <div className="min-h-screen text-gray-900 flex flex-col justify-start items-center">
         <main className="container mx-auto px-3 sm:px-6 py-5 sm:py-10 pb-32 max-w-5xl w-full">
           <div className="flex flex-col sm:flex-row sm:justify-between sm:items-end mb-2 sm:mb-6 gap-2">
             <div>
-              <h1 className="text-2xl sm:text-3xl font-bold mb-0 leading-tight font-mega">
+              <h1 className="text-2xl sm:text-3xl font-bold mb-0 leading-tight">
                 Nano Banana AIO
               </h1>
               <p className="text-sm sm:text-base text-gray-500 mt-1">
@@ -471,7 +499,7 @@ export default function Home() {
                 by{' '}
                 <a
                   className="underline"
-                  href="https://www.linkedin.com/in/prithiv-sakthi/"
+                  href="https://huggingface.co/prithivMLmods"
                   target="_blank"
                   rel="noopener noreferrer">
                   prithivsakthi-ur
@@ -480,46 +508,85 @@ export default function Home() {
             </div>
 
             <menu className="flex items-center bg-gray-300 rounded-full p-2 shadow-sm self-start sm:self-auto">
-              <div className="flex items-center bg-gray-200/80 rounded-full p-1 mr-2">
-                <button
-                  onClick={() => setMode('editor')}
-                  className={`px-3 py-1.5 rounded-full text-sm font-medium flex items-center gap-2 transition-colors ${
-                    mode === 'editor'
-                      ? 'bg-white shadow'
-                      : 'text-gray-600 hover:bg-gray-300/50'
-                  }`}
-                  aria-pressed={mode === 'editor'}>
-                  <PictureInPicture className="w-4 h-4" /> Editor
-                </button>
-                <button
-                  onClick={() => setMode('multi-img-edit')}
-                  className={`px-3 py-1.5 rounded-full text-sm font-medium flex items-center gap-2 transition-colors ${
-                    mode === 'multi-img-edit'
-                      ? 'bg-white shadow'
-                      : 'text-gray-600 hover:bg-gray-300/50'
-                  }`}
-                  aria-pressed={mode === 'multi-img-edit'}>
-                  <Library className="w-4 h-4" /> Multi-Image
-                </button>
+              <div className="flex flex-wrap justify-center items-center bg-gray-200/80 rounded-full p-1 mr-2">
+                <div className="relative" ref={dropdownRef}>
+                  <button
+                    onClick={() => setIsDropdownOpen(!isDropdownOpen)}
+                    className={`p-2 sm:px-3 sm:py-1.5 rounded-full text-sm font-medium flex items-center gap-2 transition-colors ${
+                      mode === 'editor' || mode === 'multi-img-edit'
+                        ? 'bg-white shadow'
+                        : 'text-gray-600 hover:bg-gray-300/50'
+                    }`}
+                    aria-haspopup="true"
+                    aria-expanded={isDropdownOpen}>
+                    {mode === 'multi-img-edit' ? (
+                      <>
+                        <Library className="w-4 h-4" />
+                        <span className="hidden sm:inline">Multi-Image</span>
+                      </>
+                    ) : (
+                      <>
+                        <PictureInPicture className="w-4 h-4" />
+                        <span className="hidden sm:inline">Editor</span>
+                      </>
+                    )}
+                    <ChevronDown className="w-4 h-4 opacity-70" />
+                  </button>
+                  {isDropdownOpen && (
+                    <div className="absolute top-full mt-2 w-48 bg-white rounded-lg shadow-xl z-10 border border-gray-200 py-1">
+                      <button
+                        onClick={() => {
+                          setMode('editor');
+                          setIsDropdownOpen(false);
+                        }}
+                        className={`w-full text-left px-3 py-2 text-sm flex items-center gap-3 transition-colors ${
+                          mode === 'editor'
+                            ? 'bg-gray-100 text-gray-900'
+                            : 'text-gray-700 hover:bg-gray-50'
+                        }`}
+                        aria-pressed={mode === 'editor'}>
+                        <PictureInPicture className="w-4 h-4" />
+                        <span>Single Image Edit</span>
+                      </button>
+                      <button
+                        onClick={() => {
+                          setMode('multi-img-edit');
+                          setIsDropdownOpen(false);
+                        }}
+                        className={`w-full text-left px-3 py-2 text-sm flex items-center gap-3 transition-colors ${
+                          mode === 'multi-img-edit'
+                            ? 'bg-gray-100 text-gray-900'
+                            : 'text-gray-700 hover:bg-gray-50'
+                        }`}
+                        aria-pressed={mode === 'multi-img-edit'}>
+                        <Library className="w-4 h-4" />
+                        <span>Multi-Image Edit</span>
+                      </button>
+                    </div>
+                  )}
+                </div>
+
                 <button
                   onClick={() => setMode('canvas')}
-                  className={`px-3 py-1.5 rounded-full text-sm font-medium flex items-center gap-2 transition-colors ${
+                  className={`p-2 sm:px-3 sm:py-1.5 rounded-full text-sm font-medium flex items-center gap-2 transition-colors ${
                     mode === 'canvas'
                       ? 'bg-white shadow'
                       : 'text-gray-600 hover:bg-gray-300/50'
                   }`}
                   aria-pressed={mode === 'canvas'}>
-                  <Paintbrush className="w-4 h-4" /> Canvas
+                  <Paintbrush className="w-4 h-4" />
+                  <span className="hidden sm:inline">Canvas</span>
                 </button>
                 <button
                   onClick={() => setMode('imageGen')}
-                  className={`px-3 py-1.5 rounded-full text-sm font-medium flex items-center gap-2 transition-colors ${
+                  className={`p-2 sm:px-3 sm:py-1.5 rounded-full text-sm font-medium flex items-center gap-2 transition-colors ${
                     mode === 'imageGen'
                       ? 'bg-white shadow'
                       : 'text-gray-600 hover:bg-gray-300/50'
                   }`}
                   aria-pressed={mode === 'imageGen'}>
-                  <Sparkles className="w-4 h-4" /> Image Gen
+                  <Sparkles className="w-4 h-4" />
+                  <span className="hidden sm:inline">Image Gen</span>
                 </button>
               </div>
               <button
@@ -557,36 +624,12 @@ export default function Home() {
                   onTouchStart={startDrawing}
                   onTouchMove={draw}
                   onTouchEnd={stopDrawing}
-                  className="border-2 border-black w-full sm:h-[60vh] h-[30vh] min-h-[320px] bg-white/90 touch-none"
+                  className="border-2 border-black w-full sm:h-[60vh] h-[40vh] min-h-[320px] bg-white/90 touch-none"
                   style={{
                     cursor:
                       "url('data:image/svg+xml;utf8,<svg xmlns=\"http://www.w3.org/2000/svg\" width=\"24\" height=\"24\" viewBox=\"0 0 24 24\" fill=\"none\" stroke=\"%23FF0000\" stroke-width=\"2\" stroke-linecap=\"round\" stroke-linejoin=\"round\"><path d=\"M12 5v14M5 12h14\"/></svg>') 12 12, crosshair",
                   }}
                 />
-                <div className="absolute top-2 left-2 flex gap-2">
-                  <button
-                    onClick={() => setDrawingTool('pen')}
-                    className={`p-2 rounded-md shadow transition-colors ${
-                      drawingTool === 'pen'
-                        ? 'bg-blue-200'
-                        : 'bg-white hover:bg-gray-100'
-                    }`}
-                    aria-label="Pen"
-                    aria-pressed={drawingTool === 'pen'}>
-                    <Paintbrush className="w-5 h-5" />
-                  </button>
-                  <button
-                    onClick={() => setDrawingTool('eraser')}
-                    className={`p-2 rounded-md shadow transition-colors ${
-                      drawingTool === 'eraser'
-                        ? 'bg-blue-200'
-                        : 'bg-white hover:bg-gray-100'
-                    }`}
-                    aria-label="Eraser"
-                    aria-pressed={drawingTool === 'eraser'}>
-                    <Eraser className="w-5 h-5" />
-                  </button>
-                </div>
                 <div className="absolute top-2 right-2 flex gap-2">
                   <button
                     onClick={handleUndo}
@@ -645,7 +688,7 @@ export default function Home() {
                     {multiImages.map((image, index) => (
                       <div key={index} className="relative group aspect-square">
                         <img
-                          src={image}
+                          src={image.url}
                           alt={`upload preview ${index + 1}`}
                           className="w-full h-full object-cover rounded-md"
                         />
@@ -677,8 +720,9 @@ export default function Home() {
                 )}
               </div>
             ) : (
+              // Image Gen mode display
               <div
-                className={`${baseDisplayClass} border-2 ${
+                className={`relative ${baseDisplayClass} border-2 ${
                   generatedImage ? 'border-black' : 'border-gray-400'
                 }`}>
                 {generatedImage ? (
@@ -697,6 +741,7 @@ export default function Home() {
             )}
           </div>
 
+          {/* Input form */}
           <form onSubmit={handleSubmit} className="w-full">
             <div className="relative">
               <input
@@ -706,9 +751,11 @@ export default function Home() {
                 placeholder={
                   mode === 'imageGen'
                     ? 'Describe the image you want to create...'
+                    : mode === 'multi-img-edit'
+                    ? 'Describe how to edit or combine the images...'
                     : 'Add your change...'
                 }
-                className="w-full p-3 sm:p-4 pr-12 sm:pr-14 text-sm sm:text-base border-2 border-black bg-white text-gray-800 shadow-sm focus:ring-2 focus:ring-gray-200 focus:outline-none transition-all font-mono"
+                className="w-full p-3 sm:p-4 pr-12 sm:pr-14 text-sm sm:text-base border-2 border-black bg-white text-gray-800 shadow-sm focus:ring-2 focus:ring-gray-200 focus:outline-none transition-all"
                 required
               />
               <button
@@ -730,72 +777,58 @@ export default function Home() {
             </div>
           </form>
         </main>
-        {showApiKeyModal && (
+        {/* Error Modal */}
+        {showErrorModal && (
           <div className="fixed inset-0 bg-black/50 flex items-center justify-center z-50 p-4">
             <div className="bg-white rounded-lg shadow-xl max-w-md w-full p-6">
-              <form onSubmit={handleApiKeySubmit}>
-                <div className="flex justify-between items-start mb-4">
-                  <h3 className="text-xl font-bold text-gray-700">
-                    Add Gemini API Key
-                  </h3>
-                  <button
-                    type="button"
-                    onClick={() => {
-                      setShowApiKeyModal(false);
-                      setErrorMessage('');
-                      submissionRef.current = null;
-                    }}
-                    className="text-gray-400 hover:text-gray-500"
-                    aria-label="Close">
-                    <X className="w-5 h-5" />
-                  </button>
-                </div>
-                <p className="text-gray-600 mb-4 text-sm">
-                  Add the API key to process the request.{' '}
-                  <strong className="text-gray-800">
-                    The API key will be removed if the app page is refreshed or
-                    closed.
-                  </strong>
-                </p>
-                {errorMessage && (
-                  <p className="text-red-500 text-sm mb-2 font-medium">
-                    {errorMessage}
-                  </p>
-                )}
-                <input
-                  type="password"
-                  value={tempApiKey}
-                  onChange={(e) => setTempApiKey(e.target.value)}
-                  placeholder="Enter your Gemini API Key"
-                  className="w-full p-2 mb-4 border-2 border-gray-300 rounded focus:ring-2 focus:ring-gray-400 focus:outline-none transition-all"
-                  required
-                  aria-label="Gemini API Key"
-                />
+              <div className="flex justify-between items-start mb-4">
+                <h3 className="text-xl font-bold text-gray-700">
+                  Failed to generate
+                </h3>
                 <button
-                  type="submit"
-                  className="w-full p-2 bg-black text-white rounded hover:bg-gray-800 transition-colors">
-                  Submit and Process
+                  onClick={closeErrorModal}
+                  className="text-gray-400 hover:text-gray-500">
+                  <X className="w-5 h-5" />
                 </button>
-              </form>
+              </div>
+              <p className="font-medium text-gray-600">
+                {parseError(errorMessage)}
+              </p>
             </div>
           </div>
         )}
-        {showErrorModal && (
+        {/* API Key Modal */}
+        {showApiKeyModal && (
           <div className="fixed inset-0 bg-black/50 flex items-center justify-center z-50 p-4">
             <div className="bg-white rounded-lg shadow-xl max-w-md w-full p-6">
               <div className="flex justify-between items-start mb-4">
                 <h3 className="text-xl font-bold text-gray-700">
-                  Failed to generate
+                  Add Gemini API Key
                 </h3>
                 <button
-                  onClick={closeErrorModal}
+                  onClick={() => setShowApiKeyModal(false)}
                   className="text-gray-400 hover:text-gray-500">
                   <X className="w-5 h-5" />
                 </button>
               </div>
-              <p className="font-medium text-gray-600">
-                {parseError(errorMessage)}
+              <p className="text-gray-600 mb-4">
+                Add the API key to process the request. The API key will be
+                removed if the app page is refreshed or closed.
               </p>
+              <form onSubmit={handleApiKeySubmit}>
+                <input
+                  type="password"
+                  name="apiKey"
+                  className="w-full p-2 border-2 border-gray-300 rounded-md mb-4"
+                  placeholder="Enter your Gemini API Key"
+                  required
+                />
+                <button
+                  type="submit"
+                  className="w-full bg-black text-white p-2 rounded-md hover:bg-gray-800 transition-colors">
+                  Submit
+                </button>
+              </form>
             </div>
           </div>
         )}