ModelTC-ComfyUI-Lightx2vWrapper/nodes.py

"""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", "wan2.1_audio", "wan2.1_distill", "hunyuan"], {"default": "wan2.1", "tooltip": "Model type"}),
                "model_path": ("STRING", {"default": "", "tooltip": "Model path"}),
                "task": (["t2v", "i2v"], {"default": "t2v", "tooltip": "Task type: text-to-video or image-to-video"}),
                "infer_steps": ("INT", {"default": 40, "min": 1, "max": 100, "tooltip": "Inference steps"}),
                "seed": ("INT", {"default": 42, "min": -1, "max": 2**32 - 1, "tooltip": "Random seed, -1 for random"}),
                "cfg_scale": ("FLOAT", {"default": 5.0, "min": 1.0, "max": 10.0, "step": 0.1, "tooltip": "CFG guidance strength"}),
                "sample_shift": ("INT", {"default": 5, "min": 0, "max": 10, "tooltip": "Sample shift"}),
                "height": ("INT", {"default": 480, "min": 64, "max": 2048, "step": 8, "tooltip": "Video height"}),
                "width": ("INT", {"default": 832, "min": 64, "max": 2048, "step": 8, "tooltip": "Video width"}),
                "video_length": ("INT", {"default": 81, "min": 16, "max": 120, "tooltip": "Video frame count"}),
                "fps": ("INT", {"default": 16, "min": 8, "max": 30, "tooltip": "Model output frame rate (cannot be changed)"}),
            }
        }

    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": "Enable TeaCache feature caching"}),
                "threshold": (
                    "FLOAT",
                    {
                        "default": 0.26,
                        "min": 0.0,
                        "max": 1.0,
                        "step": 0.01,
                        "tooltip": "Cache threshold, lower values provide more speedup: 0.1 ~2x speedup, 0.2 ~3x speedup",
                    },
                ),
                "use_ret_steps": ("BOOLEAN", {"default": False, "tooltip": "Only cache key steps to balance quality and speed"}),
            }
        }

    RETURN_TYPES = ("TEACACHE_CONFIG",)
    RETURN_NAMES = ("teacache_config",)
    FUNCTION = "create_config"
    CATEGORY = "LightX2V/Config"

    def create_config(self, enable, threshold, use_ret_steps):
        """Create TeaCache configuration."""
        config = {
            "enable": enable,
            "threshold": threshold,
            "use_ret_steps": use_ret_steps,
        }
        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 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_backends, {"default": quant_backends[0], "tooltip": "Quantization computation backend"}),
                "sensitive_layers_precision": (
                    ["fp32", "bf16"],
                    {"default": "fp32", "tooltip": "Sensitive layers (normalization and embedding) precision"},
                ),
            }
        }

    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": "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"}),
                "rotary_chunk_size": ("INT", {"default": 100, "min": 100, "max": 10000, "step": 100}),
                "clean_cuda_cache": ("BOOLEAN", {"default": False, "tooltip": "Clean CUDA cache promptly"}),
                # CPU offloading
                "enable_cpu_offload": ("BOOLEAN", {"default": False, "tooltip": "Enable CPU offloading"}),
                "offload_granularity": (["block", "phase"], {"default": "phase", "tooltip": "Offload granularity"}),
                "offload_ratio": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.1}),
                # Module management
                "lazy_load": ("BOOLEAN", {"default": False, "tooltip": "Lazy load model"}),
                "unload_after_inference": ("BOOLEAN", {"default": False, "tooltip": "Unload modules after inference"}),
            },
        }

    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": "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):
        """Create VAE configuration."""
        config = {
            "use_tiny_vae": use_tiny_vae,
            "use_tiling_vae": use_tiling_vae,
        }
        return (config,)


class LightX2VLoRALoader:
    """LoRA loader node that can be chained."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "lora_path": ("STRING", {"default": "", "tooltip": "Path to the LoRA file"}),
                "strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.1, "tooltip": "LoRA strength"}),
            },
            "optional": {
                "lora_chain": ("LORA_CHAIN", {"tooltip": "Previous LoRA chain to append to"}),
            },
        }

    RETURN_TYPES = ("LORA_CHAIN",)
    RETURN_NAMES = ("lora_chain",)
    FUNCTION = "load_lora"
    CATEGORY = "LightX2V/LoRA"

    def load_lora(self, lora_path, strength, lora_chain=None):
        """Load LoRA and chain with previous LoRAs."""
        # Initialize or extend the LoRA chain
        if lora_chain is None:
            lora_chain = []
        else:
            # Make a copy to avoid modifying the input
            lora_chain = lora_chain.copy()

        # Add new LoRA configuration
        if lora_path and lora_path.strip():
            lora_config = {"path": lora_path.strip(), "strength": strength}
            lora_chain.append(lora_config)

        return (lora_chain,)


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": "Basic inference configuration"}),
                "prompt": ("STRING", {"multiline": True, "default": "", "tooltip": "Generation prompt"}),
                "negative_prompt": ("STRING", {"multiline": True, "default": "", "tooltip": "Negative prompt"}),
            },
            "optional": {
                "image": ("IMAGE", {"tooltip": "Input image for i2v task"}),
                "teacache_config": ("TEACACHE_CONFIG", {"tooltip": "TeaCache configuration"}),
                "quantization_config": ("QUANT_CONFIG", {"tooltip": "Quantization configuration"}),
                "memory_config": ("MEMORY_CONFIG", {"tooltip": "Memory optimization configuration"}),
                "vae_config": ("VAE_CONFIG", {"tooltip": "VAE configuration"}),
                "lora_chain": ("LORA_CHAIN", {"tooltip": "LoRA chain configuration"}),
            },
        }

    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,
        lora_chain=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 LoRA configurations if provided
        if lora_chain:
            config.lora_configs = lora_chain

        # 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,
    "LightX2VLoRALoader": LightX2VLoRALoader,
    "LightX2VModularInference": LightX2VModularInference,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "LightX2VInferenceConfig": "LightX2V Inference Config",
    "LightX2VTeaCache": "LightX2V TeaCache",
    "LightX2VQuantization": "LightX2V Quantization",
    "LightX2VMemoryOptimization": "LightX2V Memory Optimization",
    "LightX2VLightweightVAE": "LightX2V Lightweight VAE",
    "LightX2VLoRALoader": "LightX2V LoRA Loader",
    "LightX2VModularInference": "LightX2V Modular Inference",
}