Update text2vid_torch2.py

text2vid_torch2.py

@@ -167,7 +167,8 @@ class AttnProcessor2_0:
         hidden_states = hidden_states / attn.rescale_output_factor
 
         return hidden_states
-    
+
+    """
     def get_qk(
         self, query, key):
         r"""
@@ -221,7 +222,53 @@ class AttnProcessor2_0:
     
 
         return query, key, dynamic_lambda, key1
+    """
+
+    def get_qk(
+        self, query, key):
+        r"""
+        Compute the attention scores.
+        Args:
+            query (`torch.Tensor`): The query tensor.
+            key (`torch.Tensor`): The key tensor.
+            attention_mask (`torch.Tensor`, *optional*): The attention mask to use. If `None`, no mask is applied.
+        Returns:
+            `torch.Tensor`: The attention probabilities/scores.
+        """
+        q_old = query.clone()
+        k_old = key.clone()
+        dynamic_lambda = None
+        key1 = None
+        
+        if self.use_last_attn_slice:
+            if self.last_attn_slice is not None:
+                
+                query_list = self.last_attn_slice[0]
+                key_list = self.last_attn_slice[1]
+                
+                if query.shape[1] == self.num_frames and query.shape == key.shape:
+                    key1 = key.clone()
+                    key1[:,:1,:key_list.shape[2]] = key_list[:,:1]
+                    dynamic_lambda = torch.tensor([1 + self.LAMBDA * (i/50) for i in range(self.num_frames)]).to(key.dtype).cuda()
     
+                if q_old.shape == k_old.shape and q_old.shape[1] != self.num_frames:
+                    batch_dim = query_list.shape[0] // self.bs
+                    all_dim = query.shape[0] // self.bs
+                    for i in range(self.bs):
+                        # Ensure slice dimensions match
+                        target_size = min(query[i*all_dim:(i*all_dim) + batch_dim, :query_list.shape[1], :query_list.shape[2]].size(0), 
+                                          query_list[i*batch_dim:(i+1)*batch_dim].size(0))
+                        
+                        # Assign values from query_list to query
+                        query[i*all_dim:(i*all_dim) + target_size, :query_list.shape[1], :query_list.shape[2]] = \
+                            query_list[i*batch_dim:i*batch_dim + target_size]
+    
+        if self.save_last_attn_slice:
+            self.last_attn_slice = [query, key]
+            self.save_last_attn_slice = False
+    
+        return query, key, dynamic_lambda, key1
+
 
 def init_attention_func(unet):