xinyun/ComfyUI-KJNodes
1
0
Fork 0
mirror of https://git.datalinker.icu/kijai/ComfyUI-KJNodes.git synced 2025-12-09 12:54:40 +08:00
Code Issues Packages Projects Releases Wiki Activity

TTM: Rename end_steps to steps for clarity

Browse Source
This commit is contained in:
kijai 2025-11-23 02:10:23 +02:00
parent 3b9c1b49ab
commit 8660778ea1
1 changed files with 9 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
nodes/nodes.py
View File

@ -2628,9 +2628,9 @@ class LazySwitchKJ:
from comfy.patcher_extension import WrappersMP
from comfy.sampler_helpers import prepare_mask
class TTM_SampleWrapper:
def __init__(self, mask, end_step):
def __init__(self, mask, steps):
self.mask = mask
self.end_step = end_step
self.steps = steps
def __call__(self, sampler, guider, sigmas, extra_args, callback, noise, latent_image, denoise_mask, disable_pbar):
model_options = extra_args["model_options"]
@ -2642,7 +2642,7 @@ class TTM_SampleWrapper:
motion_mask = prepare_mask(motion_mask, noise.shape, noise.device)
scale_latent_inpaint = guider.model_patcher.model.scale_latent_inpaint
w["TTM_ApplyModel_Wrapper"] = [TTM_ApplyModel_Wrapper(latent_image, noise, motion_mask, self.end_step, scale_latent_inpaint)]
w["TTM_ApplyModel_Wrapper"] = [TTM_ApplyModel_Wrapper(latent_image, noise, motion_mask, self.steps, scale_latent_inpaint)]
out = sampler(guider, sigmas, extra_args, callback, noise, latent_image, denoise_mask, disable_pbar)
@ -2650,11 +2650,11 @@ class TTM_SampleWrapper:
class TTM_ApplyModel_Wrapper:
def __init__(self, reference_samples, noise, motion_mask, end_step, scale_latent_inpaint):
def __init__(self, reference_samples, noise, motion_mask, steps, scale_latent_inpaint):
self.reference_samples = reference_samples
self.noise = noise
self.motion_mask = motion_mask
self.end_step = end_step
self.steps = steps
self.scale_latent_inpaint = scale_latent_inpaint
def __call__(self, executor, x, t, c_concat, c_crossattn, control, transformer_options, **kwargs):
@ -2669,7 +2669,7 @@ class TTM_ApplyModel_Wrapper:
next_sigma = sigmas[current_step_index + 1] if current_step_index < len(sigmas) - 1 else sigmas[current_step_index]
if current_step_index != 0 and current_step_index < self.end_step:
if current_step_index != 0 and current_step_index < self.steps:
noisy_latent = self.scale_latent_inpaint(x=x, sigma=torch.tensor([next_sigma]), noise=self.noise.to(x), latent_image=self.reference_samples.to(x))
x = x * (1-self.motion_mask).to(x) + noisy_latent * self.motion_mask.to(x)
@ -2681,7 +2681,7 @@ class LatentInpaintTTM:
def INPUT_TYPES(s):
return {"required": {
"model": ("MODEL", ),
"end_step": ("INT", {"default": 7, "min": 0, "max": 888, "step": 1}),
"steps": ("INT", {"default": 7, "min": 0, "max": 888, "step": 1, "tooltip": "Number of steps to apply TTM inpainting for."}),
},
"optional": {
"mask": ("MASK", {"tooltip": "Latent mask where white (1.0) is the area to inpaint and black (0.0) is the area to keep unchanged."}),
@ -2693,7 +2693,7 @@ class LatentInpaintTTM:
DESCRIPTION = "https://github.com/time-to-move/TTM"
CATEGORY = "KJNodes/experimental"
def patch(self, model, end_step, mask=None):
def patch(self, model, steps, mask=None):
m = model.clone()
m.add_wrapper_with_key(WrappersMP.SAMPLER_SAMPLE, "TTM_SampleWrapper", TTM_SampleWrapper(mask, end_step))
m.add_wrapper_with_key(WrappersMP.SAMPLER_SAMPLE, "TTM_SampleWrapper", TTM_SampleWrapper(mask, steps))
return (m, )