Add nodes and files

__init__.py

@@ -0,0 +1,20 @@
+from .nodes import (WanVideoEncode,
+                    TTMKSamplerSelect,
+                    AddTTMLatent,
+                    WanVideoSamplerCustomUltraAdvancedEfficient)
+
+NODE_CLASS_MAPPINGS = {
+    "WanVideoEncode": WanVideoEncode,
+    "TTMKSamplerSelect": TTMKSamplerSelect,
+    "AddTTMLatent": AddTTMLatent,
+    "WanVideoSamplerCustomUltraAdvancedEfficient": WanVideoSamplerCustomUltraAdvancedEfficient,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "WanVideoEncode": "WanVideo Encode",
+    "TTMKSamplerSelect": "TimeToMove KSampler Select",
+    "AddTTMLatent": "Add TTM Latent",
+    "WanVideoSamplerCustomUltraAdvancedEfficient": "WanVideoSampler Custom Ultra Advanced Efficient",
+}
+
+__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']

nodes.py

@@ -0,0 +1,327 @@
+import torch
+import logging
+from comfy_api.latest import io
+from comfy.utils import PROGRESS_BAR_ENABLED
+import torch.nn.functional as F
+import latent_preview
+import comfy
+from nodes import VAEDecodeTiled, PreviewImage, VAEDecode
+from comfy_extras.nodes_custom_sampler import Noise_EmptyNoise, Noise_RandomNoise
+from comfy.samplers import SAMPLER_NAMES
+from PIL import Image
+from .utils import (pil2tensor, warning, set_preview_method, sample_custom_ultra, 
+                    global_preview_method, store_ksampler_results, globals_cleanup, 
+                    add_noise_at_step, add_noise_to_reference_video)
+from .samplers import sampler_object
+
+
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+log = logging.getLogger(__name__)
+
+
+# Copied from ComfyUI Wanvideo Wrapper
+class WanVideoEncode:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {
+                    "vae": ("VAE",),
+                    "image": ("IMAGE",),
+                    "enable_vae_tiling": ("BOOLEAN", {"default": False, "tooltip": "Drastically reduces memory use but may introduce seams"}),
+                    "tile_x": ("INT", {"default": 272, "min": 64, "max": 2048, "step": 1, "tooltip": "Tile size in pixels, smaller values use less VRAM, may introduce more seams"}),
+                    "tile_y": ("INT", {"default": 272, "min": 64, "max": 2048, "step": 1, "tooltip": "Tile size in pixels, smaller values use less VRAM, may introduce more seams"}),
+                    "tile_stride_x": ("INT", {"default": 144, "min": 32, "max": 2048, "step": 32, "tooltip": "Tile stride in pixels, smaller values use less VRAM, may introduce more seams"}),
+                    "tile_stride_y": ("INT", {"default": 128, "min": 32, "max": 2048, "step": 32, "tooltip": "Tile stride in pixels, smaller values use less VRAM, may introduce more seams"}),
+                    },
+                    "optional": {
+                        "noise_aug_strength": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 10.0, "step": 0.001, "tooltip": "Strength of noise augmentation, helpful for leapfusion I2V where some noise can add motion and give sharper results"}),
+                        "latent_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001, "tooltip": "Additional latent multiplier, helpful for leapfusion I2V where lower values allow for more motion"}),
+                        "mask": ("MASK"),
+                    }
+                }
+
+    RETURN_TYPES = ("LATENT",)
+    RETURN_NAMES = ("reference_latents",)
+    FUNCTION = "encode"
+    CATEGORY = "Wan22 TimeToMove"
+
+    def encode(self, vae, image, enable_vae_tiling, tile_x, tile_y, tile_stride_x, tile_stride_y, noise_aug_strength=0.0, latent_strength=1.0, mask=None):
+        image = image.clone()
+
+        if image.shape[-1] == 4:
+            image = image[..., :3]
+
+        if noise_aug_strength > 0.0:
+            image = add_noise_to_reference_video(image, ratio=noise_aug_strength)
+        
+        if enable_vae_tiling:
+            latents = vae.encode_tiled(image[:,:,:,:3] * 2.0 - 1.0, tile_size=(tile_x//vae.upscale_ratio, tile_y//vae.upscale_ratio), tile_stride=(tile_stride_x//vae.upscale_ratio, tile_stride_y//vae.upscale_ratio))
+        else:
+            latents = vae.encode(image[:,:,:,:3] * 2.0 - 1.0)
+
+        if latent_strength != 1.0:
+            latents *= latent_strength
+
+        log.info(f"WanVideo Encode: Encoded latents shape {latents.shape}")
+ 
+        return ({"samples": latents, "noise_mask": mask},)
+    
+
+# Copied from ComfyUI Wanvideo Wrapper
+class AddTTMLatent:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {
+                    "latent": ("LATENT", {"tooltip": "wanvideo latent"}),
+                    "reference_latents": ("LATENT", {"tooltip": "Reference image to encode"}),
+                    "start_step": ("INT", {"default": 0, "min": -1, "max": 1000, "step": 1, "tooltip": "Start step for whole denoising process"}),
+                    "end_step": ("INT", {"default": 2, "min": 1, "max": 1000, "step": 1, "tooltip": "The step to stop applying TTM"}),
+                    "ref_masks": ("MASK", {"tooltip": "Reference mask to encode"}),
+                }
+        }
+
+    RETURN_TYPES = ("LATENT",)
+    RETURN_NAMES = ("latent",)
+    FUNCTION = "add"
+    CATEGORY = "Wan22 TimeToMove"
+
+    def add(self, latent, reference_latents, start_step, end_step, ref_masks):        
+        if end_step < max(0, start_step):
+            raise ValueError(f"`end_step` ({end_step}) must be >= `start_step` ({start_step}).")
+        
+        mask_sampled = ref_masks[::4]
+        mask_sampled = mask_sampled.unsqueeze(1).unsqueeze(0)  # [1, T, 1, H, W]
+        
+        vae_upscale_factor = 8
+        if reference_latents["samples"].shape[1] == 48:
+            vae_upscale_factor = 16
+        
+        # Upsample spatially to latent resolution
+        H_latent = mask_sampled.shape[-2] // vae_upscale_factor
+        W_latent = mask_sampled.shape[-1] // vae_upscale_factor
+        mask_latent = F.interpolate(
+            mask_sampled.float(),
+            size=(mask_sampled.shape[2], H_latent, W_latent),
+            mode="nearest"
+        )
+
+        latent["ttm_reference_latents"] = reference_latents["samples"].squeeze(0) # [16, T, H, W]
+        latent["ttm_mask"] =  mask_latent.squeeze(0).movedim(1, 0)  # [1, T, H, W]
+        latent["ttm_start_step"] = start_step
+        latent["ttm_end_step"] = end_step
+
+        return (latent,)
+
+
+class TTMKSamplerSelect(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="TTMKSamplerSelect",
+            category="Wan Animate End Reference",
+            inputs=[
+                io.Combo.Input("sampler_name", options=SAMPLER_NAMES, default="lcm"),
+                io.Latent.Input("latent"),
+                ],
+            outputs=[
+                io.Sampler.Output(),
+                ]
+        )
+
+    @classmethod
+    def execute(cls, sampler_name, latent) -> io.NodeOutput:
+        ttm_options = {}
+        ttm_options["ttm_reference_latents"] = latent.get("ttm_reference_latents", None)
+        ttm_options["ttm_start_step"] = latent["ttm_start_step"]
+        ttm_options["ttm_end_step"] = latent["ttm_end_step"]
+        ttm_options["latent_image"] = latent["samples"]
+        ttm_options["motion_mask"] = latent["ttm_mask"]
+
+        sampler = sampler_object(sampler_name, ttm_options)
+        return io.NodeOutput(sampler)
+
+    get_sampler = execute
+
+
+class WanVideoSamplerCustomUltraAdvancedEfficient:
+    # Image Preview code taken from jags111's efficiency-nodes (TSC_KSampler)
+    empty_image = pil2tensor(Image.new('RGBA', (1, 1), (0, 0, 0, 0)))
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required":
+                    {"model": ("MODEL",),
+                    "add_noise": ("BOOLEAN", {"default": True}),
+                    "noise_seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "control_after_generate": True}),
+                    "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}),
+                    "positive": ("CONDITIONING", ),
+                    "negative": ("CONDITIONING", ),
+                    "sampler": ("SAMPLER", ),
+                    "sigmas": ("SIGMAS", ),
+                    "latent": ("LATENT", ),
+                    "start_at_step": ("INT", {"default": 0, "min": 0, "max": 10000}),
+                    "end_at_step": ("INT", {"default": 10000, "min": 0, "max": 10000}),
+                    "return_with_leftover_noise": ("BOOLEAN", {"default": False}),
+                    "preview_method": (["auto", "latent2rgb", "taesd", "vae_decoded_only", "none"],),
+                    "vae_decode": (["true", "true (tiled)", "false"],),
+                    },
+                    "optional": {
+                        "optional_vae": ("VAE",),
+                    },
+                    "hidden": {
+                        "prompt": "PROMPT", 
+                        "extra_pnginfo": "EXTRA_PNGINFO", 
+                        "my_unique_id": "UNIQUE_ID",
+                    },
+                }
+
+    RETURN_TYPES = ("MODEL", "CONDITIONING", "CONDITIONING", "SAMPLER", "SIGMAS", "LATENT","LATENT", "IMAGE", "VAE",)
+    RETURN_NAMES = ("model", "positive", "negative", "sampler", "sigmas", "output", "denoised_output", "image", "vae", )
+    FUNCTION = "sample"
+    CATEGORY = "Wan22 TimeToMove"
+
+    def sample(self, model, add_noise, noise_seed, cfg, positive, negative, sampler, sigmas, latent, start_at_step, end_at_step, return_with_leftover_noise, preview_method, vae_decode, optional_vae=(None,), prompt=None, extra_pnginfo=None, my_unique_id=None):
+        latent_image = latent["samples"]
+        latent_image = comfy.sample.fix_empty_latent_channels(model, latent_image)
+        latent["samples"] = latent_image
+        
+        # Rename the vae variable
+        vae = optional_vae
+        # If vae is not connected, disable vae decoding
+        if vae == (None,) and vae_decode != "false":
+            print(f"{warning('Sampler Custom Ultra Advanced Warning:')} No vae input detected, proceeding as if vae_decode was false.\n")
+            vae_decode = "false"
+        
+        # ------------------------------------------------------------------------------------------------------
+        def vae_decode_latent(vae, out, vae_decode):
+            return VAEDecodeTiled().decode(vae,out,320)[0] if "tiled" in vae_decode else VAEDecode().decode(vae,out)[0]
+        # ---------------------------------------------------------------------------------------------------------------
+
+        noise_mask = None
+        if "noise_mask" in latent:
+            noise_mask = latent["noise_mask"]
+        
+        def process_latents():
+            x0_output = {}
+            # Initialize output variables
+            out = out_denoised = images = preview = previous_preview_method = None
+
+            if not add_noise:
+                noise = Noise_EmptyNoise().generate_noise(latent)
+            else:
+                noise = Noise_RandomNoise(noise_seed).generate_noise(latent)
+                
+                #Time-to-move (TTM)
+                ttm_start_step = 0
+                ttm_reference_latents = latent.get("ttm_reference_latents", None)
+                if ttm_reference_latents is not None:
+                    motion_mask = latent["ttm_mask"].to(latent_image.device, latent_image.dtype)
+                    ttm_start_step = max(latent["ttm_start_step"] - start_at_step, 0)
+                    ttm_end_step = latent["ttm_end_step"] - start_at_step
+    
+                    if ttm_start_step > end_at_step:
+                        raise ValueError("TTM start step is beyond the total number of steps")
+                    
+                    sigma = sigmas[ttm_start_step]
+
+                    if ttm_end_step > ttm_start_step:
+                        log.info("Using Time-to-move (TTM)")
+                        log.info(f"TTM reference latents shape: {ttm_reference_latents.shape}")
+                        log.info(f"TTM motion mask shape: {motion_mask.shape}")
+                        log.info(f"Applying TTM from step {ttm_start_step} to {ttm_end_step}")
+
+                        noise = add_noise_at_step(ttm_reference_latents,
+                                            noise, 
+                                            sigma
+                        ).to(latent_image.device, latent_image.dtype)
+                #--------------------------------------------------------------
+                
+            try:
+                # Change the global preview method (temporarily)
+                set_preview_method(preview_method)
+                
+                x0_output = {}
+                callback = latent_preview.prepare_callback(model, sigmas.shape[-1] - 1, x0_output)
+
+                disable_pbar = not PROGRESS_BAR_ENABLED
+        
+                disable_noise = False
+                if not add_noise:
+                    disable_noise = True
+                                
+                # Prepare noise for img specified by batch_inds
+                if disable_noise:
+                    noise = torch.zeros(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, device="cpu")
+                else:
+                    batch_inds = latent["batch_index"] if "batch_index" in latent else None
+                    noise = comfy.sample.prepare_noise(latent_image, noise_seed, batch_inds)
+                  
+                force_full_denoise = True
+                if return_with_leftover_noise:
+                    force_full_denoise = False
+        
+                device = comfy.model_management.intermediate_device()
+                model_options = model.model_options
+                start_step = start_at_step
+                last_step = end_at_step
+                denoise_mask = noise_mask
+
+                samples = sample_custom_ultra(model, device, 
+                                        noise, 
+                                        sampler, 
+                                        positive, negative, 
+                                        cfg, model_options, 
+                                        latent_image, 
+                                        start_step, last_step, 
+                                        force_full_denoise, denoise_mask, 
+                                        sigmas, 
+                                        callback, disable_pbar, noise_seed)
+                
+                samples = samples.to(comfy.model_management.intermediate_device())
+                
+                out = latent.copy()
+                out["samples"] = samples
+                if "x0" in x0_output:
+                    out_denoised = latent.copy()
+                    out_denoised["samples"] = model.model.process_latent_out(x0_output["x0"].cpu())
+                else:
+                    out_denoised = out
+    
+                previous_preview_method = global_preview_method()
+
+                # ---------------------------------------------------------------------------------------------------------------
+                # Decode image if not yet decoded
+                if "true" in vae_decode:
+                    if images is None:
+                        images = vae_decode_latent(vae, out, vae_decode)
+                        # Store decoded image as base image of no script is detected
+                        store_ksampler_results("image", my_unique_id, images)
+    
+                # Define preview images
+                if preview_method == "none" or (preview_method == "vae_decoded_only" and vae_decode == "false"):
+                    preview = {"images": list()}
+                elif images is not None:
+                    preview = PreviewImage().save_images(images, prompt=prompt, extra_pnginfo=extra_pnginfo)["ui"]
+        
+                # Define a dummy output image
+                if images is None and vae_decode == "false":
+                    images = WanVideoSamplerCustomUltraAdvancedEfficient.empty_image
+
+            finally:
+                # Restore global changes
+                set_preview_method(previous_preview_method)
+              
+            return out, out_denoised, preview, images
+        
+        # ---------------------------------------------------------------------------------------------------------------
+        # Clean globally stored objects of non-existant nodes
+        globals_cleanup(prompt)
+        # ---------------------------------------------------------------------------------------------------------------
+        out, out_denoised, preview, images = process_latents()
+
+        result = (model, positive, negative, sampler, sigmas, 
+                  out, out_denoised, images, vae,)
+
+        if preview is None:
+            return {"result": result}
+        else:
+            return {"ui": preview, "result": result}

