nodes/ModelTC-ComfyUI-Lightx2vWrapper
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

WIP: refactor nodes

Browse Source
This commit is contained in:
gaclove
2025-07-14 11:38:48 +08:00
parent 92a4886bd1
commit e234343952
8 changed files with 1512 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+22 -2
View File
@@ -44,7 +44,7 @@ your_model_dir/
### 推荐模型
- **Wan2.1-I2V-14B-480P**: 图像生成视频模型(480p分辨率)
- **Wan2.1-I2V-14B-720P**: 图像生成视频模型(720p分辨率)
- **Wan2.1-I2V-14B-720P**: 图像生成视频模型(720p分辨率)
- **Wan2.1-T2V-14B**: 文本生成视频模型
## 使用方法
@@ -83,11 +83,13 @@ T2V工作流与I2V类似，但使用 `Lightx2vWanVideoEmptyEmbeds` 替代图像
### 核心节点
#### Lightx2vWanVideoModelDir
- **功能**: 设置模型目录路径
- **输入**:
- **输入**:
- `model_dir`: 模型目录路径
#### Lightx2vWanVideoT5EncoderLoader
- **功能**: 加载T5文本编码器
- **输入**:
- `model_name`: T5模型文件名
@@ -95,6 +97,7 @@ T2V工作流与I2V类似，但使用 `Lightx2vWanVideoEmptyEmbeds` 替代图像
- `device`: 设备(cuda/cpu)
#### Lightx2vWanVideoT5Encoder
- **功能**: 编码文本提示词
- **输入**:
- `t5_encoder`: T5编码器实例
@@ -102,6 +105,7 @@ T2V工作流与I2V类似，但使用 `Lightx2vWanVideoEmptyEmbeds` 替代图像
- `negative_prompt`: 负面提示词
#### Lightx2vWanVideoVaeLoader
- **功能**: 加载VAE模型
- **输入**:
- `model_name`: VAE模型文件名
@@ -110,6 +114,7 @@ T2V工作流与I2V类似，但使用 `Lightx2vWanVideoEmptyEmbeds` 替代图像
- `parallel`: 是否并行处理
#### Lightx2vWanVideoImageEncoder
- **功能**: 编码输入图像用于I2V生成
- **输入**:
- `vae`: VAE模型实例
@@ -119,6 +124,7 @@ T2V工作流与I2V类似，但使用 `Lightx2vWanVideoEmptyEmbeds` 替代图像
- `num_frames`: 生成帧数
#### Lightx2vWanVideoModelLoader
- **功能**: 加载主要的Wan生成模型
- **输入**:
- `model_type`: 模型类型(t2v/i2v)
@@ -129,6 +135,7 @@ T2V工作流与I2V类似，但使用 `Lightx2vWanVideoEmptyEmbeds` 替代图像
- `teacache_args`: TeaCache参数(可选)
#### Lightx2vWanVideoSampler
- **功能**: 执行视频生成采样
- **输入**:
- `model`: Wan模型实例
@@ -139,6 +146,7 @@ T2V工作流与I2V类似，但使用 `Lightx2vWanVideoEmptyEmbeds` 替代图像
- `seed`: 随机种子
#### Lightx2vWanVideoVaeDecoder
- **功能**: 解码潜在表示为视频帧
- **输入**:
- `wan_vae`: VAE模型实例
@@ -147,6 +155,7 @@ T2V工作流与I2V类似，但使用 `Lightx2vWanVideoEmptyEmbeds` 替代图像
### 优化节点
#### Lightx2vTeaCache
- **功能**: TeaCache加速配置
- **输入**:
- `rel_l1_thresh`: 缓存阈值
@@ -155,6 +164,7 @@ T2V工作流与I2V类似，但使用 `Lightx2vWanVideoEmptyEmbeds` 替代图像
- `cache_device`: 缓存设备
#### Lightx2vWanVideoEmptyEmbeds
- **功能**: 为T2V任务提供空图像嵌入
- **输入**:
- `width/height`: 目标分辨率
@@ -163,20 +173,24 @@ T2V工作流与I2V类似，但使用 `Lightx2vWanVideoEmptyEmbeds` 替代图像
## 参数说明
### 分辨率设置
- **480P**: width=832, height=480
- **720P**: width=1280, height=720
- 确保宽高为8的倍数
### 帧数设置
- 推荐帧数: 81帧 (约5秒 @ 16fps)
- 帧数应为4的倍数+1 (如: 81, 85, 89等)
### 精度选择
- **bf16**: 推荐用于主模型，平衡精度和性能
- **fp16**: 适用于VAE和CLIP，节省显存
- **fp32**: 最高精度，显存占用大
### CFG引导
- **CFG Scale**: 1.0-20.0，值越高越遵循提示词
- **推荐值**: 5.0-8.0
@@ -213,3 +227,9 @@ T2V工作流与I2V类似，但使用 `Lightx2vWanVideoEmptyEmbeds` 替代图像
- 集成TeaCache加速功能
- 添加LoRA微调支持
- 优化内存管理和性能
## 参考资料
- [Develop Custom Nodes](https://docs.comfy.org/custom-nodes/walkthrough)
- [一份ComfyUI 自定义节点指南](https://developer.volcengine.com/articles/7399549896778317833)
- [Lightx2v](https://github.com/ModelTC/lightx2v)
lightx2v
+1 -1
Submodule lightx2v updated: 5658c4b410...85582c5a45
lightx2v_refactored/__init__.py
+42
View File
@@ -0,0 +1,42 @@
# Refactored LightX2V module
from .config import LightX2VConfig
from .factory import LightX2VFactory
from .models import (
LightX2VT5Encoder,
LightX2VClipVisionEncoder,
LightX2VVae,
LightX2VModel,
)
from .nodes import (
Lightx2vWanVideoModelDir,
Lightx2vWanVideoT5EncoderLoader,
Lightx2vWanVideoT5Encoder,
Lightx2vWanVideoClipVisionEncoderLoader,
Lightx2vWanVideoVaeLoader,
Lightx2vWanVideoVaeDecoder,
Lightx2vWanVideoImageEncoder,
Lightx2vWanVideoEmptyEmbeds,
Lightx2vWanVideoModelLoader,
Lightx2vWanVideoSampler,
WanVideoTeaCache,
)
__all__ = [
"LightX2VConfig",
"LightX2VFactory",
"LightX2VT5Encoder",
"LightX2VClipVisionEncoder",
"LightX2VVae",
"LightX2VModel",
"Lightx2vWanVideoModelDir",
"Lightx2vWanVideoT5EncoderLoader",
"Lightx2vWanVideoT5Encoder",
"Lightx2vWanVideoClipVisionEncoderLoader",
"Lightx2vWanVideoVaeLoader",
"Lightx2vWanVideoVaeDecoder",
"Lightx2vWanVideoImageEncoder",
"Lightx2vWanVideoEmptyEmbeds",
"Lightx2vWanVideoModelLoader",
"Lightx2vWanVideoSampler",
"WanVideoTeaCache",
]
lightx2v_refactored/config.py
+155
View File
@@ -0,0 +1,155 @@
"""Configuration management for LightX2V."""
from dataclasses import dataclass, field
from typing import Optional, Dict, Any, List, Tuple
from pathlib import Path
import torch
from easydict import EasyDict
@dataclass
class TeaCacheConfig:
"""Configuration for TeaCache optimization."""
rel_l1_thresh: float = 0.26
start_percent: float = 0.1
end_percent: float = 1.0
cache_device: torch.device = field(default_factory=lambda: torch.device("cpu"))
coefficients: List[List[float]] = field(default_factory=list)
use_ret_steps: bool = False
mode: str = "e"
@dataclass
class VideoConfig:
"""Video generation configuration."""
width: int = 832
height: int = 480
num_frames: int = 81
vae_stride: Tuple[int, int, int] = (4, 8, 8)
patch_size: Tuple[int, int, int] = (1, 2, 2)
@property
def max_area(self) -> int:
return self.height * self.width
@dataclass
class ModelConfig:
"""Model loading and inference configuration."""
model_path: Path
model_type: str = "i2v" # "t2v" or "i2v"
precision: str = "bf16" # "bf16", "fp16", "fp32"
device: str = "cuda"
attention_type: str = "flash_attn3"
cpu_offload: bool = False
offload_granularity: str = "phase" # "block" or "phase"
# Optional configurations
lora_path: Optional[Path] = None
lora_strength: float = 1.0
mm_config: Dict[str, Any] = field(default_factory=dict)
# Inference settings
steps: int = 20
shift: float = 5.0
cfg_scale: float = 5.0
seed: int = 42
def to_dtype(self) -> torch.dtype:
"""Convert precision string to torch dtype."""
dtype_map = {
"bf16": torch.bfloat16,
"fp16": torch.float16,
"fp32": torch.float32,
}
return dtype_map[self.precision]
def to_device(self) -> torch.device:
"""Get torch device."""
if self.device == "cuda":
return torch.device("cuda")
return torch.device("cpu")
@dataclass
class EncoderConfig:
"""Encoder configuration."""
model_path: Path
dtype: torch.dtype
device: torch.device
# T5 specific
text_len: int = 512
tokenizer_path: Optional[Path] = None
cpu_offload: bool = False
# CLIP specific
clip_quantized: bool = False
clip_quantized_ckpt: Optional[Path] = None
quant_scheme: Optional[str] = None
# VAE specific
z_dim: int = 16
parallel: bool = False
@dataclass
class LightX2VConfig:
"""Main configuration container for LightX2V."""
model: ModelConfig
video: VideoConfig
teacache: Optional[TeaCacheConfig] = None
@classmethod
def from_dict(cls, config_dict: Dict[str, Any]) -> "LightX2VConfig":
"""Create configuration from dictionary."""
model_config = ModelConfig(**config_dict.get("model", {}))
video_config = VideoConfig(**config_dict.get("video", {}))
teacache_config = None
if "teacache" in config_dict:
teacache_config = TeaCacheConfig(**config_dict["teacache"])
return cls(
model=model_config,
video=video_config,
teacache=teacache_config
)
def to_easydict(self) -> EasyDict:
"""Convert to EasyDict for legacy compatibility."""
config_dict = {
"model_path": str(self.model.model_path),
"task": self.model.model_type,
"dtype": self.model.to_dtype(),
"device": self.model.to_device(),
"attention_type": self.model.attention_type,
"cpu_offload": self.model.cpu_offload,
"offload_granularity": self.model.offload_granularity,
"target_height": self.video.height,
"target_width": self.video.width,
"target_video_length": self.video.num_frames,
"vae_stride": self.video.vae_stride,
"patch_size": self.video.patch_size,
"infer_steps": self.model.steps,
"sample_shift": self.model.shift,
"sample_guide_scale": self.model.cfg_scale,
"seed": self.model.seed,
"enable_cfg": self.model.cfg_scale != 1.0,
"mm_config": self.model.mm_config,
}
if self.teacache:
config_dict.update({
"feature_caching": "Tea",
"teacache_thresh": self.teacache.rel_l1_thresh,
"use_ret_steps": self.teacache.use_ret_steps,
"coefficients": self.teacache.coefficients,
})
else:
config_dict["feature_caching"] = "NoCaching"
if self.model.lora_path:
config_dict["lora_path"] = str(self.model.lora_path)
config_dict["strength_model"] = self.model.lora_strength
return EasyDict(config_dict)
lightx2v_refactored/factory.py
+243
View File
@@ -0,0 +1,243 @@
"""Factory pattern for creating LightX2V components."""
from pathlib import Path
from typing import Optional, Dict, Any, Union
import torch
import logging
from .config import EncoderConfig, ModelConfig, VideoConfig
from .models import (
LightX2VT5Encoder,
LightX2VClipVisionEncoder,
LightX2VVae,
LightX2VModel,
)
# Import original LightX2V modules
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.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
class LightX2VFactory:
"""Factory for creating LightX2V components with proper configuration."""
@staticmethod
def create_t5_encoder(
model_path: Union[str, Path],
tokenizer_path: Optional[Union[str, Path]] = None,
dtype: Optional[torch.dtype] = None,
device: Optional[torch.device] = None,
cpu_offload: bool = False,
) -> LightX2VT5Encoder:
"""Create a T5 encoder with configuration."""
model_path = Path(model_path)
# Auto-detect tokenizer path if not provided
if tokenizer_path is None:
tokenizer_path = model_path.parent / "google" / "umt5-xxl"
if not tokenizer_path.exists():
raise ValueError(f"Tokenizer not found at {tokenizer_path}")
config = EncoderConfig(
model_path=model_path,
tokenizer_path=Path(tokenizer_path),
dtype=dtype or torch.bfloat16,
device=device or torch.device("cuda"),
cpu_offload=cpu_offload,
)
# Create underlying T5 model
t5_model = T5EncoderModel(
text_len=config.text_len,
dtype=config.dtype,
device=config.device,
checkpoint_path=str(config.model_path),
tokenizer_path=str(config.tokenizer_path),
shard_fn=None,
cpu_offload=config.cpu_offload,
)
return LightX2VT5Encoder(t5_model, config)
@staticmethod
def create_clip_vision_encoder(
model_path: Union[str, Path],
dtype: Optional[torch.dtype] = None,
device: Optional[torch.device] = None,
clip_quantized: bool = False,
clip_quantized_ckpt: Optional[Union[str, Path]] = None,
quant_scheme: Optional[str] = None,
) -> LightX2VClipVisionEncoder:
"""Create a CLIP vision encoder with configuration."""
config = EncoderConfig(
model_path=Path(model_path),
dtype=dtype or torch.float16,
device=device or torch.device("cuda"),
clip_quantized=clip_quantized,
clip_quantized_ckpt=Path(clip_quantized_ckpt) if clip_quantized_ckpt else None,
quant_scheme=quant_scheme,
)
# Create underlying CLIP model
clip_model = ClipVisionModel(
dtype=config.dtype,
device=config.device,
checkpoint_path=str(config.model_path),
clip_quantized=config.clip_quantized,
clip_quantized_ckpt=str(config.clip_quantized_ckpt) if config.clip_quantized_ckpt else None,
quant_scheme=config.quant_scheme,
)
return LightX2VClipVisionEncoder(clip_model, config)
@staticmethod
def create_vae(
model_path: Union[str, Path],
dtype: Optional[torch.dtype] = None,
device: Optional[torch.device] = None,
parallel: bool = False,
z_dim: int = 16,
) -> LightX2VVae:
"""Create a VAE with configuration."""
config = EncoderConfig(
model_path=Path(model_path),
dtype=dtype or torch.float16,
device=device or torch.device("cuda"),
parallel=parallel,
z_dim=z_dim,
)
# Create underlying VAE model
vae_model = WanVAE(
z_dim=config.z_dim,
vae_pth=str(config.model_path),
dtype=config.dtype,
device=config.device,
parallel=config.parallel,
)
return LightX2VVae(vae_model, config)
@staticmethod
def create_model(
config: ModelConfig,
video_config: Optional[VideoConfig] = None,
) -> LightX2VModel:
"""Create a complete LightX2V model with configuration."""
# Load config.json if it exists
config_json_path = config.model_path / "config.json"
config_json = {}
if config_json_path.exists():
import json
with open(config_json_path, "r") as f:
config_json = json.load(f)
else:
logging.warning(f"Config file not found at {config_json_path}")
# Create model configuration dict
model_config_dict = {
"model_path": str(config.model_path),
"task": config.model_type,
"dtype": config.to_dtype(),
"device": config.to_device(),
"attention_type": config.attention_type,
"cpu_offload": config.cpu_offload,
"offload_granularity": config.offload_granularity,
"mm_config": config.mm_config,
"model_cls": "wan2.1",
"do_mm_calib": False,
"parallel_attn_type": None,
"parallel_vae": False,
"use_bfloat16": config.to_dtype() == torch.bfloat16,
}
# Add video config if provided
if video_config:
model_config_dict.update({
"target_height": video_config.height,
"target_width": video_config.width,
"target_video_length": video_config.num_frames,
"vae_stride": video_config.vae_stride,
"patch_size": video_config.patch_size,
"max_area": video_config.max_area,
})
# Merge with config.json
model_config_dict.update(config_json)
# Create EasyDict for compatibility
from easydict import EasyDict
easydict_config = EasyDict(model_config_dict)
# Create underlying model
wan_model = WanModel(
str(config.model_path),
easydict_config,
config.to_device()
)
# Apply LoRA if specified
if config.lora_path and config.lora_path.exists():
logging.info(f"Applying LoRA from {config.lora_path}")
lora_wrapper = WanLoraWrapper(wan_model)
lora_name = lora_wrapper.load_lora(str(config.lora_path))
lora_wrapper.apply_lora(lora_name, config.lora_strength)
logging.info(f"LoRA {lora_name} applied successfully")
return LightX2VModel(wan_model, config, easydict_config)
@staticmethod
def create_from_paths(
model_dir: Union[str, Path],
model_name: str,
model_type: str = "i2v",
precision: str = "bf16",
device: str = "cuda",
**kwargs
) -> Dict[str, Any]:
"""Convenience method to create all components from a model directory."""
model_dir = Path(model_dir)
# Create model config
model_config = ModelConfig(
model_path=model_dir / model_name,
model_type=model_type,
precision=precision,
device=device,
**kwargs
)
# Create components
components = {
"model": LightX2VFactory.create_model(model_config),
}
# Try to create encoders if paths exist
t5_path = model_dir / "models_t5_umt5-xxl-enc-bf16.pth"
if t5_path.exists():
components["t5_encoder"] = LightX2VFactory.create_t5_encoder(
t5_path,
dtype=model_config.to_dtype(),
device=model_config.to_device(),
)
clip_path = model_dir / "models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth"
if clip_path.exists():
components["clip_encoder"] = LightX2VFactory.create_clip_vision_encoder(
clip_path,
dtype=torch.float16, # CLIP typically uses fp16
device=model_config.to_device(),
)
vae_path = model_dir / "Wan2.1_VAE.pth"
if vae_path.exists():
components["vae"] = LightX2VFactory.create_vae(
vae_path,
dtype=torch.float16, # VAE typically uses fp16
device=model_config.to_device(),
)
return components
lightx2v_refactored/models.py
+196
View File
@@ -0,0 +1,196 @@
"""Model wrappers for LightX2V components."""
from typing import Any, Dict, List, Optional, Tuple, Union
import torch
from dataclasses import dataclass
from abc import ABC, abstractmethod
from .config import EncoderConfig, ModelConfig, VideoConfig
class BaseModel(ABC):
"""Base class for all LightX2V models."""
def __init__(self, config: Union[EncoderConfig, ModelConfig]):
self.config = config
@abstractmethod
def to(self, device: torch.device) -> "BaseModel":
"""Move model to device."""
pass
class LightX2VT5Encoder(BaseModel):
"""Wrapper for T5 text encoder."""
def __init__(self, t5_model: Any, config: EncoderConfig):
super().__init__(config)
self._model = t5_model
def encode(self, prompts: List[str]) -> Dict[str, torch.Tensor]:
"""Encode text prompts."""
context = self._model.infer(prompts)
return {"context": context}
def encode_with_negative(
self,
prompt: str,
negative_prompt: Optional[str] = None
) -> Dict[str, torch.Tensor]:
"""Encode prompt with negative prompt."""
context = self._model.infer([prompt])
context_null = self._model.infer([negative_prompt if negative_prompt else ""])
return {
"context": context,
"context_null": context_null
}
def to(self, device: torch.device) -> "LightX2VT5Encoder":
"""Move encoder to device."""
# T5 model handles device internally
return self
@property
def device(self) -> torch.device:
"""Get current device."""
return self.config.device
class LightX2VClipVisionEncoder(BaseModel):
"""Wrapper for CLIP vision encoder."""
def __init__(self, clip_model: Any, config: EncoderConfig):
super().__init__(config)
self._model = clip_model
def encode(
self,
images: torch.Tensor,
video_config: Optional[VideoConfig] = None
) -> torch.Tensor:
"""Encode images with CLIP."""
if video_config:
# Convert VideoConfig to dict format expected by CLIP
config_dict = {
"target_height": video_config.height,
"target_width": video_config.width,
"target_video_length": video_config.num_frames,
"vae_stride": video_config.vae_stride,
"patch_size": video_config.patch_size,
}
else:
config_dict = {}
# Ensure images are in correct format [B, C, T, H, W]
if images.dim() == 3: # [C, H, W]
images = images.unsqueeze(0).unsqueeze(2) # [1, C, 1, H, W]
elif images.dim() == 4: # [B, C, H, W]
images = images.unsqueeze(2) # [B, C, 1, H, W]
return self._model.visual(images, config_dict)
def to(self, device: torch.device) -> "LightX2VClipVisionEncoder":
"""Move encoder to device."""
# CLIP model handles device internally
return self
@property
def device(self) -> torch.device:
"""Get current device."""
return self.config.device
class LightX2VVae(BaseModel):
"""Wrapper for VAE."""
def __init__(self, vae_model: Any, config: EncoderConfig):
super().__init__(config)
self._model = vae_model
def encode(
self,
videos: List[torch.Tensor],
video_config: Optional[VideoConfig] = None,
cpu_offload: bool = False
) -> List[torch.Tensor]:
"""Encode videos to latent space."""
config_dict = {"cpu_offload": cpu_offload}
if video_config:
config_dict.update({
"target_height": video_config.height,
"target_width": video_config.width,
"target_video_length": video_config.num_frames,
"vae_stride": video_config.vae_stride,
"patch_size": video_config.patch_size,
})
from easydict import EasyDict
return self._model.encode(videos, EasyDict(config_dict))
def decode(
self,
latents: torch.Tensor,
generator: Optional[torch.Generator] = None,
cpu_offload: bool = False
) -> torch.Tensor:
"""Decode latents to video."""
from easydict import EasyDict
config = EasyDict({"cpu_offload": cpu_offload})
return self._model.decode(latents, generator=generator, config=config)
def to(self, device: torch.device) -> "LightX2VVae":
"""Move VAE to device."""
# VAE model handles device internally
return self
@property
def device(self) -> torch.device:
"""Get current device."""
return self.config.device
class LightX2VModel(BaseModel):
"""Wrapper for main LightX2V model."""
def __init__(self, wan_model: Any, config: ModelConfig, easydict_config: Any):
super().__init__(config)
self._model = wan_model
self._easydict_config = easydict_config
self._scheduler = None
def set_scheduler(self, scheduler: Any):
"""Set the scheduler for the model."""
self._scheduler = scheduler
self._model.set_scheduler(scheduler)
def infer(self, inputs: Dict[str, Any]):
"""Run inference."""
return self._model.infer(inputs)
def prepare_inputs(
self,
text_embeddings: Dict[str, torch.Tensor],
image_embeddings: Optional[Dict[str, Any]] = None,
) -> Dict[str, Any]:
"""Prepare inputs for inference."""
inputs = {
"text_encoder_output": text_embeddings,
"image_encoder_output": image_embeddings or {},
}
return inputs
def to(self, device: torch.device) -> "LightX2VModel":
"""Move model to device."""
# Model handles device internally
return self
@property
def device(self) -> torch.device:
"""Get current device."""
return self.config.to_device()
@property
def easydict_config(self) -> Any:
"""Get EasyDict config for compatibility."""
return self._easydict_config
lightx2v_refactored/nodes.py
+842
View File
@@ -0,0 +1,842 @@
"""Refactored ComfyUI nodes for LightX2V."""
import os
import torch
import gc
import logging
import json
import numpy as np
from pathlib import Path
from typing import Any, Dict, Optional, Tuple, cast
from easydict import EasyDict
from tqdm import tqdm
import comfy.model_management as comfy_mm
from comfy.utils import ProgressBar
from .config import LightX2VConfig, ModelConfig, VideoConfig, TeaCacheConfig
from .factory import LightX2VFactory
from .models import LightX2VT5Encoder, LightX2VClipVisionEncoder, LightX2VVae, LightX2VModel
# Import original LightX2V modules
from ..lightx2v.lightx2v.utils.profiler import ProfilingContext
from ..lightx2v.lightx2v.models.schedulers.wan.scheduler import WanScheduler
from ..lightx2v.lightx2v.models.schedulers.wan.feature_caching.scheduler import WanSchedulerTeaCaching
# Coefficient values for TeaCache
TEACACHE_COEFFICIENTS = {
"i2v-14B-480p": [
[2.57151496e05, -3.54229917e04, 1.40286849e03, -1.35890334e01, 1.32517977e-01],
[-3.02331670e02, 2.23948934e02, -5.25463970e01, 5.87348440e00, -2.01973289e-01],
],
"i2v-14B-720p": [
[8.10705460e03, 2.13393892e03, -3.72934672e02, 1.66203073e01, -4.17769401e-02],
[-114.36346466, 65.26524496, -18.82220707, 4.91518089, -0.23412683],
],
"t2v-1.3B": [
[-5.21862437e04, 9.23041404e03, -5.28275948e02, 1.36987616e01, -4.99875664e-02],
[2.39676752e03, -1.31110545e03, 2.01331979e02, -8.29855975e00, 1.37887774e-01],
],
"t2v-14B": [
[-3.03318725e05, 4.90537029e04, -2.65530556e03, 5.87365115e01, -3.15583525e-01],
[-5784.54975374, 5449.50911966, -1811.16591783, 256.27178429, -13.02252404],
],
}
class BaseNode:
"""Base class for ComfyUI nodes with common functionality."""
CATEGORY = "LightX2V"
@classmethod
def get_device(cls, device_str: str) -> torch.device:
"""Convert device string to torch device."""
if device_str == "cuda":
return comfy_mm.get_torch_device()
return torch.device("cpu")
@classmethod
def get_dtype(cls, precision_str: str) -> torch.dtype:
"""Convert precision string to torch dtype."""
dtype_map = {
"bf16": torch.bfloat16,
"fp16": torch.float16,
"fp32": torch.float32,
}
return dtype_map[precision_str]
class WanVideoTeaCache(BaseNode):
"""TeaCache configuration node."""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"rel_l1_thresh": (
"FLOAT",
{
"default": 0.26,
"min": 0.0,
"max": 10.0,
"step": 0.001,
"tooltip": "Threshold for cache application",
},
),
"start_percent": (
"FLOAT",
{
"default": 0.1,
"min": 0.0,
"max": 1.0,
"step": 0.01,
"tooltip": "Start percentage for TeaCache",
},
),
"end_percent": (
"FLOAT",
{
"default": 1.0,
"min": 0.0,
"max": 1.0,
"step": 0.01,
"tooltip": "End percentage for TeaCache",
},
),
"cache_device": (
["main_device", "offload_device"],
{"default": "offload_device", "tooltip": "Device to cache to"},
),
"coefficients": (
list(TEACACHE_COEFFICIENTS.keys()),
{
"default": "i2v-14B-720p",
"tooltip": "Coefficient preset for TeaCache",
},
),
"use_ret_steps": ("BOOLEAN", {"default": False}),
},
"optional": {
"mode": (
["e", "e0"],
{
"default": "e",
"tooltip": "Time embedding mode",
},
),
},
}
RETURN_TYPES = ("LIGHT_TEACACHEARGS",)
RETURN_NAMES = ("teacache_args",)
FUNCTION = "process"
EXPERIMENTAL = True
def process(
self,
rel_l1_thresh: float,
start_percent: float,
end_percent: float,
cache_device: str,
coefficients: str,
use_ret_steps: bool,
mode: str = "e",
) -> Tuple[TeaCacheConfig]:
"""Create TeaCache configuration."""
device = comfy_mm.get_torch_device() if cache_device == "main_device" else comfy_mm.unet_offload_device()
config = TeaCacheConfig(
rel_l1_thresh=rel_l1_thresh,
start_percent=start_percent,
end_percent=end_percent,
cache_device=device,
coefficients=TEACACHE_COEFFICIENTS[coefficients],
use_ret_steps=use_ret_steps,
mode=mode,
)
return (config,)
class Lightx2vWanVideoModelDir(BaseNode):
"""Model directory specification node."""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"model_dir": (
"STRING",
{"default": "/mnt/aigc/users/lijiaqi2/wan_model/Wan2.1-I2V-14B-480P"},
)
}
}
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("model_dir",)
FUNCTION = "process"
def process(self, model_dir: str) -> Tuple[str]:
"""Validate and return model directory."""
path = Path(model_dir)
if not path.exists():
raise ValueError(f"Model directory {model_dir} does not exist.")
return (model_dir,)
class Lightx2vWanVideoT5EncoderLoader(BaseNode):
"""T5 encoder loader node."""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"model_name": (
"STRING",
{"default": "models_t5_umt5-xxl-enc-bf16.pth"},
),
"precision": (["bf16", "fp16", "fp32"], {"default": "bf16"}),
"device": (["cuda", "cpu"], {"default": "cuda"}),
},
"optional": {
"model_dir": ("STRING", {"default": None}),
},
}
RETURN_TYPES = ("LIGHT_T5_ENCODER",)
RETURN_NAMES = ("t5_encoder",)
FUNCTION = "load_t5_encoder"
def load_t5_encoder(
self,
model_name: str,
precision: str,
device: str,
model_dir: Optional[str] = None,
) -> Tuple[LightX2VT5Encoder]:
"""Load T5 encoder."""
dtype = self.get_dtype(precision)
device_obj = self.get_device(device)
if model_dir:
model_path = Path(model_dir) / model_name
else:
model_path = Path(model_name)
if not model_path.exists():
raise ValueError(f"T5 model path {model_path} does not exist.")
encoder = LightX2VFactory.create_t5_encoder(
model_path=model_path,
dtype=dtype,
device=device_obj,
cpu_offload=(device == "cpu"),
)
return (encoder,)
class Lightx2vWanVideoT5Encoder(BaseNode):
"""T5 text encoding node."""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"t5_encoder": ("LIGHT_T5_ENCODER",),
"prompt": (
"STRING",
{
"multiline": True,
"default": "Summer beach vacation style...",
},
),
"negative_prompt": (
"STRING",
{
"multiline": True,
"default": "",
},
),
}
}
RETURN_TYPES = ("LIGHT_TEXT_EMBEDDINGS",)
RETURN_NAMES = ("text_embeddings",)
FUNCTION = "encode_text"
def encode_text(
self,
t5_encoder: LightX2VT5Encoder,
prompt: str,
negative_prompt: str = "",
) -> Tuple[Dict[str, torch.Tensor]]:
"""Encode text with T5."""
embeddings = t5_encoder.encode_with_negative(prompt, negative_prompt)
return (embeddings,)
class Lightx2vWanVideoClipVisionEncoderLoader(BaseNode):
"""CLIP vision encoder loader node."""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"model_name": (
"STRING",
{"default": "models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth"},
),
"tokenizer_path": (
"STRING",
{"default": "xlm-roberta-large"},
),
"precision": (["fp16", "fp32"], {"default": "fp16"}),
"device": (["cuda", "cpu"], {"default": "cuda"}),
},
"optional": {
"model_dir": ("STRING", {"default": None}),
},
}
RETURN_TYPES = ("LIGHT_CLIP_VISION_ENCODER",)
RETURN_NAMES = ("clip_vision_encoder",)
FUNCTION = "load_clip_vision_encoder"
def load_clip_vision_encoder(
self,
model_name: str,
tokenizer_path: str,
precision: str,
device: str,
model_dir: Optional[str] = None,
) -> Tuple[LightX2VClipVisionEncoder]:
"""Load CLIP vision encoder."""
dtype = self.get_dtype(precision)
device_obj = self.get_device(device)
if model_dir:
model_path = Path(model_dir) / model_name
else:
model_path = Path(model_name)
if not model_path.exists():
raise ValueError(f"CLIP model path {model_path} does not exist.")
encoder = LightX2VFactory.create_clip_vision_encoder(
model_path=model_path,
dtype=dtype,
device=device_obj,
)
return (encoder,)
class Lightx2vWanVideoVaeLoader(BaseNode):
"""VAE loader node."""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"model_name": (
"STRING",
{"default": "Wan2.1_VAE.pth"},
),
"precision": (["bf16", "fp16", "fp32"], {"default": "fp16"}),
"device": (["cuda", "cpu"], {"default": "cuda"}),
"parallel": ("BOOLEAN", {"default": False}),
},
"optional": {
"model_dir": ("STRING", {"default": None}),
},
}
RETURN_TYPES = ("LIGHT_WAN_VAE",)
RETURN_NAMES = ("wan_vae",)
FUNCTION = "load_vae"
def load_vae(
self,
model_name: str,
precision: str,
device: str,
parallel: bool,
model_dir: Optional[str] = None,
) -> Tuple[Dict[str, Any]]:
"""Load VAE."""
dtype = self.get_dtype(precision)
device_obj = self.get_device(device)
if model_dir:
model_path = Path(model_dir) / model_name
else:
model_path = Path(model_name)
if not model_path.exists():
raise ValueError(f"VAE model path {model_path} does not exist.")
vae = LightX2VFactory.create_vae(
model_path=model_path,
dtype=dtype,
device=device_obj,
parallel=parallel,
)
# Return in legacy format for compatibility
return ({"vae_cls": vae, "device": device},)
class Lightx2vWanVideoVaeDecoder(BaseNode):
"""VAE decoder node."""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"wan_vae": ("LIGHT_WAN_VAE",),
"latent": ("LIGHT_LATENT",),
}
}
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("images",)
FUNCTION = "decode_latent"
def decode_latent(
self,
wan_vae: Dict[str, Any],
latent: Dict[str, Any],
) -> Tuple[torch.Tensor]:
"""Decode latents to images."""
vae_instance = wan_vae["vae_cls"]
if isinstance(vae_instance, LightX2VVae):
vae_model = vae_instance
else:
# Legacy compatibility
vae_model = vae_instance
latents = latent["samples"]
generator = latent["generator"]
cpu_offload = wan_vae["device"] == "cpu"
with torch.no_grad():
with ProfilingContext("*decoded images*"):
if isinstance(vae_model, LightX2VVae):
decoded_images = vae_model.decode(latents, generator, cpu_offload)
else:
# Legacy compatibility
config = EasyDict({"cpu_offload": cpu_offload})
decoded_images = vae_model.decode(latents, generator=generator, config=config)
# Normalize from [-1, 1] to [0, 1]
images = (decoded_images + 1) / 2
# Rearrange dimensions for ComfyUI [T, H, W, C]
images = images.squeeze(0).permute(1, 2, 3, 0).cpu()
images = torch.clamp(images, 0, 1)
# Cleanup
torch.cuda.empty_cache()
gc.collect()
return (images,)
class Lightx2vWanVideoImageEncoder(BaseNode):
"""Image encoding node."""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"vae": ("LIGHT_WAN_VAE",),
"clip_vision_encoder": ("LIGHT_CLIP_VISION_ENCODER",),
"image": ("IMAGE",),
"width": (
"INT",
{
"default": 832,
"min": 64,
"max": 2048,
"step": 8,
"tooltip": "Width of the image",
},
),
"height": (
"INT",
{
"default": 480,
"min": 64,
"max": 2048,
"step": 8,
"tooltip": "Height of the image",
},
),
"num_frames": (
"INT",
{
"default": 81,
"min": 1,
"max": 10000,
"step": 4,
"tooltip": "Number of frames",
},
),
}
}
RETURN_TYPES = ("LIGHT_IMAGE_EMBEDDINGS",)
RETURN_NAMES = ("image_embeddings",)
FUNCTION = "encode_image"
def encode_image(
self,
vae: Dict[str, Any],
clip_vision_encoder: LightX2VClipVisionEncoder,
image: torch.Tensor,
width: int,
height: int,
num_frames: int,
) -> Tuple[Dict[str, Any]]:
"""Encode image with CLIP and VAE."""
vae_instance = vae["vae_cls"]
# Create video configuration
video_config = VideoConfig(
width=width,
height=height,
num_frames=num_frames,
)
# Convert image format
device = comfy_mm.get_torch_device()
img = image[0].permute(2, 0, 1).to(device)
img = img.sub_(0.5).div_(0.5) # Normalize to [-1, 1]
# Encode with CLIP
with ProfilingContext("*clip encoder*"):
if isinstance(clip_vision_encoder, LightX2VClipVisionEncoder):
clip_out = clip_vision_encoder.encode(img, video_config)
clip_out = clip_out.squeeze(0).to(torch.bfloat16)
else:
# Legacy compatibility
config_dict = video_config.__dict__ if hasattr(video_config, "__dict__") else video_config
clip_out = clip_vision_encoder.visual([img[:, None, :, :]], config_dict)
clip_out = clip_out.squeeze(0).to(torch.bfloat16)
# Calculate dimensions
h, w = img.shape[1:]
aspect_ratio = h / w
max_area = video_config.max_area
lat_h = round(np.sqrt(max_area * aspect_ratio) // video_config.vae_stride[1] // video_config.patch_size[1] * video_config.patch_size[1])
lat_w = round(np.sqrt(max_area / aspect_ratio) // video_config.vae_stride[2] // video_config.patch_size[2] * video_config.patch_size[2])
# Update config
config_dict = video_config.to_easydict() if hasattr(video_config, "to_easydict") else EasyDict(video_config.__dict__)
config_dict.lat_h = lat_h
config_dict.lat_w = lat_w
h = lat_h * video_config.vae_stride[1]
w = lat_w * video_config.vae_stride[2]
# Create mask
msk = torch.ones(1, num_frames, 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]
# Encode with VAE
with ProfilingContext("*vae encoder*"):
video_tensor = torch.concat(
[
torch.nn.functional.interpolate(img[None].cpu(), size=(h, w), mode="bicubic").transpose(0, 1),
torch.zeros(3, num_frames - 1, h, w),
],
dim=1,
).cuda()
if isinstance(vae_instance, LightX2VVae):
vae_out = vae_instance.encode([video_tensor], video_config)[0]
else:
# Legacy compatibility
vae_out = vae_instance.encode([video_tensor], config_dict)[0]
vae_out = torch.concat([msk, vae_out]).to(torch.bfloat16)
image_embeddings = {
"clip_encoder_out": clip_out,
"vae_encode_out": vae_out,
"config": config_dict,
}
logging.info(f"Image encoder outputs - CLIP: {clip_out.shape}, VAE: {vae_out.shape}")
return (image_embeddings,)
class Lightx2vWanVideoEmptyEmbeds(BaseNode):
"""Empty embeddings for T2V."""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"width": ("INT", {"default": 832, "min": 64, "max": 2048, "step": 8}),
"height": ("INT", {"default": 480, "min": 64, "max": 2048, "step": 8}),
"num_frames": ("INT", {"default": 81, "min": 1, "max": 10000, "step": 4}),
}
}
RETURN_TYPES = ("LIGHT_IMAGE_EMBEDDINGS",)
RETURN_NAMES = ("image_embeddings",)
FUNCTION = "process"
def process(self, num_frames: int, width: int, height: int) -> Tuple[Dict[str, Any]]:
"""Create empty image embeddings for T2V."""
video_config = VideoConfig(
width=width,
height=height,
num_frames=num_frames,
)
return ({"config": video_config.to_easydict()},)
class Lightx2vWanVideoModelLoader(BaseNode):
"""Main model loader node."""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"model_name": ("STRING", {"default": ""}),
"model_type": (["t2v", "i2v"], {"default": "i2v"}),
"precision": (["bf16", "fp16", "fp32"], {"default": "bf16"}),
"device": (["cuda", "cpu"], {"default": "cuda"}),
"attention_type": (
["sdpa", "flash_attn2", "flash_attn3"],
{"default": "flash_attn3"},
),
"cpu_offload": ("BOOLEAN", {"default": False}),
"offload_granularity": (["block", "phase"], {"default": "phase"}),
},
"optional": {
"mm_type": ("STRING", {"default": None}),
"teacache_args": ("LIGHT_TEACACHEARGS", {"default": None}),
"lora_path": ("STRING", {"default": None}),
"lora_strength": (
"FLOAT",
{"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.01},
),
"model_dir": (
"STRING",
{"default": "/mnt/aigc/users/lijiaqi2/wan_model/Wan2.1-I2V-14B-480P"},
),
},
}
RETURN_TYPES = ("LIGHT_WAN_MODEL",)
RETURN_NAMES = ("wan_model",)
FUNCTION = "load_model"
def load_model(
self,
model_name: str,
model_type: str,
precision: str,
device: str,
attention_type: str,
offload_granularity: str,
mm_type: Optional[str] = None,
lora_path: Optional[str] = None,
lora_strength: float = 1.0,
cpu_offload: bool = False,
teacache_args: Optional[TeaCacheConfig] = None,
model_dir: Optional[str] = None,
) -> Tuple[Dict[str, Any]]:
"""Load the main model."""
if model_dir:
model_path = Path(model_dir) / model_name
else:
model_path = Path(model_name)
if not model_path.exists():
raise ValueError(f"Model path {model_path} does not exist.")
# Parse mm_config
mm_config = {}
if mm_type:
try:
mm_config = json.loads(mm_type)
except Exception as e:
logging.error(f"Invalid mm_type config: {e}")
# Create model configuration
model_config = ModelConfig(
model_path=model_path,
model_type=model_type,
precision=precision,
device=device,
attention_type=attention_type,
cpu_offload=cpu_offload,
offload_granularity=offload_granularity,
lora_path=Path(lora_path) if lora_path and lora_path.strip() else None,
lora_strength=lora_strength,
mm_config=mm_config,
)
# Create model
model = LightX2VFactory.create_model(model_config)
# Add TeaCache config if provided
easydict_config = model.easydict_config
if teacache_args:
easydict_config.feature_caching = "Tea"
easydict_config.teacache_thresh = teacache_args.rel_l1_thresh
easydict_config.use_ret_steps = teacache_args.use_ret_steps
easydict_config.coefficients = teacache_args.coefficients
easydict_config.teacache_start_percent = teacache_args.start_percent
easydict_config.teacache_end_percent = teacache_args.end_percent
easydict_config.teacache_device = teacache_args.cache_device
easydict_config.teacache_mode = teacache_args.mode
logging.info(f"Loaded model from {model_path} with type {model_type}")
return ({"wan_model": model._model, "config": easydict_config},)
class Lightx2vWanVideoSampler(BaseNode):
"""Video sampling node."""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"model": ("LIGHT_WAN_MODEL",),
"text_embeddings": ("LIGHT_TEXT_EMBEDDINGS",),
"image_embeddings": ("LIGHT_IMAGE_EMBEDDINGS",),
"steps": ("INT", {"default": 20, "min": 1, "max": 100, "step": 1}),
"shift": ("FLOAT", {"default": 5.0}),
"cfg_scale": (
"FLOAT",
{"default": 5, "min": 1, "max": 20.0, "step": 0.1},
),
"seed": ("INT", {"default": 42, "min": 0, "max": 0xFFFFFFFFFFFFFFFF, "step": 1}),
}
}
RETURN_TYPES = ("LIGHT_LATENT",)
RETURN_NAMES = ("latent",)
FUNCTION = "sample"
def sample(
self,
model: Dict[str, Any],
text_embeddings: Dict[str, Any],
steps: int,
shift: float,
cfg_scale: float,
seed: int,
image_embeddings: Dict[str, Any],
) -> Tuple[Dict[str, Any]]:
"""Sample video latents."""
model_config = cast(EasyDict, model.get("config"))
model_config.update(image_embeddings.get("config", {}))
wan_model = model.get("wan_model")
clip_encoder_out = image_embeddings.get("clip_encoder_out", None)
vae_encode_out = image_embeddings.get("vae_encode_out", None)
if model_config.task == "i2v" and (clip_encoder_out is None or vae_encode_out is None):
raise ValueError("Image embeddings required for i2v task")
# Update config
model_config.infer_steps = steps
model_config.sample_shift = shift
model_config.sample_guide_scale = cfg_scale
model_config.seed = seed
model_config.enable_cfg = cfg_scale != 1.0
# Set target shape
num_channels_latents = model_config.get("num_channels_latents", 16)
if model_config.task == "i2v":
model_config.target_shape = (
num_channels_latents,
(model_config.target_video_length - 1) // model_config.vae_stride[0] + 1,
model_config.lat_h,
model_config.lat_w,
)
else: # t2v
model_config.target_shape = (
16,
(model_config.target_video_length - 1) // 4 + 1,
int(model_config.target_height) // model_config.vae_stride[1],
int(model_config.target_width) // model_config.vae_stride[2],
)
# Create scheduler
if model_config.feature_caching == "NoCaching":
scheduler = WanScheduler(model_config)
elif model_config.feature_caching == "Tea":
scheduler = WanSchedulerTeaCaching(model_config)
else:
raise NotImplementedError(f"Unsupported caching: {model_config.feature_caching}")
wan_model.set_scheduler(scheduler)
# Prepare inputs
inputs = {
"text_encoder_output": text_embeddings,
"image_encoder_output": image_embeddings,
}
scheduler.prepare(inputs.get("image_encoder_output"))
# Run sampling
progress = ProgressBar(steps)
for step_index in tqdm(range(scheduler.infer_steps), desc="Sampling"):
scheduler.step_pre(step_index=step_index)
with ProfilingContext("model.infer"):
wan_model.infer(inputs)
scheduler.step_post()
progress.update(1)
latents, generator = scheduler.latents, scheduler.generator
scheduler.clear()
# Cleanup
del inputs, scheduler, text_embeddings, image_embeddings
torch.cuda.empty_cache()
return ({"samples": latents, "generator": generator},)
# Node mappings
NODE_CLASS_MAPPINGS = {
"Lightx2vWanVideoModelDir": Lightx2vWanVideoModelDir,
"Lightx2vWanVideoT5EncoderLoader": Lightx2vWanVideoT5EncoderLoader,
"Lightx2vWanVideoT5Encoder": Lightx2vWanVideoT5Encoder,
"Lightx2vWanVideoClipVisionEncoderLoader": Lightx2vWanVideoClipVisionEncoderLoader,
"Lightx2vWanVideoVaeLoader": Lightx2vWanVideoVaeLoader,
"Lightx2vTeaCache": WanVideoTeaCache,
"Lightx2vWanVideoEmptyEmbeds": Lightx2vWanVideoEmptyEmbeds,
"Lightx2vWanVideoImageEncoder": Lightx2vWanVideoImageEncoder,
"Lightx2vWanVideoVaeDecoder": Lightx2vWanVideoVaeDecoder,
"Lightx2vWanVideoModelLoader": Lightx2vWanVideoModelLoader,
"Lightx2vWanVideoSampler": Lightx2vWanVideoSampler,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"Lightx2vWanVideoModelDir": "LightX2V WAN Model Directory",
"Lightx2vWanVideoT5EncoderLoader": "LightX2V WAN T5 Encoder Loader",
"Lightx2vWanVideoT5Encoder": "LightX2V WAN T5 Encoder",
"Lightx2vWanVideoClipVisionEncoderLoader": "LightX2V WAN CLIP Vision Encoder Loader",
"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",
"Lightx2vTeaCache": "LightX2V WAN Tea Cache",
"Lightx2vWanVideoEmptyEmbeds": "LightX2V WAN Video Empty Embeds",
}
nodes_refactored.py
+11
View File
@@ -0,0 +1,11 @@
"""
Refactored nodes.py that uses the new modular structure while maintaining backward compatibility.
"""
# Import refactored modules
from .lightx2v_refactored.nodes import (
NODE_CLASS_MAPPINGS,
NODE_DISPLAY_NAME_MAPPINGS,
)
# Export the mappings
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]