Update text2vid_torch2.py

text2vid_torch2.py

@@ -168,7 +168,7 @@ class AttnProcessor2_0:
 
         return hidden_states
 
-    '''
+    
     def get_qk(
         self, query, key):
         r"""
@@ -222,73 +222,6 @@ 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.
-        """
-        try:
-            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]
-    
-                    # Ensure that shapes are compatible before performing assignments
-                    if query.shape[1] == self.num_frames and query.shape == key.shape:
-                        key1 = key.clone()
-    
-                        # Safety check: ensure key1 can receive the value from key_list without causing size mismatch
-                        if key1.shape[0] >= key_list.shape[0]:
-                            key1[:, :1, :key_list.shape[2]] = key_list[:, :1]
-                        else:
-                            raise RuntimeError(f"Shape mismatch: key1 has {key1.shape[0]} batches, but key_list has {key_list.shape[0]} batches.")
-    
-                        # Dynamic lambda scaling
-                        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:
-                        # Ensure batch size division is valid
-                        batch_dim = query_list.shape[0] // self.bs
-                        all_dim = query.shape[0] // self.bs
-                        
-                        for i in range(self.bs):
-                            # Safety check for slicing indices to avoid memory access errors
-                            query_slice = query[i * all_dim:(i * all_dim) + batch_dim, :query_list.shape[1], :query_list.shape[2]]
-                            target_slice = query_list[i * batch_dim:(i + 1) * batch_dim]
-    
-                            # Validate dimensions match before assignment
-                            if query_slice.shape == target_slice.shape:
-                                query_slice[:] = target_slice
-                            else:
-                                raise RuntimeError(f"Shape mismatch during slicing: query slice shape {query_slice.shape}, target slice shape {target_slice.shape}")
-    
-            if self.save_last_attn_slice:
-                self.last_attn_slice = [query, key]
-                self.save_last_attn_slice = False
-    
-        except RuntimeError as e:
-            # If a RuntimeError happens, catch it and clean CUDA memory
-            print(f"RuntimeError occurred: {e}. Cleaning up CUDA memory...")
-            torch.cuda.empty_cache()  # Free up CUDA memory to avoid further issues
-            raise  # Re-raise the error to propagate it if needed
-    
-        return query, key, dynamic_lambda, key1
-
-
 
 
 def init_attention_func(unet):