Repository: huggingface/diffusers

Issue:
StableCascadeDecoderPipeline Error with num_images_per_prompt > 1
### Describe the bug

I'm encountering a RuntimeError when using the `StableCascadeDecoderPipeline` with `num_images_per_prompt` set to a value greater than 1. Since parallel processing functions correctly in the `StableCascadePriorPipeline`, I suspect an issue with how latents are being passed to the decoder pipeline.

```cmd
Traceback (most recent call last):
  File "E:\stable-cascade-one-click-installer\issue.py", line 18, in <module>
    decoder_output = decoder(
  File "E:\stable-cascade-one-click-installer\venv\lib\site-packages\torch\utils\_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "E:\stable-cascade-one-click-installer\venv\lib\site-packages\diffusers\pipelines\stable_cascade\pipeline_stable_cascade.py", line 429, in __call__
    predicted_latents = self.decoder(
  File "E:\stable-cascade-one-click-installer\venv\lib\site-packages\torch\nn\modules\module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "E:\stable-cascade-one-click-installer\venv\lib\site-packages\torch\nn\modules\module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "E:\stable-cascade-one-click-installer\venv\lib\site-packages\accelerate\hooks.py", line 166, in new_forward
    output = module._old_forward(*args, **kwargs)
  File "E:\stable-cascade-one-click-installer\venv\lib\site-packages\diffusers\models\unets\unet_stable_cascade.py", line 604, in forward
    level_outputs = self._down_encode(x, timestep_ratio_embed, clip)
  File "E:\stable-cascade-one-click-installer\venv\lib\site-packages\diffusers\models\unets\unet_stable_cascade.py", line 496, in _down_encode
    x = block(x, clip)
  File "E:\stable-cascade-one-click-installer\venv\lib\site-packages\torch\nn\modules\module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "E:\stable-cascade-one-click-installer\venv\lib\site-packages\torch\nn\modules\module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "E:\stable-cascade-one-click-installer\venv\lib\site-packages\diffusers\models\unets\unet_stable_cascade.py", line 109, in forward
    kv = torch.cat([norm_x.view(batch_size, channel, -1).transpose(1, 2), kv], dim=1)
RuntimeError: Sizes of tensors must match except in dimension 1. Expected size 2 but got size 1 for tensor number 1 in the list.
```

### Reproduction

```Python
import torch
from diffusers import StableCascadeDecoderPipeline, StableCascadePriorPipeline

prompt = "an image of a smiley, donning a spacesuit and helmet"

prior = StableCascadePriorPipeline.from_pretrained("stabilityai/stable-cascade-prior", variant="bf16", torch_dtype=torch.bfloat16)
decoder = StableCascadeDecoderPipeline.from_pretrained("stabilityai/stable-cascade", variant="bf16", torch_dtype=torch.float16)

prior.enable_model_cpu_offload()
prior_output = prior(
    prompt=prompt,
    height=1024,
    width=1024,
    num_images_per_prompt=2,
)

decoder.enable_model_cpu_offload()
decoder_output = decoder(
    image_embeddings=prior_output.image_embeddings.to(torch.float16),
    prompt=prompt,
    output_type="pil",
).images[0]
decoder_output.save("cascade.png")
```

### Logs

