Repository: comfyanonymous/ComfyUI

Issue:
Differential Diffusion: Giving Each Pixel Its Strength
Hello,
I would like to suggest implementing my paper: Differential Diffusion: Giving Each Pixel Its Strength.
The paper allows a user to edit a picture by a change map that describes how much each region should change.
The editing process is typically guided by textual instructions, although it can also be applied without guidance.
We support both continuous and discrete editing.
Our framework is training and fine tuning free! And has negligible penalty of the inference time.
Our implementation is diffusers-based.
We already tested it on 4 different diffusion models (Kadinsky, DeepFloyd IF, SD, SD XL).
We are confident that the framework can also be ported to other diffusion models, such as SD Turbo, Stable Cascade, and amused.
I notice that you usually stick to white==change convention, which is opposite to the convention we used in the paper.
The paper can be thought of as a generalization to some of the existing techniques.
A black map is just regular txt2img ("0"),
A map of one color (which isn't black) can be thought as img2img,
A map of two colors which one color is white can be thought as inpaint.
And the rest? It's completely new!
In the paper, we suggest some further applications such as soft inpainting and strength visualization.

Site:
https://differential-diffusion.github.io/
Paper:
https://differential-diffusion.github.io/paper.pdf
Repo:
https://github.com/exx8/differential-diffusion
Feature Request: Support Differential Diffusion for inpainting.
This is a nice alternative to standard inpainting, it allows for the mask to be a gradient for control of strength on top of denoising.


https://github.com/exx8/differential-diffusion


Discussion:
So is the primary difference between this and something like the QRCodeMonster controlnet in speed?  The controlnets can slow things down quite a bit but it has the same effect in terms of masking out which areas will change when applied on a latent.  

Here's a fun thought for you... have you considered allowing the use of normal maps as a method of describing directionality in the area with two channels and overall change with the intensity of the third?   I don't know how hard this would be but it would be iteresting if normals could be painted such that the direction of things like water flows, trees, grass, and hair would generally try to follow the specified direction.  It could probably be used to help fix things that tend to generate wrong on a regular basis (e.g. try the prompt "smoking a cigarette" on a character or human and it'll nearly always be floating but on the off chance it's in their mouth it's almost always backwards.  This would allow a bit of control.
this looks very cool, but apparently the code hasn’t been released yet. 
True, but it's laid out quite neatly in the paper. I just don't have the chops yet to implement it.
They've now released their code!
Since the time of my post, I learned of a PR that had been sitting on the automatic1111 repo to perform the same task, I don't know how it compares. https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/14208

Resolving patch:
diff --git a/comfy/samplers.py b/comfy/samplers.py
--- a/comfy/samplers.py
+++ b/comfy/samplers.py
@@ -276,6 +276,8 @@ def __init__(self, model):
         self.inner_model = model
     def forward(self, x, sigma, uncond, cond, cond_scale, denoise_mask, model_options={}, seed=None):
         if denoise_mask is not None:
+            if "denoise_mask_function" in model_options:
+                denoise_mask = model_options["denoise_mask_function"](sigma, denoise_mask)
             latent_mask = 1. - denoise_mask
             x = x * denoise_mask + (self.latent_image + self.noise * sigma.reshape([sigma.shape[0]] + [1] * (len(self.noise.shape) - 1))) * latent_mask
         out = self.inner_model(x, sigma, cond=cond, uncond=uncond, cond_scale=cond_scale, model_options=model_options, seed=seed)
diff --git a/comfy_extras/nodes_differential_diffusion.py b/comfy_extras/nodes_differential_diffusion.py
new file mode 100644
--- /dev/null
+++ b/comfy_extras/nodes_differential_diffusion.py
@@ -0,0 +1,97 @@
+# code adapted from https://github.com/exx8/differential-diffusion
+
+import torch
+import inspect
+
+class DifferentialDiffusion():
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {"model": ("MODEL", ),
+                            }}
+    RETURN_TYPES = ("MODEL",)
+    FUNCTION = "apply"
+    CATEGORY = "_for_testing"
+    INIT = False
+
+    @classmethod
+    def IS_CHANGED(s, *args, **kwargs):
+        DifferentialDiffusion.INIT = s.INIT = True
+        return ""
+
+    def __init__(self) -> None:
+        DifferentialDiffusion.INIT = False
+        self.sigmas: torch.Tensor = None
+        self.thresholds: torch.Tensor = None
+        self.mask_i = None
+        self.valid_sigmas = False
+        self.varying_sigmas_samplers = ["dpmpp_2s", "dpmpp_sde", "dpm_2", "heun", "restart"]
+
+    def apply(self, model):
+        model = model.clone()
+        model.model_options["denoise_mask_function"] = self.forward
+        return (model,)
+    
+    def init_sigmas(self, sigma: torch.Tensor, denoise_mask: torch.Tensor):
+        self.__init__()
+        self.sigmas, sampler = find_outer_instance("sigmas", callback=get_sigmas_and_sampler) or (None, "")
+        self.valid_sigmas = not ("sample_" not in sampler or any(s in sampler for s in self.varying_sigmas_samplers)) or "generic" in sampler
+        if self.sigmas is None:
+            self.sigmas = sigma[:1].repeat(2)
+            self.sigmas[-1].zero_()
+        self.sigmas_min = self.sigmas.min()
+        self.sigmas_max = self.sigmas.max()
+        self.thresholds = torch.linspace(1, 0, self.sigmas.shape[0], dtype=sigma.dtype, device=sigma.device)
+        self.thresholds_min_len = self.thresholds.shape[0] - 1
+        if self.valid_sigmas:
+            thresholds = self.thresholds[:-1].reshape(-1, 1, 1, 1, 1)
+            mask = denoise_mask.unsqueeze(0)
+            mask = (mask >= thresholds).to(denoise_mask.dtype)
+            self.mask_i = iter(mask)
+    
+    def forward(self, sigma: torch.Tensor, denoise_mask: torch.Tensor):
+        if self.sigmas is None or DifferentialDiffusion.INIT:
+            self.init_sigmas(sigma, denoise_mask)
+        if self.valid_sigmas:
+            try:
+                return next(self.mask_i)
+            except StopIteration:
+                self.valid_sigmas = False
+        if self.thresholds_min_len > 1:
+            nearest_idx = (self.sigmas - sigma[0]).abs().argmin()
+            if not self.thresholds_min_len > nearest_idx:
+                nearest_idx = -2
+            threshold = self.thresholds[nearest_idx]
+        else:
+            threshold = (sigma[0] - self.sigmas_min) / (self.sigmas_max - self.sigmas_min)
+        return (denoise_mask >= threshold).to(denoise_mask.dtype)
+
+def get_sigmas_and_sampler(frame, target):
+    found = frame.f_locals[target]
+    if isinstance(found, torch.Tensor) and found[-1] < 0.1:
+        return found, frame.f_code.co_name
+    return False
+
+def find_outer_instance(target: str, target_type=None, callback=None):
+    frame = inspect.currentframe()
+    i = 0
+    while frame and i < 100:
+        if target in frame.f_locals:
+            if callback is not None:
+                res = callback(frame, target)
+                if res:
+                    return res
+            else:
+                found = frame.f_locals[target]
+                if isinstance(found, target_type):
+                    return found
+        frame = frame.f_back
+        i += 1
+    return None
+
+    
+NODE_CLASS_MAPPINGS = {
+    "DifferentialDiffusion": DifferentialDiffusion,
+}
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "DifferentialDiffusion": "Differential Diffusion",
+}
diff --git a/nodes.py b/nodes.py
--- a/nodes.py
+++ b/nodes.py
@@ -1961,6 +1961,7 @@ def init_custom_nodes():
         "nodes_photomaker.py",
         "nodes_cond.py",
         "nodes_stable_cascade.py",
+        "nodes_differential_diffusion.py",
     ]
 
     for node_file in extras_files:


Test patch: