diff --git a/README.md b/README.md index 05e2767..52ecb97 100644 --- a/README.md +++ b/README.md @@ -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) diff --git a/lightx2v b/lightx2v index 5658c4b..85582c5 160000 --- a/lightx2v +++ b/lightx2v @@ -1 +1 @@ -Subproject commit 5658c4b41057d8553e2e961ea4c73d0452919f6d +Subproject commit 85582c5a4568d9d0f9e8320b2018ba460b2ba402 diff --git a/lightx2v_refactored/__init__.py b/lightx2v_refactored/__init__.py new file mode 100644 index 0000000..f719866 --- /dev/null +++ b/lightx2v_refactored/__init__.py @@ -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", +] \ No newline at end of file diff --git a/lightx2v_refactored/config.py b/lightx2v_refactored/config.py new file mode 100644 index 0000000..d51fc1e --- /dev/null +++ b/lightx2v_refactored/config.py @@ -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) \ No newline at end of file diff --git a/lightx2v_refactored/factory.py b/lightx2v_refactored/factory.py new file mode 100644 index 0000000..4244cae --- /dev/null +++ b/lightx2v_refactored/factory.py @@ -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 \ No newline at end of file diff --git a/lightx2v_refactored/models.py b/lightx2v_refactored/models.py new file mode 100644 index 0000000..9217953 --- /dev/null +++ b/lightx2v_refactored/models.py @@ -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 \ No newline at end of file diff --git a/lightx2v_refactored/nodes.py b/lightx2v_refactored/nodes.py new file mode 100644 index 0000000..8eef694 --- /dev/null +++ b/lightx2v_refactored/nodes.py @@ -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", +} diff --git a/nodes_refactored.py b/nodes_refactored.py new file mode 100644 index 0000000..e9be2e4 --- /dev/null +++ b/nodes_refactored.py @@ -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"] \ No newline at end of file