Support for StableDiffusionInstructPix2PixPipeline

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import spaces
-from diffusers import AutoPipelineForImage2Image
+from diffusers import AutoPipelineForImage2Image, StableDiffusionInstructPix2PixPipeline
 from loguru import logger
 from PIL import Image
 
@@ -24,8 +24,6 @@ def generate(
         f"Starting image generation: {dict(model=model, prompt=prompt, image=init_image, strength=strength)}"
     )
 
-    pipe = AutoPipelineForImage2Image.from_pretrained(model).to("cuda")
-
     # Downscale the image
     init_image.thumbnail((1024, 1024))
 
@@ -36,12 +34,23 @@ def generate(
         progress((step_index + 1, pipe.num_timesteps))
         return callback_kwargs
 
-    images = pipe(
-        prompt=prompt,
-        image=init_image,
-        strength=strength,
-        callback_on_step_end=progress_callback,
-    ).images
+    if model == "timbrooks/instruct-pix2pix":
+        pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(model).to("cuda")
+
+        images = pipe(
+            prompt=prompt,
+            image=init_image,
+            callback_on_step_end=progress_callback,
+        ).images
+    else:
+        pipe = AutoPipelineForImage2Image.from_pretrained(model).to("cuda")
+
+        images = pipe(
+            prompt=prompt,
+            image=init_image,
+            strength=strength,
+            callback_on_step_end=progress_callback,
+        ).images
     return images[0]