diff --git a/conditioning/fusers_patch.py b/conditioning/fusers_patch.py index bb92538..c61f12e 100644 --- a/conditioning/fusers_patch.py +++ b/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) diff --git a/conditioning/instance_conditioning.py b/conditioning/instance_conditioning.py index efa77c8..ceb0c00 100644 --- a/conditioning/instance_conditioning.py +++ b/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 diff --git a/modules/attention.py b/modules/attention.py index bf39a9e..54ed650 100644 --- a/modules/attention.py +++ b/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) diff --git a/nodes/instance_diffusion_tracking_prompt_node.py b/nodes/instance_diffusion_tracking_prompt_node.py index 197fe37..3689eb2 100644 --- a/nodes/instance_diffusion_tracking_prompt_node.py +++ b/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)