[Mochi-1] ensuring to compute the fourier features in FP32 in Mochi e…

…ncoder (huggingface#10031)

compute fourier features in FP32.

src/diffusers/models/autoencoders/autoencoder_kl_mochi.py

@@ -437,7 +437,8 @@ def __init__(self, start: int = 6, stop: int = 8, step: int = 1):
 
     def forward(self, inputs: torch.Tensor) -> torch.Tensor:
         r"""Forward method of the `FourierFeatures` class."""
-
+        original_dtype = inputs.dtype
+        inputs = inputs.to(torch.float32)
         num_channels = inputs.shape[1]
         num_freqs = (self.stop - self.start) // self.step
 
@@ -450,7 +451,7 @@ def forward(self, inputs: torch.Tensor) -> torch.Tensor:
         # Scale channels by frequency.
         h = w * h
 
-        return torch.cat([inputs, torch.sin(h), torch.cos(h)], dim=1)
+        return torch.cat([inputs, torch.sin(h), torch.cos(h)], dim=1).to(original_dtype)
 
 
 class MochiEncoder3D(nn.Module):