Add GroundingDINO mask injection.

__init__.py

@@ -1,4 +1,4 @@
-from .node import BatchedMaskedReferenceGen,ReferenceConditionCombine,MasksToPatches,ConditioningMatchMask,GenSubject,PreviewSubjectMask,PreRun,CrossSimMaskExtract,GenConsistent
+from .node import BatchedMaskedReferenceGen,ReferenceConditionCombine,MasksToPatches,ConditioningMatchMask,GenSubject,PreviewSubjectMask,PreRun,CrossSimMaskExtract,GenConsistent,SetSubjectMask
 NODE_CLASS_MAPPINGS = {
 
     "BatchedMaskedReferenceGen": BatchedMaskedReferenceGen,
@@ -9,7 +9,8 @@ NODE_CLASS_MAPPINGS = {
     "PreviewSubjectMask" : PreviewSubjectMask,
     "PreRun": PreRun,
     "CrossSimMaskExtract": CrossSimMaskExtract,
-    "GenConsistent": GenConsistent
+    "GenConsistent": GenConsistent,
+    "SetSubjectMask": SetSubjectMask,
 }
 NODE_DISPLAY_NAME_MAPPINGS = {
     "ReferenceConditionCombine": "Reference Condition Combine",
@@ -20,7 +21,8 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "PreviewSubjectMask": "Preview Subject Mask",
     "PreRun": "PreRun",
     "CrossSimMaskExtract": "CrossSim Mask Extract",
-    "GenConsistent": "GenConsistent"
+    "GenConsistent": "GenConsistent",
+    "SetSubjectMask": "SetSubjectMask"
 }
 
 __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']

node.py

@@ -434,6 +434,67 @@ class PreviewSubjectMask:
             combined_mask = torch.empty(0)
             
         return (combined_mask, latent, )
+
+class SetSubjectMask:
+    RETURN_TYPES = ("SUBJECT",)
+    RETURN_NAMES = ("subject",)
+    FUNCTION = "set_mask"
+    CATEGORY = "chara_consist"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "mask": ("MASK",),
+                "subject": ("SUBJECT",),
+            }
+        }
+
+    @classmethod
+    def set_mask(cls, mask: torch.Tensor, subject: list):
+        """
+        Set the global fg_mask for a subject based on an input mask.
+        
+        Args:
+            mask: Input mask tensor with shape [B, H, W]
+            subject: Subject list containing dictionaries with step data
+            
+        Returns:
+            subject: Modified subject with global fg_mask set
+        """
+        # Get the dimensions of the input mask
+        B, H, W = mask.shape
+        
+        # Calculate patch dimensions (divide by 16)
+        h_patches = H // 16
+        w_patches = W // 16
+        
+        # Reshape mask to [B, H//16, 16, W//16, 16] for processing 16x16 blocks
+        mask_reshaped = mask.view(B, h_patches, 16, w_patches, 16)
+        
+        # Check if any pixel in each 16x16 block is non-zero
+        # Take max over the 16x16 blocks, then convert to boolean
+        patch_mask = mask_reshaped.max(dim=2)[0].max(dim=3)[0] > 0
+        
+        # Convert to boolean tensor
+        patch_mask = patch_mask.bool()
+        
+        # Create a copy of the subject to avoid modifying the original
+        new_subject = copy.deepcopy(subject)
+        
+        # Set the global fg_mask in subject[0]['fg_mask']
+        # This represents the global mask of the subject, not for each step's fg_mask
+        if len(new_subject) > 0:
+            if 0 not in new_subject[0]:
+                new_subject[0] = {}
+            if 'fg_mask' not in new_subject[0][0]:
+                new_subject[0][0]['fg_mask'] = patch_mask
+            else:
+                # If fg_mask already exists, replace it with our new global mask
+                new_subject[0][0]['fg_mask'] = patch_mask
+        
+        return (new_subject,)
+
 def hook_prepare_callback():
     flux_layers.attention = orig_attention
     chroma_layers.attention = orig_attention
@@ -486,13 +547,15 @@ class GenSubject:
                 "model": ("MODEL",),
                 "fg_condition_mask": ("MASK",),
                 "kv_start_step":  ("INT",{"default":1,"min":0}),
-                "kv_end_step":  ("INT",{"default":5,"min":0}),
-                "mask_start_step": ("INT",{"default":7,"min":0}),
-                "mask_end_step": ("INT",{"default":8,"min":0}),
+                "kv_end_step":  ("INT",{"default":9,"min":0}),
+                "x_start_step":  ("INT",{"default":1,"min":0}),
+                "x_end_step":  ("INT",{"default":9,"min":0}),
+                "mask_start_step": ("INT",{"default":8,"min":0}),
+                "mask_end_step": ("INT",{"default":9,"min":0}),
             }
         }
     @classmethod
-    def go(cls, model, fg_condition_mask, kv_start_step, kv_end_step, mask_start_step, mask_end_step):
+    def go(cls, model, fg_condition_mask, kv_start_step, kv_end_step, x_start_step, x_end_step, mask_start_step, mask_end_step):
         
         subject = [{}]
         index = len(subject) - 1
@@ -549,7 +612,7 @@ class GenSubject:
                 # Reshape back to the expected output format
                 x = rearrange(x, "B H L D -> B L (H D)")
                 # Store attention result x along with k and v
-                if sampling_step >= mask_start_step and sampling_step < mask_end_step:
+                if sampling_step >= x_start_step and sampling_step < x_end_step:
                     if not layer in subject[index][sampling_step]["layers"]:
                         subject[index][sampling_step]["layers"][layer] = {}
                     subject[index][sampling_step]["layers"][layer]["x"] = F.normalize(x[:, text_len:,:], dim=-1).transpose(-2, -1).cpu()
@@ -614,7 +677,7 @@ class PreRun:
                 "subject": ("SUBJECT",),
                 "fg_condition_mask": ("MASK",),
                 "cross_sim_step": ("INT",),
-                "sim_thredhold": ("FLOAT",{"default":0.5,"min":0, "max":2}),
+                "sim_thredhold": ("FLOAT",{"default":0.5,"min":0,"max":2}),
             },
             "optional":{
                 "latent": ("LATENT",),
@@ -722,8 +785,12 @@ class PreRun:
                     all_id_fg_inds = []
                     all_curr_fg_inds = []
                     for i in range(len(fg_condition_mask)):
+                        if 0 in subject[i] and "fg_mask" in subject[i][0]:
+                            fg_mask = subject[i][0]["fg_mask"][0]
+                        else:
+                            fg_mask = subject[i][sampling_step]["fg_mask"][0]
                         id_fg_inds, curr_fg_inds = get_shared_fg_mask(
-                            subject[i][sampling_step]["fg_mask"][0].to(device), 
+                            fg_mask.to(device), 
                             bitmask[i], 
                             argmax_indices[i],
                             max_sim[i]>sim_thredhold