nodes.py

import os
import glob
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchaudio
import tempfile
import numpy as np
from loguru import logger
from typing import Optional, Tuple, Any, List, Dict
import folder_paths
import random
import urllib.request
import zipfile
import tarfile
from pathlib import Path
from datetime import datetime
import shutil
import time
import math
from tqdm import tqdm
from accelerate import init_empty_weights
from transformers import AutoTokenizer, AutoModel, ClapTextModelWithProjection
import comfy.model_management as mm
import comfy.utils
from safetensors.torch import load_file

try:
    import requests
    HAS_REQUESTS = True
except ImportError:
    HAS_REQUESTS = False

# Import ComfyUI video types
try:
    from comfy_api.input_impl import VideoFromFile
except ImportError:
    try:
        # Fallback to latest API location
        from comfy_api.latest._input_impl.video_types import VideoFromFile
    except ImportError:
        logger.warning("VideoFromFile not available, will return file paths only")
        VideoFromFile = None

# Add foley models directory to ComfyUI folder paths
foley_models_dir = os.path.join(folder_paths.models_dir, "foley")
if "foley" not in folder_paths.folder_names_and_paths:
    folder_paths.folder_names_and_paths["foley"] = ([foley_models_dir], folder_paths.supported_pt_extensions)

# Import the HunyuanVideo-Foley modules
try:
    from hunyuanvideo_foley.utils.model_utils import load_model
    from .utils import denoise_process_safely
    from .utils import feature_process_unified, extract_video_path, create_node_exit_values
    from hunyuanvideo_foley.utils.media_utils import merge_audio_video
    from .model_urls import get_model_url 
except ImportError as e:
    logger.error(f"Failed to import HunyuanVideo-Foley modules: {e}")
    logger.error("Make sure the HunyuanVideo-Foley package is installed and accessible")
    raise

