Update app.py

app.py

@@ -7,77 +7,74 @@ HF_TOKEN = os.getenv("HF_TOKEN")
 TARGET_SPACE = "selfit-camera/omni-image-editor"
 
 def remote_bridge(inputs, prompt, mode):
-    """ฟังก์ชันกลางสำหรับส่งคำสั่งไปที่ API ต้นทาง"""
     if not inputs and mode != "/t2i":
         return None, "Status: Please provide input files."
-    
     try:
         client = Client(TARGET_SPACE, hf_token=HF_TOKEN)
-        
         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")
         else:
-            # Single Edit (/predict), Upscale, Watermark
             result = client.predict(image=handle_file(inputs), prompt=prompt, api_name=mode)
-            
-        return result, f"Status: Processed via {mode} successfully."
+        return result, f"✅ Success ({mode})"
     except Exception as e:
-        return None, f"Remote Error: {str(e)}"
+        return None, f"❌ Error: {str(e)}"
 
-# สร้างหน้าจอ UI แบบแบ่งส่วนชัดเจน
-with gr.Blocks() as demo:
-    gr.HTML("<h1 style='text-align: center;'>🎨 Omni Editor 2.0 (Full Remote)</h1>")
+with gr.Blocks(css=".container { max-width: 1100px; margin: auto; }") as demo:
+    gr.HTML("<h1 style='text-align: center; margin-bottom: 0;'>🎨 Omni Editor 2.0 (Remote)</h1>")
+    gr.HTML("<p style='text-align: center; opacity: 0.8;'>Bridge Mode: Powered by 8B MM-DiT</p>")
     
     with gr.Tabs():
         # --- 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="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)
+                with gr.Column(scale=5):
+                    s_in = gr.Image(type="filepath", label="Input Image", height=400)
+                    s_prompt = gr.Textbox(label="Prompt", placeholder="Describe your edit...", lines=2)
+                    s_run = gr.Button("🎨 Execute Edit", variant="primary")
+                with gr.Column(scale=5):
+                    s_out = gr.Image(label="Result", height=400)
+                    s_status = gr.Textbox(label="System Status", interactive=False)
+                    s_use = gr.Button("🔄 Use Result 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 (แก้ให้รับได้หลายรูปตามสเปก 8B MM-DiT) ---
+        # --- TAB 2: Multi-Image (Layout ใหม่ ไม่มั่ว) ---
         with gr.TabItem("🖼️🖼️ Multi-Image Edit"):
-            gr.Markdown("### 🚀 Multi-Modal Fusion (Select multiple reference images)")
+            gr.Markdown("### 🚀 Multi-Modal Fusion (Select up to 3 images)")
             with gr.Row():
-                with gr.Column(scale=1):
-                    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")
+                with gr.Column(scale=5):
+                    m_in = gr.File(label="Upload Images (Base/Ref/Style)", file_count="multiple", file_types=["image"], height=200)
+                    m_prompt = gr.Textbox(label="Instruction", placeholder="How to fuse these images?", lines=2)
+                    m_run = gr.Button("🎯 Execute Multi-Edit", variant="primary")
+                with gr.Column(scale=5):
+                    m_out = gr.Image(label="Fused Result", height=400)
                     m_status = gr.Textbox(label="Status", interactive=False)
             m_run.click(lambda i, p: remote_bridge(i, p, "/multi_edit"), [m_in, m_prompt], [m_out, m_status])
 
         # --- TAB 3: Text to Image ---
         with gr.TabItem("✨ Text to Image"):
             with gr.Row():
-                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")
+                with gr.Column(scale=5):
+                    t_prompt = gr.Textbox(label="Description", placeholder="What do you want to create?", lines=5)
+                    t_run = gr.Button("✨ Generate Image", variant="primary")
+                with gr.Column(scale=5):
+                    t_out = gr.Image(label="AI 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 (Upscale/Watermark)"):
+        # --- TAB 4: Tools ---
+        with gr.TabItem("🔍 Tools"):
             with gr.Row():
-                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")
+                with gr.Column(scale=4):
+                    tool_mode = gr.Radio(["Upscale", "Watermark"], label="Mode", value="Upscale")
+                    tool_in = gr.Image(type="filepath", label="Source Image")
+                    tool_run = gr.Button("🚀 Process", variant="primary")
+                with gr.Column(scale=6):
+                    tool_out = gr.Image(label="Output")
                     tool_status = gr.Textbox(label="Status", interactive=False)
             
             def tool_router(mode, img):
@@ -86,8 +83,6 @@ with gr.Blocks() as demo:
             
             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>")
+    gr.HTML("<hr><p style='text-align: center; font-size: 0.8em;'>Omni Creator 2.0 | FP8 + RoPE + AMG Optimization</p>")
 
-# Launch settings
-demo.queue().launch(server_name="0.0.0.0")
+demo.queue().launch()