From ce31337813cec2e1e4e647f684550f1d5d0060b0 Mon Sep 17 00:00:00 2001
From: ThereforeGames <eric@sparknight.io>
Date: Thu, 21 Nov 2024 21:02:49 -0500
Subject: [PATCH 1/2] Update nodes_differential_diffusion.py

---
 comfy_extras/nodes_differential_diffusion.py | 28 +++++++++++++++-----
 1 file changed, 21 insertions(+), 7 deletions(-)

diff --git a/comfy_extras/nodes_differential_diffusion.py b/comfy_extras/nodes_differential_diffusion.py
index 98dbbf10..03f2a1b0 100644
--- a/comfy_extras/nodes_differential_diffusion.py
+++ b/comfy_extras/nodes_differential_diffusion.py
@@ -5,19 +5,27 @@ import torch
 class DifferentialDiffusion():
     @classmethod
     def INPUT_TYPES(s):
-        return {"required": {"model": ("MODEL", ),
-                            }}
+        return {
+            "required": {
+                "model": ("MODEL", ),
+                "strength": ("FLOAT", {
+                    "default": 1.0,
+                    "min": 0.0,
+                    "max": 1.0
+                }),
+            }
+        }
     RETURN_TYPES = ("MODEL",)
     FUNCTION = "apply"
     CATEGORY = "_for_testing"
     INIT = False
 
-    def apply(self, model):
+    def apply(self, model, strength=1.0):
         model = model.clone()
-        model.set_model_denoise_mask_function(self.forward)
-        return (model,)
+        model.set_model_denoise_mask_function(lambda *args, **kwargs: self.forward(*args, **kwargs, strength=strength))
+        return (model, )
 
-    def forward(self, sigma: torch.Tensor, denoise_mask: torch.Tensor, extra_options: dict):
+    def forward(self, sigma: torch.Tensor, denoise_mask: torch.Tensor, extra_options: dict, strength: float):
         model = extra_options["model"]
         step_sigmas = extra_options["sigmas"]
         sigma_to = model.inner_model.model_sampling.sigma_min
@@ -31,7 +39,13 @@ class DifferentialDiffusion():
 
         threshold = (current_ts - ts_to) / (ts_from - ts_to)
 
-        return (denoise_mask >= threshold).to(denoise_mask.dtype)
+        # Generate the binary mask based on the threshold
+        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
 
 
 NODE_CLASS_MAPPINGS = {

From 1a864435f611428e0304412c7b4f649161a9e9c4 Mon Sep 17 00:00:00 2001
From: Eric <eric@sparknight.io>
Date: Thu, 21 Nov 2024 23:43:17 -0500
Subject: [PATCH 2/2] Update nodes_differential_diffusion.py

---
 comfy_extras/nodes_differential_diffusion.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/comfy_extras/nodes_differential_diffusion.py b/comfy_extras/nodes_differential_diffusion.py
index 03f2a1b0..1ada3270 100644
--- a/comfy_extras/nodes_differential_diffusion.py
+++ b/comfy_extras/nodes_differential_diffusion.py
@@ -43,9 +43,11 @@ class DifferentialDiffusion():
         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 = {