Update app.py

app.py

@@ -2,82 +2,92 @@ import os
 import gradio as gr
 from gradio_client import Client, handle_file
 
-# Config
+# Configuration
 HF_TOKEN = os.getenv("HF_TOKEN")
-TARGET = "selfit-camera/omni-image-editor"
+TARGET_SPACE = "selfit-camera/omni-image-editor"
 
-def call_remote_api(images, prompt, mode_api_name):
+def remote_bridge(inputs, prompt, mode):
+    """ฟังก์ชันกลางสำหรับส่งคำสั่งไปที่ API ต้นทาง"""
+    if not inputs and mode != "/t2i":
+        return None, "Status: Please provide input files."
+    
     try:
-        client = Client(TARGET, hf_token=HF_TOKEN)
+        client = Client(TARGET_SPACE, hf_token=HF_TOKEN)
         
-        if mode_api_name == "/t2i":
+        if mode == "/multi_edit":
+            # ส่งไฟล์หลายใบ (List of files) สำหรับ Multi-modal fusion
+            file_paths = [handle_file(f.name) for f in inputs]
+            result = client.predict(images=file_paths, prompt=prompt, api_name="/multi_edit")
+        elif mode == "/t2i":
             result = client.predict(prompt=prompt, api_name="/t2i")
-        elif mode_api_name == "/multi_edit":
-            # ส่ง List ของไฟล์ภาพทั้งหมด (รองรับสูงสุด 3 ภาพตามสเปก 8B MM-DiT)
-            files = [handle_file(img.name if hasattr(img, 'name') else img) for img in images]
-            result = client.predict(images=files, prompt=prompt, api_name="/multi_edit")
         else:
-            # Single Edit, Upscale, Watermark
-            result = client.predict(image=handle_file(images), prompt=prompt, api_name=mode_api_name)
+            # Single Edit (/predict), Upscale, Watermark
+            result = client.predict(image=handle_file(inputs), prompt=prompt, api_name=mode)
             
-        return result, f"Status: Success ({mode_api_name})"
+        return result, f"Status: Processed via {mode} successfully."
     except Exception as e:
         return None, f"Remote Error: {str(e)}"
 
+# สร้างหน้าจอ UI แบบแบ่งส่วนชัดเจน
 with gr.Blocks() as demo:
-    gr.HTML("<h1 style='text-align: center;'>🎨 Omni Editor 2.0 (Remote)</h1>")
+    gr.HTML("<h1 style='text-align: center;'>🎨 Omni Editor 2.0 (Full Remote)</h1>")
     
     with gr.Tabs():
-        # 1. Single Image Edit
+        # --- TAB 1: Single Image ---
         with gr.TabItem("🖼️ Single Image Edit"):
             with gr.Row():
                 with gr.Column(scale=1):
                     s_in = gr.Image(type="filepath", label="Input Image")
-                    s_prompt = gr.Textbox(label="Prompt", placeholder="Describe your edit...")
+                    s_prompt = gr.Textbox(label="Prompt", placeholder="e.g., Change background to snow")
                     s_run = gr.Button("Execute Edit", variant="primary")
                 with gr.Column(scale=1):
                     s_out = gr.Image(label="Result")
                     s_status = gr.Textbox(label="Status", interactive=False)
-            s_run.click(lambda img, txt: call_remote_api(img, txt, "/predict"), [s_in, s_prompt], [s_out, s_status])
+            s_run.click(lambda i, p: remote_bridge(i, p, "/predict"), [s_in, s_prompt], [s_out, s_status])
 
-        # 2. Multi-Image Edit (แก้ไขจุดที่มักง่ายแล้ว: รับได้หลายรูป)
+        # --- TAB 2: Multi-Image (แก้ให้รับได้หลายรูปตามสเปก 8B MM-DiT) ---
         with gr.TabItem("🖼️🖼️ Multi-Image Edit"):
