Repository: huggingface/diffusers

Issue:
Exception raised in `LDMTextToImagePipeline`
### Describe the bug

Invoking a LDMTextToImagePipeline loaded from model `CompVis/ldm-text2im-large-256` results in an exception, in the current `main` branch.

I haven't checked the code yet, maybe I'm using the wrong model.

### Reproduction

```
from diffusers import LDMTextToImagePipeline
ldm = LDMTextToImagePipeline.from_pretrained("CompVis/ldm-text2im-large-256")
ldm("Some prompt", num_inference_steps=1)
```

### Logs

```shell
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Input In [6], in <cell line: 1>()
----> 1 ldm("Some prompt", num_inference_steps=1)
      2 print(ldm.vqvae.device, ldm.bert.device, ldm.unet.device)

File ~/miniconda3/envs/hf/lib/python3.9/site-packages/torch/autograd/grad_mode.py:27, in _DecoratorContextManager.__call__.<locals>.decorate_context(*args, **kwargs)
     24 @functools.wraps(func)
     25 def decorate_context(*args, **kwargs):
     26     with self.clone():
---> 27         return func(*args, **kwargs)

File ~/code/hf/diffusers/diffusers/src/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion.py:83, in LDMTextToImagePipeline.__call__(self, prompt, batch_size, generator, torch_device, eta, guidance_scale, num_inference_steps, output_type)
     80     context = torch.cat([uncond_embeddings, text_embeddings])
     82 # predict the noise residual
---> 83 noise_pred = self.unet(latents_input, t, encoder_hidden_states=context)["sample"]
     84 # perform guidance
     85 if guidance_scale != 1.0:

File ~/miniconda3/envs/hf/lib/python3.9/site-packages/torch/nn/modules/module.py:1130, in Module._call_impl(self, *input, **kwargs)
   1126 # If we don't have any hooks, we want to skip the rest of the logic in
   1127 # this function, and just call forward.
   1128 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1129         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1130     return forward_call(*input, **kwargs)
   1131 # Do not call functions when jit is used
   1132 full_backward_hooks, non_full_backward_hooks = [], []

File ~/code/hf/diffusers/diffusers/src/diffusers/models/unet_2d_condition.py:150, in UNet2DConditionModel.forward(self, sample, timestep, encoder_hidden_states)
    147 for downsample_block in self.down_blocks:
    149     if hasattr(downsample_block, "attentions") and downsample_block.attentions is not None:
--> 150         sample, res_samples = downsample_block(
    151             hidden_states=sample, temb=emb, encoder_hidden_states=encoder_hidden_states
    152         )
    153     else:
    154         sample, res_samples = downsample_block(hidden_states=sample, temb=emb)

File ~/miniconda3/envs/hf/lib/python3.9/site-packages/torch/nn/modules/module.py:1130, in Module._call_impl(self, *input, **kwargs)
   1126 # If we don't have any hooks, we want to skip the rest of the logic in
   1127 # this function, and just call forward.
   1128 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1129         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1130     return forward_call(*input, **kwargs)
   1131 # Do not call functions when jit is used
   1132 full_backward_hooks, non_full_backward_hooks = [], []

File ~/code/hf/diffusers/diffusers/src/diffusers/models/unet_blocks.py:505, in CrossAttnDownBlock2D.forward(self, hidden_states, temb, encoder_hidden_states)
    503 for resnet, attn in zip(self.resnets, self.attentions):
    504     hidden_states = resnet(hidden_states, temb)
--> 505     hidden_states = attn(hidden_states, context=encoder_hidden_states)
    506     output_states += (hidden_states,)
    508 if self.downsamplers is not None:

File ~/miniconda3/envs/hf/lib/python3.9/site-packages/torch/nn/modules/module.py:1130, in Module._call_impl(self, *input, **kwargs)
   1126 # If we don't have any hooks, we want to skip the rest of the logic in
   1127 # this function, and just call forward.
   1128 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1129         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1130     return forward_call(*input, **kwargs)
   1131 # Do not call functions when jit is used
   1132 full_backward_hooks, non_full_backward_hooks = [], []

File ~/code/hf/diffusers/diffusers/src/diffusers/models/attention.py:168, in SpatialTransformer.forward(self, x, context)
    166 x = x.permute(0, 2, 3, 1).reshape(b, h * w, c)
    167 for block in self.transformer_blocks:
--> 168     x = block(x, context=context)
    169 x = x.reshape(b, h, w, c).permute(0, 3, 1, 2)
    170 x = self.proj_out(x)

File ~/miniconda3/envs/hf/lib/python3.9/site-packages/torch/nn/modules/module.py:1130, in Module._call_impl(self, *input, **kwargs)
   1126 # If we don't have any hooks, we want to skip the rest of the logic in
   1127 # this function, and just call forward.
   1128 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1129         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1130     return forward_call(*input, **kwargs)
   1131 # Do not call functions when jit is used
   1132 full_backward_hooks, non_full_backward_hooks = [], []

File ~/code/hf/diffusers/diffusers/src/diffusers/models/attention.py:197, in BasicTransformerBlock.forward(self, x, context)
    195 def forward(self, x, context=None):
    196     x = self.attn1(self.norm1(x)) + x
--> 197     x = self.attn2(self.norm2(x), context=context) + x
    198     x = self.ff(self.norm3(x)) + x
    199     return x

File ~/miniconda3/envs/hf/lib/python3.9/site-packages/torch/nn/modules/module.py:1130, in Module._call_impl(self, *input, **kwargs)
   1126 # If we don't have any hooks, we want to skip the rest of the logic in
   1127 # this function, and just call forward.
   1128 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1129         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1130     return forward_call(*input, **kwargs)
   1131 # Do not call functions when jit is used
   1132 full_backward_hooks, non_full_backward_hooks = [], []

File ~/code/hf/diffusers/diffusers/src/diffusers/models/attention.py:245, in CrossAttention.forward(self, x, context, mask)
    242 k = self.reshape_heads_to_batch_dim(k)
    243 v = self.reshape_heads_to_batch_dim(v)
--> 245 sim = torch.einsum("b i d, b j d -> b i j", q, k) * self.scale
    247 if exists(mask):
    248     mask = mask.reshape(batch_size, -1)

File ~/miniconda3/envs/hf/lib/python3.9/site-packages/torch/functional.py:360, in einsum(*args)
    356     # recurse incase operands contains value that has torch function
    357     # in the original implementation this line is omitted
    358     return einsum(equation, *_operands)
--> 360 return _VF.einsum(equation, operands)

RuntimeError: einsum(): operands do not broadcast with remapped shapes [original->remapped]: [88, 1024, 40]->[88, 1024, 1, 40] [8, 77, 40]->[8, 1, 77, 40]
```


