Add InstructPix2Pix pipeline support. (#625)

* Add InstructPix2Pix pipeline support. * Replace flax pipe with torch. * Adjust annotations and docs style. * Add prerequisite on cfg. * Update docs.
huggingface · Jun 30, 2024 · 86900e7 · 86900e7
1 parent 2c524df
commit 86900e7
Show file tree

Hide file tree

Showing 8 changed files with 539 additions and 0 deletions.
diff --git a/docs/source/tutorials/stable_diffusion.mdx b/docs/source/tutorials/stable_diffusion.mdx
@@ -173,6 +173,40 @@ image.save("cat_on_bench.png")
 :-------------------------:|:-------------------------:|:-------------------------:|-------------------------:|
 <img src="https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png" alt="drawing" width="250"/> | <img src="https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png" alt="drawing" width="250"/> | ***Face of a yellow cat, high resolution, sitting on a park bench*** | <img src="https://huggingface.co/datasets/optimum/documentation-images/resolve/main/neuron/models/05-sd-inpaint.png" alt="drawing" width="250"/> |
 
+### InstructPix2Pix
+
+With the `NeuronStableDiffusionInstructPix2PixPipeline` class, you can apply instruction-based image editing using both text guidance and image guidance.
+
+```python
+import requests
+import PIL
+from io import BytesIO
+from optimum.neuron import NeuronStableDiffusionInstructPix2PixPipeline
+
+def download_image(url):
+    response = requests.get(url)
+    return PIL.Image.open(BytesIO(response.content)).convert("RGB")
+
+model_id = "timbrooks/instruct-pix2pix"
+input_shapes = {"batch_size": 1, "height": 512, "width": 512}
+
+pipe = NeuronStableDiffusionInstructPix2PixPipeline.from_pretrained(
+  model_id, export=True, dynamic_batch_size=True, **input_shapes,
+)
+pipe.save_pretrained("sd_ip2p/")
+
+img_url = "https://huggingface.co/datasets/diffusers/diffusers-images-docs/resolve/main/mountain.png"
+init_image = download_image(img_url).resize((512, 512))
+
+prompt = "Add a beautiful sunset"
+image = pipe(prompt=prompt, image=init_image).images[0]
+image.save("sunset_mountain.png")
+```
+
+`image`          | `prompt` | output |
+:-------------------------:|:-------------------------:|-------------------------:|
+<img src="https://huggingface.co/datasets/diffusers/diffusers-images-docs/resolve/main/mountain.png" alt="drawing" width="250"/> | ***Add a beautiful sunset*** | <img src="https://huggingface.co/datasets/optimum/documentation-images/resolve/main/neuron/models/11-sd-ip2p.png" alt="drawing" width="250"/> |
+
 ## Stable Diffusion XL
 
 *There is a notebook version of that tutorial [here](https://github.com/huggingface/optimum-neuron/blob/main/notebooks/stable-diffusion/stable-diffusion-xl-txt2img.ipynb)*.

diff --git a/optimum/neuron/__init__.py b/optimum/neuron/__init__.py
@@ -48,6 +48,7 @@
         "NeuronStableDiffusionPipeline",
         "NeuronStableDiffusionImg2ImgPipeline",
         "NeuronStableDiffusionInpaintPipeline",
+        "NeuronStableDiffusionInstructPix2PixPipeline",
         "NeuronLatentConsistencyModelPipeline",
         "NeuronStableDiffusionXLPipeline",
         "NeuronStableDiffusionXLImg2ImgPipeline",
@@ -88,6 +89,7 @@
         NeuronStableDiffusionControlNetPipeline,
         NeuronStableDiffusionImg2ImgPipeline,
         NeuronStableDiffusionInpaintPipeline,
+        NeuronStableDiffusionInstructPix2PixPipeline,
         NeuronStableDiffusionPipeline,
         NeuronStableDiffusionPipelineBase,
         NeuronStableDiffusionXLImg2ImgPipeline,

diff --git a/optimum/neuron/modeling_diffusion.py b/optimum/neuron/modeling_diffusion.py
@@ -94,6 +94,7 @@
         NeuronStableDiffusionControlNetPipelineMixin,
         NeuronStableDiffusionImg2ImgPipelineMixin,
         NeuronStableDiffusionInpaintPipelineMixin,
+        NeuronStableDiffusionInstructPix2PixPipelineMixin,
         NeuronStableDiffusionPipelineMixin,
         NeuronStableDiffusionXLControlNetPipelineMixin,
         NeuronStableDiffusionXLImg2ImgPipelineMixin,
@@ -1222,6 +1223,12 @@ class NeuronStableDiffusionInpaintPipeline(
     __call__ = NeuronStableDiffusionInpaintPipelineMixin.__call__
 
 
+class NeuronStableDiffusionInstructPix2PixPipeline(
+    NeuronStableDiffusionPipelineBase, NeuronStableDiffusionInstructPix2PixPipelineMixin
+):
+    __call__ = NeuronStableDiffusionInstructPix2PixPipelineMixin.__call__
+
+
 class NeuronLatentConsistencyModelPipeline(NeuronStableDiffusionPipelineBase, NeuronLatentConsistencyPipelineMixin):
     __call__ = NeuronLatentConsistencyPipelineMixin.__call__
 

diff --git a/optimum/neuron/pipelines/__init__.py b/optimum/neuron/pipelines/__init__.py
@@ -24,6 +24,7 @@
         "NeuronStableDiffusionPipelineMixin",
         "NeuronStableDiffusionImg2ImgPipelineMixin",
         "NeuronStableDiffusionInpaintPipelineMixin",
+        "NeuronStableDiffusionInstructPix2PixPipelineMixin",
         "NeuronLatentConsistencyPipelineMixin",
         "NeuronStableDiffusionControlNetPipelineMixin",
         "NeuronStableDiffusionXLPipelineMixin",
@@ -39,6 +40,7 @@
         NeuronStableDiffusionControlNetPipelineMixin,
         NeuronStableDiffusionImg2ImgPipelineMixin,
         NeuronStableDiffusionInpaintPipelineMixin,
+        NeuronStableDiffusionInstructPix2PixPipelineMixin,
         NeuronStableDiffusionPipelineMixin,
         NeuronStableDiffusionXLControlNetPipelineMixin,
         NeuronStableDiffusionXLImg2ImgPipelineMixin,

diff --git a/optimum/neuron/pipelines/diffusers/__init__.py b/optimum/neuron/pipelines/diffusers/__init__.py
@@ -19,6 +19,7 @@
 from .pipeline_stable_diffusion import NeuronStableDiffusionPipelineMixin
 from .pipeline_stable_diffusion_img2img import NeuronStableDiffusionImg2ImgPipelineMixin
 from .pipeline_stable_diffusion_inpaint import NeuronStableDiffusionInpaintPipelineMixin
+from .pipeline_stable_diffusion_instruct_pix2pix import NeuronStableDiffusionInstructPix2PixPipelineMixin
 from .pipeline_stable_diffusion_xl import NeuronStableDiffusionXLPipelineMixin
 from .pipeline_stable_diffusion_xl_img2img import NeuronStableDiffusionXLImg2ImgPipelineMixin
 from .pipeline_stable_diffusion_xl_inpaint import NeuronStableDiffusionXLInpaintPipelineMixin