#
Stable Diffusion

##
Overview

Stable Diffusion is a text-to-image latent diffusion model built upon the work of the original Stable Diffusion, and it was led by Robin Rombach and Katherine Crowson from Stability AI and LAION.

🤗 `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:

    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/

> 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`:

> 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 to better understand the options for your compilation.

    >>> 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:

    >>> 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]

##
Image-to-Image

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

    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 |
--- | --- | --- | ---
| ghibli style, a fantasy landscape with snowcapped mountains, trees, lake with detailed reflection. warm colors, 8K |  |

##
Inpaint

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

    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
--- | --- | --- | ---
|  | Face of a yellow cat, high resolution, sitting on a park bench |

##
NeuronStableDiffusionPipeline

###
class optimum.neuron.NeuronStableDiffusionPipeline

< source >

( config: dict[str, typing.Any] configs: dict[str, 'PretrainedConfig'] neuron_configs: dict[str, 'NeuronDefaultConfig'] data_parallel_mode: typing.Literal['none', 'unet', 'transformer', 'all'] scheduler: diffusers.schedulers.scheduling_utils.SchedulerMixin | None vae_decoder: torch.jit._script.ScriptModule | NeuronModelVaeDecoder text_encoder: torch.jit._script.ScriptModule | NeuronModelTextEncoder | None = None text_encoder_2: torch.jit._script.ScriptModule | NeuronModelTextEncoder | None = None unet: torch.jit._script.ScriptModule | NeuronModelUnet | None = None transformer: torch.jit._script.ScriptModule | NeuronModelTransformer | None = None vae_encoder: torch.jit._script.ScriptModule | NeuronModelVaeEncoder | None = None image_encoder: torch.jit._script.ScriptModule | None = None safety_checker: torch.jit._script.ScriptModule | None = None tokenizer: transformers.models.clip.tokenization_clip.CLIPTokenizer | transformers.models.t5.tokenization_t5.T5Tokenizer | None = None tokenizer_2: transformers.models.clip.tokenization_clip.CLIPTokenizer | None = None feature_extractor: transformers.models.clip.feature_extraction_clip.CLIPFeatureExtractor | None = None controlnet: torch.jit._script.ScriptModule | list[torch.jit._script.ScriptModule]| NeuronControlNetModel | NeuronMultiControlNetModel | None = None requires_aesthetics_score: bool = False force_zeros_for_empty_prompt: bool = True add_watermarker: bool | None = None model_save_dir: str | pathlib.Path | tempfile.TemporaryDirectory | None = None model_and_config_save_paths: dict[str, tuple[str, pathlib.Path]] | None = None )

####
__call__

< source >

( *args **kwargs )

##
NeuronStableDiffusionImg2ImgPipeline

###
class optimum.neuron.NeuronStableDiffusionImg2ImgPipeline

< source >

( config: dict[str, typing.Any] configs: dict[str, 'PretrainedConfig'] neuron_configs: dict[str, 'NeuronDefaultConfig'] data_parallel_mode: typing.Literal['none', 'unet', 'transformer', 'all'] scheduler: diffusers.schedulers.scheduling_utils.SchedulerMixin | None vae_decoder: torch.jit._script.ScriptModule | NeuronModelVaeDecoder text_encoder: torch.jit._script.ScriptModule | NeuronModelTextEncoder | None = None text_encoder_2: torch.jit._script.ScriptModule | NeuronModelTextEncoder | None = None unet: torch.jit._script.ScriptModule | NeuronModelUnet | None = None transformer: torch.jit._script.ScriptModule | NeuronModelTransformer | None = None vae_encoder: torch.jit._script.ScriptModule | NeuronModelVaeEncoder | None = None image_encoder: torch.jit._script.ScriptModule | None = None safety_checker: torch.jit._script.ScriptModule | None = None tokenizer: transformers.models.clip.tokenization_clip.CLIPTokenizer | transformers.models.t5.tokenization_t5.T5Tokenizer | None = None tokenizer_2: transformers.models.clip.tokenization_clip.CLIPTokenizer | None = None feature_extractor: transformers.models.clip.feature_extraction_clip.CLIPFeatureExtractor | None = None controlnet: torch.jit._script.ScriptModule | list[torch.jit._script.ScriptModule]| NeuronControlNetModel | NeuronMultiControlNetModel | None = None requires_aesthetics_score: bool = False force_zeros_for_empty_prompt: bool = True add_watermarker: bool | None = None model_save_dir: str | pathlib.Path | tempfile.TemporaryDirectory | None = None model_and_config_save_paths: dict[str, tuple[str, pathlib.Path]] | None = None )

####
__call__

< source >

( *args **kwargs )

##
NeuronStableDiffusionInpaintPipeline

###
class optimum.neuron.NeuronStableDiffusionInpaintPipeline

< source >

( config: dict[str, typing.Any] configs: dict[str, 'PretrainedConfig'] neuron_configs: dict[str, 'NeuronDefaultConfig'] data_parallel_mode: typing.Literal['none', 'unet', 'transformer', 'all'] scheduler: diffusers.schedulers.scheduling_utils.SchedulerMixin | None vae_decoder: torch.jit._script.ScriptModule | NeuronModelVaeDecoder text_encoder: torch.jit._script.ScriptModule | NeuronModelTextEncoder | None = None text_encoder_2: torch.jit._script.ScriptModule | NeuronModelTextEncoder | None = None unet: torch.jit._script.ScriptModule | NeuronModelUnet | None = None transformer: torch.jit._script.ScriptModule | NeuronModelTransformer | None = None vae_encoder: torch.jit._script.ScriptModule | NeuronModelVaeEncoder | None = None image_encoder: torch.jit._script.ScriptModule | None = None safety_checker: torch.jit._script.ScriptModule | None = None tokenizer: transformers.models.clip.tokenization_clip.CLIPTokenizer | transformers.models.t5.tokenization_t5.T5Tokenizer | None = None tokenizer_2: transformers.models.clip.tokenization_clip.CLIPTokenizer | None = None feature_extractor: transformers.models.clip.feature_extraction_clip.CLIPFeatureExtractor | None = None controlnet: torch.jit._script.ScriptModule | list[torch.jit._script.ScriptModule]| NeuronControlNetModel | NeuronMultiControlNetModel | None = None requires_aesthetics_score: bool = False force_zeros_for_empty_prompt: bool = True add_watermarker: bool | None = None model_save_dir: str | pathlib.Path | tempfile.TemporaryDirectory | None = None model_and_config_save_paths: dict[str, tuple[str, pathlib.Path]] | None = None )

####
__call__

< source >

( *args **kwargs )

Are there any other diffusion features that you want us to support in 🤗`Optimum-neuron`? Please file an issue to `Optimum-neuron` Github repo or discuss with us on HuggingFace’s community forum, cheers 🤗 !