```shell
DEBUG:urllib3.connectionpool:Starting new HTTPS connection (1): huggingface.co:443
DEBUG:urllib3.connectionpool:https://huggingface.co:443 "GET /api/models/stabilityai/stable-cascade-prior HTTP/1.1" 200 3491
DEBUG:urllib3.connectionpool:https://huggingface.co:443 "HEAD /stabilityai/stable-cascade-prior/resolve/main/model_index.json HTTP/1.1" 200 0
DEBUG:urllib3.connectionpool:https://huggingface.co:443 "GET /api/models/stabilityai/stable-cascade HTTP/1.1" 200 9845
DEBUG:urllib3.connectionpool:https://huggingface.co:443 "HEAD /stabilityai/stable-cascade/resolve/main/model_index.json HTTP/1.1" 200 0
DEBUG:urllib3.connectionpool:Starting new HTTPS connection (1): huggingface.co:443
DEBUG:urllib3.connectionpool:https://huggingface.co:443 "GET /api/models/stabilityai/stable-cascade-prior HTTP/1.1" 200 3491
DEBUG:urllib3.connectionpool:https://huggingface.co:443 "HEAD /stabilityai/stable-cascade-prior/resolve/main/model_index.json HTTP/1.1" 200 0
DEBUG:urllib3.connectionpool:https://huggingface.co:443 "GET /api/models/stabilityai/stable-cascade HTTP/1.1" 200 9845
DEBUG:urllib3.connectionpool:https://huggingface.co:443 "HEAD /stabilityai/stable-cascade/resolve/main/model_index.json HTTP/1.1" 200 0
ERROR:__main__:Error occurred during image generation:
Traceback (most recent call last):
  File "E:\stable-cascade-one-click-installer\issue.py", line 23, in <module>
    decoder_output = decoder(
  File "E:\stable-cascade-one-click-installer\venv\lib\site-packages\torch\utils\_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "E:\stable-cascade-one-click-installer\venv\lib\site-packages\diffusers\pipelines\stable_cascade\pipeline_stable_cascade.py", line 429, in __call__
    predicted_latents = self.decoder(
  File "E:\stable-cascade-one-click-installer\venv\lib\site-packages\torch\nn\modules\module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "E:\stable-cascade-one-click-installer\venv\lib\site-packages\torch\nn\modules\module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "E:\stable-cascade-one-click-installer\venv\lib\site-packages\accelerate\hooks.py", line 166, in new_forward
    output = module._old_forward(*args, **kwargs)
  File "E:\stable-cascade-one-click-installer\venv\lib\site-packages\diffusers\models\unets\unet_stable_cascade.py", line 604, in forward
    level_outputs = self._down_encode(x, timestep_ratio_embed, clip)
  File "E:\stable-cascade-one-click-installer\venv\lib\site-packages\diffusers\models\unets\unet_stable_cascade.py", line 496, in _down_encode
    x = block(x, clip)
  File "E:\stable-cascade-one-click-installer\venv\lib\site-packages\torch\nn\modules\module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "E:\stable-cascade-one-click-installer\venv\lib\site-packages\torch\nn\modules\module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "E:\stable-cascade-one-click-installer\venv\lib\site-packages\diffusers\models\unets\unet_stable_cascade.py", line 109, in forward
    kv = torch.cat([norm_x.view(batch_size, channel, -1).transpose(1, 2), kv], dim=1)
RuntimeError: Sizes of tensors must match except in dimension 1. Expected size 2 but got size 1 for tensor number 1 in the list.
DEBUG:urllib3.connectionpool:Starting new HTTPS connection (1): huggingface.co:443
DEBUG:urllib3.connectionpool:https://huggingface.co:443 "GET /api/models/stabilityai/stable-cascade-prior HTTP/1.1" 200 3491
DEBUG:urllib3.connectionpool:https://huggingface.co:443 "HEAD /stabilityai/stable-cascade-prior/resolve/main/model_index.json HTTP/1.1" 200 0
DEBUG:urllib3.connectionpool:https://huggingface.co:443 "GET /api/models/stabilityai/stable-cascade HTTP/1.1" 200 9845
DEBUG:urllib3.connectionpool:https://huggingface.co:443 "HEAD /stabilityai/stable-cascade/resolve/main/model_index.json HTTP/1.1" 200 0
ERROR:__main__:Error occurred during image generation:
Traceback (most recent call last):
  File "E:\stable-cascade-one-click-installer\issue.py", line 21, in <module>
    logger.debug(f"norm_x shape: {norm_x.shape}")
NameError: name 'norm_x' is not defined
DEBUG:urllib3.connectionpool:Starting new HTTPS connection (1): huggingface.co:443
DEBUG:urllib3.connectionpool:https://huggingface.co:443 "GET /api/models/stabilityai/stable-cascade-prior HTTP/1.1" 200 3491
DEBUG:urllib3.connectionpool:https://huggingface.co:443 "HEAD /stabilityai/stable-cascade-prior/resolve/main/model_index.json HTTP/1.1" 200 0
DEBUG:urllib3.connectionpool:https://huggingface.co:443 "GET /api/models/stabilityai/stable-cascade HTTP/1.1" 200 9845
DEBUG:urllib3.connectionpool:https://huggingface.co:443 "HEAD /stabilityai/stable-cascade/resolve/main/model_index.json HTTP/1.1" 200 0
ERROR:__main__:Error occurred during image generation:
Traceback (most recent call last):
  File "E:\stable-cascade-one-click-installer\issue.py", line 22, in <module>
    logger.debug(f"norm_x shape: {norm_x.shape}")
NameError: name 'norm_x' is not defined
DEBUG:urllib3.connectionpool:Starting new HTTPS connection (1): huggingface.co:443
DEBUG:urllib3.connectionpool:https://huggingface.co:443 "GET /api/models/stabilityai/stable-cascade-prior HTTP/1.1" 200 3491
DEBUG:urllib3.connectionpool:https://huggingface.co:443 "HEAD /stabilityai/stable-cascade-prior/resolve/main/model_index.json HTTP/1.1" 200 0
DEBUG:urllib3.connectionpool:https://huggingface.co:443 "GET /api/models/stabilityai/stable-cascade HTTP/1.1" 200 9845
DEBUG:urllib3.connectionpool:https://huggingface.co:443 "HEAD /stabilityai/stable-cascade/resolve/main/model_index.json HTTP/1.1" 200 0
ERROR:__main__:Error occurred during image generation:
Traceback (most recent call last):
  File "E:\stable-cascade-one-click-installer\issue.py", line 23, in <module>
    decoder_output = decoder(
  File "E:\stable-cascade-one-click-installer\venv\lib\site-packages\torch\utils\_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "E:\stable-cascade-one-click-installer\venv\lib\site-packages\diffusers\pipelines\stable_cascade\pipeline_stable_cascade.py", line 429, in __call__
    predicted_latents = self.decoder(
  File "E:\stable-cascade-one-click-installer\venv\lib\site-packages\torch\nn\modules\module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "E:\stable-cascade-one-click-installer\venv\lib\site-packages\torch\nn\modules\module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "E:\stable-cascade-one-click-installer\venv\lib\site-packages\accelerate\hooks.py", line 166, in new_forward
    output = module._old_forward(*args, **kwargs)
  File "E:\stable-cascade-one-click-installer\venv\lib\site-packages\diffusers\models\unets\unet_stable_cascade.py", line 604, in forward
    level_outputs = self._down_encode(x, timestep_ratio_embed, clip)
  File "E:\stable-cascade-one-click-installer\venv\lib\site-packages\diffusers\models\unets\unet_stable_cascade.py", line 496, in _down_encode
    x = block(x, clip)
  File "E:\stable-cascade-one-click-installer\venv\lib\site-packages\torch\nn\modules\module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "E:\stable-cascade-one-click-installer\venv\lib\site-packages\torch\nn\modules\module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "E:\stable-cascade-one-click-installer\venv\lib\site-packages\diffusers\models\unets\unet_stable_cascade.py", line 109, in forward
    kv = torch.cat([norm_x.view(batch_size, channel, -1).transpose(1, 2), kv], dim=1)
RuntimeError: Sizes of tensors must match except in dimension 1. Expected size 2 but got size 1 for tensor number 1 in the list.
```


