Repository: huggingface/diffusers
Issue category: Bug Report

Issue:
[BUG - STABLE DIFFUSION 3] Grey images generated
### Describe the bug

I'm running the SD3 model [stabilityai/stable-diffusion-3-medium](https://huggingface.co/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_


Hints:
Does this happen when you switch to torch.bfloat16? Also, was this working ever as expected and suddenly stopped working as expected? 
@sayakpaul Hi, the error persists when using BF16.
Out of the 32 generated images, the first 16 are fine, but the last 16 are all in grayscale.
Weird. 

Does it happen on other prompts?
Yes, it appears that this error occurs with any prompt.
@asomoza have you heard about this before?
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](https://huggingface.co/stabilityai/stable-diffusion-3-medium) instead of [stable-diffusion-3.5-medium](https://huggingface.co/stabilityai/stable-diffusion-3.5-medium)?
@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.
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 

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?
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.
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:

```python
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)
```

Patch:
diff --git a/src/diffusers/models/upsampling.py b/src/diffusers/models/upsampling.py
index cf07e45b0c5c..af04ae4b93cf 100644
--- a/src/diffusers/models/upsampling.py
+++ b/src/diffusers/models/upsampling.py
@@ -165,6 +165,14 @@ def forward(self, hidden_states: torch.Tensor, output_size: Optional[int] = None
         # if `output_size` is passed we force the interpolation output
         # size and do not make use of `scale_factor=2`
         if self.interpolate:
+            # upsample_nearest_nhwc also fails when the number of output elements is large
+            # https://github.com/pytorch/pytorch/issues/141831
+            scale_factor = (
+                2 if output_size is None else max([f / s for f, s in zip(output_size, hidden_states.shape[-2:])])
+            )
+            if hidden_states.numel() * scale_factor > pow(2, 31):
+                hidden_states = hidden_states.contiguous()
+
             if output_size is None:
                 hidden_states = F.interpolate(hidden_states, scale_factor=2.0, mode="nearest")
             else:


Test patch: