almost working

.gitignore

@@ -1,3 +1,4 @@
 .git
 .vscode
 **/__pycache__
+**/*.ckpt

__init__.py

@@ -1,17 +1,22 @@
+from .nodes.apply_scaleu_model_node import ApplyScaleUModelNode
+from .nodes.load_instance_scaleu_node import LoadInstanceScaleUNode
+from .nodes.load_instance_fusers_node import LoadInstanceFusersNode
+from .nodes.load_instance_positionnet_node import LoadInstancePositionNetNode
 from .nodes.instance_diffusion_tracking_prompt_node import InstanceDiffusionTrackingPromptNode
 
+
 NODE_CLASS_MAPPINGS = {
-    "LoadInstanceDiffusion": None,
-    "LoadInstancePositionNetModel": None,
-    "LoadInstanceScaleUModel": None,
-    "LoadInstanceAttentionModel": None,
+    "ApplyScaleUModelNode": ApplyScaleUModelNode,
+    "LoadInstanceScaleUNode": LoadInstanceScaleUNode,
+    "LoadInstancePositionNetModel": LoadInstancePositionNetNode,
+    "LoadInstanceFusersNode": LoadInstanceFusersNode,
     "InstanceDiffusionTrackingPrompt": InstanceDiffusionTrackingPromptNode,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
-    "LoadInstanceDiffusion": "Load Instance Diffusion",
+    "ApplyScaleUModelNode": "Apply Instance Diffusion ScaleU",
     "LoadInstancePositionNetModel": "Load Instance PositionNet Model",
     "LoadInstanceScaleUModel": "Load Instance ScaleU Model",
-    "LoadInstanceAttentionModel": "Load Instance Attention Model",
+    "LoadInstanceFusersNode": "Load Instance Fusers Model",
     "InstanceDiffusionTrackingPrompt": "Instance Diffusion Tracking Prompt"
 }

conditioning/embeddings.py

@@ -167,7 +167,7 @@ def prepare_embeddings(conds, latent_shape, idxs, use_masked_att=False):
     batch_size, latent_width, latent_height)
   if use_masked_att:
     embeddings['att_masks'] = torch.zeros(
-        idx, N_MAX_OBJECTS, latent_width, latent_height)
+        batch_size, N_MAX_OBJECTS, latent_width, latent_height)
 
   for idx in idxs:
     for cond_idx, cond in enumerate(conds):

conditioning/fusers_patch.py

@@ -0,0 +1,21 @@
+import torch
+
+
+class FusersPatch(torch.nn.Module):
+  def __init__(self, fusers_list, objs, drop_box_mask, device, fusers_batch_size=5):
+    super(FusersPatch, self).__init__()
+    self.fusers_list = fusers_list
+    self.transformer_options = {'objs': objs, 'drop_box_mask': drop_box_mask}
+    self.device = device
+    self.fusers_batch_size = fusers_batch_size
+    # call_idx traces where in the UNet this is being called from
+    # a hack around lack of Comfy support
+    self.call_idx = 0
+
+  @torch.no_grad()
+  def forward(self, x, extra_options):
+    fuser = self.fusers_list[self.call_idx %
+                             len(self.fusers_list)]
+    attn = fuser(x, self.transformer_options)
+    self.call_idx += 1
+    return attn.to(torch.float16)

conditioning/instance_conditioning.py

@@ -3,6 +3,7 @@ import torch
 import comfy.model_management
 
 from .embeddings import prepare_embeddings
+from .fusers_patch import FusersPatch
 
 
 class InstanceConditioning:
@@ -15,9 +16,11 @@ class InstanceConditioning:
     self.positionnet = positionnet['model']
     self.fusers_batch_size = fusers_batch_size
     self.conds = []
+    self.current_device = comfy.model_management.intermediate_device()
 
     # ControlNet hacks
     self.load_device = comfy.model_management.get_torch_device()
+    self.offload_device = comfy.model_management.intermediate_device()
 
     # Gligen hacks
     self.model = self
@@ -27,8 +30,8 @@ class InstanceConditioning:
     self.full_latent_length = None
     self.context_length = None
 
-  def get_fusers_patch(self, objs, drop_box_mask):
-    return
+  def get_fusers_patch(self, objs, drop_box_mask, device):
+    return FusersPatch(self.fusers_list, objs, drop_box_mask, device, self.fusers_batch_size)
 
   def set_position(self, latent_shape, _, device):
     # Called in samplers by gligen cond to return middle attention patch
@@ -36,8 +39,9 @@ class InstanceConditioning:
     idxs = self.sub_idxs if self.sub_idxs is not None else list(
       range(batch_size))
 
-    # TODO conds
     embeddings = prepare_embeddings(self.conds, latent_shape, idxs, True)
+    for key in embeddings:
+      embeddings[key] = embeddings[key].to(device)
     objs, drop_box_mask = self.positionnet.to(device)(embeddings)
     fusers_patch = self.get_fusers_patch(objs, drop_box_mask, device)
     return fusers_patch
@@ -57,6 +61,26 @@ class InstanceConditioning:
   def is_clone(self, other):
     return other == self
 
+  def model_size(self):
+    return 0
+
+  def model_patches_to(self, device_or_dtype):
+    if device_or_dtype == torch.float16 or device_or_dtype == torch.float32:
+      return
+    if device_or_dtype is None:
+      return
+    self.positionnet = self.positionnet.to(device_or_dtype)
+    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
+    return torch.float32
+
+  def patch_model(self, device_to):
+    # TODO what is this for?
+    return
+
   # def __getattribute__(self, __name: str) -> Any:
   #   if __name == 'gligen':
   #     return 'something'  # TODO

nodes/apply_scaleu_model_node.py

@@ -1,8 +1,10 @@
+import comfy.model_management
+
 from .. import constants as constants
-from ..models.prepare_scaleu import get_scaleu_patch
+from ..model_helpers.prepare_scaleu import get_scaleu_patch
 
 
-class ApplyScaleUModel:
+class ApplyScaleUModelNode:
   @classmethod
   def INPUT_TYPES(s):
     return {"required": {
@@ -17,14 +19,18 @@ class ApplyScaleUModel:
 
   def apply(self, model, scaleu):
     # Validate patches dict is setup correctly
-    if 'patches' not in model or model['patches'] is None:
-      model['patches'] = {}
+    transformer_options = model.model_options['transformer_options']
+    if 'patches' not in transformer_options:
+      transformer_options['patches'] = {}
 
-    if 'output_block_patch' not in model['patches']:
-      model['patches']['output_block_patch'] = []
+    if 'output_block_patch' not in transformer_options['patches']:
+      transformer_options['patches']['output_block_patch'] = []
 
     # Add scaleu patch to model patches
     scaleu_nets = scaleu['model_list']
-    model['patches']['output_block_patch'].append(
+    # TODO make this load in KSampler
+    for i, scaleu in enumerate(scaleu_nets):
+      scaleu_nets[i] = scaleu.to(comfy.model_management.get_torch_device())
+    transformer_options['patches']['output_block_patch'].append(
       get_scaleu_patch(scaleu_nets))
     return (model,)

nodes/instance_diffusion_tracking_prompt_node.py

@@ -7,7 +7,7 @@ class InstanceDiffusionTrackingPromptNode:
   def INPUT_TYPES(s):
     return {"required": {"conditioning": ("CONDITIONING", ),
                          "clip": ("CLIP", ),
-                         "tracking": ("TRACKING", ),
+                         #  "tracking": ("TRACKING", ),
                          "positionnet": ("POSITIONNET", ),
                          "fusers": ("FUSERS", ),
                          "fusers_batch_size": ("INT", {"default": 5, "min": 0, "max": 10000, "step": 8}),
@@ -18,25 +18,32 @@ class InstanceDiffusionTrackingPromptNode:
 
   CATEGORY = "instance/conditioning"
 
-  def append(self, conditioning, clip, tracking, fusers, positionnet, fusers_batch_size, text):
+  def append(self, conditioning, clip, fusers, positionnet, fusers_batch_size, text):
     # Get prompts and their class id and trakcer id
     prompt_pairs = extract_prompts(text)
 
     # Go through prompt pairs, encode prompts, and join with positions from tracking
-    position_conds = []
-    for tracker_id, class_id, prompt in prompt_pairs:
-      _, cond_pooled = clip.encode_from_tokens(
+    # position_conds = []
+    # for tracker_id, class_id, prompt in prompt_pairs:
+    #   _, cond_pooled = clip.encode_from_tokens(
+    #     clip.tokenize(prompt), return_pooled=True)
+    #   # A tracker_id of -1 means that it is prompting all instances of a single class
+    #   if tracker_id != -1:
+    #     position_cond = {'cond_pooled': cond_pooled, 'positions':
+    #                      tracking[class_id][tracker_id]}
+    #     position_conds.append(position_cond)
+    #   else:
+    #     for tracker_id in tracking[class_id]:
+    #       position_cond = {'cond_pooled': cond_pooled,
+    #                        'positions': tracking[class_id][tracker_id]}
+    #       position_conds.append(position_cond)
+
+    prompt = "chicken"
+    _, cond_pooled = clip.encode_from_tokens(
         clip.tokenize(prompt), return_pooled=True)
-      # A tracker_id of -1 means that it is prompting all instances of a single class
-      if tracker_id != -1:
-        position_cond = {'cond_pooled': cond_pooled, 'positions':
-                         tracking[class_id][tracker_id]}
-        position_conds.append(position_cond)
-      else:
-        for tracker_id in tracking[class_id]:
-          position_cond = {'cond_pooled': cond_pooled,
-                           'positions': tracking[class_id][tracker_id]}
-          position_conds.append(position_cond)
+
+    position_conds = [{'cond_pooled': cond_pooled,
+                       'positions': [[100, 100, 200, 200, 512, 512]]}]
 
     # Add prompts+embeddings to the input conditionings
     cond_out = []

nodes/load_instance_fusers_node.py

@@ -1,13 +1,13 @@
 from .. import constants as constants
-from ..utils.model_utils import get_model_dir, load_checkpoint
-from ..models.prepare_fusers import prepare_fusers
+from ..utils.model_utils import get_model_list, load_checkpoint
+from ..model_helpers.prepare_fusers import prepare_fusers
 
 
 class LoadInstanceFusersNode:
   @classmethod
   def INPUT_TYPES(s):
     return {"required": {
-        "model_path": (get_model_dir(constants.INSTANCE_FUSERS_DIR)),
+        "model_filename": (get_model_list(constants.INSTANCE_FUSERS_DIR),),
     }}
 
   RETURN_TYPES = ("FUSERS",)
@@ -15,8 +15,8 @@ class LoadInstanceFusersNode:
 
   CATEGORY = "instance/loaders"
 
-  def load_model(self, model_path: str):
-    checkpoint = load_checkpoint(model_path)
+  def load_model(self, model_filename: str):
+    checkpoint = load_checkpoint(constants.INSTANCE_FUSERS_DIR, model_filename)
     fusers_list = prepare_fusers(checkpoint)
     fusers = {
       'model_list': fusers_list

nodes/load_instance_positionnet_node.py

@@ -1,13 +1,13 @@
 from .. import constants as constants
-from ..utils.model_utils import get_model_dir, load_checkpoint
-from ..models.prepare_positionnet import prepare_positionnet, get_positionnet_default_params
+from ..utils.model_utils import get_model_list, load_checkpoint
+from ..model_helpers.prepare_positionnet import prepare_positionnet, get_positionnet_default_params
 
 
 class LoadInstancePositionNetNode:
   @classmethod
   def INPUT_TYPES(s):
     return {"required": {
-        "model_path": (get_model_dir(constants.INSTANCE_POSITIONNET_DIR)),
+        "model_filename": (get_model_list(constants.INSTANCE_POSITIONNET_DIR),),
     }}
 
   RETURN_TYPES = ("POSITIONNET",)
@@ -15,8 +15,9 @@ class LoadInstancePositionNetNode:
 
   CATEGORY = "instance/loaders"
 
-  def load_model(self, model_path: str):
-    checkpoint = load_checkpoint(model_path)
+  def load_model(self, model_filename: str):
+    checkpoint = load_checkpoint(
+      constants.INSTANCE_POSITIONNET_DIR, model_filename)
     params = get_positionnet_default_params()
     model = prepare_positionnet(checkpoint, params)
     positionnet = {

nodes/load_instance_scaleu_node.py

@@ -1,13 +1,13 @@
 from .. import constants as constants
-from ..utils.model_utils import get_model_dir, load_checkpoint
-from ..models.prepare_scaleu import prepare_scaleu_nets
+from ..utils.model_utils import get_model_list, load_checkpoint
+from ..model_helpers.prepare_scaleu import prepare_scaleu_nets
 
 
 class LoadInstanceScaleUNode:
   @classmethod
   def INPUT_TYPES(s):
     return {"required": {
-        "model_path": (get_model_dir(constants.INSTANCE_SCALEU_DIR)),
+        "model_filename": (get_model_list(constants.INSTANCE_SCALEU_DIR),),
     }}
 
   RETURN_TYPES = ("SCALEU",)
@@ -15,8 +15,8 @@ class LoadInstanceScaleUNode:
 
   CATEGORY = "instance/loaders"
 
-  def load_model(self, model_path: str):
-    checkpoint = load_checkpoint(model_path)
+  def load_model(self, model_filename: str):
+    checkpoint = load_checkpoint(constants.INSTANCE_SCALEU_DIR, model_filename)
     scaleu_list = prepare_scaleu_nets(checkpoint)
     scaleu = {
       'model_list': scaleu_list

utils/model_utils.py

@@ -1,18 +1,24 @@
 import os
 import folder_paths
-import torch
 
 import comfy.utils
 
-from . import constants as constants
+from .. import constants
 
 
-def get_model_dir(model_dir) -> list[str]:
-  path = os.path.join(constants.INSTANCE_MODELS_DIR,
+def get_model_dir(model_dir):
+  root_path = folder_paths.folder_names_and_paths['custom_nodes'][0][0]
+  path = os.path.join(root_path, 'ComfyUI-InstanceDiffusion', constants.INSTANCE_MODELS_DIR,
                       model_dir)
-  return folder_paths.get_filename_list(path)
+  return path
 
 
-def load_checkpoint(checkpoint_path):
+def get_model_list(model_dir) -> list[str]:
+  path = get_model_dir(model_dir)
+  return os.listdir(path)
+
+
+def load_checkpoint(model_dir, filename):
+  checkpoint_path = os.path.join(get_model_dir(model_dir), filename)
   checkpoint = comfy.utils.load_torch_file(checkpoint_path, safe_load=True)
   return checkpoint