diff --git a/.gitignore b/.gitignore index bbbe7d6..dd0c0c3 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ .git .vscode **/__pycache__ +**/*.ckpt \ No newline at end of file diff --git a/__init__.py b/__init__.py index 7a4e73e..ec86e01 100644 --- a/__init__.py +++ b/__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" } diff --git a/conditioning/embeddings.py b/conditioning/embeddings.py index aa0373d..6b120fa 100644 --- a/conditioning/embeddings.py +++ b/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): diff --git a/conditioning/fusers_patch.py b/conditioning/fusers_patch.py index e69de29..2bc8297 100644 --- a/conditioning/fusers_patch.py +++ b/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) diff --git a/conditioning/instance_conditioning.py b/conditioning/instance_conditioning.py index 13eb977..a3a3f1b 100644 --- a/conditioning/instance_conditioning.py +++ b/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 diff --git a/models/prepare_fusers.py b/model_helpers/prepare_fusers.py similarity index 100% rename from models/prepare_fusers.py rename to model_helpers/prepare_fusers.py diff --git a/models/prepare_positionnet.py b/model_helpers/prepare_positionnet.py similarity index 100% rename from models/prepare_positionnet.py rename to model_helpers/prepare_positionnet.py diff --git a/models/prepare_scaleu.py b/model_helpers/prepare_scaleu.py similarity index 100% rename from models/prepare_scaleu.py rename to model_helpers/prepare_scaleu.py diff --git a/nodes/apply_scaleu_model_node.py b/nodes/apply_scaleu_model_node.py index 7d9f26e..0aa9f5b 100644 --- a/nodes/apply_scaleu_model_node.py +++ b/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,) diff --git a/nodes/instance_diffusion_tracking_prompt_node.py b/nodes/instance_diffusion_tracking_prompt_node.py index 030d02e..352b115 100644 --- a/nodes/instance_diffusion_tracking_prompt_node.py +++ b/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 = [] diff --git a/nodes/load_instance_fusers_node.py b/nodes/load_instance_fusers_node.py index 9386800..ccf253b 100644 --- a/nodes/load_instance_fusers_node.py +++ b/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 diff --git a/nodes/load_instance_positionnet_node.py b/nodes/load_instance_positionnet_node.py index d22a9b9..ec55db6 100644 --- a/nodes/load_instance_positionnet_node.py +++ b/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 = { diff --git a/nodes/load_instance_scaleu_node.py b/nodes/load_instance_scaleu_node.py index fe0bdff..7548756 100644 --- a/nodes/load_instance_scaleu_node.py +++ b/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 diff --git a/utils/model_utils.py b/utils/model_utils.py index 69138a5..37c3e61 100644 --- a/utils/model_utils.py +++ b/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