Use float32 RoPE freqs in Wan with MPS backends (#11643)

Use float32 for RoPE on MPS in Wan

Author: hvaara

src/diffusers/models/transformers/transformer_wan.py

@@ -72,7 +72,8 @@ def __call__(
         if rotary_emb is not None:
 
             def apply_rotary_emb(hidden_states: torch.Tensor, freqs: torch.Tensor):
-                x_rotated = torch.view_as_complex(hidden_states.to(torch.float64).unflatten(3, (-1, 2)))
+                dtype = torch.float32 if hidden_states.device.type == "mps" else torch.float64
+                x_rotated = torch.view_as_complex(hidden_states.to(dtype).unflatten(3, (-1, 2)))
                 x_out = torch.view_as_real(x_rotated * freqs).flatten(3, 4)
                 return x_out.type_as(hidden_states)
 
@@ -190,9 +191,10 @@ def __init__(
         t_dim = attention_head_dim - h_dim - w_dim
 
         freqs = []
+        freqs_dtype = torch.float32 if torch.backends.mps.is_available() else torch.float64
         for dim in [t_dim, h_dim, w_dim]:
             freq = get_1d_rotary_pos_embed(
-                dim, max_seq_len, theta, use_real=False, repeat_interleave_real=False, freqs_dtype=torch.float64
+                dim, max_seq_len, theta, use_real=False, repeat_interleave_real=False, freqs_dtype=freqs_dtype
             )
             freqs.append(freq)
         self.freqs = torch.cat(freqs, dim=1)