update

app.py

@@ -2,121 +2,125 @@ import gradio as gr
 import requests
 import fal_client
 import os
-import json
-import time
 import base64
-import io
 from PIL import Image
+import io
 
 RUNPOD_API_KEY = os.getenv("RUNPOD_API_KEY")
 FAL_KEY = os.getenv("FAL_KEY")
-
-QWEN_ENDPOINT_ID = os.environ["ENDPOINTID"]
+QWEN_ENDPOINT_ID = "jzpm1xin5cprff"
 
 os.environ["FAL_KEY"] = FAL_KEY if FAL_KEY else ""
 
 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 get_as_base64(url):
-    """Helper to convert image URL to base64 for the Qwen endpoint"""
-    response = requests.get(url)
-    return base64.b64encode(response.content).decode('utf-8')
+def process_image_source(image_file, image_url):
+    """
+    Determines if we should use the uploaded file or the URL.
+    Returns base64 for Qwen and a URL (or data URI) for others.
+    """
+    if image_file is not None:
+        with open(image_file, "rb") as f:
+            file_bytes = f.read()
+        b64_str = base64.b64encode(file_bytes).decode('utf-8')
+        fal_url = fal_client.upload_file(image_file)
+        return b64_str, fal_url, f"data:image/png;base64,{b64_str}"
+    
+    elif image_url and image_url.strip():
+        response = requests.get(image_url)
+        if response.status_code == 200:
+            b64_str = base64.b64encode(response.content).decode('utf-8')
+            return b64_str, image_url, image_url
+    
+    return None, None, None
 
-def run_banana(image_url, prompt):
+def run_banana(image_source_url, prompt):
     url = "https://api.runpod.ai/v2/nano-banana-edit/runsync"
     headers = {"Content-Type": "application/json", "Authorization": f"Bearer {RUNPOD_API_KEY}"}
     payload = {
         "input": {
             "prompt": prompt,
-            "images": [image_url],
+            "images": [image_source_url],
             "enable_safety_checker": True
         }
     }
     try:
         response = requests.post(url, headers=headers, json=payload, timeout=90)
         return response.json()["output"][0] 
-    except Exception as e:
+    except:
         return None
 
-def run_qwen(image_url, prompt):
-    """Updated to use endpoint jzpm1xin5cprff with base64 logic"""
+def run_qwen(b64_img, prompt):
     url = f"https://api.runpod.ai/v2/{QWEN_ENDPOINT_ID}/runsync"
     headers = {"Content-Type": "application/json", "Authorization": f"Bearer {RUNPOD_API_KEY}"}
-    
-    try:
-        img_b64 = get_as_base64(image_url)
-        payload = {
-            "input": {
-                "image": img_b64,
-                "prompt": prompt,
-                "seed": 42,
-                "use_lightning": True,
-                "true_guidance_scale": 2.5,
-                "num_inference_steps": 4
-            }
+    payload = {
+        "input": {
+            "image": b64_img,
+            "prompt": prompt,
+            "seed": 42,
+            "use_lightning": True,
+            "true_guidance_scale": 2.5,
+            "num_inference_steps": 4
         }
+    }
+    try:
         response = requests.post(url, headers=headers, json=payload, timeout=90)
-        res_data = response.json()
-        
-        img_data = res_data["output"]["images"][0]
+        img_data = response.json()["output"]["images"][0]
         return f"data:image/png;base64,{img_data}"
-    except Exception as e:
-        print(f"Qwen Error: {e}")
+    except:
         return None
 
-def run_fal_flux(image_url, prompt):
+def run_fal_flux(fal_hosted_url, prompt):
     try:
         handler = fal_client.submit(
             "fal-ai/flux-2/edit",
-            arguments={
-                "prompt": prompt,
-                "image_urls": [image_url]
-            }
+            arguments={"prompt": prompt, "image_urls": [fal_hosted_url]}
         )
         result = handler.get()
         return result['images'][0]['url']
-    except Exception as e:
+    except:
         return None
 
-def compare_models(image_url, prompt):
-    if not image_url:
+def compare_models(image_file, image_url, prompt):
+    b64_img, fal_url, general_url = process_image_source(image_file, image_url)
+    
+    if not b64_img:
         return None, None, None
-    res_banana = run_banana(image_url, prompt)
-    res_qwen = run_qwen(image_url, prompt)
-    res_fal = run_fal_flux(image_url, prompt)
-    return res_banana, res_qwen, res_fal
 
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    res_qwen = run_qwen(b64_img, prompt)
+    res_fal = run_fal_flux(fal_url, prompt)
+    res_banana = run_banana(general_url, prompt)
+    
+    return res_qwen, res_fal, res_banana
+
+# UI
+with gr.Blocks() as demo:
     gr.Markdown("# 🛋️ Interior Design Model Arena")
-    gr.Markdown("Compare specialized staging models side-by-side.")
     
     with gr.Row():
         with gr.Column(scale=1):
-            input_url = gr.Textbox(label="Source Image URL", placeholder="Paste a link to your room image...")
-            input_prompt = gr.Textbox(label="Edit Prompt", value=DEFAULT_PROMPT, lines=5)
+            input_file = gr.Image(label="Upload Room Image", type="filepath")
+            input_url = gr.Textbox(label="OR: Source Image URL", placeholder="https://...")
+            input_prompt = gr.Textbox(label="Edit Prompt", value=DEFAULT_PROMPT, lines=4)
             run_btn = gr.Button("Generate Comparison", variant="primary")
             
     with gr.Row():
         with gr.Column():
-            gr.Markdown("### Qwen-Edit (New)")
-            gr.Markdown("**Price: $0.004 - $0.008 /req**")
+            gr.Markdown("### Qwen-Edit\n**Price: $0.004 - $0.008**")
             out_qwen = gr.Image(label="Qwen Result", interactive=False)
             
         with gr.Column():
-            gr.Markdown("### Flux-2 Edit")
-            gr.Markdown("**Price: $0.024 /req**")
+            gr.Markdown("### Flux-2 Edit\n**Price: $0.024**")
             out_fal = gr.Image(label="Flux Result", interactive=False)
 
         with gr.Column():
-            gr.Markdown("### Nano-Banana")
-            gr.Markdown("**Price: $0.039 /req**")
+            gr.Markdown("### Nano-Banana\n**Price: $0.039**")
             out_banana = gr.Image(label="Banana Result", interactive=False)
 
     run_btn.click(
         fn=compare_models,
-        inputs=[input_url, input_prompt],
-        outputs=[out_banana, out_qwen, out_fal]
+        inputs=[input_file, input_url, input_prompt],
+        outputs=[out_qwen, out_fal, out_banana]
     )
 
-if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+demo.launch(server_name="0.0.0.0", server_port=7860)