Update text2vid_torch2.py

text2vid_torch2.py

@@ -224,8 +224,9 @@ class AttnProcessor2_0:
         return query, key, dynamic_lambda, key1
     '''
 
-    def get_qk(
-        self, query, key):
+    import torch
+
+    def get_qk(self, query, key):
         r"""
         Compute the attention scores.
         Args:
@@ -240,45 +241,66 @@ class AttnProcessor2_0:
         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()
+        try:
+            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()
     
-                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))
+                        # 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]
     
-                        # Check if the target size is compatible with the query slice dimensions
-                        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
+                        dynamic_lambda = torch.tensor([1 + self.LAMBDA * (i/50) for i in range(self.num_frames)]).to(key.dtype).cuda()
     
-                        if query_slice_shape[1] != query_list_slice_shape[1]:  # Dimension mismatch check
-                            print(f"Warning: Dimension mismatch. query_slice_shape: {query_slice_shape}, query_list_slice_shape: {query_list_slice_shape}. Adjusting to compatible sizes.")
-                            # Adjust to the smaller dimension
-                            target_size = min(query_slice_shape[1], query_list_slice_shape[1])
+                    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):
+                            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
     
-                        # 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 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])
     
-        if self.save_last_attn_slice:
-            self.last_attn_slice = [query, key]
-            self.save_last_attn_slice = False
+                            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
     
+        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
+        
         return query, key, dynamic_lambda, key1
 
 
+
+
 def init_attention_func(unet):
 
     for name, module in unet.named_modules():