# Stable Diffusion

## Overview

Stable Diffusion is a text-to-image _latent diffusion_ model built upon the work of the original [Stable Diffusion](https://stability.ai/blog/stable-diffusion-public-release), and it was led by Robin Rombach and Katherine Crowson from [Stability AI](https://stability.ai/) and [LAION](https://laion.ai/).

🤗 `Optimum` extends `Diffusers` to support inference on the second generation of Neuron devices(powering Trainium and Inferentia 2). It aims at inheriting the ease of Diffusers on Neuron.

## Export to Neuron

To deploy models, you will need to compile them to TorchScript optimized for AWS Neuron. In the case of Stable Diffusion, there are four components which need to be exported to the `.neuron` format to boost the performance:

* Text encoder
* U-Net
* VAE encoder
* VAE decoder

You can either compile and export a Stable Diffusion Checkpoint via CLI or `NeuronStableDiffusionPipeline` class.

### Option 1: cli

Here is an example of exporting stable diffusion components with `Optimum` CLI:

```bash
optimum-cli export neuron --model stabilityai/stable-diffusion-2-1-base \
  --batch_size 1 \
  --height 512 `# height in pixels of generated image, eg. 512, 768` \
  --width 512 `# width in pixels of generated image, eg. 512, 768` \
  --num_images_per_prompt 1 `# number of images to generate per prompt, defaults to 1` \
  --auto_cast matmul `# cast only matrix multiplication operations` \
  --auto_cast_type bf16 `# cast operations from FP32 to BF16` \
  sd_neuron/
```

> [!TIP]
> We recommend using a `inf2.8xlarge` or a larger instance for the model compilation. You will also be able to compile the model with the Optimum CLI on a CPU-only instance (needs ~35 GB memory), and then run the pre-compiled model on `inf2.xlarge` to reduce the expenses. In this case, don't forget to disable validation of inference by adding the `--disable-validation` argument.

### Option 2: Python API

Here is an example of exporting stable diffusion components with `NeuronStableDiffusionPipeline`:

<Tip>

To apply optimized compute of Unet's attention score, please configure your environment variable with `export NEURON_FUSE_SOFTMAX=1`.

Besides, don't hesitate to tweak the compilation configuration to find the best tradeoff between performance v.s accuracy in your use case. By default, we suggest casting FP32 matrix multiplication operations to BF16 which offers good performance with moderate sacrifice of the accuracy. Check out the guide from [AWS Neuron documentation](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/general/appnotes/neuronx-cc/neuronx-cc-training-mixed-precision.html#neuronx-cc-training-mixed-precision) to better understand the options for your compilation.

</Tip>

```python
>>> from optimum.neuron import NeuronStableDiffusionPipeline

>>> model_id = "stable-diffusion-v1-5/stable-diffusion-v1-5"
>>> compiler_args = {"auto_cast": "matmul", "auto_cast_type": "bf16"}
>>> input_shapes = {"batch_size": 1, "height": 512, "width": 512}

>>> stable_diffusion = NeuronStableDiffusionPipeline.from_pretrained(model_id, export=True, **compiler_args, **input_shapes)

# Save locally or upload to the HuggingFace Hub
>>> save_directory = "sd_neuron/"
>>> stable_diffusion.save_pretrained(save_directory)
>>> stable_diffusion.push_to_hub(
...     save_directory, repository_id="my-neuron-repo"
... )
```

## Text-to-Image

`NeuronStableDiffusionPipeline` class allows you to generate images from a text prompt on neuron devices similar to the experience with `Diffusers`.

With pre-compiled Stable Diffusion models, now generate an image with a prompt on Neuron:

```python
>>> from optimum.neuron import NeuronStableDiffusionPipeline

>>> stable_diffusion = NeuronStableDiffusionPipeline.from_pretrained("sd_neuron/")
>>> prompt = "a photo of an astronaut riding a horse on mars"
>>> image = stable_diffusion(prompt).images[0]
```

<img
  src="https://raw.githubusercontent.com/huggingface/optimum-neuron/main/docs/assets/guides/models/01-sd-image.png"
  width="256"
  height="256"
  alt="stable diffusion generated image"
/>

## Image-to-Image

With the `NeuronStableDiffusionImg2ImgPipeline` class, you can generate a new image conditioned on a text prompt and an initial image.

```python
import requests
from PIL import Image
from io import BytesIO
from optimum.neuron import NeuronStableDiffusionImg2ImgPipeline

# compile & save
model_id = "nitrosocke/Ghibli-Diffusion"
input_shapes = {"batch_size": 1, "height": 512, "width": 512}
pipeline = NeuronStableDiffusionImg2ImgPipeline.from_pretrained(model_id, export=True, **input_shapes)
pipeline.save_pretrained("sd_img2img/")

url = "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg"

response = requests.get(url)
init_image = Image.open(BytesIO(response.content)).convert("RGB")
init_image = init_image.resize((512, 512))

prompt = "ghibli style, a fantasy landscape with snowcapped mountains, trees, lake with detailed reflection. sunlight and cloud in the sky, warm colors, 8K"

image = pipeline(prompt=prompt, image=init_image, strength=0.75, guidance_scale=7.5).images[0]
image.save("fantasy_landscape.png")
```

`image`          | `prompt` | output |
:-------------------------:|:-------------------------:|:-------------------------:|-------------------------:|
<img src="https://huggingface.co/datasets/optimum/documentation-images/resolve/main/neuron/models/03-sd-img2img-init.png" alt="landscape photo" width="256" height="256"/> | ***ghibli style, a fantasy landscape with snowcapped mountains, trees, lake with detailed reflection. warm colors, 8K*** | <img src="https://huggingface.co/datasets/optimum/documentation-images/resolve/main/neuron/models/04-sd-img2img.png" alt="drawing" width="250"/> |

## Inpaint

With the `NeuronStableDiffusionInpaintPipeline` class, you can edit specific parts of an image by providing a mask and a text prompt.

```python
import requests
from PIL import Image
from io import BytesIO
from optimum.neuron import NeuronStableDiffusionInpaintPipeline

model_id = "stable-diffusion-v1-5/stable-diffusion-inpainting"
input_shapes = {"batch_size": 1, "height": 512, "width": 512}
pipeline = NeuronStableDiffusionInpaintPipeline.from_pretrained(model_id, export=True, **input_shapes)
pipeline.save_pretrained("sd_inpaint/")

def download_image(url):
    response = requests.get(url)
    return Image.open(BytesIO(response.content)).convert("RGB")

img_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png"
mask_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png"

init_image = download_image(img_url).resize((512, 512))
mask_image = download_image(mask_url).resize((512, 512))

prompt = "Face of a yellow cat, high resolution, sitting on a park bench"
image = pipeline(prompt=prompt, image=init_image, mask_image=mask_image).images[0]
image.save("cat_on_bench.png")
```

`image`          | `mask_image` | `prompt` | output |
:-------------------------:|:-------------------------:|:-------------------------:|-------------------------:|
<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"/> |


## NeuronStableDiffusionPipeline[[optimum.neuron.NeuronStableDiffusionPipeline]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class optimum.neuron.NeuronStableDiffusionPipeline</name><anchor>optimum.neuron.NeuronStableDiffusionPipeline</anchor><source>https://github.com/huggingface/optimum-neuron/blob/vr_1005/optimum/neuron/modeling_diffusion.py#L1536</source><parameters>[{"name": "config", "val": ": dict[str, typing.Any]"}, {"name": "configs", "val": ": dict[str, 'PretrainedConfig']"}, {"name": "neuron_configs", "val": ": dict[str, 'NeuronDefaultConfig']"}, {"name": "data_parallel_mode", "val": ": typing.Literal['none', 'unet', 'transformer', 'all']"}, {"name": "scheduler", "val": ": diffusers.schedulers.scheduling_utils.SchedulerMixin | None"}, {"name": "vae_decoder", "val": ": torch.jit._script.ScriptModule | NeuronModelVaeDecoder"}, {"name": "text_encoder", "val": ": torch.jit._script.ScriptModule | NeuronModelTextEncoder | None = None"}, {"name": "text_encoder_2", "val": ": torch.jit._script.ScriptModule | NeuronModelTextEncoder | None = None"}, {"name": "unet", "val": ": torch.jit._script.ScriptModule | NeuronModelUnet | None = None"}, {"name": "transformer", "val": ": torch.jit._script.ScriptModule | NeuronModelTransformer | None = None"}, {"name": "vae_encoder", "val": ": torch.jit._script.ScriptModule | NeuronModelVaeEncoder | None = None"}, {"name": "image_encoder", "val": ": torch.jit._script.ScriptModule | None = None"}, {"name": "safety_checker", "val": ": torch.jit._script.ScriptModule | None = None"}, {"name": "tokenizer", "val": ": transformers.models.clip.tokenization_clip.CLIPTokenizer | transformers.models.t5.tokenization_t5.T5Tokenizer | None = None"}, {"name": "tokenizer_2", "val": ": transformers.models.clip.tokenization_clip.CLIPTokenizer | None = None"}, {"name": "feature_extractor", "val": ": transformers.models.clip.feature_extraction_clip.CLIPFeatureExtractor | None = None"}, {"name": "controlnet", "val": ": torch.jit._script.ScriptModule | list[torch.jit._script.ScriptModule]| NeuronControlNetModel | NeuronMultiControlNetModel | None = None"}, {"name": "requires_aesthetics_score", "val": ": bool = False"}, {"name": "force_zeros_for_empty_prompt", "val": ": bool = True"}, {"name": "add_watermarker", "val": ": bool | None = None"}, {"name": "model_save_dir", "val": ": str | pathlib.Path | tempfile.TemporaryDirectory | None = None"}, {"name": "model_and_config_save_paths", "val": ": dict[str, tuple[str, pathlib.Path]] | None = None"}]</parameters></docstring>



<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>__call__</name><anchor>optimum.neuron.NeuronStableDiffusionPipeline.__call__</anchor><source>https://github.com/huggingface/optimum-neuron/blob/vr_1005/optimum/neuron/modeling_diffusion.py#L1106</source><parameters>[{"name": "*args", "val": ""}, {"name": "**kwargs", "val": ""}]</parameters></docstring>


</div></div>

## NeuronStableDiffusionImg2ImgPipeline[[optimum.neuron.NeuronStableDiffusionImg2ImgPipeline]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class optimum.neuron.NeuronStableDiffusionImg2ImgPipeline</name><anchor>optimum.neuron.NeuronStableDiffusionImg2ImgPipeline</anchor><source>https://github.com/huggingface/optimum-neuron/blob/vr_1005/optimum/neuron/modeling_diffusion.py#L1549</source><parameters>[{"name": "config", "val": ": dict[str, typing.Any]"}, {"name": "configs", "val": ": dict[str, 'PretrainedConfig']"}, {"name": "neuron_configs", "val": ": dict[str, 'NeuronDefaultConfig']"}, {"name": "data_parallel_mode", "val": ": typing.Literal['none', 'unet', 'transformer', 'all']"}, {"name": "scheduler", "val": ": diffusers.schedulers.scheduling_utils.SchedulerMixin | None"}, {"name": "vae_decoder", "val": ": torch.jit._script.ScriptModule | NeuronModelVaeDecoder"}, {"name": "text_encoder", "val": ": torch.jit._script.ScriptModule | NeuronModelTextEncoder | None = None"}, {"name": "text_encoder_2", "val": ": torch.jit._script.ScriptModule | NeuronModelTextEncoder | None = None"}, {"name": "unet", "val": ": torch.jit._script.ScriptModule | NeuronModelUnet | None = None"}, {"name": "transformer", "val": ": torch.jit._script.ScriptModule | NeuronModelTransformer | None = None"}, {"name": "vae_encoder", "val": ": torch.jit._script.ScriptModule | NeuronModelVaeEncoder | None = None"}, {"name": "image_encoder", "val": ": torch.jit._script.ScriptModule | None = None"}, {"name": "safety_checker", "val": ": torch.jit._script.ScriptModule | None = None"}, {"name": "tokenizer", "val": ": transformers.models.clip.tokenization_clip.CLIPTokenizer | transformers.models.t5.tokenization_t5.T5Tokenizer | None = None"}, {"name": "tokenizer_2", "val": ": transformers.models.clip.tokenization_clip.CLIPTokenizer | None = None"}, {"name": "feature_extractor", "val": ": transformers.models.clip.feature_extraction_clip.CLIPFeatureExtractor | None = None"}, {"name": "controlnet", "val": ": torch.jit._script.ScriptModule | list[torch.jit._script.ScriptModule]| NeuronControlNetModel | NeuronMultiControlNetModel | None = None"}, {"name": "requires_aesthetics_score", "val": ": bool = False"}, {"name": "force_zeros_for_empty_prompt", "val": ": bool = True"}, {"name": "add_watermarker", "val": ": bool | None = None"}, {"name": "model_save_dir", "val": ": str | pathlib.Path | tempfile.TemporaryDirectory | None = None"}, {"name": "model_and_config_save_paths", "val": ": dict[str, tuple[str, pathlib.Path]] | None = None"}]</parameters></docstring>



<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>__call__</name><anchor>optimum.neuron.NeuronStableDiffusionImg2ImgPipeline.__call__</anchor><source>https://github.com/huggingface/optimum-neuron/blob/vr_1005/optimum/neuron/modeling_diffusion.py#L1106</source><parameters>[{"name": "*args", "val": ""}, {"name": "**kwargs", "val": ""}]</parameters></docstring>


</div></div>

## NeuronStableDiffusionInpaintPipeline[[optimum.neuron.NeuronStableDiffusionInpaintPipeline]]

<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>class optimum.neuron.NeuronStableDiffusionInpaintPipeline</name><anchor>optimum.neuron.NeuronStableDiffusionInpaintPipeline</anchor><source>https://github.com/huggingface/optimum-neuron/blob/vr_1005/optimum/neuron/modeling_diffusion.py#L1554</source><parameters>[{"name": "config", "val": ": dict[str, typing.Any]"}, {"name": "configs", "val": ": dict[str, 'PretrainedConfig']"}, {"name": "neuron_configs", "val": ": dict[str, 'NeuronDefaultConfig']"}, {"name": "data_parallel_mode", "val": ": typing.Literal['none', 'unet', 'transformer', 'all']"}, {"name": "scheduler", "val": ": diffusers.schedulers.scheduling_utils.SchedulerMixin | None"}, {"name": "vae_decoder", "val": ": torch.jit._script.ScriptModule | NeuronModelVaeDecoder"}, {"name": "text_encoder", "val": ": torch.jit._script.ScriptModule | NeuronModelTextEncoder | None = None"}, {"name": "text_encoder_2", "val": ": torch.jit._script.ScriptModule | NeuronModelTextEncoder | None = None"}, {"name": "unet", "val": ": torch.jit._script.ScriptModule | NeuronModelUnet | None = None"}, {"name": "transformer", "val": ": torch.jit._script.ScriptModule | NeuronModelTransformer | None = None"}, {"name": "vae_encoder", "val": ": torch.jit._script.ScriptModule | NeuronModelVaeEncoder | None = None"}, {"name": "image_encoder", "val": ": torch.jit._script.ScriptModule | None = None"}, {"name": "safety_checker", "val": ": torch.jit._script.ScriptModule | None = None"}, {"name": "tokenizer", "val": ": transformers.models.clip.tokenization_clip.CLIPTokenizer | transformers.models.t5.tokenization_t5.T5Tokenizer | None = None"}, {"name": "tokenizer_2", "val": ": transformers.models.clip.tokenization_clip.CLIPTokenizer | None = None"}, {"name": "feature_extractor", "val": ": transformers.models.clip.feature_extraction_clip.CLIPFeatureExtractor | None = None"}, {"name": "controlnet", "val": ": torch.jit._script.ScriptModule | list[torch.jit._script.ScriptModule]| NeuronControlNetModel | NeuronMultiControlNetModel | None = None"}, {"name": "requires_aesthetics_score", "val": ": bool = False"}, {"name": "force_zeros_for_empty_prompt", "val": ": bool = True"}, {"name": "add_watermarker", "val": ": bool | None = None"}, {"name": "model_save_dir", "val": ": str | pathlib.Path | tempfile.TemporaryDirectory | None = None"}, {"name": "model_and_config_save_paths", "val": ": dict[str, tuple[str, pathlib.Path]] | None = None"}]</parameters></docstring>



<div class="docstring border-l-2 border-t-2 pl-4 pt-3.5 border-gray-100 rounded-tl-xl mb-6 mt-8">


<docstring><name>__call__</name><anchor>optimum.neuron.NeuronStableDiffusionInpaintPipeline.__call__</anchor><source>https://github.com/huggingface/optimum-neuron/blob/vr_1005/optimum/neuron/modeling_diffusion.py#L1106</source><parameters>[{"name": "*args", "val": ""}, {"name": "**kwargs", "val": ""}]</parameters></docstring>


</div></div>

Are there any other diffusion features that you want us to support in 🤗`Optimum-neuron`? Please file an issue to [`Optimum-neuron` Github repo](https://github.com/huggingface/optimum-neuron) or discuss with us on [HuggingFace’s community forum](https://discuss.huggingface.co/c/optimum/), cheers 🤗 !