[BUG - STABLE DIFFUSION 3] Grey images generated

[hoangvictor]

Describe the bug

I'm running the SD3 model stabilityai/stable-diffusion-3-medium with the following settings:

• Height: 1024
• Width: 1024
• Inference steps: 50
• Guidance scale: 7
• Prompts length: 32 (using the same input prompt: "A men jumps from a high building")

The model generates 32 images, but half of them are grey.

Reproduction

Reproduction code:

import os

os.environ['CUDA_VISIBLE_DEVICES'] = '7'

import torch
from diffusers import StableDiffusion3Pipeline

base_dir = os.path.join(os.path.dirname(__file__), '..', '..')

def run():
    input_prompts = ["A men jumps from a high building"] * 32
    img_save_dir = f'{base_dir}/data/test_generated_img'
    os.makedirs(f'{img_save_dir}', exist_ok=True)

    pipe = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3-medium-diffusers", torch_dtype=torch.float16, cache_dir="/data0/tien/cache")
    pipe.to("cuda")

    torch.set_grad_enabled(False)
    images = pipe(
        prompt=input_prompts,
        negative_prompt="",
        height=1024,
        width=1024,
        num_inference_steps=50,
        guidance_scale=7.0
    ).images

    for j in range(len(input_prompts)):
        images[j].save(os.path.join(f'{img_save_dir}', f'{j}.jpg'))

    torch.cuda.empty_cache()
    
run()

Logs

No response

System Info

• Device: H100
• Driver Version: 550.127.05
• CUDA Version: 12.4
• Torch: 2.5.1+cu124
• OS: Ubuntu 22.04.3 LTS
• Python: 3.10.15
• diffusers: 0.31.0

Who can help?

No response

[sayakpaul]

Does this happen when you switch to torch.bfloat16? Also, was this working ever as expected and suddenly stopped working as expected?

[hoangvictor]

@sayakpaul Hi, the error persists when using BF16.

[hoangvictor]

Out of the 32 generated images, the first 16 are fine, but the last 16 are all in grayscale.

[sayakpaul]

Weird.

Does it happen on other prompts?

[hoangvictor]

Yes, it appears that this error occurs with any prompt.

[sayakpaul]

@asomoza have you heard about this before?

[asomoza]

no, but SD 3 had some other problems so probably no one else did some extensive tests or generations with it.

@hoangvictor is there a reason you're using stable-diffusion-3-medium instead of stable-diffusion-3.5-medium?

[hoangvictor]

@asomoza SD3.5 has been out since October, but when I started my project, only SD3 was available. I’ve also tried SD3.5, and the issue persists. It appears the problem lies with the VAE during the image decoding process.

[asomoza]

I did a test with an A100 80GB, I could only do 24 images or I get OOM. I could replicate your error, this is a rare test for me because I would never throw 32 times the same prompt at the model instead of doing it in a loop because I like to see the results and most people just use a batch of 4.

Anyway, after the 16th image, the rest are all grey images so I can confirm this issue cc: @sayakpaul

[sayakpaul]

Hmm, @DN6 could this be an issue with prompt encoding?

What happens when we generate 24 images for a single prompt by specifying num_images_per_prompt? Also does it happen with other pipelines too?

Very weird that it happens after we request for a certain number of images no?

[asomoza]

It also happens with num_images_per_prompt when it's more than 16.

Tested it with SD 1.5 and this issue didn't happen.

It happens with any number of images when it's more than 16.

[pcuenca]

This is a PyTorch bug / limitation in the upsample operation. I found a workaround that I'll push tomorrow. Meanwhile, you can retrieve the latents instead of the images using output_type="latent" in your pipeline invocation, and then decode them manually like this:

with torch.inference_mode():
    latents = (latents / pipe.vae.config.scaling_factor) + pipe.vae.config.shift_factor
    images_1 = pipe.vae.decode(latents[:16], return_dict=False)[0]
    images_2 = pipe.vae.decode(latents[16:], return_dict=False)[0]
    images = torch.cat([images_1, images_2])
    images = pipe.image_processor.postprocess(images)