From c732a9fb20460753b44cdb62b56ece920b4a0487 Mon Sep 17 00:00:00 2001 From: GACLove Date: Thu, 24 Apr 2025 20:06:05 +0800 Subject: [PATCH] feat: init --- .gitignore | 308 +++++++++++++++++++++++++ .gitmodules | 3 + .ruff.toml | 0 README.md | 0 __init__.py | 3 + lightx2v | 1 + nodes.py | 641 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 7 files changed, 956 insertions(+) create mode 100644 .gitignore create mode 100644 .gitmodules create mode 100644 .ruff.toml create mode 100644 README.md create mode 100644 __init__.py create mode 160000 lightx2v create mode 100644 nodes.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..31adb0f --- /dev/null +++ b/.gitignore @@ -0,0 +1,308 @@ +# Created by https://www.toptal.com/developers/gitignore/api/python,visualstudiocode,pycharm +# Edit at https://www.toptal.com/developers/gitignore?templates=python,visualstudiocode,pycharm + +### PyCharm ### +# Covers JetBrains IDEs: IntelliJ, RubyMine, PhpStorm, AppCode, PyCharm, CLion, Android Studio, WebStorm and Rider +# Reference: https://intellij-support.jetbrains.com/hc/en-us/articles/206544839 + +# User-specific stuff +.idea/**/workspace.xml +.idea/**/tasks.xml +.idea/**/usage.statistics.xml +.idea/**/dictionaries +.idea/**/shelf + +# AWS User-specific +.idea/**/aws.xml + +# Generated files +.idea/**/contentModel.xml + +# Sensitive or high-churn files +.idea/**/dataSources/ +.idea/**/dataSources.ids +.idea/**/dataSources.local.xml +.idea/**/sqlDataSources.xml +.idea/**/dynamic.xml +.idea/**/uiDesigner.xml +.idea/**/dbnavigator.xml + +# Gradle +.idea/**/gradle.xml +.idea/**/libraries + +# Gradle and Maven with auto-import +# When using Gradle or Maven with auto-import, you should exclude module files, +# since they will be recreated, and may cause churn. Uncomment if using +# auto-import. +# .idea/artifacts +# .idea/compiler.xml +# .idea/jarRepositories.xml +# .idea/modules.xml +# .idea/*.iml +# .idea/modules +# *.iml +# *.ipr + +# CMake +cmake-build-*/ + +# Mongo Explorer plugin +.idea/**/mongoSettings.xml + +# File-based project format +*.iws + +# IntelliJ +out/ + +# mpeltonen/sbt-idea plugin +.idea_modules/ + +# JIRA plugin +atlassian-ide-plugin.xml + +# Cursive Clojure plugin +.idea/replstate.xml + +# SonarLint plugin +.idea/sonarlint/ + +# Crashlytics plugin (for Android Studio and IntelliJ) +com_crashlytics_export_strings.xml +crashlytics.properties +crashlytics-build.properties +fabric.properties + +# Editor-based Rest Client +.idea/httpRequests + +# Android studio 3.1+ serialized cache file +.idea/caches/build_file_checksums.ser + +### PyCharm Patch ### +# Comment Reason: https://github.com/joeblau/gitignore.io/issues/186#issuecomment-215987721 + +# *.iml +# modules.xml +# .idea/misc.xml +# *.ipr + +# Sonarlint plugin +# https://plugins.jetbrains.com/plugin/7973-sonarlint +.idea/**/sonarlint/ + +# SonarQube Plugin +# https://plugins.jetbrains.com/plugin/7238-sonarqube-community-plugin +.idea/**/sonarIssues.xml + +# Markdown Navigator plugin +# https://plugins.jetbrains.com/plugin/7896-markdown-navigator-enhanced +.idea/**/markdown-navigator.xml +.idea/**/markdown-navigator-enh.xml +.idea/**/markdown-navigator/ + +# Cache file creation bug +# See https://youtrack.jetbrains.com/issue/JBR-2257 +.idea/$CACHE_FILE$ + +# CodeStream plugin +# https://plugins.jetbrains.com/plugin/12206-codestream +.idea/codestream.xml + +# Azure Toolkit for IntelliJ plugin +# https://plugins.jetbrains.com/plugin/8053-azure-toolkit-for-intellij +.idea/**/azureSettings.xml + +### Python ### +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ +cover/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +.pybuilder/ +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +# For a library or package, you might want to ignore these files since the code is +# intended to run in multiple environments; otherwise, check them in: +# .python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# poetry +# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control +#poetry.lock + +# pdm +# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. +#pdm.lock +# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it +# in version control. +# https://pdm.fming.dev/#use-with-ide +.pdm.toml + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static type analyzer +.pytype/ + +# Cython debug symbols +cython_debug/ + +# PyCharm +# JetBrains specific template is maintained in a separate JetBrains.gitignore that can +# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore +# and can be added to the global gitignore or merged into this file. For a more nuclear +# option (not recommended) you can uncomment the following to ignore the entire idea folder. +#.idea/ + +### Python Patch ### +# Poetry local configuration file - https://python-poetry.org/docs/configuration/#local-configuration +poetry.toml + +# ruff +.ruff_cache/ + +# LSP config files +pyrightconfig.json + +### VisualStudioCode ### +.vscode/* +# !.vscode/settings.json +!.vscode/tasks.json +!.vscode/launch.json +!.vscode/extensions.json +!.vscode/*.code-snippets + +# Local History for Visual Studio Code +.history/ + +# Built Visual Studio Code Extensions +*.vsix + +### VisualStudioCode Patch ### +# Ignore all local history of files +.history +.ionide + +# End of https://www.toptal.com/developers/gitignore/api/python,visualstudiocode,pycharm \ No newline at end of file diff --git a/.gitmodules b/.gitmodules new file mode 100644 index 0000000..3a84acf --- /dev/null +++ b/.gitmodules @@ -0,0 +1,3 @@ +[submodule "lightx2v"] + path = lightx2v + url = https://github.com/ModelTC/lightx2v.git diff --git a/.ruff.toml b/.ruff.toml new file mode 100644 index 0000000..e69de29 diff --git a/README.md b/README.md new file mode 100644 index 0000000..e69de29 diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..39a8c6b --- /dev/null +++ b/__init__.py @@ -0,0 +1,3 @@ +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/lightx2v b/lightx2v new file mode 160000 index 0000000..863b22c --- /dev/null +++ b/lightx2v @@ -0,0 +1 @@ +Subproject commit 863b22c6585c91083ce740d3bc3348c8509eff57 diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..877fb2b --- /dev/null +++ b/nodes.py @@ -0,0 +1,641 @@ +import os +import torch +import gc +from typing import cast, no_type_check, Any +import logging +import json +import numpy as np +import comfy.model_management as mm +from comfy.utils import ProgressBar +import folder_paths +from tqdm import tqdm +from easydict import EasyDict + + +# Import LightX2V modules +from .lightx2v.lightx2v.utils.profiler import ProfilingContext +from .lightx2v.lightx2v.models.input_encoders.hf.t5.model import T5EncoderModel +from .lightx2v.lightx2v.models.input_encoders.hf.xlm_roberta.model import ( + CLIPModel as ClipVisionModel, +) +from .lightx2v.lightx2v.utils.utils import seed_all +from .lightx2v.lightx2v.models.video_encoders.hf.wan.vae import WanVAE +from .lightx2v.lightx2v.models.networks.wan.model import WanModel +from .lightx2v.lightx2v.models.networks.wan.lora_adapter import WanLoraWrapper +from .lightx2v.lightx2v.models.schedulers.wan.scheduler import WanScheduler +from .lightx2v.lightx2v.models.schedulers.wan.feature_caching.scheduler import ( + WanSchedulerTeaCaching, +) + + +class Lightx2vWanVideoT5EncoderLoader: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "t5_model_path": ( + "STRING", + {"default": "models/t5/models_t5_umt5-xxl-enc-bf16.pth"}, + ), + "tokenizer_path": ("STRING", {"default": "models/t5/google/umt5-xxl"}), + "text_len": ( + "INT", + {"default": 512, "min": 64, "max": 2048, "step": 1}, + ), + "precision": (["bf16", "fp16", "fp32"], {"default": "bf16"}), + "device": (["cuda", "cpu"], {"default": "cuda"}), + "additional_param": ("STRING", {"default": "default_value"}), + } + } + + RETURN_TYPES = ("LIGHT_T5_ENCODER",) + RETURN_NAMES = ("t5_encoder",) + FUNCTION = "load_t5_encoder" + CATEGORY = "LightX2V" + + def load_t5_encoder( + self, + t5_model_path, + tokenizer_path, + text_len, + precision, + device, + additional_param, + ): + # Map precision to torch dtype + dtype_map = { + "bf16": torch.bfloat16, + "fp16": torch.float16, + "fp32": torch.float32, + } + dtype = dtype_map[precision] + + # Resolve device + if device == "cuda": + device = mm.get_torch_device() + else: + device = torch.device("cpu") + + # Load the T5 encoder + t5_encoder = T5EncoderModel( + text_len=text_len, + dtype=dtype, + device=device, # type:ignore + checkpoint_path=t5_model_path, + tokenizer_path=tokenizer_path, + shard_fn=None, + ) + + return (t5_encoder,) + + +class Lightx2vWanVideoT5Encoder: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "t5_encoder": ("LIGHT_T5_ENCODER",), + "prompt": ( + "STRING", + { + "multiline": True, + "default": "A beautiful landscape with mountains and a lake", + }, + ), + "negative_prompt": ("STRING", {"multiline": True, "default": ""}), + } + } + + RETURN_TYPES = ("LIGHT_TEXT_EMBEDDINGS",) + RETURN_NAMES = ("text_embeddings",) + FUNCTION = "encode_text" + CATEGORY = "LightX2V" + + def encode_text(self, t5_encoder, prompt, negative_prompt): + # Create a config object with required attributes + class Config: + def __init__(self, cpu_offload=False): + self.cpu_offload = cpu_offload + + # NOTE(xxx): adapt the config if cpu_offload is set, t5 model must be on cuda device + config = Config(cpu_offload=False) + + # Encode the text + context = t5_encoder.infer([prompt], config) + context_null = t5_encoder.infer( + [negative_prompt if negative_prompt else ""], config + ) + + # Create text embeddings dictionary + text_embeddings = {"context": context, "context_null": context_null} + + return (text_embeddings,) + + +class Lightx2vWanVideoClipVisionEncoderLoader: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "clip_model_path": ( + "STRING", + { + "default": "models/clip/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth" + }, + ), + "tokenizer_path": ( + "STRING", + {"default": "models/clip/xlm-roberta-large"}, + ), + "precision": (["fp16", "fp32"], {"default": "fp16"}), + "device": (["cuda", "cpu"], {"default": "cuda"}), + } + } + + RETURN_TYPES = ("LIGHT_CLIP_VISION_ENCODER",) + RETURN_NAMES = ("clip_vision_encoder",) + FUNCTION = "load_clip_vision_encoder" + CATEGORY = "LightX2V" + + def load_clip_vision_encoder( + self, clip_model_path, tokenizer_path, precision, device + ): + # Map precision to torch dtype + dtype_map = {"fp16": torch.float16, "fp32": torch.float32} + dtype = dtype_map[precision] + + # Resolve device + if device == "cuda": + device = mm.get_torch_device() + else: + device = torch.device("cpu") + + # Load the CLIP vision encoder + clip_vision_encoder = ClipVisionModel( + dtype=dtype, + device=device, + checkpoint_path=clip_model_path, + tokenizer_path=tokenizer_path, + ) + + return (clip_vision_encoder,) + + +class Lightx2vWanVideoVaeLoader: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "vae_model_path": ("STRING", {"default": "models/vae/Wan2.1_VAE.pth"}), + "precision": (["bf16", "fp16", "fp32"], {"default": "fp16"}), + "device": (["cuda", "cpu"], {"default": "cuda"}), + "parallel": ("BOOLEAN", {"default": False}), + } + } + + RETURN_TYPES = ("LIGHT_WAN_VAE",) + RETURN_NAMES = ("wan_vae",) + FUNCTION = "load_vae" + CATEGORY = "LightX2V" + + def load_vae(self, vae_model_path, precision, device, parallel): + # Map precision to torch dtype + dtype_map = { + "bf16": torch.bfloat16, + "fp16": torch.float16, + "fp32": torch.float32, + } + dtype = dtype_map[precision] + + # Resolve device + if device == "cuda": + device = mm.get_torch_device() + else: + device = torch.device("cpu") + + # Load the VAE + vae = WanVAE( + z_dim=16, + vae_pth=vae_model_path, + dtype=dtype, + device=device, # type:ignore + parallel=parallel, + ) + + return (vae,) + + +class Lightx2vWanVideoVaeDecoder: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "vae": ("LIGHT_WAN_VAE",), + "latent": ("LIGHT_LATENT",), + } + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "decode_latent" + CATEGORY = "LightX2V" + + def decode_latent(self, vae, latent): + # 创建一个带有必要属性的配置对象 + class Config: + def __init__(self, cpu_offload=False): + self.cpu_offload = cpu_offload + + config = Config(cpu_offload=False) + + # 获取潜在表示和生成器 + latents = latent["samples"] + generator = latent["generator"] + + # 使用VAE解码潜在表示 + with torch.no_grad(): + # 解码得到视频帧 + decoded_images = vae.decode(latents, generator=generator, config=config) + + # 将像素值从 [-1, 1] 归一化到 [0, 1] + images = (decoded_images + 1) / 2 + + # 重新排列维度为ComfyUI标准的图像格式 [T, H, W, C] + # 从 [1, C, T, H, W] 转换为 [T, H, W, C] + images = images.squeeze(0).permute(1, 2, 3, 0).cpu() + + # 确保像素值在有效范围内 + images = torch.clamp(images, 0, 1) + + # 清理缓存以释放GPU内存 + torch.cuda.empty_cache() + gc.collect() + + return (images,) + + +class Lightx2vWanVideoImageEncoder: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "vae": ("LIGHT_WAN_VAE",), + "image": ("IMAGE",), + "clip_vision_encoder": ("LIGHT_CLIP_VISION_ENCODER",), + "target_height": ( + "INT", + {"default": 576, "min": 256, "max": 1024, "step": 8}, + ), + "target_width": ( + "INT", + {"default": 1024, "min": 256, "max": 1024, "step": 8}, + ), + "vae_stride": ("INT", {"default": 8, "min": 1, "max": 32, "step": 1}), + "patch_size": ("INT", {"default": 2, "min": 1, "max": 16, "step": 1}), + } + } + + RETURN_TYPES = ("LIGHT_IMAGE_EMBEDDINGS",) + RETURN_NAMES = ("image_embeddings",) + FUNCTION = "encode_image" + CATEGORY = "LightX2V" + + def encode_image( + self, + vae, + image, + clip_vision_encoder, + target_height, + target_width, + vae_stride, + patch_size, + ): + # 创建配置对象 + class Config: + def __init__( + self, + target_height, + target_width, + vae_stride, + patch_size, + cpu_offload=False, + ): + self.target_height = target_height + self.target_width = target_width + self.vae_stride = [1, vae_stride, vae_stride] + self.patch_size = [1, patch_size, patch_size] + self.cpu_offload = cpu_offload + + config = Config(target_height, target_width, vae_stride, patch_size) + + # 将图像转换为期望的张量格式 + device = mm.get_torch_device() + img = image[0].permute(2, 0, 1).to(device) # [C, H, W] + img = img.sub_(0.5).div_(0.5) # 归一化到 [-1, 1] + + # 使用CLIP视觉编码器编码图像 + clip_encoder_out = clip_vision_encoder.visual( + [img[:, None, :, :]], config + ).squeeze(0) + + # 计算宽高比和尺寸 + h, w = img.shape[1:] + aspect_ratio = h / w + max_area = config.target_height * config.target_width + lat_h = round( + np.sqrt(max_area * aspect_ratio) + // config.vae_stride[1] + // config.patch_size[1] + * config.patch_size[1] + ) + lat_w = round( + np.sqrt(max_area / aspect_ratio) + // config.vae_stride[2] + // config.patch_size[2] + * config.patch_size[2] + ) + h = lat_h * config.vae_stride[1] + w = lat_w * config.vae_stride[2] + + # 创建掩码 + msk = torch.ones(1, 81, lat_h, lat_w, device=device) + msk[:, 1:] = 0 + msk = torch.concat( + [torch.repeat_interleave(msk[:, 0:1], repeats=4, dim=1), msk[:, 1:]], dim=1 + ) + msk = msk.view(1, msk.shape[1] // 4, 4, lat_h, lat_w) + msk = msk.transpose(1, 2)[0] + + # 使用VAE编码图像 + resized_img = torch.nn.functional.interpolate( + img[None].cpu(), size=(h, w), mode="bicubic" + ).transpose(0, 1) + padding = torch.zeros(3, 80, h, w, device=device) + concat_img = torch.concat([resized_img.to(device), padding], dim=1) + + vae_encode_out = vae.encode([concat_img], config)[0] + # TODO(xxx): hard code + vae_encode_out = torch.concat([msk, vae_encode_out]).to(torch.bfloat16) + + # 创建图像嵌入字典 + image_embeddings = { + "clip_encoder_out": clip_encoder_out, + "vae_encode_out": vae_encode_out, + "lat_h": lat_h, + "lat_w": lat_w, + } + + return (image_embeddings,) + + +class Lightx2vWanVideoModelLoader: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model_path": ("STRING", {"default": "models/wan"}), + "model_type": (["t2v", "i2v"], {"default": "t2v"}), + "precision": (["bf16", "fp16", "fp32"], {"default": "bf16"}), + "device": (["cuda", "cpu"], {"default": "cuda"}), + "feature_caching": (["NoCaching", "Tea"], {"default": "Tea"}), + "teacache_thresh": ( + "FLOAT", + { + "default": 0.26, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "tooltip": "Only used with Tea feature caching", + }, + ), + "attention_mode": ( + ["sdpa", "flash_attn_2", "flash_attn_3"], + {"default": "flash_attn_3"}, + ), + "cpu_offload": ("BOOLEAN", {"default": False}), + "mm_type": ("STRING", {"default": "Default"}), + }, + "optional": { + "lora_path": ("STRING", {"default": None}), + "lora_strength": ( + "FLOAT", + {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.01}, + ), + }, + } + + RETURN_TYPES = ("LIGHT_WAN_MODEL",) + RETURN_NAMES = ("wan_model",) + FUNCTION = "load_model" + CATEGORY = "LightX2V" + + def load_model( + self, + model_path, + model_type, + precision, + device, + feature_caching, + teacache_thresh, + attention_mode, + mm_type, + lora_path=None, + lora_strength=1.0, + ): + # 映射精度到torch dtype + dtype_map = { + "bf16": torch.bfloat16, + "fp16": torch.float16, + "fp32": torch.float32, + } + dtype = dtype_map[precision] + + # 解析设备 + if device == "cuda": + device = mm.get_torch_device() + else: + device = torch.device("cpu") + + # 首先在模型路径中查找config.json + model_path_dir = os.path.dirname(model_path) + config_json_path = os.path.join(model_path_dir, "config.json") + config_json = {} + if os.path.exists(config_json_path): + with open(config_json_path, "r") as f: + config_json = json.load(f) + else: + logging.error(f"Config file not found at {config_json_path}") + raise FileNotFoundError(f"Config file not found at {config_json_path}") + + # 创建配置字典 + config = { + "do_mm_calib": False, + "cpu_offload": False, + "parallel_attn_type": None, # [None, "ulysses", "ring"] + "parallel_vae": False, + "max_area": False, + "vae_stride": (4, 8, 8), + "patch_size": (1, 2, 2), + "feature_caching": feature_caching, # ["NoCaching", "TaylorSeer", "Tea"] + "teacache_thresh": teacache_thresh, + "use_ret_steps": False, + "use_bfloat16": dtype == torch.bfloat16, + "mm_config": { + "mm_type": mm_type, + "weight_auto_quant": False if mm_type == "Default" else True, + }, + "model_path": model_path, + "task": model_type, + "device": device, + "attention_type": attention_mode, + "lora_path": lora_path if lora_path and lora_path.strip() else None, + "strength_model": lora_strength, + } + # merge model dir config.json and config dict + config.update(**config_json) + # NOTE(xxx): adapt to Lightx2v + config = EasyDict(config) + + logging.info( + "Loaded config:\n" + json.dumps(config, indent=4, ensure_ascii=False) + ) + logging.info(f"Loading WanModel from {model_path} with type {model_type}") + model = WanModel(model_path, config, device) + + # 如果指定了LoRA路径，应用LoRA + if lora_path and os.path.exists(lora_path): + logging.info( + f"Applying LoRA from {lora_path} with strength {lora_strength}" + ) + lora_wrapper = WanLoraWrapper(model) + lora_name = lora_wrapper.load_lora(lora_path) + lora_wrapper.apply_lora(lora_name, lora_strength) + logging.info(f"LoRA {lora_name} applied successfully") + + wan_model = {"wan_model": model, "config": config} + + return (wan_model,) + + +class Lightx2vWanVideoSampler: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model": ("LIGHT_WAN_MODEL",), + "text_embeddings": ("LIGHT_TEXT_EMBEDDINGS",), + "steps": ("INT", {"default": 20, "min": 1, "max": 100, "step": 1}), + "cfg_scale": ( + "FLOAT", + {"default": 7.5, "min": 0.0, "max": 20.0, "step": 0.1}, + ), + "seed": ("INT", {"default": 42}), + }, + "optional": { + "image_embeddings": ("LIGHT_IMAGE_EMBEDDINGS",), + }, + } + + RETURN_TYPES = ("LIGHT_LATENT",) + RETURN_NAMES = ("latent",) + FUNCTION = "sample" + CATEGORY = "LightX2V" + + def sample( + self, + model, + text_embeddings, + steps, + cfg_scale, + seed, + image_embeddings=None, + ): + # Update model config + + config = cast(Any, model.get("config")) + model = cast(WanModel, model.get("wan_model")) + + seed_all(seed) + + # wan_runner.init_scheduler + if config.feature_caching == "NoCaching": # type:ignore + scheduler = WanScheduler(config) + elif config.feature_caching == "Tea": # type:ignore + scheduler = WanSchedulerTeaCaching(config) + else: + raise NotImplementedError( + f"Unsupported feature_caching type: {config.feature_caching}" # type:ignore + ) + + # setup scheduler + model.set_scheduler(scheduler) + + # wan_runner.set_target_shape + if config.task == "i2v": + if image_embeddings is None: + raise ValueError("image_embeddings must be provided for i2v task") + config.lat_h = image_embeddings["lat_h"] + config.lat_w = image_embeddings["lat_w"] + config.target_shape = (16, 21, config.lat_h, config.lat_w) + elif config.task == "t2v": + config.target_shape = ( + 16, + (config.target_video_length - 1) // 4 + 1, + int(config.target_height) // config.vae_stride[1], + int(config.target_width) // config.vae_stride[2], + ) + + # Set up inputs + inputs = { + "text_encoder_output": text_embeddings, + "image_encoder_output": image_embeddings, + } + + # Prepare for sampling + scheduler.prepare(inputs.get("image_encoder_output")) + + # Run sampling + progress = ProgressBar(steps) + for step_index in tqdm( + range(scheduler.infer_steps), desc="inference", unit="step" + ): + with ProfilingContext("scheduler.step_pre"): + scheduler.step_pre(step_index=step_index) + with ProfilingContext("model.infer"): + model.infer(inputs) + with ProfilingContext("scheduler.step_post"): + scheduler.step_post() + + progress.update(1) + + scheduler.clear() + + del inputs, model, text_embeddings, image_embeddings + torch.cuda.empty_cache() + + return ({"samples": scheduler.latents, "generator": scheduler.generator},) + + +# Register the nodes +NODE_CLASS_MAPPINGS = { + "Lightx2vWanVideoT5EncoderLoader": Lightx2vWanVideoT5EncoderLoader, + "Lightx2vWanVideoT5Encoder": Lightx2vWanVideoT5Encoder, + "Lightx2vWanVideoClipVisionEncoderLoader": Lightx2vWanVideoClipVisionEncoderLoader, + "Lightx2vWanVideoVaeLoader": Lightx2vWanVideoVaeLoader, + "Lightx2vWanVideoImageEncoder": Lightx2vWanVideoImageEncoder, + "Lightx2vWanVideoVaeDecoder": Lightx2vWanVideoVaeDecoder, + "Lightx2vWanVideoModelLoader": Lightx2vWanVideoModelLoader, + "Lightx2vWanVideoSampler": Lightx2vWanVideoSampler, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "Lightx2vWanVideoT5EncoderLoader": "LightX2V WAN T5 Encoder Loader", + "Lightx2vWanVideoT5Encoder": "LightX2V WAN T5 Encoder", + "Lightx2vWanVideoClipVisionEncoderLoader": "LightX2V WAN CLIP Vision Encoder Loader", + "Lightx2vWanVideoClipVisionEncoder": "LightX2V WAN CLIP Vision Encoder", + "Lightx2vWanVideoVaeLoader": "LightX2V WAN VAE Loader", + "Lightx2vWanVideoImageEncoder": "LightX2V WAN Image Encoder", + "Lightx2vWanVideoVaeDecoder": "LightX2V WAN VAE Decoder", + "Lightx2vWanVideoModelLoader": "LightX2V WAN Model Loader", + "Lightx2vWanVideoSampler": "LightX2V WAN Video Sampler", +}