Update app.py

app.py

@@ -2,46 +2,71 @@ import os
 import gradio as gr
 from gradio_client import Client, handle_file
 
-# Fetch Token from Space Secrets (Name it HF_TOKEN in Settings)
+# Configuration
 HF_TOKEN = os.getenv("HF_TOKEN")
 TARGET_SPACE = "selfit-camera/omni-image-editor"
 
-def bridge_api(image_path, prompt):
-    if not image_path:
-        return None
+def process_api(image, prompt, mode):
+    if not image:
+        return None, "Please upload an image."
     
     try:
-        # Direct connection to the target API
         client = Client(TARGET_SPACE, hf_token=HF_TOKEN)
-        
-        # API prediction call
-        # Note: Ensure parameter names match target space's client.view_api()
+        # Mapping api_name based on UI tabs if necessary
+        # Defaulting to /predict as discussed
         result = client.predict(
-            image=handle_file(image_path),
+            image=handle_file(image),
             prompt=prompt,
             api_name="/predict"
         )
-        return result
+        return result, "Success"
     except Exception as e:
-        return f"API Bridge Error: {str(e)}"
+        return None, f"Error: {str(e)}"
 
-# UI Layout
-with gr.Blocks(title="Omni Editor Proxy") as demo:
-    gr.Markdown("## 🚀 Omni Image Editor (Proxy Mode)")
-    gr.Markdown("Processing is handled by the backend API and returned here.")
+# UI Layout mimicking Omni Editor 2.0
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.HTML("<h1 style='text-align: center;'>🎨 Omni Editor 2.0</h1>")
+    gr.HTML("<p style='text-align: center; background: #6366f1; color: white; padding: 10px; border-radius: 20px;'>Please give us a ❤️ if you find it helpful. Free trials refreshed every few days.</p>")
     
-    with gr.Row():
-        with gr.Column():
-            input_img = gr.Image(type="filepath", label="Input Image")
-            input_prompt = gr.Textbox(label="Edit Prompt", placeholder="e.g., 'change background to sunset'")
-            submit_btn = gr.Button("Process", variant="primary")
-        with gr.Column():
-            output_img = gr.Image(label="Output")
+    with gr.Tabs() as tabs:
+        with gr.TabItem("🖼️ Single Image Edit", id=0):
+            with gr.Row():
+                with gr.Column():
+                    gr.Markdown("### 📥 Upload Image")
+                    input_img = gr.Image(type="filepath", label="Select image to edit")
+                    input_prompt = gr.Textbox(label="Instruction", placeholder="Describe what to edit...")
+                    submit_btn = gr.Button("Execute", variant="primary")
+                    
+                with gr.Column():
+                    gr.Markdown("### 🎯 Editing Result")
+                    output_img = gr.Image(label="Result")
+                    use_as_input = gr.Button("🔄 Use as Input")
+                    
+                    gr.Markdown("### 📋 Processing Status")
+                    status_box = gr.Textbox(label="", interactive=False)
 
+        with gr.TabItem("🖼️🖼️ Multi-Image Edit", id=1):
+            gr.Markdown("Feature coming soon...")
+        with gr.TabItem("✨ Text to Image", id=2):
+            gr.Markdown("Feature coming soon...")
+        with gr.TabItem("🔍 Image Upscale", id=3):
+            gr.Markdown("Feature coming soon...")
+        with gr.TabItem("🏷️ Remove Watermark", id=4):
+            gr.Markdown("Feature coming soon...")
+
+    # Event Logic
     submit_btn.click(
-        fn=bridge_api,
+        fn=process_api,
         inputs=[input_img, input_prompt],
-        outputs=output_img
+        outputs=[output_img, status_box]
+    )
+    
+    # Logic for 'Use as Input'
+    use_as_input.click(
+        fn=lambda x: x,
+        inputs=output_img,
+        outputs=input_img
     )
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()