diff --git a/model_helpers/prepare_fusers.py b/model_helpers/prepare_fusers.py index 8951bec..9bf6064 100644 --- a/model_helpers/prepare_fusers.py +++ b/model_helpers/prepare_fusers.py @@ -4,12 +4,13 @@ import torch from comfy.gligen import GatedSelfAttentionDense as GSAD class GatedSelfAttentionDense(GSAD): + def forward(self, x, instance_options={}): objs = instance_options['objs'] x = super().forward(x, objs) return x.to(torch.float16) -def prepare_fusers(fusers_ckpt) -> list[torch.nn.Module]: +def prepare_fusers(fusers_ckpt, fusers_scale) -> list[torch.nn.Module]: fusers_list = [] for key in fusers_ckpt['input_blocks']: fusers_ckpt['input_blocks'][key]['params']['query_dim'] = fusers_ckpt['input_blocks'][key]['params']['n_heads'] * \ @@ -17,6 +18,7 @@ def prepare_fusers(fusers_ckpt) -> list[torch.nn.Module]: fuser = GatedSelfAttentionDense( **fusers_ckpt['input_blocks'][key]['params']) fuser.load_state_dict(fusers_ckpt['input_blocks'][key]['state']) + fuser.scale = fusers_scale fusers_list.append(fuser) fusers_ckpt['middle_block']['1']['params']['query_dim'] = fusers_ckpt['middle_block']['1']['params']['n_heads'] * \ @@ -24,6 +26,7 @@ def prepare_fusers(fusers_ckpt) -> list[torch.nn.Module]: fuser = GatedSelfAttentionDense( **fusers_ckpt['middle_block']['1']['params']) fuser.load_state_dict(fusers_ckpt['middle_block']['1']['state']) + fuser.scale = fusers_scale fusers_list.append(fuser) for key in fusers_ckpt['output_blocks']: @@ -32,6 +35,7 @@ def prepare_fusers(fusers_ckpt) -> list[torch.nn.Module]: fuser = GatedSelfAttentionDense( **fusers_ckpt['output_blocks'][key]['params']) fuser.load_state_dict(fusers_ckpt['output_blocks'][key]['state']) + fuser.scale = fusers_scale fusers_list.append(fuser) return fusers_list diff --git a/nodes/download_and_load_models.py b/nodes/download_and_load_models.py index 052e129..3eb9933 100644 --- a/nodes/download_and_load_models.py +++ b/nodes/download_and_load_models.py @@ -16,6 +16,7 @@ class DownloadInstanceDiffusionModels: def INPUT_TYPES(s): return {"required": { "use_segs": ("BOOLEAN", {"default": True}), + "fusers_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}), }} RETURN_TYPES = ("POSITIONNET", "FUSERS", "SCALEU", ) @@ -23,7 +24,7 @@ class DownloadInstanceDiffusionModels: CATEGORY = "instance/loaders" - def load_model(self, use_segs: bool): + def load_model(self, use_segs: bool, fusers_scale: float): repo_id = "logtd/instance_diffusion" instance_models_folder = os.path.join(folder_paths.models_dir, constants.INSTANCE_MODELS_DIR) @@ -59,7 +60,7 @@ class DownloadInstanceDiffusionModels: } fusers_checkpoint = comfy.utils.load_torch_file(fusers_file, safe_load=True) - fusers_list = prepare_fusers(fusers_checkpoint) + fusers_list = prepare_fusers(fusers_checkpoint, fusers_scale) fusers = { 'model_list': fusers_list } diff --git a/nodes/load_instance_fusers_node.py b/nodes/load_instance_fusers_node.py index 612b216..9b1486a 100644 --- a/nodes/load_instance_fusers_node.py +++ b/nodes/load_instance_fusers_node.py @@ -8,6 +8,7 @@ class LoadInstanceFusersNode: def INPUT_TYPES(s): return {"required": { "model_filename": (get_model_list(constants.INSTANCE_FUSERS_DIR),), + "fusers_scale": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}), }} RETURN_TYPES = ("FUSERS",) @@ -15,10 +16,10 @@ class LoadInstanceFusersNode: CATEGORY = "instance/loaders" - def load_model(self, model_filename: str): + def load_model(self, model_filename: str, fusers_scale: float): checkpoint = load_checkpoint( constants.INSTANCE_FUSERS_DIR, model_filename) - fusers_list = prepare_fusers(checkpoint) + fusers_list = prepare_fusers(checkpoint, fusers_scale) fusers = { 'model_list': fusers_list }