Exception raised in `LDMTextToImagePipeline`

[pcuenca]

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

---------------------------------------------------------------------------
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

diffusers @ main (df90f0c)
PyTorch 1.12.1+cu113

[anton-l]

@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

[pcuenca]

Running pipe(prompt_str) instead of pipe([prompt_str]) is a recent feature and our tests didn't have coverage of that

Ah, makes sense!