|
|
|
|
|
import torch |
|
|
|
class DifferentialDiffusion(): |
|
@classmethod |
|
def INPUT_TYPES(s): |
|
return {"required": {"model": ("MODEL", ), |
|
}} |
|
RETURN_TYPES = ("MODEL",) |
|
FUNCTION = "apply" |
|
CATEGORY = "_for_testing" |
|
INIT = False |
|
|
|
def apply(self, model): |
|
model = model.clone() |
|
model.set_model_denoise_mask_function(self.forward) |
|
return (model,) |
|
|
|
def forward(self, sigma: torch.Tensor, denoise_mask: torch.Tensor, extra_options: dict): |
|
model = extra_options["model"] |
|
step_sigmas = extra_options["sigmas"] |
|
sigma_to = model.inner_model.model_sampling.sigma_min |
|
if step_sigmas[-1] > sigma_to: |
|
sigma_to = step_sigmas[-1] |
|
sigma_from = step_sigmas[0] |
|
|
|
ts_from = model.inner_model.model_sampling.timestep(sigma_from) |
|
ts_to = model.inner_model.model_sampling.timestep(sigma_to) |
|
current_ts = model.inner_model.model_sampling.timestep(sigma[0]) |
|
|
|
threshold = (current_ts - ts_to) / (ts_from - ts_to) |
|
|
|
return (denoise_mask >= threshold).to(denoise_mask.dtype) |
|
|
|
|
|
NODE_CLASS_MAPPINGS = { |
|
"DifferentialDiffusion": DifferentialDiffusion, |
|
} |
|
NODE_DISPLAY_NAME_MAPPINGS = { |
|
"DifferentialDiffusion": "Differential Diffusion", |
|
} |
|
|
