Update nodes_differential_diffusion.py

comfyanonymous · Nov 22, 2024 · 1a86443 · 1a86443
1 parent ce31337
commit 1a86443
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/comfy_extras/nodes_differential_diffusion.py b/comfy_extras/nodes_differential_diffusion.py
@@ -43,9 +43,11 @@ def forward(self, sigma: torch.Tensor, denoise_mask: torch.Tensor, extra_options
         binary_mask = (denoise_mask >= threshold).to(denoise_mask.dtype)
 
         # Blend binary mask with the original denoise_mask using strength
-        blended_mask = strength * binary_mask + (1 - strength) * denoise_mask
-
-        return blended_mask
+        if strength and strength < 1:
+            blended_mask = strength * binary_mask + (1 - strength) * denoise_mask
+            return blended_mask
+        else:
+            return binary_mask
 
 
 NODE_CLASS_MAPPINGS = {