if
b7fef4f116
- Fixes a TypeError in the advanced generator by removing the 'compiled_model' kwarg before the function call. - Adds automatic model downloading to the HunyuanVideoFoleyModelLoader, ensuring the advanced workflow can self-initialize. This resolves issue #22.
939 lines
41 KiB
Python
939 lines
41 KiB
Python
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)",
|
|
} |