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

feat: enhance GPU capability checks and refactor configuration management in bridge.py; add support for model class list in model_utils.py; update inference and memory optimization configurations in nodes.py

Browse Source
This commit is contained in:
gaclove
2025-09-04 22:36:37 +08:00
parent 40bbc722e7
commit 4e75a5dc0c
4 changed files with 369 additions and 356 deletions
Download Patch File Download Diff File Expand all files Collapse all files
bridge.py
+216 -230
View File
@@ -1,5 +1,3 @@
"""Modular configuration system for LightX2V ComfyUI integration."""
import copy
import importlib.util
import json
@@ -11,14 +9,30 @@ import torch
from easydict import EasyDict
def is_fp8_supported_gpu():
def get_gpu_capability():
if not torch.cuda.is_available():
return None, None
try:
return torch.cuda.get_device_capability(0)
except Exception as e:
logging.warning(f"Failed to get GPU capability: {e}")
return None, None
def is_fp8_supported_gpu():
major, minor = get_gpu_capability()
if major is None:
return False
compute_capability = torch.cuda.get_device_capability(0)
major, minor = compute_capability
return (major == 8 and minor == 9) or (major >= 9)
def is_ada_architecture_gpu():
major, minor = get_gpu_capability()
if major is None:
return False
return major == 8 and minor == 9
def is_module_installed(module_name):
try:
spec = importlib.util.find_spec(module_name)
@@ -27,76 +41,51 @@ def is_module_installed(module_name):
return False
def get_available_quant_ops():
def get_available_ops(op_mapping):
"""通用的操作可用性检查函数"""
available_ops = []
for op_name, module_name in op_mapping.items():
is_available = is_module_installed(module_name)
available_ops.append((op_name, is_available))
return available_ops
vllm_installed = is_module_installed("vllm")
if vllm_installed:
available_ops.append(("vllm", True))
else:
available_ops.append(("vllm", False))
sgl_installed = is_module_installed("sgl_kernel")
if sgl_installed:
available_ops.append(("sgl", True))
else:
available_ops.append(("sgl", False))
def get_available_quant_ops():
quant_mapping = {"sgl": "sgl_kernel", "vllm": "vllm", "q8f": "q8_kernels", "torchao": "torchao"}
q8f_installed = is_module_installed("q8_kernels")
if q8f_installed:
available_ops.append(("q8f", True))
else:
available_ops.append(("q8f", False))
available_ops = get_available_ops(quant_mapping)
torchao_installed = is_module_installed("torchao")
if torchao_installed:
available_ops.append(("torchao", True))
else:
available_ops.append(("torchao", False))
# Ada架构GPU优先使用q8f
if is_ada_architecture_gpu():
q8f_available = next((op for op in available_ops if op[0] == "q8f" and op[1]), None)
if q8f_available:
available_ops.remove(q8f_available)
available_ops.insert(0, q8f_available)
return available_ops
def get_available_attn_ops():
available_ops = []
attn_mapping = {"sage_attn2": "sageattention", "flash_attn3": "flash_attn_interface", "flash_attn2": "flash_attn", "torch_sdpa": "torch"}
vllm_installed = is_module_installed("flash_attn")
if vllm_installed:
available_ops.append(("flash_attn2", True))
else:
available_ops.append(("flash_attn2", False))
sgl_installed = is_module_installed("flash_attn_interface")
if sgl_installed:
available_ops.append(("flash_attn3", True))
else:
available_ops.append(("flash_attn3", False))
q8f_installed = is_module_installed("sageattention")
if q8f_installed:
available_ops.append(("sage_attn2", True))
else:
available_ops.append(("sage_attn2", False))
torch_installed = is_module_installed("torch")
if torch_installed:
available_ops.append(("torch_sdpa", True))
else:
available_ops.append(("torch_sdpa", False))
return available_ops
return get_available_ops(attn_mapping)
class LightX2VDefaultConfig:
"""Central default configuration for LightX2V."""
# 分组常量
DEFAULT_ATTENTION_TYPE = "flash_attn3"
DEFAULT_QUANTIZATION_SCHEMES = {"dit": "bf16", "t5": "bf16", "clip": "fp16"}
DEFAULT_VIDEO_PARAMS = {"height": 480, "width": 832, "length": 81, "fps": 16, "vae_stride": [4, 8, 8], "patch_size": [1, 2, 2]}
DEFAULT_CONFIG = {
# ========== Model Configuration ==========
# Model Configuration
"model_cls": "wan2.1",
"model_path": "",
"task": "t2v",
"mode": "infer",
# ========== Inference Parameters ==========
# Inference Parameters
"infer_steps": 40,
"seed": 42,
"sample_guide_scale": 5.0,
@@ -104,38 +93,38 @@ class LightX2VDefaultConfig:
"enable_cfg": True,
"prompt": "",
"negative_prompt": "",
# ========== Video Parameters ==========
"target_height": 480,
"target_width": 832,
"target_video_length": 81,
"fps": 16,
"vae_stride": [4, 8, 8],
"patch_size": [1, 2, 2],
# ========== Feature Caching (TeaCache) ==========
# Video Parameters
"target_height": DEFAULT_VIDEO_PARAMS["height"],
"target_width": DEFAULT_VIDEO_PARAMS["width"],
"target_video_length": DEFAULT_VIDEO_PARAMS["length"],
"fps": DEFAULT_VIDEO_PARAMS["fps"],
"vae_stride": DEFAULT_VIDEO_PARAMS["vae_stride"],
"patch_size": DEFAULT_VIDEO_PARAMS["patch_size"],
# TeaCache
"feature_caching": "NoCaching",
"teacache_thresh": 0.26,
"coefficients": None, # Auto-calculated
"coefficients": None,
"use_ret_steps": False,
# ========== Quantization ==========
"dit_quant_scheme": "bf16",
"t5_quant_scheme": "bf16",
"clip_quant_scheme": "fp16",
# Quantization
"dit_quant_scheme": DEFAULT_QUANTIZATION_SCHEMES["dit"],
"t5_quant_scheme": DEFAULT_QUANTIZATION_SCHEMES["t5"],
"clip_quant_scheme": DEFAULT_QUANTIZATION_SCHEMES["clip"],
"quant_op": "vllm",
"precision_mode": "fp32",
"dit_quantized_ckpt": None,
"t5_quantized_ckpt": None,
"clip_quantized_ckpt": None,
"mm_config": {"mm_type": "Default"},
# ========== GPU Memory Optimization ==========
# Memory Optimization
"rotary_chunk": False,
"rotary_chunk_size": 100,
"clean_cuda_cache": False,
"torch_compile": False,
"attention_type": "flash_attn3",
"self_attn_1_type": "flash_attn3",
"cross_attn_1_type": "flash_attn3",
"cross_attn_2_type": "flash_attn3",
# ========== Async Offloading ==========
"attention_type": DEFAULT_ATTENTION_TYPE,
"self_attn_1_type": DEFAULT_ATTENTION_TYPE,
"cross_attn_1_type": DEFAULT_ATTENTION_TYPE,
"cross_attn_2_type": DEFAULT_ATTENTION_TYPE,
# CPU Offloading
"cpu_offload": False,
"offload_granularity": "phase",
"offload_ratio": 1.0,
@@ -143,12 +132,12 @@ class LightX2VDefaultConfig:
"t5_offload_granularity": "model",
"lazy_load": False,
"unload_modules": False,
# ========== Lightweight VAE ==========
# VAE Settings
"use_tiny_vae": False,
"tiny_vae": False,
"tiny_vae_path": None,
"use_tiling_vae": False,
# ========== Other Settings ==========
# Other Settings
"lora_path": None,
"strength_model": 1.0,
"do_mm_calib": False,
@@ -160,7 +149,6 @@ class LightX2VDefaultConfig:
"max_area": False,
"use_prompt_enhancer": False,
"text_len": 512,
# ========== Wan2.1 Audio Settings ==========
"use_31_block": True,
}
@@ -270,76 +258,73 @@ class ModularConfigManager:
self._available_attn_ops = None
self._available_quant_ops = None
def _get_available_ops(self, ops_list: List[Tuple[str, bool]], fallback: str = None) -> List[str]:
"""从操作列表中提取可用的操作"""
available = [op_name for op_name, is_available in ops_list if is_available]
if fallback and fallback not in available:
available.append(fallback)
return available
@property
def available_attention_types(self) -> List[str]:
"""Get available attention types."""
if self._available_attn_ops is None:
self._available_attn_ops = get_available_attn_ops()
available = []
for op_name, is_available in self._available_attn_ops:
if is_available:
available.append(op_name)
if "torch_sdpa" not in available:
available.append("torch_sdpa")
return available
return self._get_available_ops(self._available_attn_ops, "torch_sdpa")
@property
def available_quant_schemes(self) -> List[str]:
"""Get available quantization schemes."""
if self._available_quant_ops is None:
self._available_quant_ops = get_available_quant_ops()
return self._get_available_ops(self._available_quant_ops)
available = []
for op_name, is_available in self._available_quant_ops:
if is_available:
available.append(op_name)
return available
def _update_from_config(self, updates: Dict, config: Dict, mappings: Dict[str, str]) -> None:
"""通用配置更新方法"""
for config_key, update_key in mappings.items():
if config_key in config:
if config_key == "seed" and config[config_key] == -1:
continue
updates[update_key] = config[config_key]
def apply_inference_config(self, config: Dict[str, Any]) -> Dict[str, Any]:
"""Apply basic inference configuration."""
updates = {}
if "model_cls" in config:
updates["model_cls"] = config["model_cls"]
if "model_path" in config:
updates["model_path"] = config["model_path"]
if "task" in config:
updates["task"] = config["task"]
# 基础映射配置
basic_mappings = {
"model_cls": "model_cls",
"model_path": "model_path",
"task": "task",
"infer_steps": "infer_steps",
"seed": "seed",
"sample_shift": "sample_shift",
"height": "target_height",
"width": "target_width",
"video_length": "target_video_length",
"fps": "fps",
"video_duration": "video_duration",
"resize_mode": "resize_mode",
"denoising_step_list": "denoising_step_list",
"use_31_block": "use_31_block",
"prev_frame_length": "prev_frame_length",
}
if "infer_steps" in config:
updates["infer_steps"] = config["infer_steps"]
if "seed" in config and config["seed"] != -1:
updates["seed"] = config["seed"]
self._update_from_config(updates, config, basic_mappings)
# CFG特殊处理
if "cfg_scale" in config:
updates["sample_guide_scale"] = config["cfg_scale"]
updates["enable_cfg"] = config["cfg_scale"] != 1.0
if "sample_shift" in config:
updates["sample_shift"] = config["sample_shift"]
if "height" in config:
updates["target_height"] = config["height"]
if "width" in config:
updates["target_width"] = config["width"]
if "video_length" in config:
updates["target_video_length"] = config["video_length"]
if "fps" in config:
updates["fps"] = config["fps"]
# 注意力类型配置
attention_type = config.get("attention_type", LightX2VDefaultConfig.DEFAULT_ATTENTION_TYPE)
for attn_key in ["attention_type", "self_attn_1_type", "cross_attn_1_type", "cross_attn_2_type"]:
updates[attn_key] = attention_type
if "video_duration" in config:
updates["video_duration"] = config["video_duration"]
if "adaptive_resize" in config:
updates["adaptive_resize"] = config["adaptive_resize"]
if "denoising_step_list" in config:
updates["denoising_step_list"] = config["denoising_step_list"]
if "use_31_block" in config:
updates["use_31_block"] = config["use_31_block"]
# TinyVAE配置
if config.get("use_tiny_vae", False):
updates.update({"use_tiny_vae": True, "tiny_vae": True, "tiny_vae_path": os.path.join(config["model_path"], "taew2_1.pth")})
return updates
@@ -366,130 +351,143 @@ class ModularConfigManager:
return updates
def _get_mm_type(self, dit_scheme: str, quant_backend: str) -> str:
"""获取mm_type配置"""
if dit_scheme == "bf16":
return "Default"
base_pattern = f"W-{dit_scheme}-channel-sym-A-{dit_scheme}-channel-sym-dynamic"
if quant_backend == "vllm":
return f"{base_pattern}-Vllm"
elif quant_backend == "sgl":
suffix = "-Sgl-ActVllm" if dit_scheme == "int8" else "-Sgl"
return f"{base_pattern}{suffix}"
elif quant_backend == "q8f":
return f"{base_pattern}-Q8F"
elif quant_backend == "torchao":
return f"{base_pattern}-Torchao"
else:
return "Default"
def apply_quantization_config(self, config: Dict[str, Any], model_path: str) -> Dict[str, Any]:
"""Apply quantization configuration."""
updates = {}
defaults = LightX2VDefaultConfig.DEFAULT_QUANTIZATION_SCHEMES
dit_scheme = config.get("dit_precision", "bf16")
updates["dit_quant_scheme"] = dit_scheme
if dit_scheme != "bf16":
# 获取量化方案
dit_scheme = config.get("dit_quant_scheme", defaults["dit"])
t5_scheme = config.get("t5_quant_scheme", defaults["t5"])
clip_scheme = config.get("clip_quant_scheme", defaults["clip"])
quant_backend = config.get("quant_op", "vllm")
updates.update(
{
"dit_quant_scheme": dit_scheme,
"t5_quant_scheme": t5_scheme,
"clip_quant_scheme": clip_scheme,
"t5_quantized": t5_scheme != defaults["t5"],
"clip_quantized": clip_scheme != defaults["clip"],
}
)
# 设置检查点路径
if dit_scheme != defaults["dit"]:
updates["dit_quantized_ckpt"] = os.path.join(model_path, dit_scheme)
t5_scheme = config.get("t5_precision", "bf16")
updates["t5_quant_scheme"] = t5_scheme
updates["t5_quantized"] = t5_scheme != "bf16"
if t5_scheme != "bf16":
if t5_scheme != defaults["t5"]:
t5_path = os.path.join(model_path, t5_scheme)
updates["t5_quantized_ckpt"] = os.path.join(t5_path, f"models_t5_umt5-xxl-enc-{t5_scheme}.pth")
clip_scheme = config.get("clip_precision", "fp16")
updates["clip_quant_scheme"] = clip_scheme
updates["clip_quantized"] = clip_scheme != "fp16"
if clip_scheme != "fp16":
if clip_scheme != defaults["clip"]:
clip_path = os.path.join(model_path, clip_scheme)
updates["clip_quantized_ckpt"] = os.path.join(clip_path, f"clip-{clip_scheme}.pth")
quant_backend = config.get("quant_backend", "vllm")
updates["quant_op"] = quant_backend
# 特殊后端处理
if quant_backend in ["q8f", "torchao"]:
backend_suffix = f"int8-{quant_backend}"
updates.update({"t5_quant_scheme": backend_suffix, "clip_quant_scheme": backend_suffix})
if dit_scheme != "bf16":
if quant_backend == "vllm":
mm_type = f"W-{dit_scheme}-channel-sym-A-{dit_scheme}-channel-sym-dynamic-Vllm"
elif quant_backend == "sgl":
if dit_scheme == "int8":
mm_type = f"W-{dit_scheme}-channel-sym-A-{dit_scheme}-channel-sym-dynamic-Sgl-ActVllm"
else:
mm_type = f"W-{dit_scheme}-channel-sym-A-{dit_scheme}-channel-sym-dynamic-Sgl"
elif quant_backend == "q8f":
mm_type = f"W-{dit_scheme}-channel-sym-A-{dit_scheme}-channel-sym-dynamic-Q8F"
updates["t5_quant_scheme"] = "int8-q8f"
updates["clip_quant_scheme"] = "int8-q8f"
elif quant_backend == "torchao":
mm_type = f"W-{dit_scheme}-channel-sym-A-{dit_scheme}-channel-sym-dynamic-Torchao"
updates["t5_quant_scheme"] = "int8-torchao"
updates["clip_quant_scheme"] = "int8-torchao"
else:
mm_type = "Default"
updates["mm_config"] = {"mm_type": mm_type}
else:
updates["mm_config"] = {"mm_type": "Default"}
updates["precision_mode"] = config.get("sensitive_layers_precision", "fp32")
# 设置mm_config
mm_type = self._get_mm_type(dit_scheme, quant_backend)
updates["mm_config"] = {"mm_type": mm_type}
return updates
def apply_memory_optimization(self, config: Dict[str, Any]) -> Dict[str, Any]:
"""Apply memory optimization settings."""
updates = {}
level = config.get("optimization_level", "none")
# GPU optimization
if config.get("enable_rotary_chunk", False) or level in ["high", "extreme"]:
updates["rotary_chunk"] = True
# 级别映射配置
level_configs = {
"medium": {"cpu_offload": True},
"high": {"cpu_offload": True, "rotary_chunk": True, "t5_cpu_offload": True, "t5_offload_granularity": "model"},
"extreme": {
"cpu_offload": True,
"rotary_chunk": True,
"clean_cuda_cache": True,
"t5_cpu_offload": True,
"t5_offload_granularity": "block",
"lazy_load": True,
"unload_modules": True,
},
}
# 应用级别配置
if level in level_configs:
updates.update(level_configs[level])
# 直接配置项映射
direct_mappings = {
"enable_rotary_chunk": "rotary_chunk",
"clean_cuda_cache": "clean_cuda_cache",
"cpu_offload": "cpu_offload",
"lazy_load": "lazy_load",
"unload_after_inference": "unload_modules",
"use_tiling_vae": "use_tiling_vae",
}
for config_key, update_key in direct_mappings.items():
if config.get(config_key, False):
updates[update_key] = True
# 附加配置
if updates.get("rotary_chunk"):
updates["rotary_chunk_size"] = config.get("rotary_chunk_size", 100)
if config.get("clean_cuda_cache", False) or level == "extreme":
updates["clean_cuda_cache"] = True
# CPU offloading
if config.get("enable_cpu_offload", False) or level in [
"medium",
"high",
"extreme",
]:
updates["cpu_offload"] = True
updates["offload_granularity"] = config.get("offload_granularity", "phase")
updates["offload_ratio"] = config.get("offload_ratio", 1.0)
# T5 offloading
if level in ["high", "extreme"]:
updates["t5_cpu_offload"] = True
updates["t5_offload_granularity"] = "block" if level == "extreme" else "model"
# Module management
if config.get("lazy_load", False) or level == "extreme":
updates["lazy_load"] = True
if config.get("unload_after_inference", False) or level == "extreme":
updates["unload_modules"] = True
# Attention type
attention_type = config.get("attention_type", "flash_attn3")
updates["attention_type"] = attention_type
updates["self_attn_1_type"] = attention_type
updates["cross_attn_1_type"] = attention_type
updates["cross_attn_2_type"] = attention_type
if updates.get("cpu_offload"):
updates.update({"offload_granularity": config.get("offload_granularity", "phase"), "offload_ratio": config.get("offload_ratio", 1.0)})
return updates
def apply_vae_config(self, config: Dict[str, Any], model_path: str) -> Dict[str, Any]:
"""Apply VAE configuration."""
updates = {}
def _load_model_config(self, model_path: str) -> Dict[str, Any]:
"""加载模型配置文件"""
config_path = os.path.join(model_path, "config.json")
if not os.path.exists(config_path):
return {}
if config.get("use_tiny_vae", False):
updates["use_tiny_vae"] = True
updates["tiny_vae"] = True
updates["tiny_vae_path"] = os.path.join(model_path, "taew2_1.pth")
if config.get("use_tiling_vae", False):
updates["use_tiling_vae"] = True
return updates
try:
with open(config_path, "r") as f:
return json.load(f)
except Exception as e:
logging.warning(f"Failed to load model config: {e}")
return {}
def build_final_config(self, configs: Dict[str, Dict[str, Any]]) -> EasyDict:
"""Build final configuration from module configs."""
final_config = copy.deepcopy(self.base_config)
if "inference" in configs:
final_config.update(self.apply_inference_config(configs["inference"]))
# 应用配置模块
config_modules = [("inference", self.apply_inference_config), ("memory", self.apply_memory_optimization)]
for module_name, apply_func in config_modules:
if module_name in configs:
final_config.update(apply_func(configs[module_name]))
# 特殊处理的模块
if "teacache" in configs:
teacache_updates = self.apply_teacache_config(
configs["teacache"],
final_config,
)
teacache_updates = self.apply_teacache_config(configs["teacache"], final_config)
final_config.update(teacache_updates)
if "quantization" in configs:
@@ -497,22 +495,10 @@ class ModularConfigManager:
quant_updates = self.apply_quantization_config(configs["quantization"], model_path)
final_config.update(quant_updates)
if "memory" in configs:
final_config.update(self.apply_memory_optimization(configs["memory"]))
if "vae" in configs:
model_path = final_config.get("model_path", "")
final_config.update(self.apply_vae_config(configs["vae"], model_path))
model_config_path = os.path.join(final_config["model_path"], "config.json")
if os.path.exists(model_config_path):
try:
with open(model_config_path, "r") as f:
model_config = json.load(f)
for key, value in model_config.items():
if key not in final_config or final_config[key] is None:
final_config[key] = value
except Exception as e:
logging.warning(f"Failed to load model config: {e}")
# 加载模型配置
model_config = self._load_model_config(final_config.get("model_path", ""))
for key, value in model_config.items():
if key not in final_config or final_config[key] is None:
final_config[key] = value
return EasyDict(final_config)
lightx2v
+1 -1
Submodule lightx2v updated: d39d61f4c5...8652b9878e
model_utils.py
+16
View File
@@ -29,6 +29,22 @@ def scan_models() -> List[str]:
return ["None"] + models if models else ["None"]
def support_model_cls_list(model_cls: str) -> List[str]:
return [
"wan2.1",
"wan2.1_distill",
"wan2.1_vace",
"cogvideox",
"seko_talk",
"wan2.2_moe",
"wan2.2",
"wan2.2_moe_audio",
"wan2.2_audio",
"wan2.2_moe_distill",
"qwen_image",
]
def scan_loras() -> List[str]:
loras = []
base_path = get_model_base_path()
nodes.py
+136 -125
View File
@@ -23,6 +23,7 @@ from .model_utils import (
get_model_full_path,
scan_loras,
scan_models,
support_model_cls_list,
)
@@ -30,11 +31,21 @@ class LightX2VInferenceConfig:
@classmethod
def INPUT_TYPES(cls):
available_models = scan_models()
support_model_classes = support_model_cls_list()
available_attn = get_available_attn_ops()
attn_types = []
for op_name, is_available in available_attn:
if is_available:
attn_types.append(op_name)
if "torch_sdpa" not in attn_types:
attn_types.append("torch_sdpa")
return {
"required": {
"model_cls": (
["wan2.1", "wan2.1_audio", "wan2.1_distill", "hunyuan"],
support_model_classes,
{"default": "wan2.1", "tooltip": "Model type"},
),
"model_name": (
@@ -47,13 +58,13 @@ class LightX2VInferenceConfig:
"task": (
["t2v", "i2v"],
{
"default": "t2v",
"default": "i2v",
"tooltip": "Task type: text-to-video or image-to-video",
},
),
"infer_steps": (
"INT",
{"default": 40, "min": 1, "max": 100, "tooltip": "Inference steps"},
{"default": 4, "min": 1, "max": 100, "tooltip": "Inference steps"},
),
"seed": (
"INT",
@@ -108,6 +119,10 @@ class LightX2VInferenceConfig:
"tooltip": "Video duration in seconds",
},
),
"attention_type": (
attn_types,
{"default": attn_types[0], "tooltip": "Attention mechanism type"},
),
},
"optional": {
"denoising_steps": (
@@ -117,11 +132,36 @@ class LightX2VInferenceConfig:
"tooltip": "Custom denoising steps for distillation models (comma-separated, e.g., '999,750,500,250'). Leave empty to use model defaults.",
},
),
"adaptive_resize": (
"resize_mode": (
["adaptive", "keep_ratio_fixed_area", "fixed_min_area", "fixed_max_area", "fixed_shape"],
{
"default": "adaptive",
"tooltip": "Adaptive resize input image to target aspect ratio",
},
),
"segment_length": (
"INT",
{
"default": 81,
"min": 16,
"max": 256,
"tooltip": "Segment length in frames for sekotalk models (target_video_length)",
},
),
"prev_frame_length": (
"INT",
{
"default": 5,
"min": 0,
"max": 16,
"tooltip": "Previous frame overlap for sekotalk models",
},
),
"use_tiny_vae": (
"BOOLEAN",
{
"default": False,
"tooltip": "Adaptive resize input image to target aspect ratio",
"tooltip": "Use lightweight VAE to accelerate decoding",
},
),
},
@@ -144,8 +184,12 @@ class LightX2VInferenceConfig:
height,
width,
duration,
attention_type,
denoising_steps="",
adaptive_resize=False,
resize_mode="adaptive",
segment_length=81,
prev_frame_length=5,
use_tiny_vae=False,
):
"""Create basic inference configuration."""
model_path = get_model_full_path(model_name)
@@ -158,6 +202,7 @@ class LightX2VInferenceConfig:
video_length = int(round(duration * fps))
if video_length < 16:
logging.warning("Video length is too short, setting to 16")
video_length = 16
remainder = (video_length - 1) % 4
@@ -166,8 +211,8 @@ class LightX2VInferenceConfig:
# TODO(xxx):
use_31_block = True
if "wan2.1_audio" in [model_cls]:
video_length = 81
if "seko" in [model_cls]:
video_length = segment_length
use_31_block = False
config = {
@@ -183,9 +228,13 @@ class LightX2VInferenceConfig:
"video_length": video_length,
"fps": fps,
"video_duration": duration,
"adaptive_resize": adaptive_resize,
"resize_mode": resize_mode,
"use_31_block": use_31_block,
"attention_type": attention_type,
"use_tiny_vae": use_tiny_vae,
}
if "seko" in [model_cls]:
config["prev_frame_length"] = prev_frame_length
if denoising_steps and denoising_steps.strip():
try:
@@ -257,35 +306,36 @@ class LightX2VQuantization:
if not quant_backends:
quant_backends = ["none"]
supported_quant_schemes = ["bf16", "fp8", "int8"]
return {
"required": {
"dit_precision": (
["bf16", "int8", "fp8"],
{"default": "bf16", "tooltip": "DIT model quantization precision"},
),
"t5_precision": (
["bf16", "int8", "fp8"],
{"default": "bf16", "tooltip": "T5 encoder quantization precision"},
),
"clip_precision": (
["fp16", "int8", "fp8"],
{
"default": "fp16",
"tooltip": "CLIP encoder quantization precision",
},
),
"quant_backend": (
"quant_op": (
quant_backends,
{
"default": quant_backends[0],
"tooltip": "Quantization computation backend",
},
),
"sensitive_layers_precision": (
["fp32", "bf16"],
"dit_quant_scheme": (
supported_quant_schemes,
{
"default": "fp32",
"tooltip": "Sensitive layers (normalization and embedding) precision",
"default": supported_quant_schemes[0],
"tooltip": "DIT model quantization precision",
},
),
"t5_quant_scheme": (
supported_quant_schemes,
{
"default": supported_quant_schemes[0],
"tooltip": "T5 encoder quantization precision",
},
),
"clip_quant_scheme": (
supported_quant_schemes,
{
"default": supported_quant_schemes[0],
"tooltip": "CLIP encoder quantization precision",
},
),
}
@@ -298,19 +348,17 @@ class LightX2VQuantization:
def create_config(
self,
dit_precision,
t5_precision,
clip_precision,
quant_backend,
sensitive_layers_precision,
quant_op,
dit_quant_scheme,
t5_quant_scheme,
clip_quant_scheme,
):
"""Create quantization configuration."""
config = {
"dit_precision": dit_precision,
"t5_precision": t5_precision,
"clip_precision": clip_precision,
"quant_backend": quant_backend,
"sensitive_layers_precision": sensitive_layers_precision,
"dit_quant_scheme": dit_quant_scheme,
"t5_quant_scheme": t5_quant_scheme,
"clip_quant_scheme": clip_quant_scheme,
"quant_op": quant_op,
}
return (config,)
@@ -320,16 +368,6 @@ class LightX2VMemoryOptimization:
@classmethod
def INPUT_TYPES(cls):
available_attn = get_available_attn_ops()
attn_types = []
for op_name, is_available in available_attn:
if is_available:
attn_types.append(op_name)
if "torch_sdpa" not in attn_types:
attn_types.append("torch_sdpa")
return {
"required": {
"optimization_level": (
@@ -338,14 +376,9 @@ class LightX2VMemoryOptimization:
"default": "none",
"tooltip": "Memory optimization level, higher levels save more memory but may affect speed",
},
),
"attention_type": (
attn_types,
{"default": attn_types[0], "tooltip": "Attention mechanism type"},
),
)
},
"optional": {
# GPU optimization
"enable_rotary_chunk": (
"BOOLEAN",
{"default": False, "tooltip": "Enable rotary encoding chunking"},
@@ -358,20 +391,42 @@ class LightX2VMemoryOptimization:
"BOOLEAN",
{"default": False, "tooltip": "Clean CUDA cache promptly"},
),
# CPU offloading
"enable_cpu_offload": (
"cpu_offload": (
"BOOLEAN",
{"default": False, "tooltip": "Enable CPU offloading"},
{"default": True, "tooltip": "Enable CPU offloading"},
),
"offload_granularity": (
["block", "phase"],
{"default": "phase", "tooltip": "Offload granularity"},
["block", "phase", "model"],
{"default": "block", "tooltip": "Offload granularity"},
),
"offload_ratio": (
"FLOAT",
{"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.1},
),
# Module management
"t5_cpu_offload": (
"BOOLEAN",
{"default": True, "tooltip": "Enable T5 CPU offloading"},
),
"t5_offload_granularity": (
["model", "block"],
{"default": "model", "tooltip": "T5 offload granularity"},
),
"audio_encoder_cpu_offload": (
"BOOLEAN",
{"default": True, "tooltip": "Enable audio encoder CPU offloading"},
),
"audio_adapter_cpu_offload": (
"BOOLEAN",
{"default": True, "tooltip": "Enable audio adapter CPU offloading"},
),
"vae_cpu_offload": (
"BOOLEAN",
{"default": True, "tooltip": "Enable VAE CPU offloading"},
),
"use_tiling_vae": (
"BOOLEAN",
{"default": True, "tooltip": "Enable VAE tiling inference"},
),
"lazy_load": (
"BOOLEAN",
{"default": False, "tooltip": "Lazy load model"},
@@ -391,66 +446,39 @@ class LightX2VMemoryOptimization:
def create_config(
self,
optimization_level,
attention_type,
enable_rotary_chunk=False,
rotary_chunk_size=100,
clean_cuda_cache=False,
enable_cpu_offload=False,
cpu_offload=False,
offload_granularity="phase",
offload_ratio=1.0,
t5_cpu_offload=True,
t5_offload_granularity="model",
audio_encoder_cpu_offload=False,
audio_adapter_cpu_offload=False,
vae_cpu_offload=False,
lazy_load=False,
unload_after_inference=False,
):
config = {
"optimization_level": optimization_level,
"attention_type": attention_type,
"enable_rotary_chunk": enable_rotary_chunk,
"rotary_chunk_size": rotary_chunk_size,
"clean_cuda_cache": clean_cuda_cache,
"enable_cpu_offload": enable_cpu_offload,
"cpu_offload": cpu_offload,
"offload_granularity": offload_granularity,
"offload_ratio": offload_ratio,
"t5_cpu_offload": t5_cpu_offload,
"t5_offload_granularity": t5_offload_granularity,
"audio_encoder_cpu_offload": audio_encoder_cpu_offload,
"audio_adapter_cpu_offload": audio_adapter_cpu_offload,
"vae_cpu_offload": vae_cpu_offload,
"lazy_load": lazy_load,
"unload_after_inference": unload_after_inference,
}
return (config,)
class LightX2VLightweightVAE:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"use_tiny_vae": (
"BOOLEAN",
{
"default": False,
"tooltip": "Use lightweight VAE to accelerate decoding",
},
),
"use_tiling_vae": (
"BOOLEAN",
{
"default": False,
"tooltip": "Use VAE tiling inference to reduce VRAM usage",
},
),
}
}
RETURN_TYPES = ("VAE_CONFIG",)
RETURN_NAMES = ("vae_config",)
FUNCTION = "create_config"
CATEGORY = "LightX2V/Config"
def create_config(self, use_tiny_vae, use_tiling_vae):
config = {
"use_tiny_vae": use_tiny_vae,
"use_tiling_vae": use_tiling_vae,
}
return (config,)
class LightX2VLoRALoader:
@classmethod
def INPUT_TYPES(cls):
@@ -530,7 +558,6 @@ class LightX2VConfigCombiner:
"MEMORY_CONFIG",
{"tooltip": "Memory optimization configuration"},
),
"vae_config": ("VAE_CONFIG", {"tooltip": "VAE configuration"}),
"lora_chain": ("LORA_CHAIN", {"tooltip": "LoRA chain configuration"}),
},
}
@@ -546,27 +573,25 @@ class LightX2VConfigCombiner:
teacache_config=None,
quantization_config=None,
memory_config=None,
vae_config=None,
lora_chain=None,
):
configs = {
"inference": inference_config,
}
if teacache_config:
configs["teacache"] = teacache_config
if quantization_config:
configs["quantization"] = quantization_config
if memory_config:
configs["memory"] = memory_config
if vae_config:
configs["vae"] = vae_config
config = self.config_manager.build_final_config(configs)
if lora_chain:
config.lora_configs = lora_chain
logging.info("lightx2v config: " + json.dumps(config, indent=2, ensure_ascii=False))
return (config,)
@@ -636,6 +661,8 @@ class LightX2VModularInference:
os.environ["ENABLE_GRAPH_MODE"] = "false"
if "ENABLE_PROFILING_DEBUG" not in os.environ:
os.environ["ENABLE_PROFILING_DEBUG"] = "true"
if "SENSITIVE_LAYER_DTYPE" not in os.environ:
os.environ["SENSITIVE_LAYER_DTYPE"] = "FP32"
config = combined_config
@@ -658,16 +685,8 @@ class LightX2VModularInference:
temp_files.append(tmp.name)
logging.info(f"Image saved to {tmp.name}")
if (
audio is not None
and hasattr(config, "model_cls")
and "audio" in config.model_cls
):
if (
isinstance(audio, dict)
and "waveform" in audio
and "sample_rate" in audio
):
if audio is not None and hasattr(config, "model_cls") and "audio" in config.model_cls:
if isinstance(audio, dict) and "waveform" in audio and "sample_rate" in audio:
waveform = audio["waveform"]
sample_rate = audio["sample_rate"]
@@ -694,9 +713,7 @@ class LightX2VModularInference:
import wave
with wave.open(tmp.name, "wb") as wav_file:
wav_file.setnchannels(
1 if waveform.ndim == 1 else waveform.shape[-1]
)
wav_file.setnchannels(1 if waveform.ndim == 1 else waveform.shape[-1])
wav_file.setsampwidth(2) # 16-bit
wav_file.setframerate(sample_rate)
if waveform.dtype != np.int16:
@@ -716,11 +733,7 @@ class LightX2VModularInference:
logging.info(f"Audio saved to {tmp.name}")
config_hash = self._get_config_hash(config)
needs_reinit = (
self._current_runner is None
or self._current_config_hash != config_hash
or getattr(config, "lazy_load", False)
)
needs_reinit = self._current_runner is None or self._current_config_hash != config_hash or getattr(config, "lazy_load", False)
if needs_reinit:
if self._current_runner is not None:
@@ -772,7 +785,6 @@ NODE_CLASS_MAPPINGS = {
"LightX2VTeaCache": LightX2VTeaCache,
"LightX2VQuantization": LightX2VQuantization,
"LightX2VMemoryOptimization": LightX2VMemoryOptimization,
"LightX2VLightweightVAE": LightX2VLightweightVAE,
"LightX2VLoRALoader": LightX2VLoRALoader,
"LightX2VConfigCombiner": LightX2VConfigCombiner,
"LightX2VModularInference": LightX2VModularInference,
@@ -783,7 +795,6 @@ NODE_DISPLAY_NAME_MAPPINGS = {
"LightX2VTeaCache": "LightX2V TeaCache",
"LightX2VQuantization": "LightX2V Quantization",
"LightX2VMemoryOptimization": "LightX2V Memory Optimization",
"LightX2VLightweightVAE": "LightX2V Lightweight VAE",
"LightX2VLoRALoader": "LightX2V LoRA Loader",
"LightX2VConfigCombiner": "LightX2V Config Combiner",
"LightX2VModularInference": "LightX2V Modular Inference",