update

app.py

@@ -5,49 +5,44 @@ import os
 import base64
 import io
 from PIL import Image
+from google import genai
+from google.genai.types import GenerateContentConfig, ImageConfig, Part
 
 RUNPOD_API_KEY = os.getenv("RUNPOD_API_KEY")
 FAL_KEY = os.getenv("FAL_KEY")
+GEMINI_API_KEY = os.getenv("GEMINI_VLM_KEY")
 QWEN_ENDPOINT_ID = "jzpm1xin5cprff"
 
 os.environ["FAL_KEY"] = FAL_KEY if FAL_KEY else ""
+gemini_client = genai.Client(api_key=GEMINI_API_KEY) if GEMINI_API_KEY else None
 
 DEFAULT_PROMPT = "Add furnishings and accessories to this room as an interior designer would do for a real estate staging. The generated image shall have the exact same dimensions as the original image and architectural details. Respect doorways and windows and make sure they are consistent with the source image and not blocked by furniture. Use cute accessories and with appropriate wall space, add smart simple graphic paintings. Use neutral colors with light colored accents to match the colors of the room. Give the area an attractive glow."
 
 def b64_to_pil(b64_str):
-    """Safely convert base64 string to PIL Image"""
-    try:
-        if not b64_str: return None
-        if "base64," in b64_str:
-            b64_str = b64_str.split("base64,")[1]
-        img_data = base64.b64decode(b64_str)
-        return Image.open(io.BytesIO(img_data))
-    except:
-        return None
+    if not b64_str: return None
+    if "base64," in b64_str:
+        b64_str = b64_str.split("base64,")[1]
+    return Image.open(io.BytesIO(base64.b64decode(b64_str)))
 
-def url_to_pil(url):
-    """Fetch URL and return PIL Image to ensure it displays in Gradio"""
-    try:
-        resp = requests.get(url, timeout=10)
-        return Image.open(io.BytesIO(resp.content))
-    except:
-        return None
+def bytes_to_pil(img_bytes):
+    return Image.open(io.BytesIO(img_bytes))
 
 def get_image_inputs(image_file, image_url):
-    """Pre-processes inputs for all models"""
+    """Returns (raw_bytes, raw_b64, fal_url)"""
     if image_file:
         with open(image_file, "rb") as f:
             raw_bytes = f.read()
         raw_b64 = base64.b64encode(raw_bytes).decode('utf-8')
         fal_url = fal_client.upload_file(image_file)
-        return raw_b64, fal_url
+        return raw_bytes, raw_b64, fal_url
     
     elif image_url:
         resp = requests.get(image_url)
-        raw_b64 = base64.b64encode(resp.content).decode('utf-8')
-        return raw_b64, image_url
+        raw_bytes = resp.content
+        raw_b64 = base64.b64encode(raw_bytes).decode('utf-8')
+        return raw_bytes, raw_b64, image_url
     
-    return None, None
+    return None, None, None
 
 def run_qwen(raw_b64, prompt):
     url = f"https://api.runpod.ai/v2/{QWEN_ENDPOINT_ID}/runsync"
@@ -63,7 +58,7 @@ def run_qwen(raw_b64, prompt):
         }
     }
     try:
-        response = requests.post(url, headers=headers, json=payload, timeout=120)
+        response = requests.post(url, headers=headers, json=payload, timeout=60)
         return b64_to_pil(response.json()["output"]["images"][0])
     except:
         return None
@@ -75,42 +70,49 @@ def run_fal_flux(image_url, prompt):
             arguments={"prompt": prompt, "image_urls": [image_url]}
         )
         result = handler.get()
-        return url_to_pil(result['images'][0]['url'])
+        resp = requests.get(result['images'][0]['url'])
+        return bytes_to_pil(resp.content)
     except:
         return None
 