### System Info

```shell
diffusers @ main (df90f0c)
PyTorch 1.12.1+cu113
```



Discussion:
@pcuenca good catch! Running `pipe(prompt_str)` instead of `pipe([prompt_str])` is a recent feature and our tests didn't have coverage of that, I'll push a fix asap :) 
cc @patil-suraj 

Resolving patch:
diff --git a/src/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion.py b/src/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion.py
--- a/src/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion.py
+++ b/src/diffusers/pipelines/latent_diffusion/pipeline_latent_diffusion.py
@@ -1,5 +1,5 @@
 import inspect
-from typing import Optional, Tuple, Union
+from typing import List, Optional, Tuple, Union
 
 import torch
 import torch.nn as nn
@@ -24,20 +24,30 @@ def __init__(self, vqvae, bert, tokenizer, unet, scheduler):
     @torch.no_grad()
     def __call__(
         self,
-        prompt,
-        batch_size=1,
-        generator=None,
-        torch_device=None,
-        eta=0.0,
-        guidance_scale=1.0,
-        num_inference_steps=50,
-        output_type="pil",
+        prompt: Union[str, List[str]],
+        height: Optional[int] = 256,
+        width: Optional[int] = 256,
+        num_inference_steps: Optional[int] = 50,
+        guidance_scale: Optional[float] = 1.0,
+        eta: Optional[float] = 0.0,
+        generator: Optional[torch.Generator] = None,
+        torch_device: Optional[Union[str, torch.device]] = None,
+        output_type: Optional[str] = "pil",
     ):
         # eta corresponds to η in paper and should be between [0, 1]
 
         if torch_device is None:
             torch_device = "cuda" if torch.cuda.is_available() else "cpu"
-        batch_size = len(prompt)
+
+        if isinstance(prompt, str):
+            batch_size = 1
+        elif isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+
+        if height % 8 != 0 or width % 8 != 0:
+            raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
 
         self.unet.to(torch_device)
         self.vqvae.to(torch_device)
@@ -53,7 +63,7 @@ def __call__(
         text_embeddings = self.bert(text_input.input_ids.to(torch_device))[0]
 
         latents = torch.randn(
-            (batch_size, self.unet.in_channels, self.unet.sample_size, self.unet.sample_size),
+            (batch_size, self.unet.in_channels, height // 8, width // 8),
             generator=generator,
         )
         latents = latents.to(torch_device)


Test patch:
diff --git a/tests/test_modeling_utils.py b/tests/test_modeling_utils.py
--- a/tests/test_modeling_utils.py
+++ b/tests/test_modeling_utils.py
@@ -854,7 +854,7 @@ def test_ldm_text2img_fast(self):
 
         prompt = "A painting of a squirrel eating a burger"
         generator = torch.manual_seed(0)
-        image = ldm([prompt], generator=generator, num_inference_steps=1, output_type="numpy")["sample"]
+        image = ldm(prompt, generator=generator, num_inference_steps=1, output_type="numpy")["sample"]
 
         image_slice = image[0, -3:, -3:, -1]