Rope in float32 for mps or npu compatibility (#12665)

rope in float32

Author: DavidBert

src/diffusers/models/transformers/transformer_prx.py

@@ -275,7 +275,12 @@ def __init__(self, dim: int, theta: int, axes_dim: List[int]):
 
     def rope(self, pos: torch.Tensor, dim: int, theta: int) -> torch.Tensor:
         assert dim % 2 == 0
-        scale = torch.arange(0, dim, 2, dtype=torch.float64, device=pos.device) / dim
+
+        is_mps = pos.device.type == "mps"
+        is_npu = pos.device.type == "npu"
+        dtype = torch.float32 if (is_mps or is_npu) else torch.float64
+
+        scale = torch.arange(0, dim, 2, dtype=dtype, device=pos.device) / dim
         omega = 1.0 / (theta**scale)
         out = pos.unsqueeze(-1) * omega.unsqueeze(0)
         out = torch.stack([torch.cos(out), -torch.sin(out), torch.sin(out), torch.cos(out)], dim=-1)