Merge branch 'ClownsharkBatwing:main' into main

drozbay · Dec 17, 2024 · f4b4583 · f4b4583
2 parents 132aff4 + 8355357
commit f4b4583
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 7 deletions.
diff --git a/conditioning.py b/conditioning.py
@@ -445,12 +445,16 @@ def __init__(self, mask: torch.Tensor, conditioning: torch.Tensor, conditioning_
         self.img_len = img_len
         self.text_len = text_len
 
-    def __call__(self, transformer_options, dtype=torch.bfloat16, *args, **kwargs):
+    def __call__(self, transformer_options, weight=0, dtype=torch.bfloat16, *args, **kwargs):
         sigma = transformer_options['sigmas'][0]
         if self.start_percent <= 1 - sigma < self.end_percent:
             if self.mask_type == "gradient":
-                return self.mask.clone().to(sigma.device).to(torch.bool)
-                return self.mask.clone().to(sigma.device).to(dtype)
+                mask = self.mask.clone().to(sigma.device)
+                mask[self.text_len:,self.text_len:] = mask[self.text_len:,self.text_len:] > 1-weight
+                #mask[self.text_len:,self.text_len:] = torch.clamp(mask[self.text_len:,self.text_len:], min=1-weight)
+
+                return mask.to(torch.bool)
+
 
         return None
 

diff --git a/flux/model.py b/flux/model.py
@@ -101,9 +101,10 @@ def forward_blocks(self, img: Tensor, img_ids: Tensor, txt: Tensor, txt_ids: Ten
             mask = None
             mask_obj = transformer_options.get('patches', {}).get('regional_conditioning_mask', None)
             if mask_obj is not None and weight >= 0:
-                mask = mask_obj[0](transformer_options)
+                mask = mask_obj[0](transformer_options, weight.item())
+                """mask = mask_obj[0](transformer_options)
                 text_len = mask_obj[0].text_len
-                mask[text_len:,text_len:] = torch.clamp(mask[text_len:,text_len:], min=1-weight.to(mask.device))
+                mask[text_len:,text_len:] = torch.clamp(mask[text_len:,text_len:], min=1-weight.to(mask.device))"""
 
             img, txt = block(img=img, txt=txt, vec=vec, pe=pe, timestep=timesteps, transformer_options=transformer_options, mask=mask) #, mask=mask)
 
@@ -119,9 +120,10 @@ def forward_blocks(self, img: Tensor, img_ids: Tensor, txt: Tensor, txt_ids: Ten
             mask = None
             mask_obj = transformer_options.get('patches', {}).get('regional_conditioning_mask', None)
             if mask_obj is not None and weight >= 0:
-                mask = mask_obj[0](transformer_options)
+                mask = mask_obj[0](transformer_options, weight.item())
+                """mask = mask_obj[0](transformer_options)
                 text_len = mask_obj[0].text_len
-                mask[text_len:,text_len:] = torch.clamp(mask[text_len:,text_len:], min=1-weight.to(mask.device))
+                mask[text_len:,text_len:] = torch.clamp(mask[text_len:,text_len:], min=1-weight.to(mask.device))"""
 
             img = block(img, vec=vec, pe=pe, timestep=timesteps, transformer_options=transformer_options, mask=mask)