(
self,
x_B_T_H_W_D: torch.Tensor,
emb_B_T_D: torch.Tensor,
adaln_lora_B_T_3D: Optional[torch.Tensor] = None,
)
| 904 | self.layer_norm.reset_parameters() |
| 905 | |
| 906 | def forward( |
| 907 | self, |
| 908 | x_B_T_H_W_D: torch.Tensor, |
| 909 | emb_B_T_D: torch.Tensor, |
| 910 | adaln_lora_B_T_3D: Optional[torch.Tensor] = None, |
| 911 | ): |
| 912 | if self.use_adaln_lora: |
| 913 | assert adaln_lora_B_T_3D is not None |
| 914 | shift_B_T_D, scale_B_T_D = ( |
| 915 | self.adaln_modulation(emb_B_T_D) + adaln_lora_B_T_3D[:, :, : 2 * self.hidden_size] |
| 916 | ).chunk(2, dim=-1) |
| 917 | else: |
| 918 | shift_B_T_D, scale_B_T_D = self.adaln_modulation(emb_B_T_D).chunk(2, dim=-1) |
| 919 | |
| 920 | shift_B_T_1_1_D, scale_B_T_1_1_D = rearrange(shift_B_T_D, "b t d -> b t 1 1 d"), rearrange( |
| 921 | scale_B_T_D, "b t d -> b t 1 1 d" |
| 922 | ) |
| 923 | |
| 924 | def _fn( |
| 925 | _x_B_T_H_W_D: torch.Tensor, |
| 926 | _norm_layer: nn.Module, |
| 927 | _scale_B_T_1_1_D: torch.Tensor, |
| 928 | _shift_B_T_1_1_D: torch.Tensor, |
| 929 | ) -> torch.Tensor: |
| 930 | return _norm_layer(_x_B_T_H_W_D) * (1 + _scale_B_T_1_1_D) + _shift_B_T_1_1_D |
| 931 | |
| 932 | x_B_T_H_W_D = _fn(x_B_T_H_W_D, self.layer_norm, scale_B_T_1_1_D, shift_B_T_1_1_D) |
| 933 | x_B_T_H_W_O = self.linear(x_B_T_H_W_D) |
| 934 | return x_B_T_H_W_O |
| 935 | |
| 936 | |
| 937 | class Block(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected