update to work with new comfy

conditioning/fusers_patch.py

@@ -30,11 +30,10 @@ block_map = {
 
 
 class FusersPatch(torch.nn.Module):
-    def __init__(self, conds, fusers_list, fusers_batch_size, positionnet, latent_shape, idxs, device):
+    def __init__(self, conds, fusers_list, positionnet, latent_shape, idxs, device):
         super(FusersPatch, self).__init__()
         self.conds = conds
         self.fusers_list = fusers_list
-        self.fusers_batch_size = fusers_batch_size
         self.positionnet = positionnet
         self.latent_shape = latent_shape
         self.idxs = idxs
@@ -64,11 +63,6 @@ class FusersPatch(torch.nn.Module):
         fuser = self.fusers_list[fuser_idx]
         attn_total = []
         idxs = self._get_idxs(x, extra_options)
-        for i in range(0, len(x), self.fusers_batch_size):
-            batch_idxs = idxs[i:i+self.fusers_batch_size]
-            x_idxs = list(range(i, i+self.fusers_batch_size))
-            attn = fuser(x[x_idxs], self._get_position_objs(batch_idxs))
-            attn_total.append(attn)
-        attn_total = torch.cat(attn_total)
-        # print(attn_total.shape)
+
+        attn_total = fuser(x, self._get_position_objs(idxs))
         return attn_total.to(torch.float16)

conditioning/instance_conditioning.py

@@ -3,7 +3,6 @@ import torch
 import comfy.model_management
 
 from .fusers_patch import FusersPatch
-from .embeddings import prepare_embeddings
 
 
 class InstanceConditioning:
@@ -11,10 +10,9 @@ class InstanceConditioning:
     This class masquerades as Gligen in order to trigger setup
     """
 
-    def __init__(self, fusers, positionnet, fusers_batch_size):
+    def __init__(self, fusers, positionnet):
         self.fusers_list = fusers['model_list']
         self.positionnet = positionnet['model']
-        self.fusers_batch_size = fusers_batch_size
         self.conds = []
         self.current_device = comfy.model_management.intermediate_device()
 
@@ -24,7 +22,7 @@ class InstanceConditioning:
         self.offload_device = comfy.model_management.intermediate_device()
 
     def get_fusers_patch(self, latent_shape, idxs, device):
-        return FusersPatch(self.conds, self.fusers_list, self.fusers_batch_size, self.positionnet, latent_shape, idxs, device)
+        return FusersPatch(self.conds, self.fusers_list, self.positionnet, latent_shape, idxs, device)
 
     def set_position(self, latent_shape, _, device):
         # Called in samplers by gligen cond to return middle attention patch
@@ -36,21 +34,21 @@ class InstanceConditioning:
     def add_conds(self, conds):
         self.conds.extend(conds)
 
-    def get_models(self) -> list[torch.nn.Module]:
+    def get_models(self, *args, **kwargs) -> list[torch.nn.Module]:
         # Used to get models for loading/offloading
         return [(None, model) for model in [*self.fusers_list, self.positionnet]]
 
-    def inference_memory_requirements(self, dtype) -> int:
+    def inference_memory_requirements(self, dtype, *args, **kwargs) -> int:
         # Used to calculate memory requirements by ControlNet
         return 0
 
-    def is_clone(self, other):
+    def is_clone(self, other, *args, **kwargs):
         return other == self
 
-    def model_size(self):
+    def model_size(self, *args, **kwargs):
         return 0
 
-    def model_patches_to(self, device_or_dtype):
+    def model_patches_to(self, device_or_dtype, *args, **kwargs):
         if device_or_dtype == torch.float16 or device_or_dtype == torch.float32:
             return
         if device_or_dtype is None:
@@ -59,10 +57,11 @@ class InstanceConditioning:
         for i, fuser in enumerate(self.fusers_list):
             self.fusers_list[i] = fuser.to(device_or_dtype)
 
-    def model_dtype(self):
-        # Fusers requires float32 so we will ignore this
+    def model_dtype(self, *args, **kwargs):
         return torch.float32
 
-    def patch_model(self, device_to):
-        # TODO what is this for?
+    def patch_model(self, device_to, *args, **kwargs):
+        return
+
+    def unpatch_weights(self, *args, **kwargs):
         return

modules/attention.py

@@ -4,6 +4,8 @@ from torch import nn
 from inspect import isfunction
 import torch.nn.functional as F
 
+from comfy.gligen import GatedSelfAttentionDense as GSAD
+
 
 def exists(val):
     return val is not None
@@ -46,7 +48,7 @@ class FeedForward(nn.Module):
         return self.net(x)
 
 
-class SelfAttention(nn.Module):
+class SelfAttentionFusers(nn.Module):
     def __init__(self, query_dim, heads=8, dim_head=64, dropout=0.):
         super().__init__()
         inner_dim = dim_head * heads
@@ -61,6 +63,7 @@ class SelfAttention(nn.Module):
             nn.Linear(inner_dim, query_dim), nn.Dropout(dropout))
 
     def forward(self, x, instance_options={}):
+        # TODO: use attention masking
         grounding_input = instance_options.get('grounding_input', None)
         drop_box_mask = instance_options.get('drop_box_mask', None)
         q = self.to_q(x)  # B*N*(H*C)
@@ -86,9 +89,9 @@ class SelfAttention(nn.Module):
             n_objs = grounding_input["att_masks"].shape[1]
 
             att_masks_reshape = None
-            if N - n_objs * 4 == 64 * 64:
-                att_masks_reshape = grounding_input["att_masks"]
-                att_key = "att_masks_selfAtt64"
+            # if N - n_objs * 4 == 64 * 64:
+            #     att_masks_reshape = grounding_input["att_masks"]
+            #     att_key = "att_masks_selfAtt64"
 
             # if the mask is all zeros or box&mask are dropped, we do not need to mask the self-attention
             if att_masks_reshape is not None and torch.sum(att_masks_reshape) > 0.0 and not drop_box_mask:
@@ -174,36 +177,8 @@ class SelfAttention(nn.Module):
         return self.to_out(out)
 
 
-class GatedSelfAttentionDense(nn.Module):
-    def __init__(self, query_dim, context_dim, n_heads, d_head):
-        super().__init__()
-        # we need a linear projection since we need cat visual feature and obj feature
-        self.linear = nn.Linear(context_dim, query_dim)
-
-        self.attn = SelfAttention(
-            query_dim=query_dim, heads=n_heads, dim_head=d_head)
-        self.ff = FeedForward(query_dim, glu=True)
-
-        self.norm1 = nn.LayerNorm(query_dim)
-        self.norm2 = nn.LayerNorm(query_dim)
-
-        self.register_parameter('alpha_attn', nn.Parameter(torch.tensor(0.)))
-        self.register_parameter('alpha_dense', nn.Parameter(torch.tensor(0.)))
-
-        # this can be useful: we can externally change magnitude of tanh(alpha)
-        # for example, when it is set to 0, then the entire model is same as original one
-        self.scale = 1
-
+class GatedSelfAttentionDense(GSAD):
     def forward(self, x, instance_options={}):
         objs = instance_options['objs']
-
-        N_visual = x.shape[1]
-        objs = self.linear(objs)
-        attention_output = self.attn(self.norm1(
-            torch.cat([x, objs], dim=1)), instance_options=instance_options)
-        x = x + self.scale * \
-            torch.tanh(self.alpha_attn) * attention_output[:, 0:N_visual, :]
-        x = x + self.scale * \
-            torch.tanh(self.alpha_dense) * self.ff(self.norm2(x))
-
+        x = super().forward(x, objs)
         return x.to(torch.float16)

nodes/instance_diffusion_tracking_prompt_node.py

@@ -11,7 +11,6 @@ class InstanceDiffusionTrackingPromptNode:
                              "tracking": ("TRACKING", ),
                              "positionnet": ("POSITIONNET", ),
                              "fusers": ("FUSERS", ),
-                             "fusers_batch_size": ("INT", {"default": 8, "min": 1, "max": 10000, "step": 1}),
                              "positive_text": ("STRING", {"multiline": True}),
                              "negative_text": ("STRING", {"multiline": True}),
                              }}
@@ -43,7 +42,7 @@ class InstanceDiffusionTrackingPromptNode:
 
         return position_conds
 
-    def _apply_position_conds(self, position_conds, conditioning, fusers, positionnet, fusers_batch_size):
+    def _apply_position_conds(self, position_conds, conditioning, fusers, positionnet):
         # Add prompts+embeddings to the input conditionings
         cond_out = []
         for t in conditioning:
@@ -52,7 +51,7 @@ class InstanceDiffusionTrackingPromptNode:
             prev = []
             has_instance = 'instance_diffusion' in cond
             instance_conditioning = conditioning['instance_diffusion'] if has_instance else InstanceConditioning(
-                fusers, positionnet, fusers_batch_size)
+                fusers, positionnet)
             cond['instance_diffusion'] = instance_conditioning
             instance_conditioning.add_conds(position_conds)
 
@@ -62,16 +61,16 @@ class InstanceDiffusionTrackingPromptNode:
 
         return cond_out
 
-    def append(self, positive, negative, clip, tracking, fusers, positionnet, fusers_batch_size, positive_text, negative_text):
+    def append(self, positive, negative, clip, tracking, fusers, positionnet, positive_text, negative_text, fusers_batch_size=None):
 
         positive_positions = self._get_position_conds(
             clip, tracking, positive_text)
         positive = self._apply_position_conds(
-            positive_positions, positive, fusers, positionnet, fusers_batch_size)
+            positive_positions, positive, fusers, positionnet)
 
         negative_positions = self._get_position_conds(
             clip, tracking, negative_text)
         negative = self._apply_position_conds(
-            negative_positions, negative, fusers, positionnet, fusers_batch_size)
+            negative_positions, negative, fusers, positionnet)
 
         return (positive, negative)