class HunyuanVideoFoleyNode:
    """
    A node for generating audio using the HunyuanVideo-Foley model
    """
    # Class-level model storage for persistence
    _model_dict = None
    _cfg = None
    _device = None
    _model_path = None
    _memory_efficient = False  # Track memory mode
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "text_prompt": ("STRING", {
                    "multiline": True,
                    "default": "footstep sound, impact, water splash",
                    "display": "textarea",
                    "placeholder": "Describe the audio you want to generate..."
                }),
                "guidance_scale": ("FLOAT", {
                    "default": 4.5,
                    "min": 1.0,
                    "max": 10.0,
                    "step": 0.1,
                    "display": "slider"
                }),
                "num_inference_steps": ("INT", {
                    "default": 50,
                    "min": 10,
                    "max": 200,
                    "step": 1,
                    "display": "slider"
                }),
                "sample_nums": ("INT", {
                    "default": 1,
                    "min": 1,
                    "max": 4,
                    "step": 1,
                    "display": "slider"
                }),
                "seed": ("INT", {
                    "default": 42,
                    "min": 0,
                    "max": 2**32 - 1,
                    "display": "number"
                }),
            },
            "optional": {
                # Hybrid input support - either VIDEO or IMAGE
                "video": ("VIDEO",),
                "images": ("IMAGE",),  # Support for frame sequences from PR
                "fps": ("FLOAT", {
                    "default": 24.0,
                    "min": 1.0,
                    "max": 120.0,
                    "step": 0.1,
                    "tooltip": "Frames per second (only used with IMAGE input)"
                }),
                # Negative prompt from PR
                "negative_prompt": ("STRING", {
                    "multiline": True,
                    "default": "",
                    "display": "textarea",
                    "placeholder": "Additional negative prompts (optional). Will be combined with built-in quality controls."
                }),
                # Output options
                "output_format": (["video_path", "frames", "both"], {
                    "default": "both",
                    "tooltip": "Choose output format: video file path only works with VIDEO input, frames only works with IMAGE input and VIDEO input"
                }),
                "output_folder": ("STRING", {
                    "default": "hunyuan_foley",
                    "multiline": False,
                    "display": "text",
                    "placeholder": "Subfolder name in ComfyUI/output/"
                }),
                "filename_prefix": ("STRING", {
                    "default": "foley_",
                    "multiline": False,
                    "display": "text",
                    "placeholder": "Prefix for output filename"
                }),
                "feature_extraction_batch_size": ("INT", {
                    "default": 4, "min": 1, "max": 128, "step": 1,
                    "tooltip": "Frames to process at once for SigLIP2. Lower to save VRAM on long videos."
                }),
                "syncformer_batch_size": ("INT", {
                    "default": 1, "min": 1, "max": 64, "step": 1,
                    "tooltip": "Internal batch size for the Syncformer model. Lower if you still get OOM errors."
                }),
                # Memory optimization options
                "memory_efficient": ("BOOLEAN", {
                    "default": False,
                    "display": "checkbox",
                    "tooltip": "Enable memory-efficient mode: unloads models after generation and uses aggressive garbage collection"
                }),
                "cpu_offload": ("BOOLEAN", {
                    "default": False,
                    "display": "checkbox",
                    "tooltip": "Offload models to CPU when not in use (slower but saves VRAM)"
                }),
                "enabled": ("BOOLEAN", {
                    "default": True,
                    "display": "checkbox",
                    "tooltip": "Enable or disable the entire audio generation process. If false, returns a silent or null audio output."
                }),
                "silent_audio": ("BOOLEAN", {
                    "default": True,
                    "display": "checkbox",
                    "tooltip": "If true, returns a silent audio clip when disabled or on failure. If false, returns None."
                }),
            }
        }
    
    RETURN_TYPES = ("STRING", "IMAGE", "AUDIO", "STRING")
    RETURN_NAMES = ("video_path", "video_frames", "audio", "status_message")
    FUNCTION = "generate_audio"
    CATEGORY = "HunyuanVideo-Foley"
    DESCRIPTION = "Generate synchronized audio for videos using HunyuanVideo-Foley model"
    
    @classmethod
    def download_models(cls, model_name: str, destination_dir: str) -> Tuple[bool, str]:
        """Download models from URLs if they don't exist."""
        model_info = get_model_url(model_name)
        if not model_info:
            return False, f"Model '{model_name}' not found in configuration."

        os.makedirs(destination_dir, exist_ok=True)

        for model_file in model_info.get("models", []):
            filename = model_file.get("filename")
            url = model_file.get("url")

            if not filename or not url:
                logger.warning("Skipping a model file due to missing filename or URL.")
                continue

            destination_path = os.path.join(destination_dir, filename)

            if os.path.exists(destination_path):
                logger.info(f"'{filename}' already exists. Skipping download.")
                continue

            logger.info(f"Downloading '{filename}' from '{url}'...")

            try:
                comfy.utils.pb_download(url, destination_path)
            except Exception as e:
                return False, f"Failed to download '{filename}': {str(e)}"

        return True, "All models downloaded or already exist."
        
    @classmethod
    def setup_device(cls, device_type: str = "auto", device_id: int = 0):
        """Setup and return the appropriate device with memory optimization"""
        # Try to import ComfyUI's model management if available
        try:
            import comfy.model_management as mm
            device = mm.get_torch_device()
            logger.info(f"Using ComfyUI device: {device}")
            return device
        except:
            pass
        
        if device_type == "auto":
            if torch.cuda.is_available():
                device = torch.device(f"cuda:{device_id}")
                # Clear cache before loading
                torch.cuda.empty_cache()
                if hasattr(torch.cuda, 'reset_peak_memory_stats'):
                    torch.cuda.reset_peak_memory_stats(device)
            else:
                device = torch.device("cpu")
        else:
            device = torch.device(device_type)
        
        logger.info(f"Using device: {device}")
        return device
    
    @classmethod
    def load_models(cls, model_path: str = "", config_path: str = "", 
                   memory_efficient: bool = False, cpu_offload: bool = False) -> Tuple[bool, str]:
        """Load models if not already loaded or if path changed"""
        try:

            # Your proactive VRAM cleanup
            if cpu_offload or memory_efficient:
                logger.info("Proactively unloading other models to make space...")
                mm.unload_all_models()
                import gc; gc.collect()
                if torch.cuda.is_available():
                    torch.cuda.empty_cache() 

            # His default path logic
            if not model_path.strip():
                foley_models_dir = folder_paths.get_folder_paths("foley")[0]
                model_path = os.path.join(foley_models_dir, "hunyuanvideo-foley-xxl")

            # His auto-downloader logic
            if not os.path.exists(model_path):
                logger.info(f"Model directory not found at '{model_path}'. Attempting to download...")
                success, message = cls.download_models("hunyuanvideo-foley-xxl", model_path)
                if not success:
                    return False, f"Auto-download failed: {message}"
            
            # Your path logic
            if not config_path.strip():
                current_dir = os.path.dirname(os.path.abspath(__file__))
                config_path = os.path.join(current_dir, "configs", "hunyuanvideo-foley-xxl.yaml")
            
            # Your caching check
            if (cls._model_dict is not None and 
                cls._cfg is not None and 
                (cls._model_path == model_path or cls._model_path == "preloaded") and
                cls._memory_efficient == memory_efficient):
                return True, "Models already loaded"
            
            # Your setup logic
            cls._device = cls.setup_device("auto", 0)
            logger.info(f"Loading models from: {model_path}")
            logger.info(f"Config: {config_path}")
            
            # Your call to the memory-safe loader in utils
            cls._model_dict, cls._cfg = load_model(model_path, config_path, cls._device)
            
            # Your state management logic
            if cls._model_path != "preloaded":
                cls._model_path = model_path
            cls._memory_efficient = memory_efficient
            
            logger.info("Models loaded successfully!")
            return True, "Models loaded successfully!"
        
        except Exception as e:
            error_msg = f"Failed to load models: {str(e)}"
            logger.error(error_msg)
            cls._model_dict = None
            cls._cfg = None
            cls._device = None
            cls._model_path = None
            return False, error_msg
    
    def set_seed(self, seed: int):
        """Set random seed for reproducibility"""
        # Clamp seed to valid range for numpy (0 to 2^32-1)
        seed = int(seed) % (2**32)
        
        torch.manual_seed(seed)
        np.random.seed(seed)
        random.seed(seed)
        if torch.cuda.is_available():
            torch.cuda.manual_seed(seed)
            torch.cuda.manual_seed_all(seed)
    

    @classmethod
    def _extract_frames_from_image_input(cls, images, fps=24.0):
        """Convert IMAGE input to video frames and create a temporary video file"""
        import cv2
        
        try:
            if images is None:
                return None, "No images provided"
            
            # Handle different image input formats
            if hasattr(images, 'shape'):
                # Tensor input [batch, height, width, channels]
                if len(images.shape) == 4:
                    frames = images.cpu().numpy()
                else:
                    return None, f"Unexpected image tensor shape: {images.shape}"
            else:
                return None, f"Unsupported image input type: {type(images)}"
            
            # Convert to uint8 if needed
            if frames.dtype != np.uint8:
                if frames.max() <= 1.0:
                    frames = (frames * 255).astype(np.uint8)
                else:
                    frames = frames.astype(np.uint8)
            
            if frames.shape[-1] == 4:
                frames = frames[..., :3]
                
            # Get dimensions
            batch_size, height, width = frames.shape[:3]
            
            # Create temporary video file
            temp_fd, temp_path = tempfile.mkstemp(suffix='.mp4')
            os.close(temp_fd)
            
            # Setup video writer
            fourcc = cv2.VideoWriter_fourcc(*'mp4v')
            out = cv2.VideoWriter(temp_path, fourcc, fps, (width, height))
            
            if not out.isOpened():
                return None, "Failed to open video writer"
            
            # Write frames
            for i in range(batch_size):
                frame = frames[i]
                
                # Handle channels - convert RGB to BGR for OpenCV
                if len(frame.shape) == 3 and frame.shape[2] == 3:
                    frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
                elif len(frame.shape) == 3 and frame.shape[2] == 4:
                    frame = cv2.cvtColor(frame, cv2.COLOR_RGBA2BGR)
                
                out.write(frame)
            
            out.release()
            
            logger.info(f"Created temporary video from {batch_size} frames at {fps} FPS: {temp_path}")
            return temp_path, "Success"
            
        except Exception as e:
            return None, f"Error converting images to video: {str(e)}"
    
    # In HunyuanVideoFoleyNode class:
    @torch.inference_mode()
    def generate_audio(self, text_prompt: str, guidance_scale: float, 
                      num_inference_steps: int, sample_nums: int, seed: int,
                      video=None, images=None, fps=24.0,
                      negative_prompt="",
                      output_format="video_path",
                      output_folder: str = "hunyuan_foley",
                      filename_prefix: str = "foley_",
                      memory_efficient: bool = False,
                      cpu_offload: bool = False,
                      enabled: bool = True,
                      silent_audio: bool = True,
                      feature_extraction_batch_size: int = 4,
                      syncformer_batch_size: int = 1):       
        """Generate audio for the input video/images with the given text prompt"""
        
        # 1. Enabled/Disabled Short-Circuit with Intelligent Passthrough
        if not enabled:
            logger.info("HunyuanVideo-Foley node is disabled. Passing through inputs.")
            return create_node_exit_values(
                silent_audio=silent_audio,
                passthrough_video=video,
                passthrough_images=images,
                message="✅ Node disabled. Skipped."
            )
        
                # In the generate_audio function:
        try:
            self.set_seed(seed)

            # Proactive VRAM Cleanup (This part is correct)
            logger.info("Performing pre-run VRAM cleanup...")
            mm.unload_all_models()
            import gc; gc.collect()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            
            # --- CORRECTED MODEL LOADING LOGIC ---
            # This is the single source of truth for loading/validating models.
            if self._model_dict is None or self._cfg is None:
                logger.info("No models loaded, attempting to load fresh...")
                success, message = self.load_models(memory_efficient=memory_efficient, cpu_offload=cpu_offload)
                if not success:
                    raise Exception(f"Model loading failed: {message}")
            else:
                logger.info("Using pre-loaded models from class state.")
            
            # Final validation check after attempting to load/use pre-loaded.
            if self._model_dict is None or self._cfg is None:
                raise Exception("Models not loaded, and fallback loading failed.")
            # --- END OF FIX ---

            # Validate inputs
            if video is None and images is None:
                raise Exception("Please provide either video or images input!")
            
            # Determine video file path
            logger.info("Processing media features...")
            visual_feats, text_feats, audio_len_in_s = feature_process_unified(
                video_input=video,
                image_input=images,
                prompt=text_prompt,
                negative_prompt=negative_prompt,
                model_dict=self._model_dict,
                cfg=self._cfg,
                fps_hint=fps,
                batch_size=feature_extraction_batch_size,
                sync_batch_size=syncformer_batch_size
            )
            
            # Ensure the core models for denoising are on the correct device before processing.
            # --- State Correction and VRAM Management for Denoising ---
            target_device = self._device
            logger.info(f"Preparing for denoising on device: {target_device}")

            # 1. Offload the large feature tensors to CPU to make space.
            visual_feats['siglip2_feat'] = visual_feats['siglip2_feat'].to("cpu")
            visual_feats['syncformer_feat'] = visual_feats['syncformer_feat'].to("cpu")
            text_feats['text_feat'] = text_feats['text_feat'].to("cpu")
            text_feats['uncond_text_feat'] = text_feats['uncond_text_feat'].to("cpu")
            logger.info("Feature tensors moved to CPU.")

            # 2. Now that VRAM is clear, move the core models to the GPU.
            self._model_dict.foley_model.to(target_device)
            self._model_dict.dac_model.to(target_device)
            logger.info("Core models moved to GPU.")
            # --- End of Fix ---

            logger.info("Generating audio...")
            audio, sample_rate = denoise_process_safely(
                visual_feats, text_feats, audio_len_in_s, self._model_dict, self._cfg,
                guidance_scale=guidance_scale, num_inference_steps=num_inference_steps, batch_size=sample_nums
            )
            
            output_dir = os.path.join(folder_paths.get_output_directory(), output_folder)
            os.makedirs(output_dir, exist_ok=True)
            
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            audio_filename = f"{filename_prefix}audio_{timestamp}_{seed}.wav"
            audio_output_path = os.path.join(output_dir, audio_filename)
            torchaudio.save(audio_output_path, audio[0], sample_rate)
            
            audio_tensor = audio[0].unsqueeze(0)
            if len(audio_tensor.shape) == 2:
                audio_tensor = audio_tensor.unsqueeze(1)
            audio_result = {"waveform": audio_tensor, "sample_rate": sample_rate}
            
            video_file = None
            temp_video_created = False
            
            if video is not None:
                video_file = extract_video_path(video)
            elif images is not None:
                # If the input was an image tensor, we must create a temporary video file
                # so that merge_audio_video has a video stream to work with.
                video_file, _ = self._extract_frames_from_image_input(images, fps)
                temp_video_created = True

            if video_file is None:
                logger.warning("No valid video file path found for output merging/frame extraction.")
                
            video_output_path = ""
            video_frames = torch.zeros((1, 1, 1, 3), dtype=torch.float32)
            
            if output_format in ["video_path", "both"]:
                video_filename = f"{filename_prefix}video_{timestamp}_{seed}.mp4"
                video_output_path = os.path.join(output_dir, video_filename)
                try:
                    merge_audio_video(audio_output_path, video_file, video_output_path)
                    logger.info(f"Created video with audio: {video_output_path}")
                except Exception as e:
                    logger.error(f"Failed to merge audio and video: {e}")
                    video_output_path = video_file
            
            if output_format in ["frames", "both"]:
                if images is not None:
                    # If the input was an image tensor, just pass it through.
                    video_frames = images
                elif video_file:
                    # If the input was a video file, extract the frames.
                    try:
                        import cv2
                        cap = cv2.VideoCapture(video_file)
                        frames_list = []
                        while True:
                            ret, frame = cap.read()
                            if not ret: break
                            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                            frames_list.append(np.array(frame_rgb, dtype=np.float32) / 255.0)
                        cap.release()
                        if frames_list:
                            video_frames = torch.from_numpy(np.stack(frames_list))
                            logger.info(f"Extracted {len(frames_list)} frames for output")
                    except Exception as e:
                        logger.warning(f"Could not extract frames for output: {e}")
            
            if temp_video_created and os.path.exists(video_file):
                try: os.remove(video_file)
                except: pass
            
            if memory_efficient:
                # Only unload completely if the model was NOT preloaded by an upstream node.
                if self._model_path != "preloaded":
                    logger.info("Memory efficient mode: Unloading models completely.")
                    self._model_dict = None
                    self._cfg = None
                    self._model_path = None
                else:
                    logger.info("Memory efficient mode: Models are preloaded, skipping complete unload.")
            
            success_msg = f"✅ Generated audio successfully"
            return (video_output_path, video_frames, audio_result, success_msg)
            
        except Exception as e:
            import traceback
            error_msg = f"❌ Generation failed: {str(e)}"
            logger.error(error_msg); logger.error(traceback.format_exc())
            return create_node_exit_values(
                silent_audio=silent_audio,
                passthrough_video=video,
                passthrough_images=images,
                message=error_msg
            )
        
        finally:
            logger.info("HunyuanVideo-Foley: Starting guaranteed cleanup...")
            if (memory_efficient or cpu_offload) and self._model_dict:
                logger.info("Offloading models to CPU...")
                for model_obj in self._model_dict.values():
                    if hasattr(model_obj, 'to'):
                        try: model_obj.to("cpu")
                        except Exception: pass
            
            if memory_efficient and self._model_path != "preloaded":
                logger.info("Memory efficient mode: Unloading models completely.")
                self._model_dict, self._cfg, self._model_path = None, None, None

            import gc; gc.collect()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            logger.info("HunyuanVideo-Foley: Cleanup complete.")
        
