Update app.py

app.py

@@ -2,72 +2,103 @@ import os
 import gradio as gr
 from gradio_client import Client, handle_file
 
-# Configuration
+# Configuration - ดึง Token จาก Secrets
 HF_TOKEN = os.getenv("HF_TOKEN")
 TARGET_SPACE = "selfit-camera/omni-image-editor"
 
 def remote_bridge(inputs, prompt, mode):
+    """ฟังก์ชันส่งคำสั่งไปที่ API ของ Omni 2.0"""
     try:
         client = Client(TARGET_SPACE, hf_token=HF_TOKEN)
+        
         if mode == "/multi_edit":
-            # กรองเอาเฉพาะช่องที่มีการอัปโหลดรูปจริง
+            # กรองรูปจากทั้ง 3 ช่อง (Base, Ref, Style)
             valid_files = [handle_file(img) for img in inputs if img is not None]
-            if not valid_files: return None, "Status: Please upload at least one image."
+            if not valid_files:
+                return None, "Status: Please upload images in the slots."
             result = client.predict(images=valid_files, prompt=prompt, api_name="/multi_edit")
+            
         elif mode == "/t2i":
             result = client.predict(prompt=prompt, api_name="/t2i")
+            
         else:
+            # Single Edit, Upscale, Watermark
+            if inputs is None: return None, "Status: No image provided."
             result = client.predict(image=handle_file(inputs), prompt=prompt, api_name=mode)
-        return result, f"✅ Success ({mode})"
+            
+        return result, f"✅ Success: {mode}"
     except Exception as e:
-        return None, f"❌ Error: {str(e)}"
+        return None, f"❌ API Error: {str(e)}"
 
+# สร้าง UI ด้วยความละเอียดรอบคอบ
 with gr.Blocks(css=".container { max-width: 1200px; margin: auto; }") as demo:
-    gr.HTML("<h1 style='text-align: center;'>🎨 Omni Editor 2.0 (Remote)</h1>")
+    gr.HTML("<h1 style='text-align: center; margin-bottom: 0;'>🎨 Omni Editor 2.0 (Full Remote)</h1>")
+    gr.HTML("<p style='text-align: center; opacity: 0.7;'>Native 8B MM-DiT Bridge | High-Fidelity Multi-Modal Results</p>")
     
     with gr.Tabs():
-        # --- TAB 1: Single Image ---
+        # --- TAB 1: Single Image Edit ---
         with gr.TabItem("🖼️ Single Image Edit"):
             with gr.Row():
                 with gr.Column(scale=5):
-                    s_in = gr.Image(type="filepath", label="Input Image")
-                    s_prompt = gr.Textbox(label="Prompt", lines=2)
-                    s_run = gr.Button("🎨 Execute", variant="primary")
+                    s_in = gr.Image(type="filepath", label="Input Image", height=420)
+                    s_prompt = gr.Textbox(label="Editing Instruction", placeholder="e.g. Add a red hat", lines=2)
+                    s_run = gr.Button("🎯 Execute Edit", variant="primary")
                 with gr.Column(scale=5):
-                    s_out = gr.Image(label="Result")
-                    s_status = gr.Textbox(label="Status", interactive=False)
+                    s_out = gr.Image(label="Output Result", height=420)
+                    s_status = gr.Textbox(label="System Status", interactive=False)
+                    s_use = gr.Button("🔄 Use as New Input")
             s_run.click(lambda i, p: remote_bridge(i, p, "/predict"), [s_in, s_prompt], [s_out, s_status])
+            s_use.click(fn=lambda x: x, inputs=s_out, outputs=s_in)
 
-        # --- TAB 2: Multi-Image (แก้ให้มี 3 ช่องแยกตามสเปก 8B MM-DiT) ---
+        # --- TAB 2: Multi-Image Edit (แยก 3 ช่องตามสเปก AMG) ---
         with gr.TabItem("🖼️🖼️ Multi-Image Edit"):
-            gr.Markdown("### 🚀 Multi-Modal Fusion (Adaptive Multi-Modal Gating)")
+            gr.Markdown("### 🚀 Multi-Modal Fusion (Unified Transformer Backbone)")
             with gr.Row():
