(
self,
hidden_states,
encoder_hidden_states=None,
timestep=None,
attention_mask=None,
video_length=None,
iter_cur=None,
save_kv=None,
source_masks=None,
target_masks=None,
long_context=None,
)
| 104 | ) |
| 105 | |
| 106 | def forward( |
| 107 | self, |
| 108 | hidden_states, |
| 109 | encoder_hidden_states=None, |
| 110 | timestep=None, |
| 111 | attention_mask=None, |
| 112 | video_length=None, |
| 113 | iter_cur=None, |
| 114 | save_kv=None, |
| 115 | source_masks=None, |
| 116 | target_masks=None, |
| 117 | long_context=None, |
| 118 | ): |
| 119 | |
| 120 | norm_hidden_states = ( |
| 121 | self.norm1(hidden_states, timestep) |
| 122 | if self.use_ada_layer_norm |
| 123 | else self.norm1(hidden_states) |
| 124 | ) |
| 125 | if self.unet_use_cross_frame_attention: |
| 126 | hidden_states = ( |
| 127 | self.attn1( |
| 128 | norm_hidden_states, |
| 129 | attention_mask=attention_mask, |
| 130 | video_length=video_length, |
| 131 | iter_cur=iter_cur, |
| 132 | save_kv=save_kv, |
| 133 | source_masks=source_masks, |
| 134 | target_masks=target_masks, |
| 135 | long_context=long_context, |
| 136 | ) |
| 137 | + hidden_states |
| 138 | ) |
| 139 | else: |
| 140 | if iter_cur is not None and save_kv is not None: |
| 141 | hidden_states = ( |
| 142 | self.attn1(norm_hidden_states, attention_mask=attention_mask, iter_cur=iter_cur,save_kv=save_kv,source_masks=source_masks, target_masks=target_masks, long_context=long_context) |
| 143 | + hidden_states |
| 144 | ) |
| 145 | else: |
| 146 | hidden_states = ( |
| 147 | self.attn1(norm_hidden_states, attention_mask=attention_mask,source_masks=source_masks, target_masks=target_masks,)+hidden_states |
| 148 | ) |
| 149 | |
| 150 | if self.attn2 is not None: |
| 151 | # Cross-Attention |
| 152 | norm_hidden_states = ( |
| 153 | self.norm2(hidden_states, timestep) |
| 154 | if self.use_ada_layer_norm |
| 155 | else self.norm2(hidden_states) |
| 156 | ) |
| 157 | hidden_states = ( |
| 158 | self.attn2( |
| 159 | norm_hidden_states, |
| 160 | encoder_hidden_states=encoder_hidden_states, |
| 161 | attention_mask=attention_mask, |
| 162 | ) |
| 163 | + hidden_states |
nothing calls this directly
no outgoing calls
no test coverage detected