Update text2vid_torch2.py

text2vid_torch2.py

@@ -224,67 +224,57 @@ class AttnProcessor2_0:
         return query, key, dynamic_lambda, key1
     '''
 
-    import torch
-
     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
-        
         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()
     
-                        # Ensure the batch dimension of key1 and key_list match
-                        batch_size_key1 = key1.shape[0]
-                        batch_size_key_list = key_list.shape[0]
-                        
-                        if batch_size_key1 != batch_size_key_list:
-                            # Handle mismatch: either pad or slice to match sizes
-                            if batch_size_key1 > batch_size_key_list:
-                                # Pad key_list if key1 batch size is larger
-                                padding = (0, 0, 0, batch_size_key1 - batch_size_key_list)  # (left, right, top, bottom)
-                                key_list = torch.nn.functional.pad(key_list, padding, "constant", 0)
-                            else:
-                                # Slice key1 if key_list batch size is larger
-                                key1 = key1[:batch_size_key_list]
-                        
-                        # Proceed with assignment after matching batch dimensions
-                        key1[:,:1,:key_list.shape[2]] = key_list[:,:1]
+                        # 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 = torch.tensor([1 + self.LAMBDA * (i/50) for i in range(self.num_frames)]).to(key.dtype).cuda()
+                        # 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):
-                            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))
-                            query_slice_shape = query[i*all_dim:(i*all_dim) + target_size, :query_list.shape[1], :query_list.shape[2]].shape
-                            query_list_slice_shape = query_list[i*batch_dim:i*batch_dim + target_size].shape
-    
-                            if query_slice_shape[1] != query_list_slice_shape[1]:
-                                print(f"Warning: Dimension mismatch. query_slice_shape: {query_slice_shape}, query_list_slice_shape: {query_list_slice_shape}. Adjusting to compatible sizes.")
-                                target_size = min(query_slice_shape[1], query_list_slice_shape[1])
+                            # 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]
     
-                            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]
+                            # 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]
@@ -293,9 +283,9 @@ class AttnProcessor2_0:
         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()
-            raise  # Re-raise the error to let the caller handle it further if needed
-        
+            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