-            gr.Markdown("### 🚀 Multi-Modal Fusion (Select up to 3 images)")
+            gr.Markdown("### 🚀 Multi-Modal Fusion (Select multiple reference images)")
             with gr.Row():
                 with gr.Column(scale=1):
-                    # ใช้ File component แบบ multiple เพื่อให้เลือกได้หลายรูปพร้อมกัน
-                    m_in = gr.File(label="Upload Images (Base/Ref/Style)", file_count="multiple", file_types=["image"])
-                    m_prompt = gr.Textbox(label="Instruction", placeholder="How should these images interact?")
+                    m_in = gr.File(label="Upload Images (Original + References)", file_count="multiple", file_types=["image"])
+                    m_prompt = gr.Textbox(label="Instruction", placeholder="How to fuse these images?")
                     m_run = gr.Button("Execute Multi-Edit", variant="primary")
                 with gr.Column(scale=1):
                     m_out = gr.Image(label="Result")
                     m_status = gr.Textbox(label="Status", interactive=False)
-            m_run.click(call_remote_api, [m_in, m_prompt, gr.State("/multi_edit")], [m_out, m_status])
+            m_run.click(lambda i, p: remote_bridge(i, p, "/multi_edit"), [m_in, m_prompt], [m_out, m_status])
 
-        # 3. Text to Image
+        # --- TAB 3: Text to Image ---
         with gr.TabItem("✨ Text to Image"):
             with gr.Row():
-                with gr.Column():
-                    t_prompt = gr.Textbox(label="Describe the image you want")
-                    t_run = gr.Button("Generate", variant="primary")
-                with gr.Column():
+                with gr.Column(scale=1):
+                    t_prompt = gr.Textbox(label="Prompt", placeholder="What do you want to generate?")
+                    t_run = gr.Button("Generate Image", variant="primary")
+                with gr.Column(scale=1):
                     t_out = gr.Image(label="Result")
-            t_run.click(lambda txt: call_remote_api(None, txt, "/t2i"), [t_prompt], [t_out, s_status])
+                    t_status = gr.Textbox(label="Status", interactive=False)
+            t_run.click(lambda p: remote_bridge(None, p, "/t2i"), [t_prompt], [t_out, t_status])
 
-        # 4. Image Upscale & Watermark
-        with gr.TabItem("🔍 Tools"):
-            tool_mode = gr.Radio(["Upscale", "Remove Watermark"], label="Select Tool", value="Upscale")
+        # --- TAB 4: Tools (Upscale/Watermark) ---
+        with gr.TabItem("🔍 Tools (Upscale/Watermark)"):
             with gr.Row():
-                with gr.Column():
-                    tool_in = gr.Image(type="filepath", label="Input")
-                    tool_run = gr.Button("Process", variant="primary")
-                with gr.Column():
+                with gr.Column(scale=1):
+                    tool_mode = gr.Radio(["Upscale", "/remove_watermark"], label="Select Tool", value="Upscale")
+                    tool_in = gr.Image(type="filepath", label="Input Image")
+                    tool_run = gr.Button("Process Tool", variant="primary")
+                with gr.Column(scale=1):
                     tool_out = gr.Image(label="Result")
+                    tool_status = gr.Textbox(label="Status", interactive=False)
             
-            def tool_handler(mode, img):
-                api_name = "/upscale" if mode == "Upscale" else "/remove_watermark"
-                return call_remote_api(img, "", api_name)
+            def tool_router(mode, img):
+                api_path = "/upscale" if mode == "Upscale" else "/remove_watermark"
+                return remote_bridge(img, "", api_path)
             
-            tool_run.click(tool_handler, [tool_mode, tool_in], [tool_out, s_status])
+            tool_run.click(tool_router, [tool_mode, tool_in], [tool_out, tool_status])
+
+    gr.Markdown("---")
+    gr.HTML("<p style='text-align: center;'>Omni Creator 2.0: 8B Unified Multi-Modal Diffusion Transformer</p>")
 
-demo.queue().launch()
+# Launch settings
+demo.queue().launch(server_name="0.0.0.0")