class LinearFP8Wrapper(nn.Module):
    """FP8 quantization wrapper for linear layers"""
    def __init__(self, original_linear, dtype="fp8_e4m3fn"):
        super().__init__()
        self.dtype = dtype
        self.weight_shape = original_linear.weight.shape
        self.bias = original_linear.bias
        
        # Quantize weights to FP8
        if dtype == "fp8_e4m3fn":
            self.weight_fp8 = original_linear.weight.to(torch.float8_e4m3fn)
        elif dtype == "fp8_e5m2":
            self.weight_fp8 = original_linear.weight.to(torch.float8_e5m2)
        else:
            self.weight_fp8 = original_linear.weight
    
    def forward(self, x):
        # Convert back to computation dtype for matmul
        weight = self.weight_fp8.to(x.dtype)
        return F.linear(x, weight, self.bias)

class HunyuanVideoFoleyModelLoader:
    """Separate model loader with FP8 quantization support"""
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "quantization": (["none", "fp8_e4m3fn", "fp8_e5m2"], {
                    "default": "none",
                    "tooltip": "FP8 weight-only quantization for VRAM savings"
                }),
                "cpu_offload": ("BOOLEAN", {
                    "default": False,
                    "tooltip": "Offload models to CPU when not in use"
                }),
            },
            "optional": {
                "model_path": ("STRING", {
                    "default": "",
                    "multiline": False,
                    "placeholder": "Path to model weights (leave empty for default)"
                }),
                "config_path": ("STRING", {
                    "default": "",
                    "multiline": False,
                    "placeholder": "Path to config file (leave empty for default)"
                }),
            }
        }
    
    RETURN_TYPES = ("FOLEY_MODEL", "STRING")
    RETURN_NAMES = ("model", "status")
    FUNCTION = "load_model"
    CATEGORY = "HunyuanVideo-Foley/Loaders"
    DESCRIPTION = "Load HunyuanVideo-Foley model with optional FP8 quantization"
    
    def load_model(self, quantization="none", cpu_offload=False, model_path="", config_path=""):
        try:
            # Set default paths
            if not model_path or not model_path.strip():
                foley_models_dir = folder_paths.folder_names_and_paths.get("foley", [None])[0]
                if foley_models_dir and len(foley_models_dir) > 0:
                    # Check for the model in the hunyuanvideo-foley-xxl subdirectory
                    model_path = os.path.join(foley_models_dir[0], "hunyuanvideo-foley-xxl")
                else:
                    current_dir = os.path.dirname(os.path.abspath(__file__))
                    model_path = os.path.join(current_dir, "pretrained_models")
            
            # Auto-download models if the directory doesn't exist
            if not os.path.exists(model_path):
                logger.info(f"Model directory not found at '{model_path}'. Attempting to download...")
                # Use the download method from the base node
                success, message = HunyuanVideoFoleyNode.download_models("hunyuanvideo-foley-xxl", model_path)
                if not success:
                    return (None, f"❌ Auto-download failed: {message}")
            
            if not config_path or not config_path.strip():
                current_dir = os.path.dirname(os.path.abspath(__file__))
                config_path = os.path.join(current_dir, "configs", "hunyuanvideo-foley-xxl.yaml")
            
            logger.info(f"Model path: {model_path}")
            logger.info(f"Config path: {config_path}")
            
            # Setup device
            device = mm.get_torch_device()
            
            logger.info(f"Loading model with quantization={quantization}")
            
            # Load model and config
            model_dict, cfg = load_model(model_path, config_path, device)
            
            # Apply FP8 quantization if requested
            if quantization != "none" and "vae" in model_dict:
                logger.info(f"Applying {quantization} quantization to VAE...")
                vae_model = model_dict["vae"]
                
                # Quantize linear layers in VAE
                for name, module in vae_model.named_modules():
                    if isinstance(module, nn.Linear):
                        # Replace with FP8 wrapper
                        parent_name = ".".join(name.split(".")[:-1]) if "." in name else ""
                        child_name = name.split(".")[-1]
                        parent = vae_model if not parent_name else dict(vae_model.named_modules())[parent_name]
                        setattr(parent, child_name, LinearFP8Wrapper(module, quantization))
                
                logger.info("FP8 quantization applied")
            
            # Package model info
            model_info = {
                "model_dict": model_dict,
                "cfg": cfg,
                "device": device,
                "quantization": quantization,
                "cpu_offload": cpu_offload
            }
            
            status = f"✅ Model loaded with {quantization} quantization" if quantization != "none" else "✅ Model loaded"
            return (model_info, status)
            
        except Exception as e:
            logger.error(f"Failed to load model: {e}")
            return (None, f"❌ Failed to load model: {str(e)}")