samplers.py

@@ -0,0 +1,109 @@
+import torch
+from comfy.samplers import Sampler
+from comfy.extra_samplers import uni_pc
+
+from .k_diffusion import sampling as k_diffusion_sampling
+
+
+class KSamplerX0Inpaint:
+    def __init__(self, model, sigmas):
+        self.inner_model = model
+        self.sigmas = sigmas
+    # Add ttm_options to extra_args
+    def __call__(self, x, sigma, denoise_mask, model_options={}, seed=None,
+                 ttm_reference_latents=None, ttm_start_step=None,
+                 ttm_end_step=None, latent_image=None, motion_mask=None):
+        
+        if denoise_mask is not None:
+            if "denoise_mask_function" in model_options:
+                denoise_mask = model_options["denoise_mask_function"](sigma, denoise_mask, extra_options={"model": self.inner_model, "sigmas": self.sigmas})
+            latent_mask = 1. - denoise_mask
+            x = x * denoise_mask + self.inner_model.inner_model.scale_latent_inpaint(x=x, sigma=sigma, noise=self.noise, latent_image=self.latent_image) * latent_mask
+        model_options["ttm_reference_latents"] = ttm_reference_latents
+        model_options["ttm_start_step"] = ttm_start_step
+        model_options["ttm_end_step"] = ttm_end_step
+        model_options["latent_image"] = latent_image
+        model_options["motion_mask"] = motion_mask
+        out = self.inner_model(x, sigma, model_options=model_options, seed=seed)
+        if denoise_mask is not None:
+            out = out * denoise_mask + self.latent_image * latent_mask
+        return out
+    
+
+class KSAMPLER(Sampler):
+    def __init__(self, sampler_function, extra_options={}, inpaint_options={}):
+        self.sampler_function = sampler_function
+        self.extra_options = extra_options
+        self.inpaint_options = inpaint_options
+
+    def sample(self, model_wrap, sigmas, extra_args, callback, noise, latent_image=None, denoise_mask=None, disable_pbar=False):
+        extra_args["denoise_mask"] = denoise_mask
+        model_k = KSamplerX0Inpaint(model_wrap, sigmas)
+        model_k.latent_image = latent_image
+        if self.inpaint_options.get("random", False): #TODO: Should this be the default?
+            generator = torch.manual_seed(extra_args.get("seed", 41) + 1)
+            model_k.noise = torch.randn(noise.shape, generator=generator, device="cpu").to(noise.dtype).to(noise.device)
+        else:
+            model_k.noise = noise
+
+        noise = model_wrap.inner_model.model_sampling.noise_scaling(sigmas[0], noise, latent_image, self.max_denoise(model_wrap, sigmas))
+
+        k_callback = None
+        total_steps = len(sigmas) - 1
+        if callback is not None:
+            k_callback = lambda x: callback(x["i"], x["denoised"], x["x"], total_steps)
+        
+        samples = self.sampler_function(model_k, noise, sigmas, extra_args=extra_args, callback=k_callback, disable=disable_pbar, **self.extra_options)
+        samples = model_wrap.inner_model.model_sampling.inverse_noise_scaling(sigmas[-1], samples)
+        return samples
+    
+
+def ksampler(sampler_name, ttm_options, extra_options={}, inpaint_options={}):
+    if sampler_name == "dpm_fast":
+        def dpm_fast_function(model, noise, sigmas, extra_args, callback, disable):
+            if len(sigmas) <= 1:
+                return noise
+
+            sigma_min = sigmas[-1]
+            if sigma_min == 0:
+                sigma_min = sigmas[-2]
+            total_steps = len(sigmas) - 1
+            return k_diffusion_sampling.sample_dpm_fast(model, noise, sigma_min, sigmas[0], total_steps, extra_args=extra_args, callback=callback, disable=disable)
+        sampler_function = dpm_fast_function
+    elif sampler_name == "dpm_adaptive":
+        def dpm_adaptive_function(model, noise, sigmas, extra_args, callback, disable, **extra_options):
+            if len(sigmas) <= 1:
+                return noise
+
+            sigma_min = sigmas[-1]
+            if sigma_min == 0:
+                sigma_min = sigmas[-2]
+            return k_diffusion_sampling.sample_dpm_adaptive(model, noise, sigma_min, sigmas[0], extra_args=extra_args, callback=callback, disable=disable, **extra_options)
+        sampler_function = dpm_adaptive_function
+    elif sampler_name == "lcm":
+        def lcm_function(model, noise, sigmas, extra_args, callback, disable, **extra_options):
+            extra_args["ttm_reference_latents"] = ttm_options["ttm_reference_latents"]
+            extra_args["ttm_start_step"] = ttm_options["ttm_start_step"]
+            extra_args["ttm_end_step"] = ttm_options["ttm_end_step"]
+            extra_args["latent_image"] = ttm_options["latent_image"]
+            extra_args["motion_mask"] = ttm_options["motion_mask"]
+            return k_diffusion_sampling.sample_lcm(model, noise, sigmas, extra_args=extra_args, callback=callback, disable=disable, **extra_options)
+        sampler_function = lcm_function
+    else:
+        sampler_function = getattr(k_diffusion_sampling, "sample_{}".format(sampler_name))
+
+    return KSAMPLER(sampler_function, extra_options, inpaint_options)
+
+
+def sampler_object(name, ttm_options):
+    if name == "uni_pc":
+        sampler = KSAMPLER(uni_pc.sample_unipc)
+    elif name == "uni_pc_bh2":
+        sampler = KSAMPLER(uni_pc.sample_unipc_bh2)
+    elif name == "ddim":
+        sampler = ksampler("euler", inpaint_options={"random": True})
+    elif name == "lcm":
+        sampler = ksampler(name, ttm_options)
+    else:
+        sampler = ksampler(name)
+    return sampler