-def run_banana(image_url, prompt):
-    url = "https://api.runpod.ai/v2/nano-banana-edit/runsync"
-    headers = {"Authorization": f"Bearer {RUNPOD_API_KEY}", "Content-Type": "application/json"}
-    payload = {"input": {"prompt": prompt, "images": [image_url], "enable_safety_checker": True}}
+def run_gemini(image_bytes, prompt):
+    if not gemini_client: return None
     try:
-        response = requests.post(url, json=payload, headers=headers, timeout=120)
-        res_data = response.json()
-        output_val = res_data.get("output", [])
-        img_src = output_val[0] if isinstance(output_val, list) else output_val
-        
-        if img_src.startswith("http"):
-            return url_to_pil(img_src)
-        return b64_to_pil(img_src)
-    except:
+        response = gemini_client.models.generate_content(
+            model="gemini-2.5-flash-image", 
+            contents=[
+                Part.from_bytes(data=image_bytes, mime_type="image/jpeg"),
+                prompt
+            ],
+            config=GenerateContentConfig(
+                response_modalities=["IMAGE"],
+                image_config=ImageConfig(aspect_ratio="4:3"), 
+                candidate_count=1,
+            ),
+        )
+        for part in response.candidates[0].content.parts:
+            if part.inline_data:
+                return bytes_to_pil(part.inline_data.data)
+    except Exception as e:
+        print(f"Gemini Error: {e}")
         return None
 
 def compare_all(image_file, image_url, prompt):
-    raw_b64, web_url = get_image_inputs(image_file, image_url)
-    if not raw_b64:
+    raw_bytes, raw_b64, web_url = get_image_inputs(image_file, image_url)
+    if not raw_bytes:
         yield None, None, None
         return
 
-    qwen_res, flux_res, banana_res = None, None, None
+    qwen_img, flux_img, gemini_img = None, None, None
 
-    qwen_res = run_qwen(raw_b64, prompt)
-    yield qwen_res, flux_res, banana_res
+    qwen_img = run_qwen(raw_b64, prompt)
+    yield qwen_img, flux_img, gemini_img
 
-    flux_res = run_fal_flux(web_url, prompt)
-    yield qwen_res, flux_res, banana_res
+    flux_img = run_fal_flux(web_url, prompt)
+    yield qwen_img, flux_img, gemini_img
 
-    banana_res = run_banana(web_url, prompt)
-    yield qwen_res, flux_res, banana_res
+    gemini_img = run_gemini(raw_bytes, prompt)
+    yield qwen_img, flux_img, gemini_img
 
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown("# 🛋️ Interior Design Model Arena")
@@ -120,7 +122,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             input_file = gr.Image(label="Upload Image", type="filepath")
             input_url = gr.Textbox(label="OR: Image URL")
             input_prompt = gr.Textbox(label="Prompt", value=DEFAULT_PROMPT, lines=4)
-            run_btn = gr.Button("Generate Side-by-Side", variant="primary")
+            run_btn = gr.Button("Generate Comparison", variant="primary")
             
     with gr.Row():
         with gr.Column():
@@ -132,13 +134,13 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             out_fal = gr.Image(label="Flux Result", type="pil")
 
         with gr.Column():
-            gr.Markdown("### Nano-Banana\n**$0.039 /req**")
-            out_banana = gr.Image(label="Banana Result", type="pil")
+            gr.Markdown("### Gemini 2.5 Flash\n**$0.039 /req**")
+            out_gemini = gr.Image(label="Gemini Result", type="pil")
 
     run_btn.click(
         fn=compare_all,
         inputs=[input_file, input_url, input_prompt],
-        outputs=[out_qwen, out_fal, out_banana]
+        outputs=[out_qwen, out_fal, out_gemini]
     )
 
 demo.launch(server_name="0.0.0.0", server_port=7860)

requirements.txt

@@ -2,4 +2,5 @@ gradio
 requests
 fal-client
 pillow
-runpod
+runpod
+google-genai