Add OmniRoute Integration for Image Editing

app.py

@@ -1,6 +1,10 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
 import spaces
+from openai import OpenAI
+import io
+import requests
+from PIL import Image
 
 @spaces.GPU
 def generate_image(prompt, hf_token, progress=gr.Progress(track_tqdm=True)):
@@ -18,7 +22,46 @@ def generate_image(prompt, hf_token, progress=gr.Progress(track_tqdm=True)):
     
     return image
 
-# UI
+def edit_image(image, prompt, omni_url, omni_key, omni_model, progress=gr.Progress(track_tqdm=True)):
+    if not omni_url or not omni_key:
+        raise gr.Error("Please provide both your OmniRoute URL and API Key!")
+    if not image:
+        raise gr.Error("Please upload an image to edit!")
+    if not prompt:
+        raise gr.Error("Please enter a prompt for editing!")
+        
+    client = OpenAI(
+        base_url=omni_url.strip(),
+        api_key=omni_key.strip()
+    )
+    
+    img_byte_arr = io.BytesIO()
+    image.save(img_byte_arr, format='PNG')
+    img_byte_arr = img_byte_arr.getvalue()
+    
+    try:
+        response = client.images.edit(
+            image=img_byte_arr,
+            prompt=prompt,
+            model=omni_model.strip()
+        )
+        
+        if response.data and len(response.data) > 0:
+            url = response.data[0].url
+            if url:
+                img_response = requests.get(url)
+                edited = Image.open(io.BytesIO(img_response.content))
+                return edited
+            elif response.data[0].b64_json:
+                import base64
+                img_bytes = base64.b64decode(response.data[0].b64_json)
+                edited = Image.open(io.BytesIO(img_bytes))
+                return edited
+                
+        raise gr.Error("OmniRoute returned an empty response.")
+    except Exception as e:
+        raise gr.Error(f"OmniRoute API Error: {str(e)}")
+
 custom_theme = gr.themes.Soft(
     primary_hue="blue",
     secondary_hue="indigo",
@@ -29,41 +72,80 @@ with gr.Blocks() as demo:
     gr.Markdown(
         """
         # ⚡ FLUX.1 Image Studio (Grok Quality)
-        Generate ultra-fast images from text using the real-time FLUX.1-schnell model via Hugging Face Serverless API.
-        *No local GPU Required! Generates in the cloud.*
+        Generate ultra-fast images from text using the real-time FLUX.1-schnell model via Hugging Face Serverless API. *No local GPU Required!*
         """
     )
     
-    with gr.Row():
-        with gr.Column(scale=1):
-            hf_token = gr.Textbox(
-                label="🔑 Hugging Face Access Token",
-                placeholder="hf_...",
-                type="password",
-                info="Paste your Hugging Face Token here"
+    with gr.Tabs():
+        with gr.Tab("✨ Generate New Image"):
+            with gr.Row():
+                with gr.Column(scale=1):
+                    hf_token = gr.Textbox(
+                        label="🔑 Hugging Face Access Token",
+                        placeholder="hf_...",
+                        type="password",
+                        info="Paste your Hugging Face Token here"
+                    )
+                    prompt = gr.Textbox(
+                        label="✨ Prompt", 
+                        lines=3, 
+                        placeholder="e.g. A futuristic cyberpunk city at night...", 
+                        autofocus=True
+                    )
+                    
+                    generate_btn = gr.Button("🎨 Generate Image", variant="primary", size="lg")
+                    
+                with gr.Column(scale=1):
+                    output_image = gr.Image(label="Result", type="pil", interactive=False)
+
+            generate_btn.click(
+                fn=generate_image,
+                inputs=[prompt, hf_token],
+                outputs=[output_image]
             )
-            prompt = gr.Textbox(
-                label="✨ Prompt", 
-                lines=3, 
-                placeholder="e.g. A futuristic cyberpunk city at night...", 
-                autofocus=True
+            prompt.submit(
+                fn=generate_image,
+                inputs=[prompt, hf_token],
+                outputs=[output_image]
             )
             
-            generate_btn = gr.Button("🎨 Generate Image", variant="primary", size="lg")
+        with gr.Tab("🖌️ Edit Existing Image (OmniRoute)"):
+            with gr.Row():
+                with gr.Column(scale=1):
+                    with gr.Accordion("⚙️ OmniRoute Configuration", open=True):
+                        omni_url = gr.Textbox(
+                            label="OmniRoute Base URL",
+                            value="http://localhost:20128/v1",
+                            info="The endpoint URL of your OmniRoute gateway"
+                        )
+                        omni_key = gr.Textbox(
+                            label="🔑 OmniRoute API Key",
+                            placeholder="Enter your key...",
+                            type="password"
+                        )
+                        omni_model = gr.Textbox(
+                            label="Model ID",
+                            value="black-forest-labs/FLUX.1-schnell",
+                            info="The underlying model to route to"
+                        )
+                        
+                    input_image = gr.Image(label="Image to Edit", type="pil")
+                    edit_prompt = gr.Textbox(
+                        label="✨ Edit Prompt", 
+                        lines=2, 
+                        placeholder="e.g. change the sky to red..."
+                    )
+                    
+                    edit_btn = gr.Button("🖌️ Edit Image", variant="primary", size="lg")
+                    
+                with gr.Column(scale=1):
+                    edit_output = gr.Image(label="Edited Result", type="pil", interactive=False)
             
-        with gr.Column(scale=1):
-            output_image = gr.Image(label="Result", type="pil", interactive=False)
-
-    generate_btn.click(
-        fn=generate_image,
-        inputs=[prompt, hf_token],
-        outputs=[output_image]
-    )
-    prompt.submit(
-        fn=generate_image,
-        inputs=[prompt, hf_token],
-        outputs=[output_image]
-    )
+            edit_btn.click(
+                fn=edit_image,
+                inputs=[input_image, edit_prompt, omni_url, omni_key, omni_model],
+                outputs=[edit_output]
+            )
 
 if __name__ == "__main__":
     demo.launch(theme=custom_theme)

requirements.txt

@@ -1,3 +1,5 @@
 gradio
 huggingface_hub
-spaces
+spaces
+openai
+pillow