utils.py

@@ -0,0 +1,139 @@
+import torch
+from PIL import Image
+import numpy as np
+import latent_preview
+from comfy.cli_args import args
+from comfy.samplers import sample
+
+
+# Convert PIL to Tensor (grabbed from WAS Suite)
+def pil2tensor(image: Image.Image) -> torch.Tensor:
+    return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
+
+def format_message(text, color_code):
+    RESET_COLOR = "\033[0m"
+    return f"{color_code}{text}{RESET_COLOR}"
+
+WARNING_COLOR = "\033[93m"  # Yellow
+
+def warning(text):
+    return format_message(text, WARNING_COLOR)
+
+
+# Set global preview_method
+def set_preview_method(method):
+    if method == 'auto' or method == 'LatentPreviewMethod.Auto':
+        args.preview_method = latent_preview.LatentPreviewMethod.Auto
+    elif method == 'latent2rgb' or method == 'LatentPreviewMethod.Latent2RGB':
+        args.preview_method = latent_preview.LatentPreviewMethod.Latent2RGB
+    elif method == 'taesd' or method == 'LatentPreviewMethod.TAESD':
+        args.preview_method = latent_preview.LatentPreviewMethod.TAESD
+    else:
+        args.preview_method = latent_preview.LatentPreviewMethod.NoPreviews
+
+
+def sample_custom_ultra(model, device, noise, sampler, positive, negative, cfg, model_options={}, latent_image=None, start_step=None, last_step=None, force_full_denoise=False, denoise_mask=None, sigmas=None, callback=None, disable_pbar=False, seed=None):
+    if last_step is not None and last_step < (len(sigmas) - 1):
+        sigmas = sigmas[:last_step + 1]
+        if force_full_denoise:
+            sigmas[-1] = 0
+
+    if start_step is not None:
+        if start_step < (len(sigmas) - 1):
+            sigmas = sigmas[start_step:]
+        else:
+            if latent_image is not None:
+                return latent_image
+            else:
+                return torch.zeros_like(noise)
+    
+    return sample(model, noise, positive, negative, cfg, device, sampler, sigmas, model_options, latent_image=latent_image, denoise_mask=denoise_mask, callback=callback, disable_pbar=disable_pbar, seed=seed)
+
+
+# Extract global preview_method
+def global_preview_method():
+    return args.preview_method
+
+
+# Cache for Efficiency Node models
+loaded_objects = {
+    "ckpt": [], # (ckpt_name, ckpt_model, clip, bvae, [id])
+    "refn": [], # (ckpt_name, ckpt_model, clip, bvae, [id])
+    "vae": [],  # (vae_name, vae, [id])
+    "lora": []  # ([(lora_name, strength_model, strength_clip)], ckpt_name, lora_model, clip_lora, [id])
+}
+
+# Cache for Efficient Ksamplers
+last_helds = {
+    "latent": [],   # (latent, [parameters], id)    # Base sampling latent results
+    "image": [],    # (image, id)                   # Base sampling image results
+    "cnet_img": []  # (cnet_img, [parameters], id)  # HiRes-Fix control net preprocessor image results
+}
+
+def store_ksampler_results(key: str, my_unique_id, value, parameters_list=None):
+    global last_helds
+
+    for i, data in enumerate(last_helds[key]):
+        id_ = data[-1]  # ID will always be the last in the tuple
+        if id_ == my_unique_id:
+            # Check if parameters_list is provided or not
+            updated_data = (value, parameters_list, id_) if parameters_list is not None else (value, id_)
+            last_helds[key][i] = updated_data
+            return True
+
+    # If parameters_list is given
+    if parameters_list is not None:
+        last_helds[key].append((value, parameters_list, my_unique_id))
+    else:
+        last_helds[key].append((value, my_unique_id))
+    return True
+
+
+# This function cleans global variables associated with nodes that are no longer detected on UI
+def globals_cleanup(prompt):
+    global loaded_objects
+    global last_helds
+
+    # Step 1: Clean up last_helds
+    for key in list(last_helds.keys()):
+        original_length = len(last_helds[key])
+        last_helds[key] = [
+            (*values, id_)
+            for *values, id_ in last_helds[key]
+            if str(id_) in prompt.keys()
+        ]
+
+    # Step 2: Clean up loaded_objects
+    for key in list(loaded_objects.keys()):
+        for i, tup in enumerate(list(loaded_objects[key])):
+            # Remove ids from id array in each tuple that don't exist in prompt
+            id_array = [id for id in tup[-1] if str(id) in prompt.keys()]
+            if len(id_array) != len(tup[-1]):
+                if id_array:
+                    loaded_objects[key][i] = tup[:-1] + (id_array,)
+                    #print(f'Updated tuple at index {i} in {key} in loaded_objects: {loaded_objects[key][i]}')
+                else:
+                    # If id array becomes empty, delete the corresponding tuple
+                    loaded_objects[key].remove(tup)
+                    #print(f'Deleted tuple at index {i} in {key} in loaded_objects because its id array became empty.')
+
+
+# Copied from ComfyUI Wanvideo Wrapper
+def add_noise_to_reference_video(image, ratio=None):
+    sigma = torch.ones((image.shape[0],)).to(image.device, image.dtype) * ratio
+    image_noise = torch.randn_like(image) * sigma[:, None, None, None]
+    image_noise = torch.where(image==-1, torch.zeros_like(image), image_noise)
+    image = image + image_noise
+    return image
+
+
+# Copied from Kijai Wanvideo Wrapper
+def add_noise_at_step(
+    original_samples: torch.FloatTensor,
+    noise: torch.FloatTensor,
+    sigma: torch.IntTensor,
+) -> torch.FloatTensor:
+    
+    sigma = sigma.view(sigma.shape + (1,) * (len(noise.shape)-1))
+
+    return (1 - sigma) * original_samples + sigma * noise