Update app.py

app.py

@@ -2,63 +2,88 @@ import os
 import gradio as gr
 from gradio_client import Client, handle_file
 
-# Configuration
+# Config
 HF_TOKEN = os.getenv("HF_TOKEN")
-TARGET_SPACE = "selfit-camera/omni-image-editor"
+TARGET = "selfit-camera/omni-image-editor"
 
-def process_via_api(image_path, prompt):
-    if not image_path:
-        return None, "Status: Please upload an image."
+def call_remote_api(image, prompt, mode_api_name):
+    if image is None and mode_api_name != "/t2i": # T2I ไม่ต้องใช้รูป
+        return None, "Error: Please provide necessary input."
     
     try:
-        # Connect to the target space
-        client = Client(TARGET_SPACE, hf_token=HF_TOKEN)
+        client = Client(TARGET, hf_token=HF_TOKEN)
         
-        # Call the prediction API
-        result = client.predict(
-            image=handle_file(image_path),
-            prompt=prompt,
-            api_name="/predict"
-        )
-        return result, "Status: Processed successfully via API"
+        # จัดการ input ตามแต่ละโหมด
+        if mode_api_name == "/t2i":
+            result = client.predict(prompt=prompt, api_name="/t2i")
+        elif mode_api_name == "/multi_edit":
+            # สำหรับ Multi-edit ต้นทางอาจต้องการ list ของไฟล์
+            result = client.predict(images=[handle_file(image)], prompt=prompt, api_name="/multi_edit")
+        else:
+            # โหมดทั่วไป: Single Edit, Upscale, Watermark
+            result = client.predict(image=handle_file(image), prompt=prompt, api_name=mode_api_name)
+            
+        return result, f"Status: Processed via {mode_api_name}"
     except Exception as e:
-        return None, f"Status Error: {str(e)}"
+        return None, f"Remote Error: {str(e)}"
 
-# UI Construction (Minimalist & Stable for Python 3.13)
+# UI Layout แบบไม่มั่ว (Row/Column สัดส่วนชัดเจน)
 with gr.Blocks() as demo:
-    gr.HTML("<h1 style='text-align: center;'>🎨 Omni Editor 2.0 (Remote)</h1>")
-    gr.HTML("<div style='text-align: center; background: #6366f1; color: white; padding: 10px; border-radius: 10px; margin-bottom: 20px;'>Powered by Omni-Image-Editor API | Running on Free Tier</div>")
+    gr.HTML("<h1 style='text-align: center;'>🎨 Omni Editor 2.0 (Full Remote)</h1>")
     
-    with gr.Tabs():
-        with gr.TabItem("🖼️ Image Editing"):
+    with gr.Tabs() as tabs:
+        # 1. Single Image Edit
+        with gr.TabItem("🖼️ Single Image Edit"):
             with gr.Row():
                 with gr.Column(scale=1):
-                    input_img = gr.Image(type="filepath", label="Input Image")
-                    input_prompt = gr.Textbox(
-                        label="Editing Instruction", 
-                        placeholder="e.g. 'add a red hat to the dog'",
-                        lines=3
-                    )
-                    submit_btn = gr.Button("Execute", variant="primary")
-                
+                    s_in = gr.Image(type="filepath", label="Input Image")
+                    s_prompt = gr.Textbox(label="Prompt", placeholder="e.g. Change background to Mars")
+                    s_run = gr.Button("Execute Edit", variant="primary")
                 with gr.Column(scale=1):
-                    output_img = gr.Image(label="Output Result")
-                    status_box = gr.Textbox(label="System Status", interactive=False)
-                    use_as_input = gr.Button("🔄 Use Result as New Input")
+                    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])
 
-    # Event Handlers
-    submit_btn.click(
-        fn=process_via_api,
-        inputs=[input_img, input_prompt],
-        outputs=[output_img, status_box]
-    )
-    
-    # Loop back output to input
-    use_as_input.click(
-        fn=lambda x: x,
-        inputs=output_img,
-        outputs=input_img
-    )
+        # 2. Multi-Image Edit
+        with gr.TabItem("🖼️🖼️ Multi-Image Edit"):
+            gr.Markdown("Select multiple images to edit (Handled by 8B MM-DiT)")
+            with gr.Row():
+                with gr.Column():
+                    m_in = gr.Image(type="filepath", label="Reference Image")
+                    m_prompt = gr.Textbox(label="Instruction")
+                    m_run = gr.Button("Execute Multi-Edit", variant="primary")
+                with gr.Column():
+                    m_out = gr.Image(label="Result")
+            m_run.click(lambda img, txt: call_remote_api(img, txt, "/multi_edit"), [m_in, m_prompt], [m_out, s_status])
+
+        # 3. Text to Image
+        with gr.TabItem("✨ Text to Image"):
+            with gr.Row():
+                with gr.Column():
+                    t_prompt = gr.Textbox(label="What do you want to create?")
+                    t_run = gr.Button("Generate Image", variant="primary")
+                with gr.Column():
+                    t_out = gr.Image(label="Generated Result")
+            t_run.click(lambda txt: call_remote_api(None, txt, "/t2i"), [t_prompt], [t_out, s_status])
+
+        # 4. Image Upscale & Watermark
+        with gr.TabItem("🔍 Tools (Upscale/Watermark)"):
+            tool_mode = gr.Radio(["Upscale", "Remove Watermark"], label="Select Tool", value="Upscale")
+            with gr.Row():
+                with gr.Column():
+                    tool_in = gr.Image(type="filepath", label="Input")
+                    tool_run = gr.Button("Process", variant="primary")
+                with gr.Column():
+                    tool_out = gr.Image(label="Processed Result")
+            
+            # เลือก API name ตาม Radio
+            def tool_handler(mode, img):
+                api_name = "/upscale" if mode == "Upscale" else "/remove_watermark"
+                return call_remote_api(img, "", api_name)
+            
+            tool_run.click(tool_handler, [tool_mode, tool_in], [tool_out, s_status])
+
+    gr.Markdown("---")
+    gr.Markdown("🤖 **Omni Creator 2.0:** 8B Unified Multi-Modal Diffusion Transformer")
 
-if __name__ == "__main__":
-    demo.queue(max_size=20).launch(server_name="0.0.0.0")
+demo.queue().launch(server_name="0.0.0.0")