if@impactframes
fdc735cd05
- Fix AttributeDict handling in TorchCompile node to prevent .copy() errors - Fix duplicate model loading in GeneratorAdvanced causing memory allocation failures - Update model loader to properly handle preloaded models from pipeline - Fix DAC VAE compilation (was looking for 'vae' key instead of 'dac_model') - Make model_path and config_path optional parameters with proper defaults - Add better error handling and validation throughout the pipeline - Prevent unnecessary model reloading when using pre-compiled models The advanced pipeline now works correctly: ModelLoader → DependenciesLoader → TorchCompile → GeneratorAdvanced Co-Authored-By: if@impactframes <if@impactframes.ai>
852 lines
36 KiB
Python
852 lines
36 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, denoise_process
|
|
from hunyuanvideo_foley.utils.feature_utils import feature_process
|
|
from hunyuanvideo_foley.utils.media_utils import merge_audio_video
|
|
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"
|
|
}),
|
|
# 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)"
|
|
}),
|
|
}
|
|
}
|
|
|
|
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 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:
|
|
# Set default paths if empty
|
|
if not model_path.strip():
|
|
# Try ComfyUI foley models directory first
|
|
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:
|
|
# Fallback to custom node directory
|
|
current_dir = os.path.dirname(os.path.abspath(__file__))
|
|
model_path = os.path.join(current_dir, "pretrained_models")
|
|
|
|
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")
|
|
|
|
# Check if models are already loaded with the same path and memory mode
|
|
# Also check for "preloaded" which means models came from pipeline
|
|
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"
|
|
|
|
# Setup device
|
|
cls._device = cls.setup_device("auto", 0)
|
|
|
|
logger.info(f"Loading models from: {model_path}")
|
|
logger.info(f"Config: {config_path}")
|
|
|
|
# Load models
|
|
cls._model_dict, cls._cfg = load_model(model_path, config_path, cls._device)
|
|
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)
|
|
|
|
@staticmethod
|
|
def _extract_video_path(video):
|
|
"""Extract a file path from various potential video input types"""
|
|
# Handle ComfyUI VideoFromFile object
|
|
if hasattr(video, '__class__') and 'VideoFromFile' in video.__class__.__name__:
|
|
if hasattr(video, '_VideoFromFile__file'):
|
|
return getattr(video, '_VideoFromFile__file')
|
|
for attr in ['file', 'path', 'filename']:
|
|
if hasattr(video, attr):
|
|
value = getattr(video, attr)
|
|
if isinstance(value, str):
|
|
return value
|
|
|
|
# Direct string path
|
|
if isinstance(video, str):
|
|
return video
|
|
elif isinstance(video, dict) and 'path' in video:
|
|
return video['path']
|
|
|
|
return None
|
|
|
|
@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)
|
|
|
|
# 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)}"
|
|
|
|
@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):
|
|
"""Generate audio for the input video/images with the given text prompt"""
|
|
try:
|
|
# Set seed for reproducibility
|
|
self.set_seed(seed)
|
|
|
|
# Check if models are already loaded (from pipeline or previous run)
|
|
if self._model_dict is None or self._cfg is None:
|
|
# Load models if needed
|
|
success, message = self.load_models("", "", memory_efficient, cpu_offload)
|
|
if not success:
|
|
logger.error(f"Model loading failed: {message}")
|
|
empty_audio = {"waveform": torch.zeros((1, 1, 48000)), "sample_rate": 48000}
|
|
empty_frames = torch.zeros((1, 64, 64, 3), dtype=torch.float32)
|
|
return ("", empty_frames, empty_audio, f"❌ {message}")
|
|
else:
|
|
logger.info("Using already loaded models")
|
|
|
|
# Validate that models are loaded
|
|
if self._model_dict is None or self._cfg is None:
|
|
error_msg = "Models not loaded"
|
|
logger.error(error_msg)
|
|
empty_audio = {"waveform": torch.zeros((1, 1, 48000)), "sample_rate": 48000}
|
|
empty_frames = torch.zeros((1, 64, 64, 3), dtype=torch.float32)
|
|
return ("", empty_frames, empty_audio, f"❌ {error_msg}")
|
|
|
|
# Validate inputs
|
|
if video is None and images is None:
|
|
empty_audio = {"waveform": torch.zeros((1, 1, 48000)), "sample_rate": 48000}
|
|
empty_frames = torch.zeros((1, 64, 64, 3), dtype=torch.float32)
|
|
return ("", empty_frames, empty_audio, "❌ Please provide either video or images input!")
|
|
|
|
# Determine video file path
|
|
video_file = None
|
|
temp_video_created = False
|
|
|
|
if video is not None:
|
|
video_file = self._extract_video_path(video)
|
|
elif images is not None:
|
|
video_file, convert_msg = self._extract_frames_from_image_input(images, fps)
|
|
temp_video_created = True
|
|
if video_file is None:
|
|
empty_audio = {"waveform": torch.zeros((1, 1, 48000)), "sample_rate": 48000}
|
|
empty_frames = torch.zeros((1, 64, 64, 3), dtype=torch.float32)
|
|
return ("", empty_frames, empty_audio, f"❌ {convert_msg}")
|
|
|
|
if video_file is None or not os.path.exists(video_file):
|
|
empty_audio = {"waveform": torch.zeros((1, 1, 48000)), "sample_rate": 48000}
|
|
empty_frames = torch.zeros((1, 64, 64, 3), dtype=torch.float32)
|
|
return ("", empty_frames, empty_audio, "❌ Video file not found")
|
|
|
|
# Process features
|
|
logger.info("Processing video features...")
|
|
visual_feats, text_feats, audio_len_in_s = feature_process(
|
|
video_file,
|
|
text_prompt,
|
|
self._model_dict,
|
|
self._cfg
|
|
)
|
|
|
|
# Generate audio
|
|
logger.info("Generating audio...")
|
|
audio, sample_rate = denoise_process(
|
|
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
|
|
)
|
|
|
|
# Create output directory
|
|
output_dir = os.path.join(folder_paths.get_output_directory(), output_folder)
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
|
|
# Generate timestamp for unique filename
|
|
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
|
|
|
# Save audio file
|
|
audio_filename = f"{filename_prefix}audio_{timestamp}_{seed}.wav"
|
|
audio_output = os.path.join(output_dir, audio_filename)
|
|
torchaudio.save(audio_output, audio[0], sample_rate)
|
|
|
|
# Create audio result dict
|
|
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}
|
|
|
|
# Handle output formats
|
|
video_output_path = ""
|
|
video_frames = torch.zeros((1, 64, 64, 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, 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"]:
|
|
# Extract frames for output
|
|
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)
|
|
frame_normalized = np.array(frame_rgb, dtype=np.float32) / 255.0
|
|
frames_list.append(frame_normalized)
|
|
|
|
cap.release()
|
|
|
|
if frames_list:
|
|
video_frames = torch.from_numpy(np.stack(frames_list))
|
|
logger.info(f"Extracted {len(frames_list)} frames")
|
|
except Exception as e:
|
|
logger.warning(f"Could not extract frames: {e}")
|
|
|
|
# Cleanup
|
|
if temp_video_created and os.path.exists(video_file):
|
|
try:
|
|
os.remove(video_file)
|
|
except:
|
|
pass
|
|
|
|
# Memory cleanup if requested
|
|
if memory_efficient:
|
|
# Clear intermediate variables
|
|
del visual_feats, text_feats, audio
|
|
|
|
# Force garbage collection
|
|
import gc
|
|
gc.collect()
|
|
|
|
# Clear CUDA cache
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
|
|
success_msg = f"✅ Generated audio successfully"
|
|
return (video_output_path, video_frames, audio_result, success_msg)
|
|
|
|
except Exception as e:
|
|
error_msg = f"❌ Generation failed: {str(e)}"
|
|
logger.error(error_msg)
|
|
empty_audio = {"waveform": torch.zeros((1, 1, 48000)), "sample_rate": 48000}
|
|
empty_frames = torch.zeros((1, 64, 64, 3), dtype=torch.float32)
|
|
return ("", empty_frames, empty_audio, error_msg)
|
|
|
|
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")
|
|
|
|
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",)
|
|
|
|
return base_inputs
|
|
|
|
FUNCTION = "generate_audio_advanced"
|
|
CATEGORY = "HunyuanVideo-Foley"
|
|
DESCRIPTION = "Generate audio with optional pre-loaded/optimized models"
|
|
|
|
def generate_audio_advanced(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="hunyuan_foley",
|
|
filename_prefix="foley_",
|
|
memory_efficient=False,
|
|
cpu_offload=False,
|
|
compiled_model=None):
|
|
"""Generate audio using either compiled model or loading fresh"""
|
|
|
|
# If compiled model provided, use it
|
|
if compiled_model is not None:
|
|
# Check if compiled_model is a tuple (from node output)
|
|
if isinstance(compiled_model, tuple):
|
|
compiled_model = compiled_model[0]
|
|
|
|
if compiled_model and isinstance(compiled_model, dict):
|
|
# Set class-level variables to prevent reloading
|
|
self.__class__._model_dict = compiled_model["model_dict"]
|
|
self.__class__._cfg = compiled_model["cfg"]
|
|
self.__class__._device = compiled_model["device"]
|
|
self.__class__._model_path = "preloaded"
|
|
self.__class__._memory_efficient = memory_efficient
|
|
|
|
logger.info("Using pre-loaded/compiled model")
|
|
else:
|
|
logger.warning("Invalid compiled model provided, will load fresh")
|
|
|
|
# Call parent generate_audio
|
|
return self.generate_audio(
|
|
text_prompt, guidance_scale, num_inference_steps, sample_nums, seed,
|
|
video, images, fps, negative_prompt, output_format,
|
|
output_folder, filename_prefix, memory_efficient, cpu_offload
|
|
)
|
|
|
|
# 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)",
|
|
} |