Update app.py

app.py

@@ -1,9 +1,11 @@
 import torch
 import spaces
 import gradio as gr
-from diffusers import DiffusionPipeline, FluxImg2ImgPipeline
+from diffusers import DiffusionPipeline, FluxImg2ImgPipeline, StableDiffusionInstructPix2PixPipeline
+from PIL import Image
 
 print("Loading pipelines...")
+
 # Text to image pipeline
 pipe_t2i = DiffusionPipeline.from_pretrained(
     "Tongyi-MAI/Z-Image-Turbo",
@@ -12,15 +14,34 @@ pipe_t2i = DiffusionPipeline.from_pretrained(
 )
 pipe_t2i.to("cuda")
 
-# Image to image pipeline
-pipe_i2i = FluxImg2ImgPipeline.from_pretrained(
+# Image to image - FLUX
+pipe_flux = FluxImg2ImgPipeline.from_pretrained(
     "black-forest-labs/FLUX.1-schnell",
     torch_dtype=torch.bfloat16,
 )
-pipe_i2i.to("cuda")
+pipe_flux.to("cuda")
+
+# Image to image - InstructPix2Pix
+pipe_ip2p = StableDiffusionInstructPix2PixPipeline.from_pretrained(
+    "timbrooks/instruct-pix2pix",
+    torch_dtype=torch.float16,
+    safety_checker=None,
+)
+pipe_ip2p.to("cuda")
 
 print("Pipelines loaded!")
 
+def resize_image(image, max_size=1024):
+    orig_w, orig_h = image.size
+    scale = min(max_size / orig_w, max_size / orig_h)
+    if scale < 1:
+        new_w = round(orig_w * scale / 64) * 64
+        new_h = round(orig_h * scale / 64) * 64
+    else:
+        new_w = round(orig_w / 64) * 64
+        new_h = round(orig_h / 64) * 64
+    return image.resize((new_w, new_h))
+
 @spaces.GPU
 def generate_t2i(prompt, height, width, num_inference_steps, seed, randomize_seed, progress=gr.Progress(track_tqdm=True)):
     if randomize_seed:
@@ -37,32 +58,35 @@ def generate_t2i(prompt, height, width, num_inference_steps, seed, randomize_see
     return image, seed
 
 @spaces.GPU
-def generate_i2i(input_image, prompt, strength, num_inference_steps, seed, randomize_seed, progress=gr.Progress(track_tqdm=True)):
+def generate_i2i(model_choice, input_image, prompt, strength, num_inference_steps, seed, randomize_seed, progress=gr.Progress(track_tqdm=True)):
+    if input_image is None:
+        raise gr.Error("Please upload an image first.")
     if randomize_seed:
         seed = torch.randint(0, 2**32 - 1, (1,)).item()
     generator = torch.Generator("cuda").manual_seed(int(seed))
-    
-    # Preserve original aspect ratio
-    orig_w, orig_h = input_image.size
-    # Round to nearest multiple of 64
-    new_w = round(orig_w / 64) * 64
-    new_h = round(orig_h / 64) * 64
-    # Cap at 1024 on longest side
-    scale = min(1024 / new_w, 1024 / new_h)
-    if scale < 1:
-        new_w = round(new_w * scale / 64) * 64
-        new_h = round(new_h * scale / 64) * 64
-    input_image = input_image.resize((new_w, new_h))
-    
-    image = pipe_i2i(
-        prompt=prompt,
-        image=input_image,
-        strength=float(strength),
-        num_inference_steps=int(num_inference_steps),
-        generator=generator,
-        width=new_w,
-        height=new_h,
-    ).images[0]
+    input_image = resize_image(input_image)
+
+    if model_choice == "FLUX.1-schnell (Creative, high change)":
+        image = pipe_flux(
+            prompt=prompt,
+            image=input_image,
+            strength=float(strength),
+            num_inference_steps=int(num_inference_steps),
+            generator=generator,
+            width=input_image.width,
+            height=input_image.height,
+        ).images[0]
+
+    elif model_choice == "InstructPix2Pix (Precise, preserves identity)":
+        image = pipe_ip2p(
+            prompt=prompt,
+            image=input_image,
+            num_inference_steps=int(num_inference_steps),
+            image_guidance_scale=1.5,
+            guidance_scale=7.5,
+            generator=generator,
+        ).images[0]
+
     return image, seed
 
 examples_t2i = [
@@ -98,16 +122,14 @@ with gr.Blocks(fill_height=True) as demo:
 
     with gr.Tabs():
 
-        # ── Tab 1: Text to Image ──────────────────────────────────────
+        # ── Tab 1: Text to Image ──────────────────────────────────
         with gr.Tab("✨ Text to Image"):
             with gr.Row(equal_height=False):
                 with gr.Column(scale=1, min_width=320):
                     t2i_prompt = gr.Textbox(
                         label="✨ Your Prompt",
                         placeholder="Describe the image you want to create...",
-                        lines=5,
-                        max_lines=10,
-                        autofocus=True,
+                        lines=5, max_lines=10, autofocus=True,
                     )
                     with gr.Accordion("⚙️ Advanced Settings", open=False):
                         with gr.Row():
@@ -131,19 +153,33 @@ with gr.Blocks(fill_height=True) as demo:
             t2i_btn.click(generate_t2i, [t2i_prompt, t2i_height, t2i_width, t2i_steps, t2i_seed, t2i_randomize], [t2i_output, t2i_used_seed])
             t2i_prompt.submit(generate_t2i, [t2i_prompt, t2i_height, t2i_width, t2i_steps, t2i_seed, t2i_randomize], [t2i_output, t2i_used_seed])
 
-        # ── Tab 2: Image to Image ─────────────────────────────────────
+        # ── Tab 2: Image to Image ─────────────────────────────────
         with gr.Tab("🖼️ Image to Image"):
             with gr.Row(equal_height=False):
                 with gr.Column(scale=1, min_width=320):
+
+                    model_choice = gr.Radio(
+                        choices=[
+                            "FLUX.1-schnell (Creative, high change)",
+                            "InstructPix2Pix (Precise, preserves identity)",
+                        ],
+                        value="InstructPix2Pix (Precise, preserves identity)",
+                        label="🤖 Model",
+                        info="InstructPix2Pix is better for targeted edits like changing colours or styles while keeping the person. FLUX is better for creative transformations."
+                    )
+
                     i2i_input = gr.Image(label="Upload Image", type="pil")
                     i2i_prompt = gr.Textbox(
                         label="✨ Edit Instruction",
-                        placeholder="Describe how you want to edit the image...",
+                        placeholder="e.g. 'change the dress to blue' or 'make it a sunset'",
                         lines=4,
                     )
+
                     with gr.Accordion("⚙️ Advanced Settings", open=False):
-                        i2i_strength = gr.Slider(0.1, 1.0, value=0.85, step=0.05, label="Strength", info="Higher = more change")
-                        i2i_steps = gr.Slider(1, 8, value=4, step=1, label="Inference Steps")
+                        i2i_strength = gr.Slider(0.1, 1.0, value=0.75, step=0.05,
+                            label="Strength (FLUX only)",
+                            info="Higher = more change. Not used by InstructPix2Pix.")
+                        i2i_steps = gr.Slider(1, 50, value=20, step=1, label="Inference Steps")
                         with gr.Row():
                             i2i_randomize = gr.Checkbox(label="🎲 Random Seed", value=True)
                             i2i_seed = gr.Number(label="Seed", value=42, precision=0, visible=False)
@@ -151,22 +187,27 @@ with gr.Blocks(fill_height=True) as demo:
                             lambda r: gr.Number(visible=not r),
                             inputs=[i2i_randomize], outputs=[i2i_seed]
                         )
+
                     i2i_btn = gr.Button("🚀 Edit Image", variant="primary", size="lg")
 
                 with gr.Column(scale=1, min_width=320):
                     i2i_output = gr.Image(label="Result", type="pil", format="png", show_label=False, height=600)
                     i2i_used_seed = gr.Number(label="🎲 Seed Used", interactive=False)
 
-            i2i_btn.click(generate_i2i, [i2i_input, i2i_prompt, i2i_strength, i2i_steps, i2i_seed, i2i_randomize], [i2i_output, i2i_used_seed])
+            i2i_btn.click(
+                generate_i2i,
+                [model_choice, i2i_input, i2i_prompt, i2i_strength, i2i_steps, i2i_seed, i2i_randomize],
+                [i2i_output, i2i_used_seed]
+            )
 
     gr.Markdown(
         """
         ---
         <div style="text-align: center; opacity: 0.7; font-size: 0.9em;">
-        <strong>T2I Model:</strong> Tongyi-MAI/Z-Image-Turbo • 
-        <strong>I2I Model:</strong> FLUX.1-schnell
+        <strong>T2I:</strong> Tongyi-MAI/Z-Image-Turbo • 
+        <strong>I2I:</strong> FLUX.1-schnell + InstructPix2Pix
         </div>
-        """,
+        """
     )
 
 if __name__ == "__main__":
@@ -183,7 +224,7 @@ if __name__ == "__main__":
         }
         .header-text p { font-size: 1.1rem !important; color: #64748b !important; }
         .gradio-container { max-width: 1400px !important; margin: 0 auto !important; }
-        button, .gr-button { transition: all 0.2s ease !important; }
+        button { transition: all 0.2s ease !important; }
         button:hover { transform: translateY(-1px); box-shadow: 0 4px 12px rgba(0,0,0,0.15) !important; }
         """,
         mcp_server=True