kanshan/diffusers
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[Mochi-1] ensuring to compute the fourier features in FP32 in Mochi encoder (#10031)

Browse Source 
compute fourier features in FP32.
This commit is contained in:
Sayak Paul 2024-11-29 14:15:00 +05:30 committed by GitHub
parent 6b288ec44d
commit c96bfa5c80
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
src/diffusers/models/autoencoders/autoencoder_kl_mochi.py
View File

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