Update app.py

app.py

@@ -1,7 +1,12 @@
 import torch
 import spaces
 import gradio as gr
-from diffusers import DiffusionPipeline, FluxImg2ImgPipeline, StableDiffusionInstructPix2PixPipeline
+from diffusers import (
+    DiffusionPipeline,
+    FluxImg2ImgPipeline,
+    StableDiffusionInstructPix2PixPipeline,
+    StableDiffusionXLImg2ImgPipeline,
+)
 from PIL import Image
 
 print("Loading pipelines...")
@@ -29,6 +34,15 @@ pipe_ip2p = StableDiffusionInstructPix2PixPipeline.from_pretrained(
 )
 pipe_ip2p.to("cuda")
 
+# Image to image - SDXL
+pipe_sdxl = StableDiffusionXLImg2ImgPipeline.from_pretrained(
+    "stabilityai/stable-diffusion-xl-base-1.0",
+    torch_dtype=torch.float16,
+    variant="fp16",
+    use_safetensors=True,
+)
+pipe_sdxl.to("cuda")
+
 print("Pipelines loaded!")
 
 def resize_image(image, max_size=1024):
@@ -87,6 +101,15 @@ def generate_i2i(model_choice, input_image, prompt, strength, num_inference_step
             generator=generator,
         ).images[0]
 
+    elif model_choice == "SDXL (Balanced quality & control)":
+        image = pipe_sdxl(
+            prompt=prompt,
+            image=input_image,
+            strength=float(strength),
+            num_inference_steps=int(num_inference_steps),
+            generator=generator,
+        ).images[0]
+
     return image, seed
 
 examples_t2i = [
@@ -162,10 +185,11 @@ with gr.Blocks(fill_height=True) as demo:
                         choices=[
                             "FLUX.1-schnell (Creative, high change)",
                             "InstructPix2Pix (Precise, preserves identity)",
+                            "SDXL (Balanced quality & control)",
                         ],
                         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."
+                        info="InstructPix2Pix: best for targeted edits like colours/styles. SDXL: balanced quality and control. FLUX: creative transformations."
                     )
 
                     i2i_input = gr.Image(label="Upload Image", type="pil")
@@ -177,7 +201,7 @@ with gr.Blocks(fill_height=True) as demo:
 
                     with gr.Accordion("⚙️ Advanced Settings", open=False):
                         i2i_strength = gr.Slider(0.1, 1.0, value=0.75, step=0.05,
-                            label="Strength (FLUX only)",
+                            label="Strength (FLUX & SDXL 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():
@@ -205,7 +229,7 @@ with gr.Blocks(fill_height=True) as demo:
         ---
         <div style="text-align: center; opacity: 0.7; font-size: 0.9em;">
         <strong>T2I:</strong> Tongyi-MAI/Z-Image-Turbo • 
-        <strong>I2I:</strong> FLUX.1-schnell + InstructPix2Pix
+        <strong>I2I:</strong> FLUX.1-schnell + InstructPix2Pix + SDXL
         </div>
         """
     )