From 1dc3a6cedcbb87d4ccda0c0e66e3b97f839f4dcd Mon Sep 17 00:00:00 2001 From: gaclove Date: Mon, 14 Jul 2025 17:45:05 +0800 Subject: [PATCH] feat: introduce UniversalLightX2V node for seamless integration between LightX2V and ComfyUI, enhancing video generation capabilities --- README.md | 2 +- lightx2v_nodes/universal_bridge.py | 326 +++++++++++++++++++++++++++++ 2 files changed, 327 insertions(+), 1 deletion(-) create mode 100644 lightx2v_nodes/universal_bridge.py diff --git a/README.md b/README.md index 52ecb97..97ea5a0 100644 --- a/README.md +++ b/README.md @@ -30,7 +30,7 @@ pip install -r lightx2v/requirements.txt # Install dependencies for lightx2v 确保你的模型目录结构如下： -``` +```text your_model_dir/ ├── config.json # 模型配置文件 ├── models_t5_umt5-xxl-enc-bf16.pth # T5文本编码器 diff --git a/lightx2v_nodes/universal_bridge.py b/lightx2v_nodes/universal_bridge.py new file mode 100644 index 0000000..1751487 --- /dev/null +++ b/lightx2v_nodes/universal_bridge.py @@ -0,0 +1,326 @@ +"""Universal bridge between LightX2V and ComfyUI for automatic adaptation.""" + +import json +import os +import torch +import numpy as np +from pathlib import Path +from typing import Any, Dict, Optional, List, Union +from PIL import Image +import tempfile +from easydict import EasyDict + +# Import LightX2V modules +from ..lightx2v.utils.set_config import set_config, get_default_config +from ..lightx2v.utils.registry_factory import RUNNER_REGISTER +from ..lightx2v.infer import init_runner + + +class LightX2VBridge: + """Universal bridge for LightX2V and ComfyUI integration.""" + + def __init__(self): + self._model_registry = None + self._config_registry = None + self._runners = {} # Cache initialized runners + + @property + def model_registry(self): + """Lazy load model registry.""" + if self._model_registry is None: + self._model_registry = self._discover_models() + return self._model_registry + + @property + def config_registry(self): + """Lazy load config registry.""" + if self._config_registry is None: + self._config_registry = self._discover_configs() + return self._config_registry + + def _discover_models(self) -> List[str]: + """Discover available model classes from LightX2V.""" + # These are the known model classes from the infer.py + return ["wan2.1", "hunyuan", "wan2.1_distill", "wan2.1_causvid", "wan2.1_skyreels_v2_df", "cogvideox", "wan2.1_audio"] + + def _discover_configs(self) -> Dict[str, Path]: + """Discover all available config files.""" + configs = {} + base_dir = Path(__file__).parent.parent / "lightx2v" / "configs" + + if base_dir.exists(): + for config_file in base_dir.rglob("*.json"): + # Create a descriptive key + relative_path = config_file.relative_to(base_dir) + key = str(relative_path).replace("/", "_").replace(".json", "") + configs[key] = config_file + + return configs + + def load_config(self, config_path: Union[str, Path]) -> Dict: + """Load a config file.""" + with open(config_path, "r") as f: + return json.load(f) + + def get_runner(self, config: EasyDict): + """Get or create a runner for the given config.""" + # Create a unique key for this configuration + key = f"{config.model_cls}_{config.model_path}_{hash(str(config))}" + + if key not in self._runners: + self._runners[key] = init_runner(config) + + return self._runners[key] + + def clear_cache(self): + """Clear cached runners to free memory.""" + self._runners.clear() + torch.cuda.empty_cache() + + +class InputOutputConverter: + """Convert between ComfyUI and LightX2V formats.""" + + @staticmethod + def comfy_image_to_pil(comfy_image: torch.Tensor) -> Image.Image: + """Convert ComfyUI image tensor to PIL Image. + + ComfyUI format: [B, H, W, C] float32 0-1 + """ + # Take first image from batch + image_np = (comfy_image[0].cpu().numpy() * 255).astype(np.uint8) + return Image.fromarray(image_np) + + @staticmethod + def comfy_audio_to_path(comfy_audio: Dict) -> str: + """Convert ComfyUI audio to file path.""" + # ComfyUI audio format: {"waveform": tensor, "sample_rate": int} + waveform = comfy_audio["waveform"] + sample_rate = comfy_audio["sample_rate"] + + # Save to temporary file + import torchaudio + + with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp: + torchaudio.save(tmp.name, waveform, sample_rate) + return tmp.name + + @staticmethod + def video_to_latent(video_path: str) -> Dict[str, Any]: + """Convert video path to ComfyUI latent format.""" + # For now, return the path wrapped in latent format + # This allows downstream nodes to handle the video + return {"samples": video_path, "type": "video", "format": "path"} + + @staticmethod + def tensor_to_latent(video_tensor: torch.Tensor) -> Dict[str, Any]: + """Convert video tensor to ComfyUI latent format.""" + return {"samples": video_tensor, "type": "video", "format": "tensor"} + + +class ConfigManager: + """Manage configuration merging and parameter mapping.""" + + # Mapping from ComfyUI parameter names to LightX2V config keys + PARAM_MAPPING = { + "steps": "infer_steps", + "cfg_scale": "sample_guide_scale", + "seed": "seed", + "height": "target_height", + "width": "target_width", + "video_length": "target_video_length", + } + + @classmethod + def create_config(cls, model_cls: str, model_path: str, task: str, base_config: Dict, overrides: Dict) -> EasyDict: + """Create a complete config for LightX2V runner.""" + # Start with default config + config = get_default_config() + + # Add required fields + config.update( + { + "model_cls": model_cls, + "model_path": model_path, + "task": task, + "mode": "infer", + } + ) + + # Apply base config from file + config.update(base_config) + + # Apply ComfyUI overrides + for comfy_key, value in overrides.items(): + if value is not None and value != -1: # -1 means use default + if comfy_key in cls.PARAM_MAPPING: + lightx2v_key = cls.PARAM_MAPPING[comfy_key] + config[lightx2v_key] = value + else: + # Direct mapping for unknown parameters + config[comfy_key] = value + + return EasyDict(config) + + @classmethod + def parse_custom_config(cls, custom_config_str: str) -> Dict: + """Parse custom config string.""" + if not custom_config_str.strip(): + return {} + + try: + return json.loads(custom_config_str) + except json.JSONDecodeError as e: + print(f"Failed to parse custom config: {e}") + return {} + + +class UniversalLightX2VNode: + """Universal ComfyUI node for all LightX2V models.""" + + def __init__(self): + self.bridge = LightX2VBridge() + self.converter = InputOutputConverter() + + @classmethod + def INPUT_TYPES(cls): + """Define inputs for the node.""" + bridge = LightX2VBridge() + + # Get available models and configs + model_choices = bridge.model_registry + config_choices = ["custom"] + list(bridge.config_registry.keys()) + + return { + "required": { + "model_cls": (model_choices, {"tooltip": "Model class to use"}), + "model_path": ("STRING", {"default": "", "tooltip": "Path to model weights"}), + "task": (["t2v", "i2v"], {"default": "t2v", "tooltip": "Task type: text-to-video or image-to-video"}), + "config_preset": (config_choices, {"default": "custom", "tooltip": "Configuration preset or 'custom'"}), + "prompt": ("STRING", {"multiline": True, "default": "", "tooltip": "Text prompt for generation"}), + "negative_prompt": ("STRING", {"multiline": True, "default": "", "tooltip": "Negative prompt"}), + "seed": ("INT", {"default": -1, "min": -1, "max": 2**32 - 1, "tooltip": "Random seed (-1 for random)"}), + "steps": ("INT", {"default": -1, "min": -1, "max": 200, "tooltip": "Inference steps (-1 for default)"}), + "cfg_scale": ("FLOAT", {"default": -1, "min": -1, "max": 30, "step": 0.1, "tooltip": "Guidance scale (-1 for default)"}), + }, + "optional": { + "image": ("IMAGE", {"tooltip": "Input image for i2v task"}), + "audio": ("AUDIO", {"tooltip": "Input audio for audio-driven generation"}), + "custom_config": ("STRING", {"multiline": True, "default": "{}", "tooltip": "Custom configuration in JSON format"}), + "lora_path": ("STRING", {"default": "", "tooltip": "Path to LoRA weights"}), + "strength_model": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.1, "tooltip": "Model strength for LoRA"}), + "height": ("INT", {"default": -1, "min": -1, "max": 2048, "step": 8, "tooltip": "Video height (-1 for default)"}), + "width": ("INT", {"default": -1, "min": -1, "max": 2048, "step": 8, "tooltip": "Video width (-1 for default)"}), + "video_length": ("INT", {"default": -1, "min": -1, "max": 300, "tooltip": "Number of frames (-1 for default)"}), + }, + } + + RETURN_TYPES = ("LATENT",) + RETURN_NAMES = ("video",) + FUNCTION = "generate" + CATEGORY = "LightX2V/Universal" + + def generate( + self, + model_cls, + model_path, + task, + config_preset, + prompt, + negative_prompt, + seed, + steps, + cfg_scale, + image=None, + audio=None, + custom_config="{}", + lora_path="", + strength_model=1.0, + height=-1, + width=-1, + video_length=-1, + **kwargs, + ): + """Generate video using LightX2V.""" + + # Load base configuration + if config_preset == "custom": + base_config = {} + else: + config_path = self.bridge.config_registry.get(config_preset) + if config_path: + base_config = self.bridge.load_config(config_path) + else: + raise ValueError(f"Config preset '{config_preset}' not found") + + # Parse custom config + custom_cfg = ConfigManager.parse_custom_config(custom_config) + base_config.update(custom_cfg) + + # Prepare overrides + overrides = { + "seed": seed if seed != -1 else None, + "steps": steps, + "cfg_scale": cfg_scale, + "height": height, + "width": width, + "video_length": video_length, + "lora_path": lora_path if lora_path else None, + "strength_model": strength_model, + "prompt": prompt, + "negative_prompt": negative_prompt, + } + + # Handle image input for i2v + if task == "i2v" and image is not None: + # Save image to temporary file + pil_image = self.converter.comfy_image_to_pil(image) + with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp: + pil_image.save(tmp.name) + overrides["image_path"] = tmp.name + + # Handle audio input + if audio is not None and "audio" in model_cls: + audio_path = self.converter.comfy_audio_to_path(audio) + overrides["audio_path"] = audio_path + + # Create final config + config = ConfigManager.create_config(model_cls, model_path, task, base_config, overrides) + + # Create a temporary config file (required by set_config) + with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as tmp: + json.dump(base_config, tmp) + config.config_json = tmp.name + + # Get or create runner + try: + runner = self.bridge.get_runner(config) + + # Run generation + result = runner.run_pipeline() + + # Convert output + if hasattr(result, "save_video_path"): + return (self.converter.video_to_latent(result.save_video_path),) + else: + # Assume tensor output + return (self.converter.tensor_to_latent(result),) + + finally: + # Cleanup temporary files + if "image_path" in overrides and os.path.exists(overrides["image_path"]): + os.unlink(overrides["image_path"]) + if "audio_path" in overrides and os.path.exists(overrides["audio_path"]): + os.unlink(overrides["audio_path"]) + if hasattr(config, "config_json") and os.path.exists(config.config_json): + os.unlink(config.config_json) + + +# Node class mapping +NODE_CLASS_MAPPINGS = { + "UniversalLightX2V": UniversalLightX2VNode, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "UniversalLightX2V": "LightX2V Universal Generator", +}