Update app.py

app.py

@@ -3,7 +3,7 @@ import numpy as np
 import random
 
 import spaces #[uncomment to use ZeroGPU]
-from diffusers import StableDiffusionXLPipeline, AutoencoderKL
+from diffusers import StableDiffusionXLPipeline, AutoencoderKL, StableDiffusionXLImg2ImgPipeline
 import torch
 from typing import Tuple
 
@@ -27,6 +27,17 @@ pipe = StableDiffusionXLPipeline.from_pretrained(
     )
 pipe.to(device)
 
+pipe_img2img = StableDiffusionXLImg2ImgPipeline.from_pretrained(
+    "RunDiffusion/Juggernaut-XL-v9",
+    vae=vae,
+    torch_dtype=torch.float16,
+    custom_pipeline="lpw_stable_diffusion_xl",
+    use_safetensors=True,
+    add_watermarker=False,
+    variant="fp16",
+)
+pipe_img2img.to(device)
+
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 4096
 
@@ -104,6 +115,8 @@ def infer(
     height,
     guidance_scale,
     num_inference_steps,
+    input_image=None,  # New parameter for input image
+    strength=0.8,  # New parameter for img2img strength
     progress=gr.Progress(track_tqdm=True),
 ):
     if randomize_seed:
@@ -111,15 +124,28 @@ def infer(
     prompt, negative_prompt = apply_style(style, prompt, negative_prompt)
     generator = torch.Generator().manual_seed(seed)
 
-    image = pipe(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
-        width=width,
-        height=height,
-        generator=generator,
-    ).images[0]
+    if input_image is not None:
+        # Use img2img pipeline if an image is provided
+        image = pipe_img2img(
+            prompt=prompt,
+            image=input_image,  # Pass the input image
+            strength=strength,   # Control how much the image is changed
+            negative_prompt=negative_prompt,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_inference_steps,
+            generator=generator,
+        ).images[0]
+    else:
+        # Use text2img pipeline otherwise
+        image = pipe(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_inference_steps,
+            width=width,
+            height=height,
+            generator=generator,
+        ).images[0]
 
     return image, seed
 
@@ -140,7 +166,6 @@ css = """
 with gr.Blocks(css=css) as demo:
     with gr.Column(elem_id="col-container"):
         gr.Markdown(" # ImageGen, the fastest and most precise image generator")
-
         with gr.Row():
             prompt = gr.Text(
                 label="Prompt",
@@ -149,11 +174,22 @@ with gr.Blocks(css=css) as demo:
                 placeholder="Enter your prompt",
                 container=False,
             )
-
             run_button = gr.Button("Run", scale=0, variant="primary")
-
         result = gr.Image(label="Result", show_label=False)
 
+        # Add image input and strength slider
+        with gr.Row():
+            input_image = gr.Image(type="pil", label="Input Image (Optional)", show_label=True, height=200)
+            with gr.Column():
+                strength = gr.Slider(
+                    label="Image Strength",
+                    minimum=0.0,
+                    maximum=1.0,
+                    step=0.01,
+                    value=0.8, # Default strength for img2img
+                    visible=True, # Make it visible if you want it always there, or toggle visibility with JS
+                )
+
         with gr.Row(visible=True):
             style_selection = gr.Radio(
                 show_label=True,
@@ -230,6 +266,8 @@ with gr.Blocks(css=css) as demo:
             height,
             guidance_scale,
             num_inference_steps,
+            input_image,  # Add input_image to inputs
+            strength,     # Add strength to inputs
         ],
         outputs=[result, seed],
     )