class HunyuanVideoFoleyDependenciesLoader:
    """Load text encoder and feature extractors separately"""
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model": ("FOLEY_MODEL",),
                "load_text_encoder": ("BOOLEAN", {
                    "default": True,
                    "tooltip": "Load CLAP text encoder"
                }),
                "load_feature_extractor": ("BOOLEAN", {
                    "default": True,
                    "tooltip": "Load visual feature extractor"
                }),
            }
        }
    
    RETURN_TYPES = ("FOLEY_DEPS", "STRING")
    RETURN_NAMES = ("dependencies", "status")
    FUNCTION = "load_dependencies"
    CATEGORY = "HunyuanVideo-Foley/Loaders"
    DESCRIPTION = "Load model dependencies (text encoder, feature extractors)"
    
    def load_dependencies(self, model, load_text_encoder=True, load_feature_extractor=True):
        try:
            if model is None:
                return (None, "❌ No model provided - check ModelLoader output")
            
            # Check if model is a tuple (from failed load)
            if isinstance(model, tuple):
                model = model[0]
            
            if model is None or not isinstance(model, dict):
                return (None, "❌ Invalid model input - ModelLoader may have failed")
            
            deps = {
                "model_dict": model["model_dict"],
                "cfg": model["cfg"],
                "device": model["device"]
            }
            
            status_parts = []
            
            if load_text_encoder:
                logger.info("Loading text encoder...")
                # Text encoder is already in model_dict
                status_parts.append("text encoder")
            
            if load_feature_extractor:
                logger.info("Loading feature extractor...")
                # Feature extractor is already in model_dict
                status_parts.append("feature extractor")
            
            status = f"✅ Loaded: {', '.join(status_parts)}" if status_parts else "✅ Dependencies ready"
            return (deps, status)
            
        except Exception as e:
            logger.error(f"Failed to load dependencies: {e}")
            return (None, f"❌ Failed: {str(e)}")

