diff --git a/.ruff.toml b/.ruff.toml index 2e4c07b..099b44d 100644 --- a/.ruff.toml +++ b/.ruff.toml @@ -1,4 +1,4 @@ line-length = 150 indent-width = 4 -extend-select = ["I"] +lint.extend-select = ["I"] diff --git a/bridge.py b/bridge.py index 5d0a2cd..2a4bc92 100644 --- a/bridge.py +++ b/bridge.py @@ -137,6 +137,7 @@ class LightX2VDefaultConfig: "seq_parallel": False, "cfg_parallel": False, "audio_sr": 16000, + "return_video": True, } diff --git a/config_builder.py b/config_builder.py new file mode 100644 index 0000000..f137ab0 --- /dev/null +++ b/config_builder.py @@ -0,0 +1,374 @@ +import hashlib +import json +import logging +import uuid +from typing import Any, Dict, List, Optional, Tuple + +import numpy as np +from easydict import EasyDict + +from .bridge import ( + ModularConfigManager, +) +from .data_models import ( + CombinedConfig, + InferenceConfig, + LoRAConfig, + MemoryOptimizationConfig, + QuantizationConfig, + TalkObject, + TalkObjectsConfig, + TeaCacheConfig, +) +from .file_handlers import ( + AudioFileHandler, + ComfyUIFileResolver, + MaskFileHandler, + TempFileManager, +) +from .model_utils import get_lora_full_path, get_model_full_path + + +class ConfigValidator: + """Validator for configuration parameters.""" + + @staticmethod + def validate_inference_config(config: InferenceConfig) -> InferenceConfig: + """Validate and adjust inference configuration.""" + # Ensure video length is valid + if config.video_length < 16: + logging.warning("Video length is too short, setting to 16") + config.video_length = 16 + + # Adjust video length to be compatible with model requirements + remainder = (config.video_length - 1) % 4 + if remainder != 0: + config.video_length = config.video_length + (4 - remainder) + + # Set FPS based on model class + if config.model_cls == "hunyuan": + config.fps = 24 + else: + config.fps = 16 + + # Handle random seed + if config.seed == -1: + config.seed = np.random.randint(0, 2**32 - 1) + + return config + + @staticmethod + def validate_dimensions(height: int, width: int) -> Tuple[int, int]: + """Validate and adjust video dimensions.""" + # Ensure dimensions are multiples of 8 + height = (height // 8) * 8 + width = (width // 8) * 8 + + # Ensure minimum dimensions + height = max(64, height) + width = max(64, width) + + # Ensure maximum dimensions + height = min(2048, height) + width = min(2048, width) + + return height, width + + +class InferenceConfigBuilder: + """Builder for inference configuration.""" + + def __init__(self): + self.validator = ConfigValidator() + + def build( + self, + model_cls: str, + model_name: str, + task: str, + infer_steps: int, + seed: int, + cfg_scale: float, + cfg_scale2: float, + sample_shift: int, + height: int, + width: int, + duration: float, + attention_type: str, + **optional_params, + ) -> InferenceConfig: + """Build inference configuration from parameters.""" + # Get model path + model_path = get_model_full_path(model_name) + + # Calculate video length from duration + fps = 24 if model_cls == "hunyuan" else 16 + video_length = int(round(duration * fps)) + + # Validate dimensions + height, width = self.validator.validate_dimensions(height, width) + + # Create base config + config = InferenceConfig( + model_cls=model_cls, + model_path=model_path, + task=task, + infer_steps=infer_steps, + seed=seed, + cfg_scale=cfg_scale, + cfg_scale2=cfg_scale2, + sample_shift=sample_shift, + height=height, + width=width, + video_length=video_length, + fps=fps, + video_duration=duration, + attention_type=attention_type, + ) + + # Handle optional parameters + self._apply_optional_params(config, optional_params) + + # Validate final config + config = self.validator.validate_inference_config(config) + + return config + + def _apply_optional_params(self, config: InferenceConfig, optional_params: Dict[str, Any]): + """Apply optional parameters to config.""" + # Handle denoising steps + if "denoising_steps" in optional_params: + steps_str = optional_params["denoising_steps"] + if steps_str and steps_str.strip(): + try: + steps_list = [int(s.strip()) for s in steps_str.split(",")] + config.denoising_step_list = steps_list + config.infer_steps = len(steps_list) + except ValueError: + logging.warning(f"Invalid denoising steps: {steps_str}") + + # Handle other optional params + for param in ["resize_mode", "fixed_area", "segment_length", "prev_frame_length", "use_tiny_vae"]: + if param in optional_params: + setattr(config, param, optional_params[param]) + + # Special handling for seko models + if "seko" in config.model_cls: + config.video_length = optional_params.get("segment_length", 81) + config.use_31_block = False + if "prev_frame_length" in optional_params: + config.prev_frame_length = optional_params["prev_frame_length"] + + +class TalkObjectConfigBuilder: + """Builder for talk object configurations.""" + + def __init__(self): + self.audio_handler = AudioFileHandler() + self.mask_handler = MaskFileHandler() + self.resolver = ComfyUIFileResolver() + + def build_from_input(self, name: str, audio: Optional[Any] = None, mask: Optional[Any] = None, save_to_input: bool = True) -> TalkObject: + """Build talk object from input data.""" + if audio is None: + return None + + talk_object = TalkObject(name=name) + + # Process audio + if save_to_input and audio is not None: + audio_path = self._save_audio_to_input(name, audio) + if audio_path: + talk_object.audio = audio_path + talk_object.source_type = "file" + else: + talk_object.audio = audio + talk_object.source_type = "data" + + # Process mask + if mask is not None: + if save_to_input: + mask_path = self._save_mask_to_input(name, mask) + if mask_path: + talk_object.mask = mask_path + else: + talk_object.mask = mask + + return talk_object + + def build_from_json(self, json_config: str) -> Optional[TalkObjectsConfig]: + """Build talk objects configuration from JSON.""" + try: + objects_data = json.loads(json_config) + if not isinstance(objects_data, list): + logging.error("JSON config must be a list") + return None + + config = TalkObjectsConfig() + + for obj_data in objects_data: + if not isinstance(obj_data, dict) or "audio" not in obj_data: + continue + + talk_obj = TalkObject(name=obj_data.get("name", "unknown"), audio=obj_data["audio"], mask=obj_data.get("mask"), source_type="path") + config.add_object(talk_obj) + + return config if config.talk_objects else None + + except json.JSONDecodeError as e: + logging.error(f"Failed to parse JSON: {e}") + + def build_from_files(self, audio_files: str, mask_files: str = "", names: str = "") -> Optional[TalkObjectsConfig]: + """Build talk objects configuration from file lists.""" + audio_list = [f.strip() for f in audio_files.split("\n") if f.strip()] + if not audio_list: + return None + + mask_list = [f.strip() for f in mask_files.split("\n") if f.strip()] if mask_files else [] + name_list = [n.strip() for n in names.split("\n") if n.strip()] if names else [] + + config = TalkObjectsConfig() + + for i, audio_file in enumerate(audio_list): + talk_obj = TalkObject( + name=name_list[i] if i < len(name_list) else f"person_{i + 1}", + audio=audio_file, + mask=mask_list[i] if i < len(mask_list) else None, + source_type="file", + ) + config.add_object(talk_obj) + + return config + + def _save_audio_to_input(self, name: str, audio_data: Any) -> Optional[str]: + """Save audio data to input directory.""" + try: + filename = f"{name}_audio_{uuid.uuid4().hex[:8]}.wav" + return self.resolver.save_to_input(audio_data, filename, self.audio_handler) + except Exception as e: + logging.error(f"Failed to save audio: {e}") + return None + + def _save_mask_to_input(self, name: str, mask_data: Any) -> Optional[str]: + """Save mask data to input directory.""" + try: + filename = f"{name}_mask_{uuid.uuid4().hex[:8]}.png" + return self.resolver.save_to_input(mask_data, filename, self.mask_handler) + except Exception as e: + logging.error(f"Failed to save mask: {e}") + return None + + +class ConfigBuilder: + """Main configuration builder that combines all configs.""" + + def __init__(self): + self.manager = ModularConfigManager() + self.temp_manager = TempFileManager() + + def combine_configs( + self, + inference_config: InferenceConfig, + teacache_config: Optional[TeaCacheConfig] = None, + quantization_config: Optional[QuantizationConfig] = None, + memory_config: Optional[MemoryOptimizationConfig] = None, + lora_chain: Optional[List[Dict[str, Any]]] = None, + talk_objects_config: Optional[TalkObjectsConfig] = None, + ) -> EasyDict: + """Combine all configurations into final config.""" + # Create combined config object + combined = CombinedConfig( + inference=inference_config, + teacache=teacache_config, + quantization=quantization_config, + memory=memory_config, + talk_objects=talk_objects_config, + ) + + # Add LoRA configs + if lora_chain: + for lora_dict in lora_chain: + lora_config = LoRAConfig(path=lora_dict["path"], strength=lora_dict.get("strength", 1.0)) + combined.lora_configs.append(lora_config) + + # Build final config using existing manager + configs_dict = { + "inference": inference_config.to_dict() if inference_config else {}, + "teacache": teacache_config.to_dict() if teacache_config else None, + "quantization": quantization_config.to_dict() if quantization_config else None, + "memory": memory_config.to_dict() if memory_config else None, + } + + # Filter out None values + configs_dict = {k: v for k, v in configs_dict.items() if v is not None} + + # Use existing manager to build config + final_config = self.manager.build_final_config(configs_dict) + + # Add additional configs + if lora_chain: + final_config.lora_configs = lora_chain + + if talk_objects_config: + final_config.update(talk_objects_config.to_dict()) + + return final_config + + def get_config_hash(self, config: EasyDict) -> str: + """Generate hash for configuration to detect changes.""" + relevant_configs = { + "model_cls": getattr(config, "model_cls", None), + "model_path": getattr(config, "model_path", None), + "task": getattr(config, "task", None), + "t5_quantized": getattr(config, "t5_quantized", False), + "clip_quantized": getattr(config, "clip_quantized", False), + "lora_configs": getattr(config, "lora_configs", None), + "mm_config": getattr(config, "mm_config", None), + "cross_attn_1_type": getattr(config, "cross_attn_1_type", None), + "cross_attn_2_type": getattr(config, "cross_attn_2_type", None), + "self_attn_1_type": getattr(config, "self_attn_1_type", None), + "self_attn_2_type": getattr(config, "self_attn_2_type", None), + "cpu_offload": getattr(config, "cpu_offload", False), + "offload_granularity": getattr(config, "offload_granularity", None), + "offload_ratio": getattr(config, "offload_ratio", None), + "t5_cpu_offload": getattr(config, "t5_cpu_offload", False), + "t5_offload_granularity": getattr(config, "t5_offload_granularity", None), + "audio_encoder_cpu_offload": getattr(config, "audio_encoder_cpu_offload", False), + "audio_adapter_cpu_offload": getattr(config, "audio_adapter_cpu_offload", False), + "vae_cpu_offload": getattr(config, "vae_cpu_offload", False), + "use_tiling_vae": getattr(config, "use_tiling_vae", False), + "unload_after_inference": getattr(config, "unload_after_inference", False), + "enable_rotary_chunk": getattr(config, "enable_rotary_chunk", False), + "rotary_chunk_size": getattr(config, "rotary_chunk_size", None), + "clean_cuda_cache": getattr(config, "clean_cuda_cache", False), + "torch_compile": getattr(config, "torch_compile", False), + "threshold": getattr(config, "threshold", None), + "use_ret_steps": getattr(config, "use_ret_steps", False), + "t5_quant_scheme": getattr(config, "t5_quant_scheme", None), + "clip_quant_scheme": getattr(config, "clip_quant_scheme", None), + "adapter_quant_scheme": getattr(config, "adapter_quant_scheme", None), + "adapter_quantized": getattr(config, "adapter_quantized", False), + "feature_caching": getattr(config, "feature_caching", None), + } + + config_str = json.dumps(relevant_configs, sort_keys=True) + return hashlib.md5(config_str.encode()).hexdigest() + + +class LoRAChainBuilder: + """Builder for LoRA chain configurations.""" + + @staticmethod + def build_chain(lora_name: str, strength: float, existing_chain: Optional[List[Dict]] = None) -> List[Dict]: + """Build or extend a LoRA chain.""" + if existing_chain is None: + chain = [] + else: + chain = existing_chain.copy() + + lora_path = get_lora_full_path(lora_name) + if lora_path: + lora_config = {"path": lora_path, "strength": strength} + chain.append(lora_config) + + return chain diff --git a/data_models.py b/data_models.py new file mode 100644 index 0000000..ab2416b --- /dev/null +++ b/data_models.py @@ -0,0 +1,198 @@ +"""Data models for LightX2V ComfyUI wrapper.""" + +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional, Union + +import numpy as np +import torch + + +@dataclass +class TalkObject: + """Single talk object containing audio and optional mask.""" + + name: str + audio: Optional[Union[str, Dict[str, Any], torch.Tensor, np.ndarray]] = None + mask: Optional[Union[str, torch.Tensor, np.ndarray]] = None + source_type: str = "data" # "data", "file", or "path" + + def to_dict(self) -> Dict[str, Any]: + """Convert to dictionary for pipeline.""" + result = {"name": self.name} + + if isinstance(self.audio, str): + result["audio"] = self.audio + elif self.audio is not None: + result["audio_data"] = self.audio + + if isinstance(self.mask, str): + result["mask"] = self.mask + elif self.mask is not None: + result["mask_data"] = self.mask + + result["source_type"] = self.source_type + return result + + +@dataclass +class InferenceConfig: + """Basic inference configuration.""" + + model_cls: str = "wan2.1" + model_path: str = "" + task: str = "i2v" + infer_steps: int = 4 + seed: int = 42 + cfg_scale: float = 5.0 + cfg_scale2: float = 5.0 + sample_shift: int = 5 + height: int = 1280 + width: int = 720 + video_length: int = 81 + fps: int = 16 + video_duration: float = 5.0 + attention_type: str = "torch_sdpa" + use_31_block: bool = True + + # Optional parameters + denoising_step_list: Optional[List[int]] = None + resize_mode: str = "adaptive" + fixed_area: str = "720p" + segment_length: int = 81 + prev_frame_length: int = 5 + use_tiny_vae: bool = False + + # Runtime parameters + prompt: str = "" + negative_prompt: str = "" + image_path: Optional[str] = None + audio_path: Optional[str] = None + + def to_dict(self) -> Dict[str, Any]: + """Convert to dictionary, excluding None values.""" + result = {} + for key, value in self.__dict__.items(): + if value is not None: + result[key] = value + return result + + +@dataclass +class TeaCacheConfig: + """TeaCache configuration.""" + + enable: bool = False + threshold: float = 0.26 + use_ret_steps: bool = False + + def to_dict(self) -> Dict[str, Any]: + return {"enable": self.enable, "threshold": self.threshold, "use_ret_steps": self.use_ret_steps} + + +@dataclass +class QuantizationConfig: + """Quantization configuration.""" + + quant_op: str = "none" + dit_quant_scheme: str = "bf16" + t5_quant_scheme: str = "bf16" + clip_quant_scheme: str = "fp16" + adapter_quant_scheme: str = "bf16" + + def to_dict(self) -> Dict[str, Any]: + return { + "quant_op": self.quant_op, + "dit_quant_scheme": self.dit_quant_scheme, + "t5_quant_scheme": self.t5_quant_scheme, + "clip_quant_scheme": self.clip_quant_scheme, + "adapter_quant_scheme": self.adapter_quant_scheme, + } + + +@dataclass +class MemoryOptimizationConfig: + """Memory optimization configuration.""" + + enable_rotary_chunk: bool = False + rotary_chunk_size: int = 100 + clean_cuda_cache: bool = False + cpu_offload: bool = True + offload_granularity: str = "block" + offload_ratio: float = 1.0 + t5_cpu_offload: bool = True + t5_offload_granularity: str = "model" + audio_encoder_cpu_offload: bool = True + audio_adapter_cpu_offload: bool = True + vae_cpu_offload: bool = True + use_tiling_vae: bool = True + lazy_load: bool = False + unload_after_inference: bool = False + + def to_dict(self) -> Dict[str, Any]: + return self.__dict__.copy() + + +@dataclass +class LoRAConfig: + """LoRA configuration.""" + + path: str + strength: float = 1.0 + + def to_dict(self) -> Dict[str, Any]: + return {"path": self.path, "strength": self.strength} + + +@dataclass +class TalkObjectsConfig: + """Configuration for multiple talk objects.""" + + talk_objects: List[TalkObject] = field(default_factory=list) + + def add_object(self, talk_object: TalkObject): + """Add a talk object to the configuration.""" + self.talk_objects.append(talk_object) + + def to_dict(self) -> Dict[str, Any]: + """Convert to dictionary for pipeline.""" + return {"talk_objects": [obj for obj in self.talk_objects]} + + def to_list(self) -> List[Dict[str, Any]]: + """Get list of talk object dictionaries.""" + return [obj for obj in self.talk_objects] + + +@dataclass +class CombinedConfig: + """Combined configuration for all modules.""" + + inference: Optional[InferenceConfig] = None + teacache: Optional[TeaCacheConfig] = None + quantization: Optional[QuantizationConfig] = None + memory: Optional[MemoryOptimizationConfig] = None + lora_configs: List[LoRAConfig] = field(default_factory=list) + talk_objects: Optional[TalkObjectsConfig] = None + + def to_dict(self) -> Dict[str, Any]: + """Convert to dictionary for pipeline.""" + result = {} + + if self.inference: + result.update(self.inference.to_dict()) + + if self.teacache: + result["teacache"] = self.teacache.to_dict() + + if self.quantization: + result["quantization"] = self.quantization.to_dict() + + if self.memory: + result["memory"] = self.memory.to_dict() + + if self.lora_configs: + result["lora_configs"] = [lora.to_dict() for lora in self.lora_configs] + + if self.talk_objects: + result["talk_objects"] = self.talk_objects.to_list() + + return result diff --git a/file_handlers.py b/file_handlers.py new file mode 100644 index 0000000..ddf29ef --- /dev/null +++ b/file_handlers.py @@ -0,0 +1,361 @@ +"""Unified file handlers for LightX2V ComfyUI wrapper.""" + +import logging +import os +import tempfile +from abc import ABC, abstractmethod +from contextlib import contextmanager +from typing import Any, Dict, List, Optional, Tuple, Union + +import numpy as np +import scipy.io.wavfile as wavfile +import torch +from PIL import Image + + +class FileHandler(ABC): + """Base class for file handling.""" + + @abstractmethod + def save(self, data: Any, path: str) -> str: + """Save data to file.""" + pass + + @abstractmethod + def load(self, path: str) -> Any: + """Load data from file.""" + pass + + +class AudioFileHandler(FileHandler): + """Handler for audio files.""" + + def __init__(self): + self.supported_formats = [".wav", ".mp3", ".flac", ".m4a"] + + def save(self, audio_data: Union[Dict, torch.Tensor, np.ndarray, Tuple], path: str, sample_rate: Optional[int] = None) -> str: + """Save audio data to file. + + Args: + audio_data: Audio data in various formats + path: Output file path + sample_rate: Sample rate (required if not in audio_data) + + Returns: + Path to saved file + """ + + waveform, sr = self._extract_audio_data(audio_data, sample_rate) + + # Ensure waveform is in correct shape + waveform = self._normalize_waveform(waveform) + + # Always save as WAV format + ext = os.path.splitext(path)[1].lower() + if ext != ".wav": + path = path.rsplit(".", 1)[0] + ".wav" + logging.info(f"Audio will be saved as WAV format: {path}") + + # Ensure waveform is in int16 format for WAV + if waveform.dtype != np.int16: + # Normalize to [-1, 1] range if not already + if waveform.dtype == np.float32 or waveform.dtype == np.float64: + # Clip to [-1, 1] to avoid overflow + waveform = np.clip(waveform, -1.0, 1.0) + waveform = (waveform * 32767).astype(np.int16) + else: + # Assume uint8 or other integer type + waveform = waveform.astype(np.int16) + + wavfile.write(path, sr, waveform) + logging.info(f"Audio saved to {path}") + return path + + def load(self, path: str) -> Tuple[np.ndarray, int]: + """Load audio from file. + + Returns: + Tuple of (waveform, sample_rate) + """ + sample_rate, waveform = wavfile.read(path) + return waveform, sample_rate + + def _extract_audio_data(self, audio_data: Any, sample_rate: Optional[int] = None) -> Tuple[np.ndarray, int]: + """Extract waveform and sample rate from various audio formats. + + Handles three main sources: + 1. ComfyUI LoadAudio output: {"waveform": tensor, "sample_rate": int} + 2. Tuple format: (waveform, sample_rate) + 3. Raw waveform with separate sample_rate + """ + if isinstance(audio_data, dict): + if "waveform" in audio_data and "sample_rate" in audio_data: + waveform = audio_data["waveform"] + sr = audio_data["sample_rate"] + + # Handle ComfyUI LoadAudio format specifically + # ComfyUI returns waveform with shape [batch, channels, samples] + if isinstance(waveform, torch.Tensor): + if waveform.dim() == 3: # [batch, channels, samples] + waveform = waveform[0] # Take first batch + if waveform.dim() == 2 and waveform.shape[0] <= 2: # [channels, samples] + waveform = waveform.transpose(0, 1) # -> [samples, channels] + waveform = waveform.cpu().numpy() + else: + raise ValueError("Audio dict must contain 'waveform' and 'sample_rate'") + elif isinstance(audio_data, tuple) and len(audio_data) == 2: + waveform, sr = audio_data + if isinstance(waveform, torch.Tensor): + waveform = waveform.cpu().numpy() + elif sample_rate is not None: + waveform = audio_data + sr = sample_rate + if isinstance(waveform, torch.Tensor): + waveform = waveform.cpu().numpy() + else: + raise ValueError("Sample rate must be provided for raw audio data") + + return waveform, sr + + def _normalize_waveform(self, waveform: np.ndarray) -> np.ndarray: + """Normalize waveform shape for saving. + + Ensures waveform is in shape [samples, channels] or [samples] for mono. + """ + # Already converted to numpy in _extract_audio_data + + # Handle different shapes + if waveform.ndim == 3: # Shouldn't happen, but handle it + waveform = waveform[0] + + if waveform.ndim == 2: + # Check if it's [channels, samples] format (channels < samples typically) + if waveform.shape[0] <= 2 and waveform.shape[0] < waveform.shape[1]: + waveform = waveform.T # -> [samples, channels] + # If mono with extra dimension, squeeze it + if waveform.shape[1] == 1: + waveform = waveform.squeeze() + + return waveform + + def _save_with_wave(self, path: str, waveform: np.ndarray, sample_rate: int): + """Save audio using wave module.""" + import wave + + with wave.open(path, "wb") as wav_file: + 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: + waveform = (waveform * 32767).astype(np.int16) + + wav_file.writeframes(waveform.tobytes()) + + def _load_with_wave(self, path: str) -> Tuple[np.ndarray, int]: + """Load audio using wave module.""" + import wave + + with wave.open(path, "rb") as wav_file: + sample_rate = wav_file.getframerate() + frames = wav_file.readframes(wav_file.getnframes()) + waveform = np.frombuffer(frames, dtype=np.int16) + + if wav_file.getnchannels() > 1: + waveform = waveform.reshape(-1, wav_file.getnchannels()) + + return waveform, sample_rate + + +class ImageFileHandler(FileHandler): + """Handler for image files.""" + + def __init__(self): + self.supported_formats = [".png", ".jpg", ".jpeg", ".bmp", ".tiff"] + + def save(self, image_data: Union[torch.Tensor, np.ndarray, Image.Image], path: str) -> str: + """Save image data to file. + + Args: + image_data: Image data in various formats + path: Output file path + + Returns: + Path to saved file + """ + if isinstance(image_data, torch.Tensor): + # Convert from tensor [H, W, C] or [C, H, W] + if image_data.dim() == 4: # [batch, ...] + image_data = image_data[0] + + image_np = image_data.cpu().numpy() + + # Handle channel ordering + if image_np.shape[0] in [1, 3, 4]: # [C, H, W] + image_np = np.transpose(image_np, (1, 2, 0)) + + # Convert to uint8 + if image_np.dtype != np.uint8: + image_np = (image_np * 255).astype(np.uint8) + + image = Image.fromarray(image_np.squeeze()) + + elif isinstance(image_data, np.ndarray): + if image_data.dtype != np.uint8: + image_data = (image_data * 255).astype(np.uint8) + image = Image.fromarray(image_data.squeeze()) + + elif isinstance(image_data, Image.Image): + image = image_data + else: + raise ValueError(f"Unsupported image format: {type(image_data)}") + + image.save(path) + logging.info(f"Image saved to {path}") + return path + + def load(self, path: str) -> Image.Image: + """Load image from file.""" + return Image.open(path) + + +class MaskFileHandler(ImageFileHandler): + """Handler specifically for mask files.""" + + def save(self, mask_data: Union[torch.Tensor, np.ndarray], path: str) -> str: + """Save mask data to file. + + Args: + mask_data: Mask data (single channel) + path: Output file path + + Returns: + Path to saved file + """ + if isinstance(mask_data, torch.Tensor): + if mask_data.dim() == 3: # [batch, H, W] + mask_data = mask_data[0] + mask_np = (mask_data.cpu().numpy() * 255).astype(np.uint8) + elif isinstance(mask_data, np.ndarray): + mask_np = (mask_data * 255).astype(np.uint8) + else: + mask_np = mask_data + + mask_image = Image.fromarray(mask_np) + mask_image.save(path) + logging.info(f"Mask saved to {path}") + return path + + +class TempFileManager: + """Manager for temporary files with automatic cleanup.""" + + def __init__(self): + self.temp_files: List[str] = [] + + @contextmanager + def temp_file(self, suffix: str = "", prefix: str = "lightx2v_", delete: bool = True): + """Context manager for temporary file creation. + + Args: + suffix: File suffix + prefix: File prefix + delete: Whether to delete on exit + + Yields: + Path to temporary file + """ + temp_file = tempfile.NamedTemporaryFile(suffix=suffix, prefix=prefix, delete=False) + temp_path = temp_file.name + temp_file.close() + + self.temp_files.append(temp_path) + + try: + yield temp_path + finally: + if delete: + self.cleanup_file(temp_path) + + def create_temp_file(self, suffix: str = "", prefix: str = "lightx2v_") -> str: + """Create a temporary file that will be tracked for cleanup. + + Args: + suffix: File suffix + prefix: File prefix + + Returns: + Path to temporary file + """ + with tempfile.NamedTemporaryFile(suffix=suffix, prefix=prefix, delete=False) as tmp: + temp_path = tmp.name + + self.temp_files.append(temp_path) + return temp_path + + def cleanup_file(self, path: str): + """Clean up a specific file.""" + if path in self.temp_files: + self.temp_files.remove(path) + + if os.path.exists(path): + try: + os.unlink(path) + logging.debug(f"Cleaned up temp file: {path}") + except Exception as e: + logging.warning(f"Failed to clean up {path}: {e}") + + def cleanup_all(self): + """Clean up all tracked temporary files.""" + for temp_file in self.temp_files[:]: + self.cleanup_file(temp_file) + + self.temp_files.clear() + + def __del__(self): + """Cleanup on deletion.""" + self.cleanup_all() + + +class ComfyUIFileResolver: + """Resolve file paths for ComfyUI input/output directories.""" + + @staticmethod + def get_input_directory() -> str: + """Get ComfyUI input directory.""" + try: + import folder_paths + + return folder_paths.get_input_directory() + except ImportError: + # Fallback if not in ComfyUI environment + return "input" + + @staticmethod + def resolve_input_path(filename: str) -> str: + """Resolve a filename to full path in input directory.""" + if os.path.isabs(filename): + return filename + + input_dir = ComfyUIFileResolver.get_input_directory() + return os.path.join(input_dir, filename) + + @staticmethod + def save_to_input(data: Any, filename: str, handler: FileHandler) -> str: + """Save data to ComfyUI input directory. + + Args: + data: Data to save + filename: Target filename + handler: File handler to use + + Returns: + Full path to saved file + """ + input_dir = ComfyUIFileResolver.get_input_directory() + full_path = os.path.join(input_dir, filename) + + # Create directory if needed + os.makedirs(input_dir, exist_ok=True) + + return handler.save(data, full_path) diff --git a/lightx2v b/lightx2v index f7eb4f1..749f8c6 160000 --- a/lightx2v +++ b/lightx2v @@ -1 +1 @@ -Subproject commit f7eb4f12e6d54c588d2a178f93b8e9be78107a6e +Subproject commit 749f8c64ef0494a06f0c3b2ea5e30ac0cfa7be89 diff --git a/nodes.py b/nodes.py index b9d298e..ab006c7 100644 --- a/nodes.py +++ b/nodes.py @@ -1,30 +1,37 @@ """Modular ComfyUI nodes for LightX2V without presets.""" import gc -import hashlib import json import logging import os -import tempfile 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 .bridge import get_available_attn_ops, get_available_quant_ops +from .config_builder import ( + ConfigBuilder, + InferenceConfigBuilder, + LoRAChainBuilder, + TalkObjectConfigBuilder, +) +from .data_models import ( + InferenceConfig, + MemoryOptimizationConfig, + QuantizationConfig, + TalkObjectsConfig, + TeaCacheConfig, +) +from .file_handlers import ( + AudioFileHandler, + ComfyUIFileResolver, + ImageFileHandler, + TempFileManager, ) from .lightx2v.lightx2v.infer import init_runner -from .model_utils import ( - get_lora_full_path, - get_model_full_path, - scan_loras, - scan_models, - support_model_cls_list, -) +from .model_utils import scan_loras, scan_models, support_model_cls_list class LightX2VInferenceConfig: @@ -211,61 +218,30 @@ class LightX2VInferenceConfig: use_tiny_vae=False, ): """Create basic inference configuration.""" - model_path = get_model_full_path(model_name) + builder = InferenceConfigBuilder() - if model_cls == "hunyuan": - fps = 24 - else: - fps = 16 + config = builder.build( + model_cls=model_cls, + model_name=model_name, + task=task, + infer_steps=infer_steps, + seed=seed, + cfg_scale=cfg_scale, + cfg_scale2=cfg_scale2, + sample_shift=sample_shift, + height=height, + width=width, + duration=duration, + attention_type=attention_type, + denoising_steps=denoising_steps, + resize_mode=resize_mode, + fixed_area=fixed_area, + segment_length=segment_length, + prev_frame_length=prev_frame_length, + use_tiny_vae=use_tiny_vae, + ) - 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 - if remainder != 0: - video_length = video_length + (4 - remainder) - - # TODO(xxx): - use_31_block = True - if "seko" in model_cls: - video_length = segment_length - use_31_block = False - - 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, - "cfg_scale2": cfg_scale2, - "sample_shift": sample_shift, - "height": height, - "width": width, - "video_length": video_length, - "fps": fps, - "video_duration": duration, - "resize_mode": resize_mode, - "fixed_area": fixed_area, - "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: - steps_list = [int(s.strip()) for s in denoising_steps.split(",")] - config["denoising_step_list"] = steps_list - config["infer_steps"] = len(steps_list) - except ValueError: - pass - - return (config,) + return (config.to_dict(),) class LightX2VTeaCache: @@ -305,12 +281,9 @@ class LightX2VTeaCache: CATEGORY = "LightX2V/Config" def create_config(self, enable, threshold, use_ret_steps): - config = { - "enable": enable, - "threshold": threshold, - "use_ret_steps": use_ret_steps, - } - return (config,) + """Create TeaCache configuration.""" + config = TeaCacheConfig(enable=enable, threshold=threshold, use_ret_steps=use_ret_steps) + return (config.to_dict(),) class LightX2VQuantization: @@ -383,14 +356,14 @@ class LightX2VQuantization: adapter_quant_scheme, ): """Create quantization configuration.""" - config = { - "dit_quant_scheme": dit_quant_scheme, - "t5_quant_scheme": t5_quant_scheme, - "clip_quant_scheme": clip_quant_scheme, - "adapter_quant_scheme": adapter_quant_scheme, - "quant_op": quant_op, - } - return (config,) + config = QuantizationConfig( + quant_op=quant_op, + dit_quant_scheme=dit_quant_scheme, + t5_quant_scheme=t5_quant_scheme, + clip_quant_scheme=clip_quant_scheme, + adapter_quant_scheme=adapter_quant_scheme, + ) + return (config.to_dict(),) class LightX2VMemoryOptimization: @@ -481,23 +454,24 @@ class LightX2VMemoryOptimization: lazy_load=False, unload_after_inference=False, ): - config = { - "enable_rotary_chunk": enable_rotary_chunk, - "rotary_chunk_size": rotary_chunk_size, - "clean_cuda_cache": clean_cuda_cache, - "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, - "use_tiling_vae": use_tiling_vae, - "lazy_load": lazy_load, - "unload_after_inference": unload_after_inference, - } - return (config,) + """Create memory optimization configuration.""" + config = MemoryOptimizationConfig( + enable_rotary_chunk=enable_rotary_chunk, + rotary_chunk_size=rotary_chunk_size, + clean_cuda_cache=clean_cuda_cache, + 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, + use_tiling_vae=use_tiling_vae, + lazy_load=lazy_load, + unload_after_inference=unload_after_inference, + ) + return (config.to_dict(),) class LightX2VLoRALoader: @@ -539,35 +513,22 @@ class LightX2VLoRALoader: CATEGORY = "LightX2V/LoRA" def load_lora(self, lora_name, strength, lora_chain=None): - if lora_chain is None: - lora_chain = [] - else: - lora_chain = lora_chain.copy() - - lora_path = get_lora_full_path(lora_name) - - if lora_path: - lora_config = {"path": lora_path, "strength": strength} - lora_chain.append(lora_config) - - return (lora_chain,) + """Load and chain LoRA configurations.""" + chain = LoRAChainBuilder.build_chain(lora_name=lora_name, strength=strength, existing_chain=lora_chain) + return (chain,) class TalkObjectInput: - """单个谈话对象（音频+遮罩）输入节点""" - @classmethod def INPUT_TYPES(cls): return { "required": { - "audio_source": (["upload", "url"], {"default": "upload", "tooltip": "音频输入方式"}), - "mask_source": (["upload", "url", "none"], {"default": "upload", "tooltip": "遮罩输入方式"}), + "name": ("STRING", {"default": "person_1", "tooltip": "说话人名称标识"}), }, "optional": { "audio": ("AUDIO", {"tooltip": "上传的音频文件"}), - "audio_url": ("STRING", {"default": "", "tooltip": "音频文件URL路径"}), - "mask": ("MASK", {"tooltip": "上传的遮罩图像"}), - "mask_url": ("STRING", {"default": "", "tooltip": "遮罩图像URL路径"}), + "mask": ("MASK", {"tooltip": "上传的遮罩图像（可选）"}), + "save_to_input": ("BOOLEAN", {"default": True, "tooltip": "是否保存到input文件夹"}), }, } @@ -576,76 +537,104 @@ class TalkObjectInput: FUNCTION = "create_talk_object" CATEGORY = "LightX2V/Audio" - def create_talk_object(self, audio_source, mask_source, audio=None, audio_url="", mask=None, mask_url=""): - """创建单个谈话对象配置""" - talk_object = {} + def create_talk_object(self, name, audio=None, mask=None, save_to_input=True): + """Create a talk object from input data.""" + builder = TalkObjectConfigBuilder() - # 处理音频 - if audio_source == "upload" and audio is not None: - # 音频将在推理时转换为临时文件 - talk_object["audio_data"] = audio - talk_object["audio_type"] = "upload" - elif audio_source == "url" and audio_url: - talk_object["audio"] = audio_url - talk_object["audio_type"] = "url" - else: - return (None,) # 无效的音频输入 + talk_object = builder.build_from_input(name=name, audio=audio, mask=mask, save_to_input=save_to_input) - # 处理遮罩 - if mask_source == "upload" and mask is not None: - talk_object["mask_data"] = mask - talk_object["mask_type"] = "upload" - elif mask_source == "url" and mask_url: - talk_object["mask"] = mask_url - talk_object["mask_type"] = "url" - elif mask_source == "none": - # 不使用遮罩 - talk_object["mask"] = None - talk_object["mask_type"] = "none" - - return (talk_object,) + if talk_object: + return (talk_object.to_dict(),) + return (None,) -class TalkObjectsBuilder: - """构建多人对话配置的节点""" +class TalkObjectsCombiner: + """组合多个谈话对象为配置""" + @classmethod + def INPUT_TYPES(cls): + inputs = {"required": {}, "optional": {}} + + # 预定义10个TALK_OBJECT输入槽 + for i in range(1, 11): + inputs["optional"][f"talk_object_{i}"] = ("TALK_OBJECT", {"tooltip": f"谈话对象{i}"}) + + return inputs + + RETURN_TYPES = ("TALK_OBJECTS_CONFIG",) + RETURN_NAMES = ("talk_objects_config",) + FUNCTION = "combine_talk_objects" + CATEGORY = "LightX2V/Audio" + + def combine_talk_objects(self, **kwargs): + config = TalkObjectsConfig() + + for i in range(1, 11): + talk_obj = kwargs.get(f"talk_object_{i}") + + if talk_obj is not None: + config.add_object(talk_obj) + + if not config.talk_objects: + return (None,) + + return (config,) + + +class TalkObjectsFromJSON: @classmethod def INPUT_TYPES(cls): return { "required": { - "number": ("INT", {"default": 1, "min": 0, "max": 10, "tooltip": "谈话对象数量，0表示无，1表示单人，2+表示多人"}), - }, - "optional": { - # 支持最多10个谈话对象 - **{f"talk_object_{i}": ("TALK_OBJECT",) for i in range(1, 11)} + "json_config": ( + "STRING", + { + "multiline": True, + "default": '[{"name": "person1", "audio": "/path/to/audio1.wav", "mask": "/path/to/mask1.png"}]', + "tooltip": "JSON格式的谈话对象配置", + }, + ), }, } RETURN_TYPES = ("TALK_OBJECTS_CONFIG",) RETURN_NAMES = ("talk_objects_config",) - FUNCTION = "build_talk_objects" + FUNCTION = "parse_json_config" CATEGORY = "LightX2V/Audio" - def build_talk_objects(self, number, **kwargs): - """构建谈话对象列表""" - if number == 0: - return (None,) + def parse_json_config(self, json_config): + builder = TalkObjectConfigBuilder() + talk_objects_config = builder.build_from_json(json_config) + return (talk_objects_config,) - talk_objects = [] - for i in range(1, number + 1): - talk_obj = kwargs.get(f"talk_object_{i}") - if talk_obj is not None: - talk_objects.append(talk_obj) - if not talk_objects: - return (None,) +class TalkObjectsFromFiles: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "audio_files": ("STRING", {"multiline": True, "default": "audio1.wav\naudio2.wav", "tooltip": "音频文件列表（每行一个）"}), + }, + "optional": { + "mask_files": ("STRING", {"multiline": True, "default": "mask1.png\nmask2.png", "tooltip": "遮罩文件列表（每行一个，可选）"}), + "names": ("STRING", {"multiline": True, "default": "person1\nperson2", "tooltip": "人物名称列表（每行一个，可选）"}), + }, + } - return ({"talk_objects": talk_objects},) + RETURN_TYPES = ("TALK_OBJECTS_CONFIG",) + RETURN_NAMES = ("talk_objects_config",) + FUNCTION = "build_from_files" + CATEGORY = "LightX2V/Audio" + + def build_from_files(self, audio_files, mask_files="", names=""): + builder = TalkObjectConfigBuilder() + talk_objects_config = builder.build_from_files(audio_files, mask_files, names) + return (talk_objects_config,) class LightX2VConfigCombiner: def __init__(self): - self.config_manager = ModularConfigManager() + self.config_builder = ConfigBuilder() @classmethod def INPUT_TYPES(cls): @@ -688,36 +677,37 @@ class LightX2VConfigCombiner: lora_chain=None, talk_objects_config=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 + """Combine multiple configurations into final config.""" + # Convert dict configs back to objects if needed - config = self.config_manager.build_final_config(configs) + # Create objects from dicts + inf_config = InferenceConfig(**inference_config) if isinstance(inference_config, dict) else None + tea_config = TeaCacheConfig(**teacache_config) if teacache_config and isinstance(teacache_config, dict) else None + quant_config = QuantizationConfig(**quantization_config) if quantization_config and isinstance(quantization_config, dict) else None + mem_config = MemoryOptimizationConfig(**memory_config) if memory_config and isinstance(memory_config, dict) else None - if lora_chain: - config.lora_configs = lora_chain - - if talk_objects_config: - config.talk_objects_config = talk_objects_config - - logging.info("lightx2v config: " + json.dumps(config, indent=2, ensure_ascii=False)) + config = self.config_builder.combine_configs( + inference_config=inf_config, + teacache_config=tea_config, + quantization_config=quant_config, + memory_config=mem_config, + lora_chain=lora_chain, + talk_objects_config=talk_objects_config, + ) return (config,) class LightX2VModularInference: - # 类变量，所有实例共享 _current_runner = None _current_config_hash = None def __init__(self): - pass + self.config_builder = ConfigBuilder() + self.temp_manager = TempFileManager() + self.image_handler = ImageFileHandler() + self.audio_handler = AudioFileHandler() + self.resolver = ComfyUIFileResolver() @classmethod def INPUT_TYPES(cls): @@ -751,43 +741,8 @@ class LightX2VModularInference: CATEGORY = "LightX2V/Inference" def _get_config_hash(self, config) -> str: - relevant_configs = { - "model_cls": getattr(config, "model_cls", None), - "model_path": getattr(config, "model_path", None), - "task": getattr(config, "task", None), - "t5_quantized": getattr(config, "t5_quantized", False), - "clip_quantized": getattr(config, "clip_quantized", False), - "lora_configs": getattr(config, "lora_configs", None), - "mm_config": getattr(config, "mm_config", None), - "cross_attn_1_type": getattr(config, "cross_attn_1_type", None), - "cross_attn_2_type": getattr(config, "cross_attn_2_type", None), - "self_attn_1_type": getattr(config, "self_attn_1_type", None), - "self_attn_2_type": getattr(config, "self_attn_2_type", None), - "cpu_offload": getattr(config, "cpu_offload", False), - "offload_granularity": getattr(config, "offload_granularity", None), - "offload_ratio": getattr(config, "offload_ratio", None), - "t5_cpu_offload": getattr(config, "t5_cpu_offload", False), - "t5_offload_granularity": getattr(config, "t5_offload_granularity", None), - "audio_encoder_cpu_offload": getattr(config, "audio_encoder_cpu_offload", False), - "audio_adapter_cpu_offload": getattr(config, "audio_adapter_cpu_offload", False), - "vae_cpu_offload": getattr(config, "vae_cpu_offload", False), - "use_tiling_vae": getattr(config, "use_tiling_vae", False), - "unload_after_inference": getattr(config, "unload_after_inference", False), - "enable_rotary_chunk": getattr(config, "enable_rotary_chunk", False), - "rotary_chunk_size": getattr(config, "rotary_chunk_size", None), - "clean_cuda_cache": getattr(config, "clean_cuda_cache", False), - "torch_compile": getattr(config, "torch_compile", False), - "threshold": getattr(config, "threshold", None), - "use_ret_steps": getattr(config, "use_ret_steps", False), - "t5_quant_scheme": getattr(config, "t5_quant_scheme", None), - "clip_quant_scheme": getattr(config, "clip_quant_scheme", None), - "adapter_quant_scheme": getattr(config, "adapter_quant_scheme", None), - "adapter_quantized": getattr(config, "adapter_quantized", False), - "feature_caching": getattr(config, "feature_caching", None), - } - - config_str = json.dumps(relevant_configs, sort_keys=True) - return hashlib.md5(config_str.encode()).hexdigest() + """Get hash of configuration to detect changes.""" + return self.config_builder.get_config_hash(config) def generate( self, @@ -798,157 +753,73 @@ class LightX2VModularInference: audio=None, **kwargs, ): + # config type is EasyDict config = combined_config - config.prompt = prompt config.negative_prompt = negative_prompt if config.task == "i2v" and image is None: raise ValueError("i2v task requires input image") - temp_files = [] - try: + # Handle image input if config.task == "i2v" and image is not None: image_np = (image[0].cpu().numpy() * 255).astype(np.uint8) pil_image = Image.fromarray(image_np) - with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp: - pil_image.save(tmp.name) - config.image_path = tmp.name - temp_files.append(tmp.name) - logging.info(f"Image saved to {tmp.name}") + temp_path = self.temp_manager.create_temp_file(suffix=".png") + pil_image.save(temp_path) + config.image_path = temp_path + logging.info(f"Image saved to {temp_path}") + # Handle audio input for seko models if audio is not None and hasattr(config, "model_cls") and "seko" 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"] + temp_path = self.temp_manager.create_temp_file(suffix=".wav") + self.audio_handler.save(audio, temp_path) + config.audio_path = temp_path + logging.info(f"Audio saved to {temp_path}") - # Handle different waveform shapes - if isinstance(waveform, torch.Tensor): - if waveform.dim() == 3: # [batch, channels, samples] - waveform = waveform[0] # Take first batch - if waveform.dim() == 2: # [channels, samples] - # Convert to [samples, channels] for wav file - waveform = waveform.transpose(0, 1) - waveform = waveform.cpu().numpy() - elif isinstance(audio, tuple) and len(audio) == 2: - # Legacy format support - waveform, sample_rate = audio - if isinstance(waveform, torch.Tensor): - waveform = waveform.cpu().numpy() - else: - raise ValueError(f"Unsupported audio format: {type(audio)}") - - with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp: - try: - import scipy.io.wavfile as wavfile - except ImportError: - 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.setsampwidth(2) # 16-bit - wav_file.setframerate(sample_rate) - if waveform.dtype != np.int16: - waveform = (waveform * 32767).astype(np.int16) - wav_file.writeframes(waveform.tobytes()) - else: - if waveform.ndim == 1: - wavfile.write(tmp.name, sample_rate, waveform) - else: - if waveform.shape[0] < waveform.shape[1]: - waveform = waveform.T - wavfile.write(tmp.name, sample_rate, waveform) - - config.audio_path = tmp.name - temp_files.append(tmp.name) - - logging.info(f"Audio saved to {tmp.name}") - - # 处理多人对话对象 - if hasattr(config, "talk_objects_config") and config.talk_objects_config: - talk_objects = config.talk_objects_config.get("talk_objects", []) + # Handle talk objects + if hasattr(config, "talk_objects") and config.talk_objects: + talk_objects = config.talk_objects processed_talk_objects = [] for talk_obj in talk_objects: processed_obj = {} - # 处理音频 - if talk_obj.get("audio_type") == "upload" and "audio_data" in talk_obj: - audio_data = talk_obj["audio_data"] + if "audio" in talk_obj: + processed_obj["audio"] = talk_obj["audio"] - # 处理音频数据，转换为临时文件 - if isinstance(audio_data, dict) and "waveform" in audio_data and "sample_rate" in audio_data: - waveform = audio_data["waveform"] - sample_rate = audio_data["sample_rate"] - - if isinstance(waveform, torch.Tensor): - if waveform.dim() == 3: - waveform = waveform[0] - if waveform.dim() == 2: - waveform = waveform.transpose(0, 1) - waveform = waveform.cpu().numpy() - - with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp: - try: - import scipy.io.wavfile as wavfile - - if waveform.ndim == 1: - wavfile.write(tmp.name, sample_rate, waveform) - else: - if waveform.shape[0] < waveform.shape[1]: - waveform = waveform.T - wavfile.write(tmp.name, sample_rate, waveform) - except ImportError: - 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.setsampwidth(2) - wav_file.setframerate(sample_rate) - if waveform.dtype != np.int16: - waveform = (waveform * 32767).astype(np.int16) - wav_file.writeframes(waveform.tobytes()) - - processed_obj["audio"] = tmp.name - temp_files.append(tmp.name) - logging.info(f"Talk object audio saved to {tmp.name}") - elif talk_obj.get("audio_type") == "url": - processed_obj["audio"] = talk_obj.get("audio") - - # 处理遮罩 - if talk_obj.get("mask_type") == "upload" and "mask_data" in talk_obj: - mask_data = talk_obj["mask_data"] - - # 假设 mask_data 是一个 tensor 或 numpy array - if isinstance(mask_data, torch.Tensor): - mask_np = (mask_data[0].cpu().numpy() * 255).astype(np.uint8) - elif isinstance(mask_data, np.ndarray): - mask_np = (mask_data * 255).astype(np.uint8) - else: - mask_np = mask_data - - mask_image = Image.fromarray(mask_np) - - with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp: - mask_image.save(tmp.name) - processed_obj["mask"] = tmp.name - temp_files.append(tmp.name) - logging.info(f"Talk object mask saved to {tmp.name}") - elif talk_obj.get("mask_type") == "url": - processed_obj["mask"] = talk_obj.get("mask") - elif talk_obj.get("mask_type") == "none": - processed_obj["mask"] = None + if "mask" in talk_obj: + processed_obj["mask"] = talk_obj["mask"] if "audio" in processed_obj: processed_talk_objects.append(processed_obj) - # 设置 talk_objects 到 config + # Resolve paths + for obj in processed_talk_objects: + if "audio" in obj and obj["audio"]: + audio_path = obj["audio"] + if not os.path.isabs(audio_path) and not audio_path.startswith("/tmp"): + obj["audio"] = self.resolver.resolve_input_path(audio_path) + logging.info(f"Resolved audio path: {audio_path} -> {obj['audio']}") + if not os.path.exists(obj["audio"]): + logging.warning(f"Audio file not found: {obj['audio']}") + + if "mask" in obj and obj["mask"]: + mask_path = obj["mask"] + if not os.path.isabs(mask_path) and not mask_path.startswith("/tmp"): + obj["mask"] = self.resolver.resolve_input_path(mask_path) + logging.info(f"Resolved mask path: {mask_path} -> {obj['mask']}") + if not os.path.exists(obj["mask"]): + logging.warning(f"Mask file not found: {obj['mask']}") + if processed_talk_objects: config.talk_objects = processed_talk_objects logging.info(f"Processed {len(processed_talk_objects)} talk objects") + logging.info("lightx2v config: " + json.dumps(config, indent=2, ensure_ascii=False)) + config_hash = self._get_config_hash(config) needs_reinit = ( self.__class__._current_runner is None or self.__class__._current_config_hash != config_hash or getattr(config, "lazy_load", False) @@ -970,13 +841,13 @@ class LightX2VModularInference: progress = ProgressBar(100) - def update_progress(current_step, total): + def update_progress(current_step, _total): progress.update_absolute(current_step) if hasattr(self.__class__._current_runner, "set_progress_callback"): self.__class__._current_runner.set_progress_callback(update_progress) - result_dict = self.__class__._current_runner.run_pipeline(save_video=False) + result_dict = self.__class__._current_runner.run_pipeline() images = result_dict.get("video", None) audio = result_dict.get("audio", None) @@ -995,12 +866,8 @@ class LightX2VModularInference: raise finally: - for temp_file in temp_files: - if os.path.exists(temp_file): - try: - os.unlink(temp_file) - except Exception: - pass + # Cleanup is handled by TempFileManager destructor + pass NODE_CLASS_MAPPINGS = { @@ -1011,8 +878,10 @@ NODE_CLASS_MAPPINGS = { "LightX2VLoRALoader": LightX2VLoRALoader, "LightX2VConfigCombiner": LightX2VConfigCombiner, "LightX2VModularInference": LightX2VModularInference, - "TalkObjectInput": TalkObjectInput, - "TalkObjectsBuilder": TalkObjectsBuilder, + "LightX2VTalkObjectInput": TalkObjectInput, + "LightX2VTalkObjectsCombiner": TalkObjectsCombiner, + "LightX2VTalkObjectsFromJSON": TalkObjectsFromJSON, + "LightX2VTalkObjectsFromFiles": TalkObjectsFromFiles, } NODE_DISPLAY_NAME_MAPPINGS = { @@ -1023,6 +892,8 @@ NODE_DISPLAY_NAME_MAPPINGS = { "LightX2VLoRALoader": "LightX2V LoRA Loader", "LightX2VConfigCombiner": "LightX2V Config Combiner", "LightX2VModularInference": "LightX2V Modular Inference", - "TalkObjectInput": "Talk Object Input", - "TalkObjectsBuilder": "Talk Objects Builder", + "LightX2VTalkObjectInput": "LightX2V Talk Object Input (Single)", + "LightX2VTalkObjectsCombiner": "LightX2V Talk Objects Combiner", + "LightX2VTalkObjectsFromFiles": "LightX2V Talk Objects From Files", + "LightX2VTalkObjectsFromJSON": "LightX2V Talk Objects From JSON (API)", }