Update app.py

app.py

@@ -1,20 +1,20 @@
 import gradio as gr
 import torch
-from diffusers import AutoPipelineForImage2Image
+from diffusers import StableDiffusionInstructPix2PixPipeline
 from PIL import Image
 
 # ---------------------------------------------------------------------------
-# Why SD-Turbo?
-# Hugging Face Free Spaces only have 2 vCPUs. Standard Image models take 
-# 3-5 minutes per image on CPU because they require 30-50 steps.
-# SD-Turbo only requires 1 to 3 steps! It is incredibly fast and perfect 
-# for a free deployment.
+# Instruct-Pix2Pix
+# This model is specifically trained to edit images while preserving 
+# their exact structure. It takes much longer on CPU (approx 2-3 mins),
+# but the text and layout will look significantly more accurate.
 # ---------------------------------------------------------------------------
 
-print("Loading SD-Turbo Model... (This may take a minute on boot)")
-pipe = AutoPipelineForImage2Image.from_pretrained(
-    "stabilityai/sd-turbo", 
-    torch_dtype=torch.float32 # Use float32 for CPU compatibility
+print("Loading Instruct-Pix2Pix Model... (This takes a moment)")
+pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(
+    "timbrooks/instruct-pix2pix", 
+    torch_dtype=torch.float32,
+    safety_checker=None
 )
 
 def process_image(init_image, prompt, strength, steps):
@@ -22,40 +22,36 @@ def process_image(init_image, prompt, strength, steps):
         return None
         
     print(f"Received request: '{prompt}'")
-    
-    # Resize image to SD-Turbo's preferred 512x512 resolution
-    # Maintaining aspect ratio by cropping or padding would be better, 
-    # but exact 512x512 prevents memory spikes on the free CPU tier.
     init_image = init_image.convert("RGB")
-    init_image = init_image.resize((512, 512))
     
-    # Run the pipeline
+    # Resize keeping aspect ratio to max 512 for CPU memory limits
+    init_image.thumbnail((512, 512))
+    
+    # InstructPix2Pix uses image_guidance_scale. 
+    # 1.5 strictly preserves original image structure.
     image = pipe(
         prompt=prompt, 
         image=init_image, 
-        num_inference_steps=int(steps), 
-        strength=float(strength), 
-        guidance_scale=0.0 # Turbo mathematically requires guidance_scale=0.0
+        num_inference_steps=20, # Higher quality, but takes 2 mins on free CPU
+        image_guidance_scale=1.5, 
+        guidance_scale=7.5
     ).images[0]
     
     return image
 
-# Define the Gradio Interface
 with gr.Blocks(theme=gr.themes.Monochrome()) as demo:
-    gr.Markdown("# 🪄 WiggleAgent // Free Img2Img Backend")
-    gr.Markdown("Powered by SD-Turbo (Optimized for Free CPU Tiers)")
+    gr.Markdown("# 🪄 WiggleAgent // Instruct-Pix2Pix Backend")
+    gr.Markdown("Powered by Instruct-Pix2Pix (High Structural Preservation)")
     
     with gr.Row():
         with gr.Column():
             input_image = gr.Image(type="pil", label="Input Image (Your Screenshot)")
-            prompt = gr.Textbox(label="Prompt", value="cyberpunk style, dark neon city, glowing interface")
-            
-            # Strength determines how much of the original image is preserved.
-            # 0.1 = Almost no change. 1.0 = Completely new image.
-            strength = gr.Slider(minimum=0.1, maximum=1.0, value=0.6, step=0.05, label="Transformation Strength")
+            prompt = gr.Textbox(label="Prompt", value="cyberpunk style, dark neon city")
             
-            # Steps determines quality. 2 is the sweet spot for SD-Turbo.
-            steps = gr.Slider(minimum=1, maximum=4, value=2, step=1, label="Inference Steps (Keep low for CPU)")
+            # These sliders exist so the Gradio Client in main.py doesn't crash,
+            # but the actual function hardcodes the best values for Instruct-Pix2Pix.
+            strength = gr.Slider(minimum=0.1, maximum=1.0, value=0.6, label="Ignored by Backend")
+            steps = gr.Slider(minimum=1, maximum=4, value=2, label="Ignored by Backend")
             
             btn = gr.Button("Generate", variant="primary")
         
@@ -66,7 +62,7 @@ with gr.Blocks(theme=gr.themes.Monochrome()) as demo:
         fn=process_image, 
         inputs=[input_image, prompt, strength, steps], 
         outputs=output_image,
-        api_name="predict" # Exposes this function to our gradio_client in WiggleAgent!
+        api_name="/predict"
     )
 
 demo.launch()