class HunyuanVideoFoleyTorchCompile:
    """Apply torch.compile optimization to the model"""
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "dependencies": ("FOLEY_DEPS",),
                "compile_vae": ("BOOLEAN", {
                    "default": True,
                    "tooltip": "Compile VAE with torch.compile for ~30% speedup"
                }),
                "compile_mode": (["default", "reduce-overhead", "max-autotune"], {
                    "default": "default",
                    "tooltip": "Compilation mode (default is fastest compile time)"
                }),
                "backend": (["inductor", "cudagraphs", "eager"], {
                    "default": "inductor",
                    "tooltip": "Compilation backend"
                }),
            }
        }
    
    RETURN_TYPES = ("FOLEY_COMPILED", "STRING")
    RETURN_NAMES = ("compiled_model", "status")
    FUNCTION = "compile_model"
    CATEGORY = "HunyuanVideo-Foley/Optimization"
    DESCRIPTION = "Optimize model with torch.compile for faster inference"
    
    def compile_model(self, dependencies, compile_vae=True, compile_mode="default", backend="inductor"):
        try:
            if dependencies is None:
                return (None, "❌ No dependencies provided - check DependenciesLoader output")
            
            # Check if dependencies is a tuple (from failed load)
            if isinstance(dependencies, tuple):
                dependencies = dependencies[0]
            
            if dependencies is None or not isinstance(dependencies, dict):
                return (None, "❌ Invalid dependencies - previous node may have failed")
            
            # Handle AttributeDict or regular dict
            model_dict = dependencies["model_dict"]
            # Don't copy AttributeDict, just reference it
            if hasattr(model_dict, '__class__') and 'AttributeDict' in str(model_dict.__class__):
                model_dict_ref = model_dict
            else:
                model_dict_ref = model_dict.copy() if hasattr(model_dict, 'copy') else model_dict
            
            compiled = {
                "model_dict": model_dict_ref,
                "cfg": dependencies["cfg"],
                "device": dependencies["device"],
                "compiled": False
            }
            
            # Check for dac_model (the actual VAE used in HunyuanVideo-Foley)
            model_dict = compiled["model_dict"]
            has_dac = hasattr(model_dict, 'dac_model') or (isinstance(model_dict, dict) and "dac_model" in model_dict)
            
            if compile_vae and has_dac:
                logger.info(f"Compiling DAC VAE with mode={compile_mode}, backend={backend}...")
                
                import torch._dynamo as dynamo
                dynamo.config.suppress_errors = True
                
                # Only compile if backend is not eager
                if backend != "eager":
                    # Access dac_model correctly whether it's dict or AttributeDict
                    if hasattr(model_dict, 'dac_model'):
                        dac_model = model_dict.dac_model
                    else:
                        dac_model = model_dict["dac_model"]
                    
                    compiled_dac = torch.compile(
                        dac_model,
                        mode=compile_mode,
                        backend=backend
                    )
                    
                    # Set the compiled model back
                    if hasattr(model_dict, 'dac_model'):
                        model_dict.dac_model = compiled_dac
                    else:
                        model_dict["dac_model"] = compiled_dac
                    
                    compiled["compiled"] = True
                    status = f"✅ DAC VAE compiled with {compile_mode}/{backend}"
                else:
                    status = "✅ Model ready (eager mode - no compilation)"
            else:
                status = "✅ Model ready (no VAE compilation)"
            
            return (compiled, status)
            
        except Exception as e:
            logger.error(f"Failed to compile model: {e}")
            # Return uncompiled model on failure
            if dependencies:
                return (dependencies, f"⚠️ Compilation failed, using uncompiled: {str(e)}")
            return (None, f"❌ Failed: {str(e)}")

