Skip to content

Stable Diffusion 3 (Image Generation)

This case study runs Stable Diffusion 3 with RBLN Optimum and analyzes the execution of submodules, including the text encoders, transformer, and VAE, placed across multiple ATOM™ devices.

Stable Diffusion 3 (SD3) is a Multimodal Diffusion Transformer (MMDiT) model, which outperforms various state-of-the-art text-to-image generation models.

Some models, such as SD3, consist of multiple submodules, and loading all submodules onto a single device may occasionally exceed the device's memory capacity. To optimize memory utilization, devices can be assigned individually to the runtime created for each submodule.

The RBLN Compiler and the RBLN Profiler support multi-device compilation and the generation of profiled data for models with multiple submodules. For example, the RBLN Profiler can generate multiple profiled data files for each submodule in SD3, such as the text encoders, the diffusion transformer, and the VAE (Variational AutoEncoder). It can also generate a single profiled data file that includes all submodules of SD3.

Preliminaries

Installation

1
2
$ cd rbln-model-zoo/ATOM/huggingface/diffusers/text-to-image/stable_diffusion_3_t2i
$ pip3 install -r requirements.txt

Compile Model and Activate RBLN Profiler

1
2
3
# Sample Text: "a photo of a cat holding a sign that says hello world”
$ python3 compile.py
$ RBLN_PROFILER=1 python3 inference.py

Or, set "activate_profiler": True to enable the RBLN Profiler in each module:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
# rbln-model-zoo/ATOM/huggingface/diffusers/text-to-image/stable_diffusion_3_t2i/inference.py
...
pipe = RBLNStableDiffusion3Pipeline.from_pretrained(
    ...
    rbln_config={
        ...
        "text_encoder": {"device": 0, "activate_profiler": True},
        "text_encoder_2": {"device": 0, "activate_profiler": True},
        "text_encoder_3": {"device": 1, "activate_profiler": True},
        "transformer": {"device": 0, "activate_profiler": True},
        "vae": {"device": 0, "activate_profiler": True},
    },
)
...

To profile only a subset of submodules, set "activate_profiler": False for the submodules to exclude from profiling.

Extract Profiled Data with Inference

Recommendation

The RBLN Profiler captures all iterations of the diffusion process, which may result in extended profiling times. To obtain quicker profiling results, set num_inference_steps=1 in inference.py. 

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
# rbln-model-zoo/ATOM/huggingface/diffusers/text-to-image/stable_diffusion_3_t2i/inference.py
...
# Generate image
image = pipe(
    # prompt, num_inference_steps=28, height=1024, width=1024, guidance_scale=7.0
    prompt,
    num_inference_steps=1,
    height=1024,
    width=1024,
    guidance_scale=7.0,
).images[0]
...

Run inference to obtain profiling data:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
# rbln-model-zoo/ATOM/huggingface/diffusers/text-to-image/stable_diffusion_3_t2i/inference.py
...
# Generate image
image = pipe(
    # prompt, num_inference_steps=28, height=1024, width=1024, guidance_scale=7.0
    prompt,
    num_inference_steps=1,
    height=1024,
    width=1024,
    guidance_scale=7.0,
).images[0]
...

Alternatively, use scoped profiling methods as shown below:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
# rbln-model-zoo/ATOM/huggingface/diffusers/text-to-image/stable_diffusion_3_t2i/inference.py
...
from rebel.profiler import profile

with profile(output_dir="./profile_out"):
    # Generate image
    image = pipe(
        # prompt, num_inference_steps=28, height=1024, width=1024, guidance_scale=7.0
        prompt,
        num_inference_steps=1,
        height=1024,
        width=1024,
        guidance_scale=7.0,
    ).images[0]
...

Analysis of Profiled Data from SD3 with RBLN Profiler

Profiling result

Perfetto timeline over two ATOM devices with the pipeline stages colour-boxed, one filling most of it. 

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
# https://github.com/huggingface/diffusers/blob/78bc824729f76a14ff2f211fc7f9a31e5500a41e/src/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion_3.py#L343
class StableDiffusion3Pipeline(DiffusionPipeline, SD3LoraLoaderMixin, FromSingleFileMixin):
...
    def encode_prompt(
        ...
        ):
      ...
      prompt_embed, pooled_prompt_embed = self._get_clip_prompt_embeds(
          prompt=prompt,
          ...
      )
      prompt_2_embed, pooled_prompt_2_embed = self._get_clip_prompt_embeds(
          prompt=prompt_2,
          ...
      )
      ...
      t5_prompt_embed = self._get_t5_prompt_embeds(
          prompt=prompt_3,
          ...
      )
      ...
      negative_prompt_embed, negative_pooled_prompt_embed = self._get_clip_prompt_embeds(
          negative_prompt,
          ...
      )
      negative_prompt_2_embed, negative_pooled_prompt_2_embed = self._get_clip_prompt_embeds(
          negative_prompt_2,
          ...
      )
      t5_negative_prompt_embed = self._get_t5_prompt_embeds(
          prompt=negative_prompt_3,
          ...
      )
      ...
      return prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds

SD3 has three text encoders (text_encoder, text_encoder_2, and text_encoder_3). According to the encode_prompt function in the pipeline_stable_diffusion_3.py code, each text encoder is called twice to generate prompts and negative prompts. Since the RBLN Profiler generates a single *.pb file when each text encoder is called, this encode_prompt process results in the generation of six *.pb files. Additionally, since the diffusion transformer is called for a specified number of iterations (num_inference_steps), one *.pb file is generated for each iteration. Finally, when the VAE decoder produces the final image, a single *.pb file is generated. A full log *.pb file is also generated to represent the entire tracing information, displaying the sequential profiling results of all submodules.

Note

Loading all submodules of SD3 on a single device may exceed the device's memory capacity. To address this, text_encoder_3 is executed on a separate device (ATOM™ 1). Furthermore, since the multi-device setup is mounted on a single host, Host commands are generated only on the root device (ATOM™ 0).

The comments in the example's inference.py explain this device placement in detail.

Tiling large operations

A run of seven identically named commands outlined across two tracks.

The RBLN Compiler efficiently optimizes large operations by applying the tiling method. As seen in the profiling results above, the tiled operations share the same name but have different indices.

References