diff --git a/nodes/nodes.py b/nodes/nodes.py index d04053c..083528c 100644 --- a/nodes/nodes.py +++ b/nodes/nodes.py @@ -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, ) \ No newline at end of file