Expose fuser scale (GatedSelfAttentionDense scale)

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

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
         }

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
         }