Update text2vid_torch2.py

text2vid_torch2.py

@@ -258,7 +258,16 @@ class AttnProcessor2_0:
                         # 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))
-                        
+    
+                        # 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
+    
+                        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])
+    
                         # 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]