Cody Schultz
b1ac358ddc
re-wrote codebase to utilize ComfyUIs model patcher functions so it is easier to use and more robust. this change also makes it so positive values add detail and negative values remove detail.
109 lines
4.3 KiB
Python
109 lines
4.3 KiB
Python
import torch
|
|
import torch.nn.functional as F
|
|
import comfy.model_patcher
|
|
import comfy.samplers
|
|
import math
|
|
|
|
class Detailer:
|
|
"""
|
|
A ComfyUI node that uses the `set_model_unet_function_wrapper` hook to precisely modify
|
|
the raw noise prediction output from the UNet model. This allows for fine-grained
|
|
control over detail before or during the CFG calculation.
|
|
|
|
- A positive `scaling_factor` amplifies details by increasing the residual noise in the output.
|
|
- A negative `scaling_factor` smoothens details by dampening the residual noise in the output.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"scaling_factor": ("FLOAT", {
|
|
"default": 0.05,
|
|
"min": -1.0,
|
|
"max": 1.0,
|
|
"step": 0.005,
|
|
"display": "slider",
|
|
"tooltip": "Positive values increase detail, negative values reduce detail."
|
|
}),
|
|
"scaling_narrowness": ("FLOAT", {
|
|
"default": 1.0,
|
|
"min": 0.1,
|
|
"max": 10.0,
|
|
"step": 0.1,
|
|
"display": "slider",
|
|
"tooltip": "Controls the focus of the scaling effect within the sampling process (higher = narrower peak)."
|
|
}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
FUNCTION = "patch"
|
|
CATEGORY = "model_patches/details" # This defines where it appears in the ComfyUI menu
|
|
|
|
def patch(self, model, scaling_factor, scaling_narrowness):
|
|
"""
|
|
Applies a patch to the model's UNet function to modify noise predictions
|
|
based on a sine wave schedule and a scaling factor.
|
|
"""
|
|
|
|
# we scale the user input to make it an acceptable sigma modification so the user doesn't have to work with very small numbers'
|
|
actual_scaling_factor = scaling_factor * 0.05
|
|
|
|
def sin_schedule_for_unet(timestep_tensor, current_sigmas):
|
|
"""
|
|
Calculates a scaling multiplier based on the current timestep (noise level).
|
|
This function creates a sine wave curve across the sampling steps,
|
|
allowing the effect to be strongest in the middle of the process.
|
|
"""
|
|
# ComfyUI's `model.apply_model`
|
|
# typically passes the current 'sigma' value as the 't' (timestep_tensor) argument.
|
|
boolean_mask = (timestep_tensor.item() == current_sigmas)
|
|
steps = len(current_sigmas) - 1
|
|
step = torch.argwhere(boolean_mask)
|
|
current_sigmas = timestep_tensor
|
|
|
|
normalized_progress = step / steps
|
|
|
|
# Clamp the normalized sigma to ensure it stays within [0, 1] just in case my code is too jank
|
|
normalized_progress_clamped = torch.clamp(normalized_progress, 0.0, 1.0)
|
|
|
|
# Apply the sine wave schedule. The `scaling_narrowness`
|
|
# controls the sharpness of the peak of the sine wave.
|
|
return torch.sin(normalized_progress_clamped * math.pi) ** scaling_narrowness
|
|
|
|
def apply_detailer(apply_model, args):
|
|
input_x = args["input"]
|
|
timestep = args["timestep"]
|
|
cond_or_uncond = args["cond_or_uncond"]
|
|
c = args["c"]
|
|
|
|
|
|
mod_strength = sin_schedule_for_unet(timestep[0], c["transformer_options"]["sample_sigmas"])
|
|
timestep[0] = torch.mul(timestep[0], (1 + (mod_strength * actual_scaling_factor)))
|
|
|
|
noise_pred = apply_model(input_x, timestep, **c)
|
|
|
|
return noise_pred
|
|
|
|
# Clone the model to ensure that this patch only affects the current branch
|
|
# of the ComfyUI workflow and does not modify the original model in place. idk if this is really necessary but I don't think it can hurt.
|
|
model_clone = model.clone()
|
|
|
|
# Apply the custom `apply_detailer` method using ComfyUI's patching mechanism.
|
|
# This hooks into the model's internal UNet forward pass.
|
|
model_clone.set_model_unet_function_wrapper(apply_detailer)
|
|
|
|
return (model_clone,)
|
|
|
|
# Node registration for ComfyUI.
|
|
# These dictionaries tell ComfyUI how to display and interact with your custom node.
|
|
NODE_CLASS_MAPPINGS = {
|
|
"Detailer": Detailer,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"Detailer": "Detailer (UNet Patch)",
|
|
}
|