class HunyuanVideoFoleyGeneratorAdvanced(HunyuanVideoFoleyNode):
    """Enhanced generator that can use separately loaded models"""
    
    @classmethod
    def INPUT_TYPES(cls):
        base_inputs = super().INPUT_TYPES()
        
        # Add optional compiled model input
        base_inputs["optional"]["compiled_model"] = ("FOLEY_COMPILED",)
        # --- QoL FIELDS HERE ---
        base_inputs["optional"]["enabled"] = ("BOOLEAN", {
            "default": True,
            "display": "checkbox",
            "tooltip": "Enable or disable this node. If disabled, it will pass through silent audio."
        })
        base_inputs["optional"]["silent_audio"] = ("BOOLEAN", {
            "default": True,
            "display": "checkbox",
            "tooltip": "If true, returns a silent audio clip when disabled or on failure. If false, returns None."
        })
        # -----------------------------

        return base_inputs
    
    FUNCTION = "generate_audio_advanced"
    CATEGORY = "HunyuanVideo-Foley"
    DESCRIPTION = "Generate audio with optional pre-loaded/optimized models"
    
    def generate_audio_advanced(self, **kwargs):
    
        compiled_model = kwargs.get("compiled_model")

        if compiled_model:
            if isinstance(compiled_model, tuple): compiled_model = compiled_model[0]
            if compiled_model and isinstance(compiled_model, dict):
                self.__class__._model_dict = compiled_model.get("model_dict")
                self.__class__._cfg = compiled_model.get("cfg")
                self.__class__._device = compiled_model.get("device")
                self.__class__._model_path = "preloaded"
                self.__class__._memory_efficient = kwargs.get("memory_efficient", False)
                logger.info("Using pre-loaded/compiled model from input.")
            else:
                logger.warning("Invalid compiled_model input. Clearing state to force reload.")
                self.__class__._model_dict, self.__class__._cfg, self.__class__._model_path = None, None, None
        else:
            logger.info("No pre-loaded model detected. Clearing state to force fresh load.")
            self.__class__._model_dict, self.__class__._cfg, self.__class__._model_path = None, None, None
            
        kwargs.pop('compiled_model', None)
        return self.generate_audio(**kwargs)

# Node mappings for ComfyUI
NODE_CLASS_MAPPINGS = {
    "HunyuanVideoFoley": HunyuanVideoFoleyNode,
    "HunyuanVideoFoleyModelLoader": HunyuanVideoFoleyModelLoader,
    "HunyuanVideoFoleyDependenciesLoader": HunyuanVideoFoleyDependenciesLoader,
    "HunyuanVideoFoleyTorchCompile": HunyuanVideoFoleyTorchCompile,
    "HunyuanVideoFoleyGeneratorAdvanced": HunyuanVideoFoleyGeneratorAdvanced,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "HunyuanVideoFoley": "HunyuanVideo-Foley Generator",
    "HunyuanVideoFoleyModelLoader": "HunyuanVideo-Foley Model Loader (FP8)",
    "HunyuanVideoFoleyDependenciesLoader": "HunyuanVideo-Foley Dependencies",
    "HunyuanVideoFoleyTorchCompile": "HunyuanVideo-Foley Torch Compile",
    "HunyuanVideoFoleyGeneratorAdvanced": "HunyuanVideo-Foley Generator (Advanced)",
}