LattePipeline fails with dtype mismatch

[vladmandic]

Describe the bug

LattePipeline fails with dtype mismatch in transformer block during attention phase if enable_temporal_attentions is enabled (which is the default). Disabling temporal attention skips the affected block of code, so model works.

Reproduction

load diffusers.LattePipeline with specific torch_dtype, it fails for both torch.float16 and torch.bfloat16

Logs

│ /home/vlado/dev/sdnext/venv/lib/python3.12/site-packages/diffusers/pipelines/latte/pipeline_latte.py:819 in __call__ │
│ │
│ 818 │ │ │ │ # predict noise model_output │
│ ❱ 819 │ │ │ │ noise_pred = self.transformer( │
│ 820 │ │ │ │ │ hidden_states=latent_model_input, │
│ │
│ /home/vlado/dev/sdnext/venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1739 in _wrapped_call_impl │
│ │
│ 1738 │ │ else: │
│ ❱ 1739 │ │ │ return self._call_impl(*args, **kwargs) │
│ 1740 │
│ │
│ /home/vlado/dev/sdnext/venv/lib/python3.12/site-packages/torch/nn/modules/module.py:1750 in _call_impl │
│ │
│ 1749 │ │ │ │ or _global_forward_hooks or _global_forward_pre_hooks): │
│ ❱ 1750 │ │ │ return forward_call(*args, **kwargs) │
│ 1751 │
│ │
│ /home/vlado/dev/sdnext/venv/lib/python3.12/site-packages/accelerate/hooks.py:176 in new_forward │
│ │
│ 175 │ │ else: │
│ ❱ 176 │ │ │ output = module._old_forward(*args, **kwargs) │
│ 177 │ │ return module._hf_hook.post_forward(module, output) │
│ │
│ /home/vlado/dev/sdnext/venv/lib/python3.12/site-packages/diffusers/models/transformers/latte_transformer_3d.py:291 in forward │
│ │
│ 290 │ │ │ │ else: │
│ ❱ 291 │ │ │ │ │ hidden_states = temp_block( │
│ 292 │ │ │ │ │ │ hidden_states, 
...
│ /home/vlado/dev/sdnext/venv/lib/python3.12/site-packages/torch/nn/modules/linear.py:125 in forward │
│ │
│ 124 │ def forward(self, input: Tensor) -> Tensor: │
│ ❱ 125 │ │ return F.linear(input, self.weight, self.bias) │
│ 126 │
╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯
RuntimeError: mat1 and mat2 must have the same dtype, but got Float and Half

System Info

diffusers==main

Who can help?

@a-r-r-o-w @DN6