### System Info

- `diffusers` version: 0.27.0
- Platform: Windows-10-10.0.22631-SP0
- Python version: 3.10.9
- PyTorch version (GPU?): 2.2.1+cu121 (True)
- Huggingface_hub version: 0.21.4
- Transformers version: 4.38.2
- Accelerate version: 0.28.0
- xFormers version: not installed
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: No

### Who can help?

@DN6 @yiyi


Discussion:

Resolving patch:
diff --git a/src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade.py b/src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade.py
--- a/src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade.py
+++ b/src/diffusers/pipelines/stable_cascade/pipeline_stable_cascade.py
@@ -100,8 +100,10 @@ def __init__(
         )
         self.register_to_config(latent_dim_scale=latent_dim_scale)
 
-    def prepare_latents(self, image_embeddings, num_images_per_prompt, dtype, device, generator, latents, scheduler):
-        batch_size, channels, height, width = image_embeddings.shape
+    def prepare_latents(
+        self, batch_size, image_embeddings, num_images_per_prompt, dtype, device, generator, latents, scheduler
+    ):
+        _, channels, height, width = image_embeddings.shape
         latents_shape = (
             batch_size * num_images_per_prompt,
             4,
@@ -383,7 +385,19 @@ def __call__(
         )
         if isinstance(image_embeddings, list):
             image_embeddings = torch.cat(image_embeddings, dim=0)
-        batch_size = image_embeddings.shape[0]
+
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        # Compute the effective number of images per prompt
+        # We must account for the fact that the image embeddings from the prior can be generated with num_images_per_prompt > 1
+        # This results in a case where a single prompt is associated with multiple image embeddings
+        # Divide the number of image embeddings by the batch size to determine if this is the case.
+        num_images_per_prompt = num_images_per_prompt * (image_embeddings.shape[0] // batch_size)
 
         # 2. Encode caption
         if prompt_embeds is None and negative_prompt_embeds is None:
@@ -417,7 +431,7 @@ def __call__(
 
         # 5. Prepare latents
         latents = self.prepare_latents(
-            image_embeddings, num_images_per_prompt, dtype, device, generator, latents, self.scheduler
+            batch_size, image_embeddings, num_images_per_prompt, dtype, device, generator, latents, self.scheduler
         )
 
         # 6. Run denoising loop


Test patch:
diff --git a/tests/pipelines/stable_cascade/test_stable_cascade_decoder.py b/tests/pipelines/stable_cascade/test_stable_cascade_decoder.py
--- a/tests/pipelines/stable_cascade/test_stable_cascade_decoder.py
+++ b/tests/pipelines/stable_cascade/test_stable_cascade_decoder.py
@@ -33,6 +33,7 @@
     slow,
     torch_device,
 )
+from diffusers.utils.torch_utils import randn_tensor
 
 from ..test_pipelines_common import PipelineTesterMixin
 
@@ -246,6 +247,66 @@ def test_stable_cascade_decoder_prompt_embeds(self):
 
         assert np.abs(decoder_output_prompt.images - decoder_output_prompt_embeds.images).max() < 1e-5
 
+    def test_stable_cascade_decoder_single_prompt_multiple_image_embeddings(self):
+        device = "cpu"
+        components = self.get_dummy_components()
+
+        pipe = StableCascadeDecoderPipeline(**components)
+        pipe.set_progress_bar_config(disable=None)
+
+        prior_num_images_per_prompt = 2
+        decoder_num_images_per_prompt = 2
+        prompt = ["a cat"]
+        batch_size = len(prompt)
+
+        generator = torch.Generator(device)
+        image_embeddings = randn_tensor(
+            (batch_size * prior_num_images_per_prompt, 4, 4, 4), generator=generator.manual_seed(0)
+        )
+        decoder_output = pipe(
+            image_embeddings=image_embeddings,
+            prompt=prompt,
+            num_inference_steps=1,
+            output_type="np",
+            guidance_scale=0.0,
+            generator=generator.manual_seed(0),
+            num_images_per_prompt=decoder_num_images_per_prompt,
+        )
+
+        assert decoder_output.images.shape[0] == (
+            batch_size * prior_num_images_per_prompt * decoder_num_images_per_prompt
+        )
+
+    def test_stable_cascade_decoder_single_prompt_multiple_image_embeddings_with_guidance(self):
+        device = "cpu"
+        components = self.get_dummy_components()
+
+        pipe = StableCascadeDecoderPipeline(**components)
+        pipe.set_progress_bar_config(disable=None)
+
+        prior_num_images_per_prompt = 2
+        decoder_num_images_per_prompt = 2
+        prompt = ["a cat"]
+        batch_size = len(prompt)
+
+        generator = torch.Generator(device)
+        image_embeddings = randn_tensor(
+            (batch_size * prior_num_images_per_prompt, 4, 4, 4), generator=generator.manual_seed(0)
+        )
+        decoder_output = pipe(
+            image_embeddings=image_embeddings,
+            prompt=prompt,
+            num_inference_steps=1,
+            output_type="np",
+            guidance_scale=2.0,
+            generator=generator.manual_seed(0),
+            num_images_per_prompt=decoder_num_images_per_prompt,
+        )
+
+        assert decoder_output.images.shape[0] == (
+            batch_size * prior_num_images_per_prompt * decoder_num_images_per_prompt
+        )
+
 
 @slow
 @require_torch_gpu