diff --git a/CrossAttentionPatch.py b/CrossAttentionPatch.py
index 06f153e..e0dd5c5 100644
--- a/CrossAttentionPatch.py
+++ b/CrossAttentionPatch.py
@@ -82,31 +82,67 @@ class CrossAttentionPatch:
                         cond = cond_alt[t_idx]
                         del cond_alt
 
-                if weight == 0:
-                    continue
+                #if isinstance(weight, torch.Tensor):
+                #    weight = tensor_to_size(weight, batch_prompt)
+                #    weight = weight.repeat(len(cond_or_uncond), 1, 1)
+                #elif weight == 0:
+                #    continue
 
                 if unfold_batch and cond.shape[0] > 1:
                     # Check AnimateDiff context window
                     if ad_params is not None and ad_params["sub_idxs"] is not None:
                         # if image length matches or exceeds full_length get sub_idx images
                         if cond.shape[0] >= ad_params["full_length"]:
+                            if isinstance(weight, torch.Tensor):
+                                weight = torch.Tensor(weight["sub_idxs"])
+                                weight = weight.repeat(len(cond_or_uncond), 1, 1) # repeat for cond and uncond
+                                if torch.all(weight == 0):
+                                    continue
+                            elif weight == 0:
+                                continue
+
                             cond = torch.Tensor(cond[ad_params["sub_idxs"]])
                             uncond = torch.Tensor(uncond[ad_params["sub_idxs"]])
                         # otherwise get sub_idxs images
                         else:
+                            if isinstance(weight, torch.Tensor):
+                                weight = tensor_to_size(weight, ad_params["full_length"])
+                                weight = weight[ad_params["sub_idxs"]]
+                                weight = weight.repeat(len(cond_or_uncond), 1, 1) # repeat for cond and uncond
+                                if torch.all(weight == 0):
+                                    continue
+                            elif weight == 0:
+                                continue
                             cond = tensor_to_size(cond, ad_params["full_length"])
                             uncond = tensor_to_size(uncond, ad_params["full_length"])
                             cond = cond[ad_params["sub_idxs"]]
                             uncond = uncond[ad_params["sub_idxs"]]
+                    else:
+                        cond = tensor_to_size(cond, batch_prompt)
+                        uncond = tensor_to_size(uncond, batch_prompt)
 
-                    cond = tensor_to_size(cond, batch_prompt)
-                    uncond = tensor_to_size(uncond, batch_prompt)
+                        if isinstance(weight, torch.Tensor):
+                            weight = tensor_to_size(weight, batch_prompt)
+                            weight = weight.repeat(len(cond_or_uncond), 1, 1)
+                            if torch.all(weight == 0):
+                                continue
+                        elif weight == 0:
+                            continue
 
                     k_cond = ipadapter.ip_layers.to_kvs[self.k_key](cond)
                     k_uncond = ipadapter.ip_layers.to_kvs[self.k_key](uncond)
                     v_cond = ipadapter.ip_layers.to_kvs[self.v_key](cond)
                     v_uncond = ipadapter.ip_layers.to_kvs[self.v_key](uncond)
                 else:
+                    # TODO: should we always convert the weights to a tensor?
+                    if isinstance(weight, torch.Tensor):
+                        weight = tensor_to_size(weight, batch_prompt)
+                        weight = weight.repeat(len(cond_or_uncond), 1, 1)
+                        if torch.all(weight == 0):
+                            continue
+                    elif weight == 0:
+                        continue
+
                     k_cond = ipadapter.ip_layers.to_kvs[self.k_key](cond).repeat(batch_prompt, 1, 1)
                     k_uncond = ipadapter.ip_layers.to_kvs[self.k_key](uncond).repeat(batch_prompt, 1, 1)
                     v_cond = ipadapter.ip_layers.to_kvs[self.v_key](cond).repeat(batch_prompt, 1, 1)