-                with gr.Column(scale=5):
-                    with gr.Row(): # แบ่งช่องใส่รูปย่อยๆ
-                        m_in1 = gr.Image(type="filepath", label="Base Image (A)", height=200)
-                        m_in2 = gr.Image(type="filepath", label="Ref Image (B)", height=200)
-                        m_in3 = gr.Image(type="filepath", label="Style Image (C)", height=200)
-                    m_prompt = gr.Textbox(label="Fusion Instruction", placeholder="e.g., Apply style C to object in A using B as reference", lines=2)
-                    m_run = gr.Button("🎯 Execute Multi-Fusion", variant="primary")
-                with gr.Column(scale=5):
-                    m_out = gr.Image(label="Fused Result", height=450)
+                with gr.Column(scale=6):
+                    with gr.Row(): # แถวช่องใส่ภาพ 3 ใบแยกกัน
+                        m_in1 = gr.Image(type="filepath", label="Base Image (A)", height=180)
+                        m_in2 = gr.Image(type="filepath", label="Ref Image (B)", height=180)
+                        m_in3 = gr.Image(type="filepath", label="Style Image (C)", height=180)
+                    m_prompt = gr.Textbox(label="Fusion Prompt", placeholder="Describe how images should interact...", lines=2)
+                    m_run = gr.Button("🔗 Execute Multi-Fusion", variant="primary")
+                with gr.Column(scale=4):
+                    m_out = gr.Image(label="Fused Result", height=400)
                     m_status = gr.Textbox(label="Status", interactive=False)
-            
-            # ส่ง List ของทั้ง 3 ช่องไปที่ฟังก์ชัน
             m_run.click(
-                fn=lambda i1, i2, i3, p: remote_bridge([i1, i2, i3], p, "/multi_edit"), 
-                inputs=[m_in1, m_in2, m_in3, m_prompt], 
+                fn=lambda i1, i2, i3, p: remote_bridge([i1, i2, i3], p, "/multi_edit"),
+                inputs=[m_in1, m_in2, m_in3, m_prompt],
                 outputs=[m_out, m_status]
             )
 
-        # --- TAB 3: Text to Image / Tools --- (ส่วนที่เหลือคงเดิม)
+        # --- TAB 3: Text to Image ---
         with gr.TabItem("✨ Text to Image"):
             with gr.Row():
                 with gr.Column(scale=5):
-                    t_prompt = gr.Textbox(label="Description", lines=4)
-                    t_run = gr.Button("✨ Generate", variant="primary")
+                    t_prompt = gr.Textbox(label="Prompt", placeholder="What do you want to see?", lines=5)
+                    t_run = gr.Button("✨ Generate Image", variant="primary")
                 with gr.Column(scale=5):
-                    t_out = gr.Image(label="AI Result")
-            t_run.click(lambda p: remote_bridge(None, p, "/t2i"), [t_prompt], [t_out, s_status])
+                    t_out = gr.Image(label="Generated Result", height=400)
+                    t_status = gr.Textbox(label="Status", interactive=False)
+            t_run.click(lambda p: remote_bridge(None, p, "/t2i"), [t_prompt], [t_out, t_status])
+
+        # --- TAB 4: Tools (Upscale / Watermark) ---
+        with gr.TabItem("🔍 Tools"):
+            with gr.Row():
+                with gr.Column(scale=5):
+                    tool_mode = gr.Radio(["Upscale", "Remove Watermark"], label="Select Function", value="Upscale")
+                    tool_in = gr.Image(type="filepath", label="Input Image")
+                    tool_run = gr.Button("🚀 Process", variant="primary")
+                with gr.Column(scale=5):
+                    tool_out = gr.Image(label="Result")
+                    tool_status = gr.Textbox(label="Status", interactive=False)
+            
+            def tool_router(mode, img):
+                api_path = "/upscale" if mode == "Upscale" else "/remove_watermark"
+                return remote_bridge(img, "", api_path)
+            
+            tool_run.click(tool_router, [tool_mode, tool_in], [tool_out, tool_status])
+
+    gr.HTML("<hr><p style='text-align: center; font-size: 0.8em;'>Omni Creator 2.0 | FP8 + RoPE + AMG Optimization</p>")
 
-demo.queue().launch()
+# Launch
+demo.queue().launch(server_name="0.0.0.0")