Update app.py

app.py

@@ -19,14 +19,14 @@ pipe_t2i = DiffusionPipeline.from_pretrained(
 )
 pipe_t2i.to("cuda")
 
-# Image to image - FLUX
+# FLUX
 pipe_flux = FluxImg2ImgPipeline.from_pretrained(
     "black-forest-labs/FLUX.1-schnell",
     torch_dtype=torch.bfloat16,
 )
 pipe_flux.to("cuda")
 
-# Image to image - InstructPix2Pix
+# InstructPix2Pix
 pipe_ip2p = StableDiffusionInstructPix2PixPipeline.from_pretrained(
     "timbrooks/instruct-pix2pix",
     torch_dtype=torch.float16,
@@ -34,16 +34,52 @@ 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! Portrait models load on first use.")
+
+# Portrait models — lazy loaded on first selection to save memory
+portrait_pipes = {}
 
-print("Pipelines loaded!")
+PORTRAIT_MODELS = {
+    "SDXL Base (Balanced quality & control)": {
+        "repo": "stabilityai/stable-diffusion-xl-base-1.0",
+        "variant": "fp16",
+    },
+    "Juggernaut XL (Photorealistic portraits)": {
+        "repo": "RunDiffusion/Juggernaut-XL-v9",
+        "variant": None,
+    },
+    "DreamShaper XL (Creative portraits)": {
+        "repo": "Lykon/dreamshaper-xl-v2-turbo",
+        "variant": None,
+    },
+    "Realistic Vision XL (Portrait photography)": {
+        "repo": "SG161222/RealVisXL_V4.0",
+        "variant": "fp16",
+    },
+    "CyberRealistic XL (Detailed skin & faces)": {
+        "repo": "stablediffusionapi/cyberrealistic-xl",
+        "variant": None,
+    },
+    "DreamShaper (SD1.5, fast portraits)": {
+        "repo": "Lykon/DreamShaper",
+        "variant": None,
+        "sd15": True,
+    },
+}
+
+def get_portrait_pipe(model_name):
+    """Lazy load portrait models on first use."""
+    if model_name not in portrait_pipes:
+        print(f"Loading {model_name} for the first time...")
+        cfg = PORTRAIT_MODELS[model_name]
+        kwargs = dict(torch_dtype=torch.float16, use_safetensors=True)
+        if cfg.get("variant"):
+            kwargs["variant"] = cfg["variant"]
+        pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(cfg["repo"], **kwargs)
+        pipe.to("cuda")
+        portrait_pipes[model_name] = pipe
+        print(f"{model_name} loaded!")
+    return portrait_pipes[model_name]
 
 def resize_image(image, max_size=1024):
     orig_w, orig_h = image.size
@@ -101,8 +137,10 @@ 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(
+    else:
+        # Portrait models (lazy loaded)
+        pipe = get_portrait_pipe(model_choice)
+        image = pipe(
             prompt=prompt,
             image=input_image,
             strength=float(strength),
@@ -134,6 +172,11 @@ custom_theme = gr.themes.Soft(
     block_title_text_weight="600",
 )
 
+I2I_MODELS = [
+    "FLUX.1-schnell (Creative, high change)",
+    "InstructPix2Pix (Precise, preserves identity)",
+] + list(PORTRAIT_MODELS.keys())
+
 with gr.Blocks(fill_height=True) as demo:
     gr.Markdown(
         """
@@ -180,30 +223,23 @@ with gr.Blocks(fill_height=True) as demo:
         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)",
-                            "SDXL (Balanced quality & control)",
-                        ],
+                        choices=I2I_MODELS,
                         value="InstructPix2Pix (Precise, preserves identity)",
                         label="🤖 Model",
-                        info="InstructPix2Pix: best for targeted edits like colours/styles. SDXL: balanced quality and control. FLUX: creative transformations."
+                        info="InstructPix2Pix: best for targeted edits. Portrait models: high quality people & faces. FLUX: creative transformations. Note: portrait models load on first use."
                     )
-
                     i2i_input = gr.Image(label="Upload Image", type="pil")
                     i2i_prompt = gr.Textbox(
                         label="✨ Edit Instruction",
-                        placeholder="e.g. 'change the dress to blue' or 'make it a sunset'",
+                        placeholder="e.g. 'woman in blue dress, photorealistic portrait' or 'make it a sunset'",
                         lines=4,
                     )
-
                     with gr.Accordion("⚙️ Advanced Settings", open=False):
-                        i2i_strength = gr.Slider(0.1, 1.0, value=0.75, step=0.05,
-                            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")
+                        i2i_strength = gr.Slider(0.1, 1.0, value=0.65, step=0.05,
+                            label="Strength (all except InstructPix2Pix)",
+                            info="Lower = more faithful to original. Higher = more creative change.")
+                        i2i_steps = gr.Slider(1, 50, value=25, 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)
@@ -211,7 +247,6 @@ 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):
@@ -229,7 +264,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 + SDXL
+        <strong>I2I:</strong> FLUX.1-schnell + InstructPix2Pix + Juggernaut XL + DreamShaper XL + RealVisXL + CyberRealistic XL
         </div>
         """
     )