(
self,
x_B_T_H_W_D: torch.Tensor,
emb_B_T_D: torch.Tensor,
crossattn_emb: torch.Tensor,
rope_emb_L_1_1_D: Optional[torch.Tensor] = None,
adaln_lora_B_T_3D: Optional[torch.Tensor] = None,
extra_per_block_pos_emb: Optional[torch.Tensor] = None,
)
| 1033 | self.mlp.init_weights() |
| 1034 | |
| 1035 | def forward( |
| 1036 | self, |
| 1037 | x_B_T_H_W_D: torch.Tensor, |
| 1038 | emb_B_T_D: torch.Tensor, |
| 1039 | crossattn_emb: torch.Tensor, |
| 1040 | rope_emb_L_1_1_D: Optional[torch.Tensor] = None, |
| 1041 | adaln_lora_B_T_3D: Optional[torch.Tensor] = None, |
| 1042 | extra_per_block_pos_emb: Optional[torch.Tensor] = None, |
| 1043 | ) -> torch.Tensor: |
| 1044 | if extra_per_block_pos_emb is not None: |
| 1045 | x_B_T_H_W_D = x_B_T_H_W_D + extra_per_block_pos_emb |
| 1046 | |
| 1047 | if self.use_adaln_lora: |
| 1048 | shift_self_attn_B_T_D, scale_self_attn_B_T_D, gate_self_attn_B_T_D = ( |
| 1049 | self.adaln_modulation_self_attn(emb_B_T_D) + adaln_lora_B_T_3D |
| 1050 | ).chunk(3, dim=-1) |
| 1051 | shift_cross_attn_B_T_D, scale_cross_attn_B_T_D, gate_cross_attn_B_T_D = ( |
| 1052 | self.adaln_modulation_cross_attn(emb_B_T_D) + adaln_lora_B_T_3D |
| 1053 | ).chunk(3, dim=-1) |
| 1054 | shift_mlp_B_T_D, scale_mlp_B_T_D, gate_mlp_B_T_D = ( |
| 1055 | self.adaln_modulation_mlp(emb_B_T_D) + adaln_lora_B_T_3D |
| 1056 | ).chunk(3, dim=-1) |
| 1057 | else: |
| 1058 | shift_self_attn_B_T_D, scale_self_attn_B_T_D, gate_self_attn_B_T_D = self.adaln_modulation_self_attn( |
| 1059 | emb_B_T_D |
| 1060 | ).chunk(3, dim=-1) |
| 1061 | shift_cross_attn_B_T_D, scale_cross_attn_B_T_D, gate_cross_attn_B_T_D = self.adaln_modulation_cross_attn( |
| 1062 | emb_B_T_D |
| 1063 | ).chunk(3, dim=-1) |
| 1064 | shift_mlp_B_T_D, scale_mlp_B_T_D, gate_mlp_B_T_D = self.adaln_modulation_mlp(emb_B_T_D).chunk(3, dim=-1) |
| 1065 | |
| 1066 | # Reshape tensors from (B, T, D) to (B, T, 1, 1, D) for broadcasting |
| 1067 | shift_self_attn_B_T_1_1_D = rearrange(shift_self_attn_B_T_D, "b t d -> b t 1 1 d") |
| 1068 | scale_self_attn_B_T_1_1_D = rearrange(scale_self_attn_B_T_D, "b t d -> b t 1 1 d") |
| 1069 | gate_self_attn_B_T_1_1_D = rearrange(gate_self_attn_B_T_D, "b t d -> b t 1 1 d") |
| 1070 | |
| 1071 | shift_cross_attn_B_T_1_1_D = rearrange(shift_cross_attn_B_T_D, "b t d -> b t 1 1 d") |
| 1072 | scale_cross_attn_B_T_1_1_D = rearrange(scale_cross_attn_B_T_D, "b t d -> b t 1 1 d") |
| 1073 | gate_cross_attn_B_T_1_1_D = rearrange(gate_cross_attn_B_T_D, "b t d -> b t 1 1 d") |
| 1074 | |
| 1075 | shift_mlp_B_T_1_1_D = rearrange(shift_mlp_B_T_D, "b t d -> b t 1 1 d") |
| 1076 | scale_mlp_B_T_1_1_D = rearrange(scale_mlp_B_T_D, "b t d -> b t 1 1 d") |
| 1077 | gate_mlp_B_T_1_1_D = rearrange(gate_mlp_B_T_D, "b t d -> b t 1 1 d") |
| 1078 | |
| 1079 | B, T, H, W, D = x_B_T_H_W_D.shape |
| 1080 | |
| 1081 | def _fn(_x_B_T_H_W_D, _norm_layer, _scale_B_T_1_1_D, _shift_B_T_1_1_D): |
| 1082 | return _norm_layer(_x_B_T_H_W_D) * (1 + _scale_B_T_1_1_D) + _shift_B_T_1_1_D |
| 1083 | |
| 1084 | normalized_x_B_T_H_W_D = _fn( |
| 1085 | x_B_T_H_W_D, |
| 1086 | self.layer_norm_self_attn, |
| 1087 | scale_self_attn_B_T_1_1_D, |
| 1088 | shift_self_attn_B_T_1_1_D, |
| 1089 | ) |
| 1090 | result_B_T_H_W_D = rearrange( |
| 1091 | self.self_attn( |
| 1092 | # normalized_x_B_T_HW_D, |
nothing calls this directly
no outgoing calls
no test coverage detected