import os import torch import gc from .utils import log, print_memory import numpy as np import math from tqdm import tqdm from .wanvideo.modules.clip import CLIPModel from .wanvideo.modules.model import WanModel, rope_params from .wanvideo.modules.t5 import T5EncoderModel from .wanvideo.utils.fm_solvers import (FlowDPMSolverMultistepScheduler, get_sampling_sigmas, retrieve_timesteps) from .wanvideo.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler from diffusers.schedulers import FlowMatchEulerDiscreteScheduler from .enhance_a_video.globals import enable_enhance, disable_enhance, set_enhance_weight, set_num_frames from accelerate import init_empty_weights from accelerate.utils import set_module_tensor_to_device import folder_paths import comfy.model_management as mm from comfy.utils import load_torch_file, save_torch_file, ProgressBar, common_upscale import comfy.model_base import comfy.latent_formats from comfy.clip_vision import clip_preprocess, ClipVisionModel script_directory = os.path.dirname(os.path.abspath(__file__)) 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 class WindowTrackerSE: def __init__(self, verbose=False): self.window_map = {} # Maps frame sequence to persistent ID self.next_id = 0 self.teacache_states = {} # Maps persistent ID to teacache state self.verbose = verbose def get_window_id(self, frames): key = tuple(sorted(frames)) # Order-independent frame sequence if key not in self.window_map: self.window_map[key] = self.next_id if self.verbose: log.info(f"New window pattern {key} -> ID {self.next_id}") self.next_id += 1 return self.window_map[key] def get_teacache(self, window_id, base_state): if window_id not in self.teacache_states: if self.verbose: log.info(f"Initializing persistent teacache for window {window_id}") self.teacache_states[window_id] = base_state.copy() return self.teacache_states[window_id] class WanVideoSEModel(comfy.model_base.BaseModel): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.pipeline = {} def __getitem__(self, k): return self.pipeline[k] def __setitem__(self, k, v): self.pipeline[k] = v try: from comfy.latent_formats import Wan21 latent_format = Wan21 except: #for backwards compatibility log.warning("Wan21 latent format not found, update ComfyUI for better livepreview") from comfy.latent_formats import HunyuanVideo latent_format = HunyuanVideo class WanVideoSEModelConfig: def __init__(self, dtype): self.unet_config = {} self.unet_extra_config = {} self.latent_format = latent_format self.latent_format.latent_channels = 16 self.manual_cast_dtype = dtype self.sampling_settings = {"multiplier": 1.0} # Don't know what this is. Value taken from ComfyUI Mochi model. self.memory_usage_factor = 2.0 # denoiser is handled by extension self.unet_config["disable_unet_model_creation"] = True def filter_state_dict_by_blocks(state_dict, blocks_mapping): filtered_dict = {} for key in state_dict: if 'blocks.' in key: block_pattern = key.split('diffusion_model.')[1].split('.', 2)[0:2] block_key = f'{block_pattern[0]}.{block_pattern[1]}.' if block_key in blocks_mapping: filtered_dict[key] = state_dict[key] return filtered_dict def standardize_lora_key_format(lora_sd): new_sd = {} for k, v in lora_sd.items(): # Diffusers format if k.startswith('transformer.'): k = k.replace('transformer.', 'diffusion_model.') # from finetrainer format if '.attn1.' in k: k = k.replace('.attn1.', '.cross_attn.') k = k.replace('.to_k.', '.k.') k = k.replace('.to_q.', '.q.') k = k.replace('.to_v.', '.v.') k = k.replace('.to_out.0.', '.o.') elif '.attn2.' in k: k = k.replace('.attn2.', '.cross_attn.') k = k.replace('.to_k.', '.k.') k = k.replace('.to_q.', '.q.') k = k.replace('.to_v.', '.v.') k = k.replace('.to_out.0.', '.o.') if "img_attn.proj" in k: k = k.replace("img_attn.proj", "img_attn_proj") if "img_attn.qkv" in k: k = k.replace("img_attn.qkv", "img_attn_qkv") if "txt_attn.proj" in k: k = k.replace("txt_attn.proj ", "txt_attn_proj") if "txt_attn.qkv" in k: k = k.replace("txt_attn.qkv", "txt_attn_qkv") new_sd[k] = v return new_sd #region Model loading class WanVideoSEModelLoader: @classmethod def INPUT_TYPES(s): return { "required": { "model": (folder_paths.get_filename_list("diffusion_models"), {"tooltip": "These models are loaded from the 'ComfyUI/models/diffusion_models' -folder",}), "base_precision": (["fp32", "bf16", "fp16", "fp16_fast"], {"default": "bf16"}), "quantization": (['disabled', 'fp8_e4m3fn', 'fp8_e4m3fn_fast', 'fp8_e5m2', 'torchao_fp8dq', "torchao_fp8dqrow", "torchao_int8dq", "torchao_fp6", "torchao_int4", "torchao_int8"], {"default": 'disabled', "tooltip": "optional quantization method"}), "load_device": (["main_device", "offload_device"], {"default": "main_device"}), }, "optional": { "attention_mode": ([ "sdpa", "flash_attn_2", "flash_attn_3", "sageattn", "spargeattn", "spargeattn_tune", ], {"default": "sdpa"}), "compile_args": ("WANCOMPILEARGS", ), "block_swap_args": ("BLOCKSWAPARGS", ), "lora": ("WANVIDLORA", {"default": None}), "vram_management_args": ("VRAM_MANAGEMENTARGS", {"default": None, "tooltip": "Alternative offloading method from DiffSynth-Studio, more aggressive in reducing memory use than block swapping, but can be slower"}), } } RETURN_TYPES = ("WANVIDEOMODEL",) RETURN_NAMES = ("model", ) FUNCTION = "loadmodel" CATEGORY = "WanVideoStartEndFrame" def loadmodel(self, model, base_precision, load_device, quantization, compile_args=None, attention_mode="sdpa", block_swap_args=None, lora=None, vram_management_args=None): assert not (vram_management_args is not None and block_swap_args is not None), "Can't use both block_swap_args and vram_management_args at the same time" transformer = None mm.unload_all_models() mm.cleanup_models() mm.soft_empty_cache() manual_offloading = True if "sage" in attention_mode: try: from sageattention import sageattn except Exception as e: raise ValueError(f"Can't import SageAttention: {str(e)}") device = mm.get_torch_device() offload_device = mm.unet_offload_device() manual_offloading = True transformer_load_device = device if load_device == "main_device" else offload_device base_dtype = {"fp8_e4m3fn": torch.float8_e4m3fn, "fp8_e4m3fn_fast": torch.float8_e4m3fn, "bf16": torch.bfloat16, "fp16": torch.float16, "fp16_fast": torch.float16, "fp32": torch.float32}[base_precision] if base_precision == "fp16_fast": if hasattr(torch.backends.cuda.matmul, "allow_fp16_accumulation"): torch.backends.cuda.matmul.allow_fp16_accumulation = True else: raise ValueError("torch.backends.cuda.matmul.allow_fp16_accumulation is not available in this version of torch, requires torch 2.7.0.dev2025 02 26 nightly minimum currently") else: try: if hasattr(torch.backends.cuda.matmul, "allow_fp16_accumulation"): torch.backends.cuda.matmul.allow_fp16_accumulation = False except: pass model_path = folder_paths.get_full_path_or_raise("diffusion_models", model) sd = load_torch_file(model_path, device=transformer_load_device, safe_load=True) first_key = next(iter(sd)) if first_key.startswith("model.diffusion_model."): new_sd = {} for key, value in sd.items(): new_key = key.replace("model.diffusion_model.", "", 1) new_sd[new_key] = value sd = new_sd dim = sd["patch_embedding.weight"].shape[0] in_channels = sd["patch_embedding.weight"].shape[1] print("in_channels: ", in_channels) ffn_dim = sd["blocks.0.ffn.0.bias"].shape[0] model_type = "i2v" if in_channels == 36 else "t2v" num_heads = 40 if dim == 5120 else 12 num_layers = 40 if dim == 5120 else 30 log.info(f"Model type: {model_type}, num_heads: {num_heads}, num_layers: {num_layers}") teacache_coefficients_map = { "1_3B": [2.39676752e+03, -1.31110545e+03, 2.01331979e+02, -8.29855975e+00, 1.37887774e-01], "14B": [-5784.54975374, 5449.50911966, -1811.16591783, 256.27178429, -13.02252404], "i2v_480": [-3.02331670e+02, 2.23948934e+02, -5.25463970e+01, 5.87348440e+00, -2.01973289e-01], "i2v_720": [-114.36346466, 65.26524496, -18.82220707, 4.91518089, -0.23412683], } if model_type == "i2v": model_variant = "i2v_480" if "480" in model else "i2v_720" elif model_type == "t2v": model_variant = "14B" if dim == 5120 else "1_3B" log.info(f"Model variant detected: {model_variant}") TRANSFORMER_CONFIG= { "dim": dim, "ffn_dim": ffn_dim, "eps": 1e-06, "freq_dim": 256, "in_dim": in_channels, "model_type": model_type, "out_dim": 16, "text_len": 512, "num_heads": num_heads, "num_layers": num_layers, "attention_mode": attention_mode, "main_device": device, "offload_device": offload_device, "teacache_coefficients": teacache_coefficients_map[model_variant], } with init_empty_weights(): transformer = WanModel(**TRANSFORMER_CONFIG) transformer.eval() comfy_model = WanVideoSEModel( WanVideoSEModelConfig(base_dtype), model_type=comfy.model_base.ModelType.FLOW, device=device, ) if not "torchao" in quantization: log.info("Using accelerate to load and assign model weights to device...") if quantization == "fp8_e4m3fn" or quantization == "fp8_e4m3fn_fast" or quantization == "fp8_scaled": dtype = torch.float8_e4m3fn elif quantization == "fp8_e5m2": dtype = torch.float8_e5m2 else: dtype = base_dtype params_to_keep = {"norm", "head", "bias", "time_in", "vector_in", "patch_embedding", "time_", "img_emb", "modulation"} for name, param in transformer.named_parameters(): #print("Assigning Parameter name: ", name) dtype_to_use = base_dtype if any(keyword in name for keyword in params_to_keep) else dtype set_module_tensor_to_device(transformer, name, device=transformer_load_device, dtype=dtype_to_use, value=sd[name]) comfy_model.diffusion_model = transformer comfy_model.load_device = transformer_load_device patcher = comfy.model_patcher.ModelPatcher(comfy_model, device, offload_device) if lora is not None: from comfy.sd import load_lora_for_models for l in lora: log.info(f"Loading LoRA: {l['name']} with strength: {l['strength']}") lora_path = l["path"] lora_strength = l["strength"] lora_sd = load_torch_file(lora_path, safe_load=True) lora_sd = standardize_lora_key_format(lora_sd) if l["blocks"]: lora_sd = filter_state_dict_by_blocks(lora_sd, l["blocks"]) patcher, _ = load_lora_for_models(patcher, None, lora_sd, lora_strength, 0) del lora_sd patcher.patch_model(device) del sd gc.collect() mm.soft_empty_cache() if load_device == "offload_device": patcher.model.diffusion_model.to(offload_device) if quantization == "fp8_e4m3fn_fast": from .fp8_optimization import convert_fp8_linear #params_to_keep.update({"ffn"}) print(params_to_keep) convert_fp8_linear(patcher.model.diffusion_model, base_dtype, params_to_keep=params_to_keep) if vram_management_args is not None: from .diffsynth.vram_management import enable_vram_management, AutoWrappedModule, AutoWrappedLinear from .wanvideo.modules.model import WanLayerNorm, WanRMSNorm total_params_in_model = sum(p.numel() for p in patcher.model.diffusion_model.parameters()) log.info(f"Total number of parameters in the loaded model: {total_params_in_model}") offload_percent = vram_management_args["offload_percent"] offload_params = int(total_params_in_model * offload_percent) params_to_keep = total_params_in_model - offload_params log.info(f"Selected params to offload: {offload_params}") enable_vram_management( patcher.model.diffusion_model, module_map = { torch.nn.Linear: AutoWrappedLinear, torch.nn.Conv3d: AutoWrappedModule, torch.nn.LayerNorm: AutoWrappedModule, WanLayerNorm: AutoWrappedModule, WanRMSNorm: AutoWrappedModule, }, module_config = dict( offload_dtype=dtype, offload_device=offload_device, onload_dtype=dtype, onload_device=device, computation_dtype=base_dtype, computation_device=device, ), max_num_param=params_to_keep, overflow_module_config = dict( offload_dtype=dtype, offload_device=offload_device, onload_dtype=dtype, onload_device=offload_device, computation_dtype=base_dtype, computation_device=device, ), ) #compile if compile_args is not None: torch._dynamo.config.cache_size_limit = compile_args["dynamo_cache_size_limit"] if compile_args["compile_transformer_blocks_only"]: for i, block in enumerate(patcher.model.diffusion_model.blocks): patcher.model.diffusion_model.blocks[i] = torch.compile(block, fullgraph=compile_args["fullgraph"], dynamic=compile_args["dynamic"], backend=compile_args["backend"], mode=compile_args["mode"]) else: patcher.model.diffusion_model = torch.compile(patcher.model.diffusion_model, fullgraph=compile_args["fullgraph"], dynamic=compile_args["dynamic"], backend=compile_args["backend"], mode=compile_args["mode"]) elif "torchao" in quantization: try: from torchao.quantization import ( quantize_, fpx_weight_only, float8_dynamic_activation_float8_weight, int8_dynamic_activation_int8_weight, int8_weight_only, int4_weight_only ) except: raise ImportError("torchao is not installed") # def filter_fn(module: nn.Module, fqn: str) -> bool: # target_submodules = {'attn1', 'ff'} # avoid norm layers, 1.5 at least won't work with quantized norm1 #todo: test other models # if any(sub in fqn for sub in target_submodules): # return isinstance(module, nn.Linear) # return False if "fp6" in quantization: quant_func = fpx_weight_only(3, 2) elif "int4" in quantization: quant_func = int4_weight_only() elif "int8" in quantization: quant_func = int8_weight_only() elif "fp8dq" in quantization: quant_func = float8_dynamic_activation_float8_weight() elif 'fp8dqrow' in quantization: from torchao.quantization.quant_api import PerRow quant_func = float8_dynamic_activation_float8_weight(granularity=PerRow()) elif 'int8dq' in quantization: quant_func = int8_dynamic_activation_int8_weight() log.info(f"Quantizing model with {quant_func}") comfy_model.diffusion_model = transformer patcher = comfy.model_patcher.ModelPatcher(comfy_model, device, offload_device) for i, block in enumerate(patcher.model.diffusion_model.blocks): log.info(f"Quantizing block {i}") for name, _ in block.named_parameters(prefix=f"blocks.{i}"): #print(f"Parameter name: {name}") set_module_tensor_to_device(patcher.model.diffusion_model, name, device=transformer_load_device, dtype=base_dtype, value=sd[name]) if compile_args is not None: patcher.model.diffusion_model.blocks[i] = torch.compile(block, fullgraph=compile_args["fullgraph"], dynamic=compile_args["dynamic"], backend=compile_args["backend"], mode=compile_args["mode"]) quantize_(block, quant_func) print(block) #block.to(offload_device) for name, param in patcher.model.diffusion_model.named_parameters(): if "blocks" not in name: set_module_tensor_to_device(patcher.model.diffusion_model, name, device=transformer_load_device, dtype=base_dtype, value=sd[name]) manual_offloading = False # to disable manual .to(device) calls log.info(f"Quantized transformer blocks to {quantization}") for name, param in patcher.model.diffusion_model.named_parameters(): print(name, param.dtype) #param.data = param.data.to(self.vae_dtype).to(device) del sd mm.soft_empty_cache() patcher.model["dtype"] = base_dtype patcher.model["base_path"] = model_path patcher.model["model_name"] = model patcher.model["manual_offloading"] = manual_offloading patcher.model["quantization"] = "disabled" patcher.model["block_swap_args"] = block_swap_args patcher.model["auto_cpu_offload"] = True if vram_management_args is not None else False for model in mm.current_loaded_models: if model._model() == patcher: mm.current_loaded_models.remove(model) return (patcher,) #region load VAE class WanVideoSEVAELoader: @classmethod def INPUT_TYPES(s): return { "required": { "model_name": (folder_paths.get_filename_list("vae"), {"tooltip": "These models are loaded from 'ComfyUI/models/vae'"}), }, "optional": { "precision": (["fp16", "fp32", "bf16"], {"default": "bf16"} ), } } RETURN_TYPES = ("WANVAE",) RETURN_NAMES = ("vae", ) FUNCTION = "loadmodel" CATEGORY = "WanVideoStartEndFrame" DESCRIPTION = "Loads Hunyuan VAE model from 'ComfyUI/models/vae'" def loadmodel(self, model_name, precision): from .wanvideo.wan_video_vae_SE import WanVideoVAE device = mm.get_torch_device() offload_device = mm.unet_offload_device() dtype = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}[precision] #with open(os.path.join(script_directory, 'configs', 'hy_vae_config.json')) as f: # vae_config = json.load(f) model_path = folder_paths.get_full_path("vae", model_name) vae_sd = load_torch_file(model_path, safe_load=True) has_model_prefix = any(k.startswith("model.") for k in vae_sd.keys()) if not has_model_prefix: vae_sd = {f"model.{k}": v for k, v in vae_sd.items()} vae = WanVideoVAE(dtype=dtype) vae.load_state_dict(vae_sd) vae.eval() vae.to(device = offload_device, dtype = dtype) return (vae,) class WanVideoSEImageClipEncode: @classmethod def INPUT_TYPES(s): return {"required": { "clip_vision": ("CLIP_VISION",), "start": ("IMAGE", {"tooltip": "Image to encode"}), "end": ("IMAGE", {"tooltip": "Image to encode"}), "vae": ("WANVAE",), "generation_width": ( "INT", {"default": 832, "min": 64, "max": 2048, "step": 8, "tooltip": "Width of the image to encode"}), "generation_height": ( "INT", {"default": 480, "min": 64, "max": 29048, "step": 8, "tooltip": "Height of the image to encode"}), "num_frames": ( "INT", {"default": 81, "min": 1, "max": 10000, "step": 4, "tooltip": "Number of frames to encode"}), }, "optional": { "force_offload": ("BOOLEAN", {"default": True}), "noise_aug_strength": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 10.0, "step": 0.001, "tooltip": "Strength of noise augmentation, helpful for 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 I2V where lower values allow for more motion"}), "clip_embed_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001, "tooltip": "Additional clip embed multiplier"}), "adjust_resolution": ("BOOLEAN", {"default": True, "tooltip": "Performs the same resolution adjustment as in the original code"}), "start_frame_weight": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01, "tooltip": "Weight for the start frame. Higher values (>1.0) make the starting image more influential throughout the video, resulting in stronger adherence to the start frame's features. Lower values (<1.0) reduce its influence, allowing more creative freedom. Set to 0 to completely ignore the start frame."}), "end_frame_weight": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01, "tooltip": "Weight for the end frame. Higher values (>1.0) make the ending image more influential, causing the video to transition more strongly toward the end frame's characteristics. Lower values (<1.0) reduce its influence. Balancing with start_frame_weight allows control over transition speed and style."}), } } RETURN_TYPES = ("WANVIDIMAGE_EMBEDS",) RETURN_NAMES = ("image_embeds",) FUNCTION = "process" CATEGORY = "WanVideoStartEndFrame" def process(self, clip_vision, vae, start, end, num_frames, generation_width, generation_height, force_offload=True, noise_aug_strength=0.0, latent_strength=1.0, clip_embed_strength=1.0, adjust_resolution=True, start_frame_weight=1.0, end_frame_weight=1.0): device = mm.get_torch_device() offload_device = mm.unet_offload_device() self.image_mean = [0.48145466, 0.4578275, 0.40821073] self.image_std = [0.26862954, 0.26130258, 0.27577711] patch_size = (1, 2, 2) vae_stride = (4, 8, 8) H, W = start.shape[1], end.shape[2] max_area = generation_width * generation_height print(clip_vision) clip_vision.model.to(device) if isinstance(clip_vision, ClipVisionModel): clip_context = clip_vision.encode_image(start).last_hidden_state.to(device) else: pixel_values = clip_preprocess(start.to(device), size=224, mean=self.image_mean, std=self.image_std, crop=True).float() clip_context = clip_vision.visual(pixel_values) if clip_embed_strength != 1.0: clip_context *= clip_embed_strength if force_offload: clip_vision.model.to(offload_device) mm.soft_empty_cache() if adjust_resolution: aspect_ratio = H / W lat_h = round( np.sqrt(max_area * aspect_ratio) // vae_stride[1] // patch_size[1] * patch_size[1]) lat_w = round( np.sqrt(max_area / aspect_ratio) // vae_stride[2] // patch_size[2] * patch_size[2]) h = lat_h * vae_stride[1] w = lat_w * vae_stride[2] else: h = generation_height w = generation_width lat_h = h // 8 lat_w = w // 8 # Step 1: Create initial mask with ones for first frame, zeros for others mask = torch.ones(1, num_frames+1, lat_h, lat_w, device=device) mask[:, 1:-1] = 0 # No longer applying weights to the mask - this should improve transition smoothness # mask[:, 0] *= start_frame_weight # mask[:, -1] *= end_frame_weight # Step 2: Repeat first frame 4 times and concatenate with remaining frames first_frame_repeated = torch.repeat_interleave(mask[:, 0:1], repeats=4, dim=1) last_frame_repeated = torch.repeat_interleave(mask[:, -1:], repeats=4, dim=1) mask = torch.concat([first_frame_repeated, mask[:, 1:-1], last_frame_repeated], dim=1) # Step 3: Reshape mask into groups of 4 frames mask = mask.view(1, mask.shape[1] // 4, 4, lat_h, lat_w) # Step 4: Transpose dimensions and select first batch mask = mask.transpose(1, 2)[0] # Calculate maximum sequence length frames_per_stride = (num_frames - 1) // vae_stride[0] + 2 patches_per_frame = lat_h * lat_w // (patch_size[1] * patch_size[2]) max_seq_len = frames_per_stride * patches_per_frame vae.to(device) # Step 1: Resize and rearrange the input image dimensions # resized_image = image.permute(0, 3, 1, 2) # Rearrange dimensions to (B, C, H, W) # resized_image = torch.nn.functional.interpolate(resized_image, size=(h, w), mode='bicubic') resized_start_image = common_upscale(start.movedim(-1, 1), w, h, "lanczos", "disabled") resized_start_image = resized_start_image.transpose(0, 1) # Transpose to match required format resized_start_image = resized_start_image * 2 - 1 resized_end_image = common_upscale(end.movedim(-1, 1), w, h, "lanczos", "disabled") resized_end_image = resized_end_image.transpose(0, 1) # Transpose to match required format resized_end_image = resized_end_image * 2 - 1 if noise_aug_strength > 0.0: resized_start_image = add_noise_to_reference_video(resized_start_image, ratio=noise_aug_strength) resized_end_image = add_noise_to_reference_video(resized_end_image, ratio=noise_aug_strength) # Apply start and end frame weights to image data resized_start_image = resized_start_image * start_frame_weight resized_end_image = resized_end_image * end_frame_weight # Step 2: Create zero padding frames zero_frames = torch.zeros(3, num_frames - 1, h, w, device=device) # Step 3: Concatenate image with zero frames concatenated = torch.concat([resized_start_image.to(device), zero_frames, resized_end_image.to(device)], dim=1).to( device=device, dtype=vae.dtype) concatenated *= latent_strength y = vae.encode([concatenated], device,end_=True)[0] # y:[4+c,t,h,w] y = torch.concat([mask, y]) vae.model.clear_cache() vae.to(offload_device) # clip_context contains the encoding of a single image image_embeds = { "image_embeds": y, "clip_context": clip_context, "max_seq_len": max_seq_len, "num_frames": num_frames, "lat_h": lat_h, "lat_w": lat_w, } return (image_embeds,) #region Sampler class WanVideoSESampler: @classmethod def INPUT_TYPES(s): return { "required": { "model": ("WANVIDEOMODEL",), "text_embeds": ("WANVIDEOTEXTEMBEDS",), "image_embeds": ("WANVIDIMAGE_EMBEDS",), "steps": ("INT", {"default": 30, "min": 1}), "cfg": ("FLOAT", {"default": 6.0, "min": 0.0, "max": 30.0, "step": 0.01}), "shift": ("FLOAT", {"default": 5.0, "min": 0.0, "max": 1000.0, "step": 0.01}), "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}), "force_offload": ( "BOOLEAN", {"default": True, "tooltip": "Moves the model to the offload device after sampling"}), "scheduler": (["unipc", "dpm++", "dpm++_sde", "euler"], { "default": 'unipc' }), "riflex_freq_index": ("INT", {"default": 0, "min": 0, "max": 1000, "step": 1, "tooltip": "Frequency index for RIFLEX, disabled when 0, default 6. Allows for new frames to be generated after without looping"}), }, "optional": { "samples": ("LATENT", {"tooltip": "init Latents to use for video2video process"}), "denoise_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), "feta_args": ("FETAARGS",), "context_options": ("WANVIDCONTEXT",), "teacache_args": ("TEACACHEARGS",), "flowedit_args": ("FLOWEDITARGS",), "slg_args": ("SLGARGS",), } } RETURN_TYPES = ("LATENT",) RETURN_NAMES = ("samples",) FUNCTION = "process" CATEGORY = "WanVideoStartEndFrame" def process(self, model, text_embeds, image_embeds, shift, steps, cfg, seed, scheduler, riflex_freq_index, force_offload=True, samples=None, feta_args=None, denoise_strength=1.0, context_options=None, teacache_args=None, flowedit_args=None,slg_args=None,): # assert not (context_options and teacache_args), "Context options cannot currently be used together with teacache." patcher = model model = model.model transformer = model.diffusion_model device = mm.get_torch_device() offload_device = mm.unet_offload_device() steps = int(steps / denoise_strength) if scheduler == 'unipc': sample_scheduler = FlowUniPCMultistepScheduler( num_train_timesteps=1000, shift=shift, use_dynamic_shifting=False) sample_scheduler.set_timesteps( steps, device=device, shift=shift) timesteps = sample_scheduler.timesteps elif scheduler == 'euler': sample_scheduler = FlowMatchEulerDiscreteScheduler( num_train_timesteps=1000, shift=shift, use_dynamic_shifting=False) sampling_sigmas = get_sampling_sigmas(steps, shift) timesteps, _ = retrieve_timesteps( sample_scheduler, device=device, sigmas=sampling_sigmas) elif 'dpm++' in scheduler: if scheduler == 'dpm++_sde': algorithm_type = "sde-dpmsolver++" else: algorithm_type = "dpmsolver++" sample_scheduler = FlowDPMSolverMultistepScheduler( num_train_timesteps=1000, shift=shift, use_dynamic_shifting=False, algorithm_type=algorithm_type) sampling_sigmas = get_sampling_sigmas(steps, shift) timesteps, _ = retrieve_timesteps( sample_scheduler, device=device, sigmas=sampling_sigmas) else: raise NotImplementedError("Unsupported solver.") if denoise_strength < 1.0: steps = int(steps * denoise_strength) timesteps = timesteps[-(steps + 1):] seed_g = torch.Generator(device=torch.device("cpu")) seed_g.manual_seed(seed) image_cond = None clip_fea = None if transformer.model_type == "i2v": lat_h = image_embeds.get("lat_h", None) lat_w = image_embeds.get("lat_w", None) if lat_h is None or lat_w is None: raise ValueError("Clip encoded image embeds must be provided for I2V (Image to Video) model") noise = torch.randn( 16, (image_embeds["num_frames"] - 1) // 4 + 2, lat_h, lat_w, dtype=torch.float32, generator=seed_g, device=torch.device("cpu")) seq_len = image_embeds["max_seq_len"] image_cond = image_embeds.get("image_embeds", None) clip_fea = image_embeds.get("clip_context", None) else: # t2v target_shape = image_embeds.get("target_shape", None) if target_shape is None: raise ValueError("Empty image embeds must be provided for T2V (Text to Video") seq_len = image_embeds["max_seq_len"] noise = torch.randn( target_shape[0], target_shape[1], target_shape[2], target_shape[3], dtype=torch.float32, device=torch.device("cpu"), generator=seed_g) latent_video_length = noise.shape[1] if context_options is not None: def create_window_mask(noise_pred_context, c, latent_video_length, context_overlap, looped=False): window_mask = torch.ones_like(noise_pred_context) # Apply left-side blending for all except first chunk (or always in loop mode) if min(c) > 0 or (looped and max(c) == latent_video_length - 1): ramp_up = torch.linspace(0, 1, context_overlap, device=noise_pred_context.device) ramp_up = ramp_up.view(1, -1, 1, 1) window_mask[:, :context_overlap] = ramp_up # Apply right-side blending for all except last chunk (or always in loop mode) if max(c) < latent_video_length - 1 or (looped and min(c) == 0): ramp_down = torch.linspace(1, 0, context_overlap, device=noise_pred_context.device) ramp_down = ramp_down.view(1, -1, 1, 1) window_mask[:, -context_overlap:] = ramp_down return window_mask context_schedule = context_options["context_schedule"] context_frames = (context_options["context_frames"] - 1) // 4 + 2 context_stride = context_options["context_stride"] // 4 context_overlap = context_options["context_overlap"] // 4 self.window_tracker = WindowTrackerSE(verbose=context_options["verbose"]) # Get total number of prompts num_prompts = len(text_embeds["prompt_embeds"]) log.info(f"Number of prompts: {num_prompts}") # Calculate which section this context window belongs to section_size = latent_video_length / num_prompts log.info(f"Section size: {section_size}") is_looped = context_schedule == "uniform_looped" seq_len = math.ceil((noise.shape[2] * noise.shape[3]) / 4 * context_frames) if context_options["freenoise"]: log.info("Applying FreeNoise") # code from AnimateDiff-Evolved by Kosinkadink (https://github.com/Kosinkadink/ComfyUI-AnimateDiff-Evolved) delta = context_frames - context_overlap for start_idx in range(0, latent_video_length - context_frames, delta): place_idx = start_idx + context_frames if place_idx >= latent_video_length: break end_idx = place_idx - 1 if end_idx + delta >= latent_video_length: final_delta = latent_video_length - place_idx list_idx = torch.tensor(list(range(start_idx, start_idx + final_delta)), device=torch.device("cpu"), dtype=torch.long) list_idx = list_idx[torch.randperm(final_delta, generator=seed_g)] noise[:, place_idx:place_idx + final_delta, :, :] = noise[:, list_idx, :, :] break list_idx = torch.tensor(list(range(start_idx, start_idx + delta)), device=torch.device("cpu"), dtype=torch.long) list_idx = list_idx[torch.randperm(delta, generator=seed_g)] noise[:, place_idx:place_idx + delta, :, :] = noise[:, list_idx, :, :] log.info( f"Context schedule enabled: {context_frames} frames, {context_stride} stride, {context_overlap} overlap") from .context import get_context_scheduler context = get_context_scheduler(context_schedule) if samples is not None and denoise_strength < 1.0: latent_timestep = timesteps[:1].to(noise) noise = noise * latent_timestep / 1000 + (1 - latent_timestep / 1000) * samples["samples"].squeeze(0).to( noise) if samples is not None: original_image = samples["samples"].clone().squeeze(0).to(device) mask = samples.get("mask", None) # latent初始化为一段噪声 latent = noise.to(device) d = transformer.dim // transformer.num_heads freqs = torch.cat([ rope_params(1024, d - 4 * (d // 6), L_test=latent_video_length, k=riflex_freq_index), rope_params(1024, 2 * (d // 6)), rope_params(1024, 2 * (d // 6)) ], dim=1) if not isinstance(cfg, list): cfg = [cfg] * (steps + 1) print("Seq len:", seq_len) pbar = ProgressBar(steps) from latent_preview import prepare_callback callback = prepare_callback(patcher, steps) # blockswap init if model["block_swap_args"] is not None: for name, param in transformer.named_parameters(): if "block" not in name: param.data = param.data.to(device) elif model["block_swap_args"]["offload_txt_emb"] and "txt_emb" in name: param.data = param.data.to(offload_device) elif model["block_swap_args"]["offload_img_emb"] and "img_emb" in name: param.data = param.data.to(offload_device) transformer.block_swap( model["block_swap_args"]["blocks_to_swap"] - 1, model["block_swap_args"]["offload_txt_emb"], model["block_swap_args"]["offload_img_emb"], ) elif model["auto_cpu_offload"]: for module in transformer.modules(): if hasattr(module, "offload"): module.offload() if hasattr(module, "onload"): module.onload() elif model["manual_offloading"]: transformer.to(device) # feta if feta_args is not None: set_enhance_weight(feta_args["weight"]) feta_start_percent = feta_args["start_percent"] feta_end_percent = feta_args["end_percent"] if context_options is not None: set_num_frames(context_frames) else: set_num_frames(latent_video_length) enable_enhance() else: disable_enhance() # Initialize TeaCache if enabled if teacache_args is not None: transformer.enable_teacache = True transformer.rel_l1_thresh = teacache_args["rel_l1_thresh"] transformer.teacache_start_step = teacache_args["start_step"] transformer.teacache_cache_device = teacache_args["cache_device"] transformer.teacache_end_step = len(timesteps) - 1 if teacache_args["end_step"] == -1 else teacache_args[ "end_step"] transformer.teacache_use_coefficients = teacache_args["use_coefficients"] else: transformer.enable_teacache = False if slg_args is not None: #assert batched_cfg is not None, "Batched cfg is not supported with SLG" transformer.slg_blocks = slg_args["blocks"] transformer.slg_start_percent = slg_args["start_percent"] transformer.slg_end_percent = slg_args["end_percent"] else: transformer.slg_blocks = None mm.unload_all_models() mm.soft_empty_cache() gc.collect() self.teacache_state = [None, None] self.teacache_state_source = [None, None] self.teacache_states_context = [] if "sparge" in transformer.attention_mode: from spas_sage_attn.autotune import ( SparseAttentionMeansim, extract_sparse_attention_state_dict, load_sparse_attention_state_dict, ) for idx, block in enumerate(transformer.blocks): block.self_attn.verbose = True block.self_attn.inner_attention = SparseAttentionMeansim(l1=0.06, pv_l1=0.065) if transformer.attention_mode == "spargeattn": try: saved_state_dict = torch.load("sparge_wan.pt") except: raise ValueError("No saved parameters found for sparse attention, tuning is required first") load_sparse_attention_state_dict(transformer, saved_state_dict, verbose=True) if flowedit_args is not None: source_embeds = flowedit_args["source_embeds"] source_image_embeds = flowedit_args.get("source_image_embeds", image_embeds) source_image_cond = source_image_embeds.get("image_embeds", None) source_clip_fea = source_image_embeds["clip_context"] skip_steps = flowedit_args["skip_steps"] drift_steps = flowedit_args["drift_steps"] source_cfg = flowedit_args["source_cfg"] if not isinstance(source_cfg, list): source_cfg = [source_cfg] * (steps + 1) drift_cfg = flowedit_args["drift_cfg"] if not isinstance(drift_cfg, list): drift_cfg = [drift_cfg] * (steps + 1) x_init = samples["samples"].clone().squeeze(0).to(device) x_tgt = samples["samples"].squeeze(0).to(device) sample_scheduler = FlowMatchEulerDiscreteScheduler( num_train_timesteps=1000, shift=flowedit_args["drift_flow_shift"], use_dynamic_shifting=False) sampling_sigmas = get_sampling_sigmas(steps, flowedit_args["drift_flow_shift"]) drift_timesteps, _ = retrieve_timesteps( sample_scheduler, device=device, sigmas=sampling_sigmas) if drift_steps > 0: drift_timesteps = torch.cat([drift_timesteps, torch.tensor([0]).to(drift_timesteps.device)]).to( drift_timesteps.device) timesteps[-drift_steps:] = drift_timesteps[-drift_steps:] def predict_with_cfg(z, cfg_scale, positive_embeds, negative_embeds, timestep, idx, image_cond=None, clip_fea=None, teacache_state=None): with torch.autocast(device_type=mm.get_autocast_device(device), dtype=model["dtype"], enabled=True): base_params = { 'clip_fea': clip_fea, 'seq_len': seq_len, 'device': device, 'freqs': freqs, 't': timestep, 'current_step': idx, 'y': image_cond, } # Get conditional prediction noise_pred_cond, teacache_state_cond = transformer( z, context=[positive_embeds], pred_id=teacache_state[0] if teacache_state else None, is_uncond=False, current_step_percentage=current_step_percentage, **base_params ) noise_pred_cond = noise_pred_cond.to(intermediate_device) # If cfg_scale is 1.0, return conditional prediction directly if math.isclose(cfg_scale, 1.0): return noise_pred_cond, [teacache_state_cond] # Get unconditional prediction and apply cfg noise_pred_uncond, teacache_state_uncond = transformer( z, context=negative_embeds, pred_id=teacache_state[1] if teacache_state else None, is_uncond=True, current_step_percentage=current_step_percentage, **base_params ) noise_pred_uncond = noise_pred_uncond.to(intermediate_device) return noise_pred_uncond + cfg_scale * (noise_pred_cond - noise_pred_uncond), [teacache_state_cond, teacache_state_uncond] try: torch.cuda.reset_peak_memory_stats(device) except: pass log.info( f"Sampling {(latent_video_length - 1) * 4 + 2} frames at {latent.shape[3] * 8}x{latent.shape[2] * 8} with {steps} steps") intermediate_device = device # diff diff prep masks = None if samples is not None and mask is not None: mask = 1 - mask thresholds = torch.arange(len(timesteps), dtype=original_image.dtype) / len(timesteps) thresholds = thresholds.unsqueeze(1).unsqueeze(1).unsqueeze(1).unsqueeze(1).to(device) masks = mask.repeat(len(timesteps), 1, 1, 1, 1).to(device) masks = masks > thresholds for idx, t in enumerate(tqdm(timesteps)): if flowedit_args is not None: if idx < skip_steps: continue # diff diff if masks is not None: if idx < len(timesteps) - 1: noise_timestep = timesteps[idx + 1] image_latent = sample_scheduler.scale_noise( original_image, torch.tensor([noise_timestep]), noise.to(device) ) mask = masks[idx] mask = mask.to(latent) latent = image_latent * mask + latent * (1 - mask) # end diff diff latent_model_input = latent.to(device) timestep = torch.tensor([t]).to(device) current_step_percentage = idx / len(timesteps) # enhance-a-video if feta_args is not None: if feta_start_percent <= current_step_percentage <= feta_end_percent: enable_enhance() else: disable_enhance() # flow-edit if flowedit_args is not None: sigma = t / 1000.0 sigma_prev = (timesteps[idx + 1] if idx < len(timesteps) - 1 else timesteps[-1]) / 1000.0 noise = torch.randn(x_init.shape, generator=seed_g, device=torch.device("cpu")) if idx < len(timesteps) - drift_steps: cfg = drift_cfg zt_src = (1 - sigma) * x_init + sigma * noise.to(t) zt_tgt = x_tgt + zt_src - x_init # source if idx < len(timesteps) - drift_steps: if context_options is not None: counter = torch.zeros_like(zt_src, device=intermediate_device) vt_src = torch.zeros_like(zt_src, device=intermediate_device) context_queue = list( context(idx, steps, latent_video_length, context_frames, context_stride, context_overlap)) for c in context_queue: window_id = self.window_tracker.get_window_id(c) if teacache_args is not None: current_teacache = self.window_tracker.get_teacache(window_id, self.teacache_state) else: current_teacache = None prompt_index = min(int(max(c) / section_size), num_prompts - 1) if context_options["verbose"]: log.info(f"Prompt index: {prompt_index}") positive = source_embeds["prompt_embeds"][prompt_index] partial_img_emb = None if source_image_cond is not None: partial_img_emb = source_image_cond[:, c, :, :] partial_img_emb[:, 0, :, :] = source_image_cond[:, 0, :, :].to(intermediate_device) partial_zt_src = zt_src[:, c, :, :] vt_src_context, new_teacache = predict_with_cfg( partial_zt_src, cfg[idx], positive, source_embeds["negative_prompt_embeds"], timestep, idx, partial_img_emb, source_clip_fea, current_teacache) if teacache_args is not None: self.window_tracker.teacache_states[window_id] = new_teacache window_mask = create_window_mask(vt_src_context, c, latent_video_length, context_overlap) vt_src[:, c, :, :] += vt_src_context * window_mask counter[:, c, :, :] += window_mask vt_src /= counter else: vt_src, self.teacache_state_source = predict_with_cfg( zt_src, cfg[idx], source_embeds["prompt_embeds"][0], source_embeds["negative_prompt_embeds"], timestep, idx, source_image_cond, source_clip_fea, teacache_state=self.teacache_state_source) else: if idx == len(timesteps) - drift_steps: x_tgt = zt_tgt zt_tgt = x_tgt vt_src = 0 # target if context_options is not None: counter = torch.zeros_like(zt_tgt, device=intermediate_device) vt_tgt = torch.zeros_like(zt_tgt, device=intermediate_device) context_queue = list( context(idx, steps, latent_video_length, context_frames, context_stride, context_overlap)) for c in context_queue: window_id = self.window_tracker.get_window_id(c) if teacache_args is not None: current_teacache = self.window_tracker.get_teacache(window_id, self.teacache_state) else: current_teacache = None prompt_index = min(int(max(c) / section_size), num_prompts - 1) if context_options["verbose"]: log.info(f"Prompt index: {prompt_index}") positive = text_embeds["prompt_embeds"][prompt_index] partial_img_emb = None if image_cond is not None: partial_img_emb = image_cond[:, c, :, :] partial_img_emb[:, 0, :, :] = image_cond[:, 0, :, :].to(intermediate_device) partial_zt_tgt = zt_tgt[:, c, :, :] vt_tgt_context, new_teacache = predict_with_cfg( partial_zt_tgt, cfg[idx], positive, text_embeds["negative_prompt_embeds"], timestep, idx, partial_img_emb, clip_fea, current_teacache) if teacache_args is not None: self.window_tracker.teacache_states[window_id] = new_teacache window_mask = create_window_mask(vt_tgt_context, c, latent_video_length, context_overlap) vt_tgt[:, c, :, :] += vt_tgt_context * window_mask counter[:, c, :, :] += window_mask vt_tgt /= counter else: vt_tgt, self.teacache_state = predict_with_cfg( zt_tgt, cfg[idx], text_embeds["prompt_embeds"][0], text_embeds["negative_prompt_embeds"], timestep, idx, image_cond, clip_fea, teacache_state=self.teacache_state) v_delta = vt_tgt - vt_src x_tgt = x_tgt.to(torch.float32) v_delta = v_delta.to(torch.float32) x_tgt = x_tgt + (sigma_prev - sigma) * v_delta x0 = x_tgt # context windowing elif context_options is not None: counter = torch.zeros_like(latent_model_input, device=intermediate_device) noise_pred = torch.zeros_like(latent_model_input, device=intermediate_device) context_queue = list( context(idx, steps, latent_video_length, context_frames, context_stride, context_overlap)) for c in context_queue: window_id = self.window_tracker.get_window_id(c) if teacache_args is not None: current_teacache = self.window_tracker.get_teacache(window_id, self.teacache_state) else: current_teacache = None prompt_index = min(int(max(c) / section_size), num_prompts - 1) if context_options["verbose"]: log.info(f"Prompt index: {prompt_index}") # Use the appropriate prompt for this section positive = text_embeds["prompt_embeds"][prompt_index] partial_img_emb = None if image_cond is not None: partial_img_emb = image_cond[:, c, :, :] partial_img_emb[:, 0, :, :] = image_cond[:, 0, :, :].to(intermediate_device) partial_latent_model_input = latent_model_input[:, c, :, :] noise_pred_context, new_teacache = predict_with_cfg( partial_latent_model_input, cfg[idx], positive, text_embeds["negative_prompt_embeds"], timestep, idx, partial_img_emb, clip_fea, current_teacache) if teacache_args is not None: self.window_tracker.teacache_states[window_id] = new_teacache window_mask = create_window_mask(noise_pred_context, c, latent_video_length, context_overlap, looped=is_looped) noise_pred[:, c, :, :] += noise_pred_context * window_mask counter[:, c, :, :] += window_mask noise_pred /= counter # normal inference else: noise_pred, self.teacache_state = predict_with_cfg( latent_model_input, cfg[idx], text_embeds["prompt_embeds"][0], text_embeds["negative_prompt_embeds"], timestep, idx, image_cond, clip_fea, teacache_state=self.teacache_state) if flowedit_args is None: latent = latent.to(intermediate_device) temp_x0 = sample_scheduler.step( noise_pred.unsqueeze(0), t, latent.unsqueeze(0), return_dict=False, generator=seed_g)[0] latent = temp_x0.squeeze(0) x0 = latent.to(device) if callback is not None: callback_latent = (latent_model_input - noise_pred.to(t.device) * t / 1000).detach().permute(1, 0, 2, 3) callback(idx, callback_latent, None, steps) else: pbar.update(1) del latent_model_input, timestep else: if callback is not None: callback_latent = (zt_tgt - vt_tgt.to(t.device) * t / 1000).detach().permute(1, 0, 2, 3) callback(idx, callback_latent, None, steps) else: pbar.update(1) if teacache_args is not None: states = transformer.teacache_state.states state_names = { 0: "conditional", 1: "unconditional" } for pred_id, state in states.items(): name = state_names.get(pred_id, f"prediction_{pred_id}") if 'skipped_steps' in state: log.info(f"TeaCache skipped: {state['skipped_steps']} {name} steps") transformer.teacache_state.clear_all() if transformer.attention_mode == "spargeattn_tune": saved_state_dict = extract_sparse_attention_state_dict(transformer) torch.save(saved_state_dict, "sparge_wan.pt") save_torch_file(saved_state_dict, "sparge_wan.safetensors") if force_offload: if model["manual_offloading"]: transformer.to(offload_device) mm.soft_empty_cache() gc.collect() try: print_memory(device) torch.cuda.reset_peak_memory_stats(device) except: pass return ({ "samples": x0.unsqueeze(0).cpu() },) class WanVideoSEDecode: @classmethod def INPUT_TYPES(s): return {"required": { "vae": ("WANVAE",), "samples": ("LATENT",), "start": ("IMAGE",), "end": ("IMAGE",), "enable_vae_tiling": ("BOOLEAN", {"default": True, "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"}), }, } RETURN_TYPES = ("IMAGE",) RETURN_NAMES = ("images",) FUNCTION = "decode" CATEGORY = "WanVideoStartEndFrame" def decode(self, vae, samples, start, end, enable_vae_tiling, tile_x, tile_y, tile_stride_x, tile_stride_y): device = mm.get_torch_device() offload_device = mm.unet_offload_device() mm.soft_empty_cache() latents = samples["samples"] vae.to(device) latents = latents.to(device = device, dtype = vae.dtype) mm.soft_empty_cache() image = vae.decode(latents, device=device, end_=True, tiled=enable_vae_tiling, tile_size=(tile_x, tile_y), tile_stride=(tile_stride_x, tile_stride_y))[0] vae.to(offload_device) vae.model.clear_cache() mm.soft_empty_cache() image = (image - image.min()) / (image.max() - image.min()) image = torch.clamp(image, 0.0, 1.0) h,w=image.shape[2],image.shape[3] end = common_upscale(end.movedim(-1, 1), w, h, "lanczos", "disabled").squeeze(0).unsqueeze(1) start = common_upscale(start.movedim(-1, 1), w, h, "lanczos", "disabled").squeeze(0).unsqueeze(1) image=torch.cat([start,image[:,1:-4,:,:],end],dim=1) image = image.permute(1, 2, 3, 0).cpu().float() #print(image.shape,end.shape) return (image,) NODE_CLASS_MAPPINGS = { "WanVideoSEModelLoader": WanVideoSEModelLoader, "WanVideoSEVAELoader": WanVideoSEVAELoader, "WanVideoSEImageClipEncode": WanVideoSEImageClipEncode, "WanVideoSEDecode": WanVideoSEDecode, "WanVideoSESampler":WanVideoSESampler, } NODE_DISPLAY_NAME_MAPPINGS = { "WanVideoSEModelLoader": "WanVideo Model Loader(SE)", "WanVideoSEVAELoader": "WanVideo VAE Loader(SE)", "WanVideoSEImageClipEncode":"WanVideo SEI Clip Encode", "WanVideoSEDecode": "Wan Video SEI Decode", "WanVideoSESampler": "Wan Video Sampler(SE)", }