| 83 | |
| 84 | |
| 85 | class SanaModulatedNorm(nn.Module): |
| 86 | def __init__(self, dim: int, elementwise_affine: bool = False, eps: float = 1e-6): |
| 87 | super().__init__() |
| 88 | self.norm = nn.LayerNorm(dim, elementwise_affine=elementwise_affine, eps=eps) |
| 89 | |
| 90 | def forward( |
| 91 | self, hidden_states: torch.Tensor, temb: torch.Tensor, scale_shift_table: torch.Tensor |
| 92 | ) -> torch.Tensor: |
| 93 | hidden_states = self.norm(hidden_states) |
| 94 | shift, scale = (scale_shift_table[None] + temb[:, None].to(scale_shift_table.device)).chunk(2, dim=1) |
| 95 | hidden_states = hidden_states * (1 + scale) + shift |
| 96 | return hidden_states |
| 97 | |
| 98 | |
| 99 | class SanaCombinedTimestepGuidanceEmbeddings(nn.Module): |