427 lines
16 KiB
Python
427 lines
16 KiB
Python
"""Modular ComfyUI nodes for LightX2V without presets."""
|
|
|
|
import asyncio
|
|
import gc
|
|
import logging
|
|
import os
|
|
import tempfile
|
|
from typing import Any, Dict
|
|
|
|
import numpy as np
|
|
import torch
|
|
from comfy.utils import ProgressBar
|
|
from PIL import Image
|
|
|
|
from .bridge import ModularConfigManager, get_available_attn_ops, get_available_quant_ops
|
|
from .lightx2v.lightx2v.infer import init_runner
|
|
|
|
|
|
class LightX2VInferenceConfig:
|
|
"""Basic inference configuration node."""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model_cls": (["wan2.1", "hunyuan"], {"default": "wan2.1", "tooltip": "模型类型"}),
|
|
"model_path": ("STRING", {"default": "", "tooltip": "模型路径"}),
|
|
"task": (["t2v", "i2v"], {"default": "t2v", "tooltip": "任务类型：文本到视频或图像到视频"}),
|
|
"infer_steps": ("INT", {"default": 40, "min": 1, "max": 100, "tooltip": "推理步数"}),
|
|
"seed": ("INT", {"default": 42, "min": -1, "max": 2**32 - 1, "tooltip": "随机种子，-1为随机"}),
|
|
"cfg_scale": ("FLOAT", {"default": 5.0, "min": 1.0, "max": 10.0, "step": 0.1, "tooltip": "CFG引导强度"}),
|
|
"sample_shift": ("INT", {"default": 5, "min": 0, "max": 10, "tooltip": "采样偏移"}),
|
|
"height": ("INT", {"default": 480, "min": 64, "max": 2048, "step": 8, "tooltip": "视频高度"}),
|
|
"width": ("INT", {"default": 832, "min": 64, "max": 2048, "step": 8, "tooltip": "视频宽度"}),
|
|
"video_length": ("INT", {"default": 81, "min": 16, "max": 120, "tooltip": "视频帧数"}),
|
|
"fps": ("INT", {"default": 16, "min": 8, "max": 30, "tooltip": "每秒帧数"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("INFERENCE_CONFIG",)
|
|
RETURN_NAMES = ("inference_config",)
|
|
FUNCTION = "create_config"
|
|
CATEGORY = "LightX2V/Config"
|
|
|
|
def create_config(self, model_cls, model_path, task, infer_steps, seed, cfg_scale, sample_shift, height, width, video_length, fps):
|
|
"""Create basic inference configuration."""
|
|
config = {
|
|
"model_cls": model_cls,
|
|
"model_path": model_path,
|
|
"task": task,
|
|
"infer_steps": infer_steps,
|
|
"seed": seed if seed != -1 else np.random.randint(0, 2**32 - 1),
|
|
"cfg_scale": cfg_scale,
|
|
"sample_shift": sample_shift,
|
|
"height": height,
|
|
"width": width,
|
|
"video_length": video_length,
|
|
"fps": fps,
|
|
}
|
|
return (config,)
|
|
|
|
|
|
class LightX2VTeaCache:
|
|
"""TeaCache configuration node."""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"enable": ("BOOLEAN", {"default": False, "tooltip": "启用TeaCache特征缓存"}),
|
|
"threshold": (
|
|
"FLOAT",
|
|
{"default": 0.26, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "缓存阈值，越低加速越多：0.1约2倍加速，0.2约3倍加速"},
|
|
),
|
|
"cache_key_steps_only": ("BOOLEAN", {"default": False, "tooltip": "只缓存关键步骤以平衡质量和速度"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("TEACACHE_CONFIG",)
|
|
RETURN_NAMES = ("teacache_config",)
|
|
FUNCTION = "create_config"
|
|
CATEGORY = "LightX2V/Config"
|
|
|
|
def create_config(self, enable, threshold, cache_key_steps_only):
|
|
"""Create TeaCache configuration."""
|
|
config = {
|
|
"enable": enable,
|
|
"threshold": threshold,
|
|
"cache_key_steps_only": cache_key_steps_only,
|
|
}
|
|
return (config,)
|
|
|
|
|
|
class LightX2VQuantization:
|
|
"""Quantization configuration node."""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
# Get available quantization backends
|
|
available_ops = get_available_quant_ops()
|
|
quant_backends = []
|
|
|
|
for op_name, is_available in available_ops:
|
|
if is_available:
|
|
quant_backends.append(op_name)
|
|
|
|
# Always have at least one option
|
|
if not quant_backends:
|
|
quant_backends = ["none"]
|
|
|
|
return {
|
|
"required": {
|
|
"dit_precision": (["bf16", "int8", "fp8"], {"default": "bf16", "tooltip": "DIT模型量化精度"}),
|
|
"t5_precision": (["bf16", "int8", "fp8"], {"default": "bf16", "tooltip": "T5编码器量化精度"}),
|
|
"clip_precision": (["fp16", "int8", "fp8"], {"default": "fp16", "tooltip": "CLIP编码器量化精度"}),
|
|
"quant_backend": (quant_backends, {"default": quant_backends[0], "tooltip": "量化计算后端"}),
|
|
"sensitive_layers_precision": (["fp32", "bf16"], {"default": "fp32", "tooltip": "敏感层（归一化和嵌入层）精度"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("QUANT_CONFIG",)
|
|
RETURN_NAMES = ("quantization_config",)
|
|
FUNCTION = "create_config"
|
|
CATEGORY = "LightX2V/Config"
|
|
|
|
def create_config(self, dit_precision, t5_precision, clip_precision, quant_backend, sensitive_layers_precision):
|
|
"""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,
|
|
}
|
|
return (config,)
|
|
|
|
|
|
class LightX2VMemoryOptimization:
|
|
"""Memory optimization configuration node."""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
# Get available attention types
|
|
available_attn = get_available_attn_ops()
|
|
attn_types = []
|
|
|
|
for op_name, is_available in available_attn:
|
|
if is_available:
|
|
attn_types.append(op_name)
|
|
|
|
# Always include fallback
|
|
if "torch_sdpa" not in attn_types:
|
|
attn_types.append("torch_sdpa")
|
|
|
|
return {
|
|
"required": {
|
|
"optimization_level": (
|
|
["none", "low", "medium", "high", "extreme"],
|
|
{"default": "none", "tooltip": "内存优化级别，越高越省内存但可能影响速度"},
|
|
),
|
|
"attention_type": (attn_types, {"default": attn_types[0], "tooltip": "注意力机制类型"}),
|
|
},
|
|
"optional": {
|
|
# GPU optimization
|
|
"enable_rotary_chunk": ("BOOLEAN", {"default": False, "tooltip": "启用旋转编码分块"}),
|
|
"rotary_chunk_size": ("INT", {"default": 100, "min": 100, "max": 10000, "step": 100}),
|
|
"clean_cuda_cache": ("BOOLEAN", {"default": False, "tooltip": "及时清理CUDA缓存"}),
|
|
# CPU offloading
|
|
"enable_cpu_offload": ("BOOLEAN", {"default": False, "tooltip": "启用CPU卸载"}),
|
|
"offload_granularity": (["block", "phase"], {"default": "phase", "tooltip": "卸载粒度"}),
|
|
"offload_ratio": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.1}),
|
|
# Module management
|
|
"lazy_load": ("BOOLEAN", {"default": False, "tooltip": "延迟加载模型"}),
|
|
"unload_after_inference": ("BOOLEAN", {"default": False, "tooltip": "推理后卸载模块"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("MEMORY_CONFIG",)
|
|
RETURN_NAMES = ("memory_config",)
|
|
FUNCTION = "create_config"
|
|
CATEGORY = "LightX2V/Config"
|
|
|
|
def create_config(
|
|
self,
|
|
optimization_level,
|
|
attention_type,
|
|
enable_rotary_chunk=False,
|
|
rotary_chunk_size=100,
|
|
clean_cuda_cache=False,
|
|
enable_cpu_offload=False,
|
|
offload_granularity="phase",
|
|
offload_ratio=1.0,
|
|
lazy_load=False,
|
|
unload_after_inference=False,
|
|
):
|
|
"""Create memory optimization configuration."""
|
|
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,
|
|
"offload_granularity": offload_granularity,
|
|
"offload_ratio": offload_ratio,
|
|
"lazy_load": lazy_load,
|
|
"unload_after_inference": unload_after_inference,
|
|
}
|
|
return (config,)
|
|
|
|
|
|
class LightX2VLightweightVAE:
|
|
"""Lightweight VAE configuration node."""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"use_tiny_vae": ("BOOLEAN", {"default": False, "tooltip": "使用轻量级VAE加速解码"}),
|
|
"use_tiling_vae": ("BOOLEAN", {"default": False, "tooltip": "使用VAE分块推理减少显存"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("VAE_CONFIG",)
|
|
RETURN_NAMES = ("vae_config",)
|
|
FUNCTION = "create_config"
|
|
CATEGORY = "LightX2V/Config"
|
|
|
|
def create_config(self, use_tiny_vae, use_tiling_vae):
|
|
"""Create VAE configuration."""
|
|
config = {
|
|
"use_tiny_vae": use_tiny_vae,
|
|
"use_tiling_vae": use_tiling_vae,
|
|
}
|
|
return (config,)
|
|
|
|
|
|
class LightX2VModularInference:
|
|
"""Modular inference node that combines all configurations."""
|
|
|
|
def __init__(self):
|
|
self.config_manager = ModularConfigManager()
|
|
self._current_runner = None
|
|
self._current_config_hash = None
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"inference_config": ("INFERENCE_CONFIG", {"tooltip": "基础推理配置"}),
|
|
"prompt": ("STRING", {"multiline": True, "default": "", "tooltip": "生成提示词"}),
|
|
"negative_prompt": ("STRING", {"multiline": True, "default": "", "tooltip": "负面提示词"}),
|
|
},
|
|
"optional": {
|
|
"image": ("IMAGE", {"tooltip": "i2v任务的输入图像"}),
|
|
"teacache_config": ("TEACACHE_CONFIG", {"tooltip": "TeaCache配置"}),
|
|
"quantization_config": ("QUANT_CONFIG", {"tooltip": "量化配置"}),
|
|
"memory_config": ("MEMORY_CONFIG", {"tooltip": "内存优化配置"}),
|
|
"vae_config": ("VAE_CONFIG", {"tooltip": "VAE配置"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("images",)
|
|
FUNCTION = "generate"
|
|
CATEGORY = "LightX2V/Inference"
|
|
|
|
def _get_config_hash(self, configs: Dict[str, Any]) -> str:
|
|
"""Generate a hash for configuration to detect changes."""
|
|
import hashlib
|
|
import json
|
|
|
|
# Only hash model-related configs
|
|
relevant_configs = {
|
|
"model_cls": configs.get("inference", {}).get("model_cls"),
|
|
"model_path": configs.get("inference", {}).get("model_path"),
|
|
"quantization": configs.get("quantization"),
|
|
"memory_lazy_load": configs.get("memory", {}).get("lazy_load"),
|
|
}
|
|
|
|
config_str = json.dumps(relevant_configs, sort_keys=True)
|
|
return hashlib.md5(config_str.encode()).hexdigest()
|
|
|
|
def generate(
|
|
self,
|
|
inference_config,
|
|
prompt,
|
|
negative_prompt,
|
|
image=None,
|
|
teacache_config=None,
|
|
quantization_config=None,
|
|
memory_config=None,
|
|
vae_config=None,
|
|
**kwargs,
|
|
):
|
|
"""Generate video using modular configuration."""
|
|
|
|
# Set environment variables
|
|
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
|
if "DTYPE" not in os.environ:
|
|
os.environ["DTYPE"] = "BF16"
|
|
if "ENABLE_GRAPH_MODE" not in os.environ:
|
|
os.environ["ENABLE_GRAPH_MODE"] = "false"
|
|
if "ENABLE_PROFILING_DEBUG" not in os.environ:
|
|
os.environ["ENABLE_PROFILING_DEBUG"] = "true"
|
|
|
|
# Collect all configurations
|
|
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
|
|
|
|
# Build final configuration
|
|
config = self.config_manager.build_final_config(configs)
|
|
|
|
# Add prompt and negative prompt
|
|
config.prompt = prompt
|
|
config.negative_prompt = negative_prompt
|
|
|
|
# Check if task requires image
|
|
if config.task == "i2v" and image is None:
|
|
raise ValueError("i2v task requires input image")
|
|
|
|
temp_files = []
|
|
|
|
try:
|
|
# Handle image input for i2v
|
|
if config.task == "i2v" and image is not None:
|
|
# Convert ComfyUI image to PIL
|
|
image_np = (image[0].cpu().numpy() * 255).astype(np.uint8)
|
|
pil_image = Image.fromarray(image_np)
|
|
|
|
# Save to temporary file
|
|
with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
|
|
pil_image.save(tmp.name)
|
|
config.image_path = tmp.name
|
|
temp_files.append(tmp.name)
|
|
|
|
# Check if we need to reinitialize runner
|
|
config_hash = self._get_config_hash(configs)
|
|
needs_reinit = (
|
|
self._current_runner is None or self._current_config_hash != config_hash or configs.get("memory", {}).get("lazy_load", False)
|
|
)
|
|
|
|
if needs_reinit:
|
|
# Clear old runner
|
|
if self._current_runner is not None:
|
|
del self._current_runner
|
|
torch.cuda.empty_cache()
|
|
gc.collect()
|
|
|
|
# Initialize new runner
|
|
self._current_runner = init_runner(config)
|
|
self._current_config_hash = config_hash
|
|
else:
|
|
# Update config for existing runner
|
|
self._current_runner.config = config
|
|
|
|
# Set up progress callback
|
|
total_steps = config.get("infer_steps", 40)
|
|
progress = ProgressBar(total_steps)
|
|
|
|
def update_progress(current_step, total):
|
|
progress.update_absolute(current_step)
|
|
|
|
self._current_runner.set_progress_callback(update_progress)
|
|
|
|
# Run inference
|
|
images = asyncio.run(self._current_runner.run_pipeline(save_video=False))
|
|
|
|
# Clean up if requested
|
|
if configs.get("memory", {}).get("unload_after_inference", False):
|
|
del self._current_runner
|
|
self._current_runner = None
|
|
self._current_config_hash = None
|
|
|
|
torch.cuda.empty_cache()
|
|
gc.collect()
|
|
|
|
# Convert output to ComfyUI format
|
|
images = (images + 1) / 2
|
|
images = images.squeeze(0).permute(1, 2, 3, 0).cpu()
|
|
images = torch.clamp(images, 0, 1)
|
|
|
|
return (images,)
|
|
|
|
except Exception as e:
|
|
logging.error(f"Error during inference: {e}")
|
|
raise
|
|
|
|
finally:
|
|
# Clean up temporary files
|
|
for temp_file in temp_files:
|
|
if os.path.exists(temp_file):
|
|
try:
|
|
os.unlink(temp_file)
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
# Node mappings
|
|
NODE_CLASS_MAPPINGS = {
|
|
"LightX2VInferenceConfig": LightX2VInferenceConfig,
|
|
"LightX2VTeaCache": LightX2VTeaCache,
|
|
"LightX2VQuantization": LightX2VQuantization,
|
|
"LightX2VMemoryOptimization": LightX2VMemoryOptimization,
|
|
"LightX2VLightweightVAE": LightX2VLightweightVAE,
|
|
"LightX2VModularInference": LightX2VModularInference,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"LightX2VInferenceConfig": "LightX2V 推理配置",
|
|
"LightX2VTeaCache": "LightX2V TeaCache缓存",
|
|
"LightX2VQuantization": "LightX2V 低精度量化",
|
|
"LightX2VMemoryOptimization": "LightX2V 内存优化",
|
|
"LightX2VLightweightVAE": "LightX2V 轻量VAE",
|
|
"LightX2VModularInference": "LightX2V 模块化推理",
|
|
}
|