""" Utility functions for ComfyUI HunyuanVideo-Foley custom node """ import os import tempfile import torch import numpy as np from typing import Union, Optional, Tuple from loguru import logger import decord from tqdm import tqdm from PIL import Image from einops import rearrange import comfy.model_management as mm # We need to import the original library functions that our safe wrappers will call. from hunyuanvideo_foley.utils.feature_utils import encode_text_feat, encode_video_with_siglip2, encode_video_with_sync from hunyuanvideo_foley.utils.config_utils import AttributeDict import time class SimpleProfiler: """A simple profiler for timing and CUDA memory logging.""" def __init__(self, name, enabled=True): self.name = name self.enabled = enabled self.start_time = None self.device = get_optimal_device() def __enter__(self): if not self.enabled: return self self.start_time = time.time() if self.device.type == 'cuda': torch.cuda.synchronize(self.device) start_mem = torch.cuda.memory_allocated(self.device) / 1024**2 logger.info(f"[Profiler:{self.name}] Start VRAM: {start_mem:.2f} MB") return self def __exit__(self, exc_type, exc_val, exc_tb): if not self.enabled: return elapsed = time.time() - self.start_time if self.device.type == 'cuda': torch.cuda.synchronize(self.device) end_mem = torch.cuda.memory_allocated(self.device) / 1024**2 logger.info(f"[Profiler:{self.name}] Time: {elapsed:.4f}s. End VRAM: {end_mem:.2f} MB") else: logger.info(f"[Profiler:{self.name}] Time: {elapsed:.4f}s. (CPU)") def tensor_to_video(video_tensor: torch.Tensor, output_path: str, fps: int = 30) -> str: """ Convert a video tensor to a video file """ try: import cv2 video_np = video_tensor.detach().cpu().numpy() if isinstance(video_tensor, torch.Tensor) else np.array(video_tensor) if video_np.ndim == 4 and video_np.shape[1] == 3: video_np = np.transpose(video_np, (0, 2, 3, 1)) elif video_np.ndim == 5 and video_np.shape[1] == 3: video_np = np.transpose(video_np[0], (0, 2, 3, 1)) video_np = (video_np * 255).astype(np.uint8) if video_np.max() <= 1.0 else video_np.astype(np.uint8) frames, height, width, channels = video_np.shape fourcc = cv2.VideoWriter_fourcc(*'mp4v') out = cv2.VideoWriter(output_path, fourcc, fps, (width, height)) for i in range(frames): frame = video_np[i] if channels == 3: frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR) out.write(frame) out.release() logger.info(f"Video saved to: {output_path}") return output_path except Exception as e: logger.error(f"Failed to convert tensor to video: {e}") raise def get_video_info(video_path: str) -> dict: """ Get information about a video file """ try: import cv2 cap = cv2.VideoCapture(video_path) info = { 'fps': cap.get(cv2.CAP_PROP_FPS), 'frame_count': int(cap.get(cv2.CAP_PROP_FRAME_COUNT)), 'width': int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), 'height': int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)), 'duration': int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) / cap.get(cv2.CAP_PROP_FPS) if cap.get(cv2.CAP_PROP_FPS) > 0 else 0 } cap.release() return info except Exception as e: logger.error(f"Failed to get video info: {e}") return {} def ensure_video_file(video_input: Union[str, torch.Tensor, np.ndarray]) -> str: """ Ensure the video input is converted to a file path """ if isinstance(video_input, str): if os.path.exists(video_input): return video_input else: raise FileNotFoundError(f"Video file not found: {video_input}") elif isinstance(video_input, (torch.Tensor, np.ndarray)): temp_dir = tempfile.mkdtemp() output_path = os.path.join(temp_dir, "input_video.mp4") return tensor_to_video(video_input, output_path) else: raise ValueError(f"Unsupported video input type: {type(video_input)}") def validate_model_files(model_path: str) -> Tuple[bool, str]: """ Validate that all required model files exist """ required_files = ["hunyuanvideo_foley.pth", "vae_128d_48k.pth", "synchformer_state_dict.pth"] missing_files = [f for f in required_files if not os.path.exists(os.path.join(model_path, f))] if missing_files: return False, f"Missing model files: {', '.join(missing_files)}" return True, "All required model files found" def get_optimal_device() -> torch.device: """ Get the optimal device for model execution """ if torch.cuda.is_available(): device = torch.device(f"cuda:{torch.cuda.current_device()}") props = torch.cuda.get_device_properties(device) logger.info(f"Using CUDA device: {props.name} with {props.total_memory / 1e9:.1f}GB memory") return device elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available(): logger.info("Using MPS device (Apple Silicon)") return torch.device("mps") else: logger.info("Using CPU device") return torch.device("cpu") def check_memory_requirements(device: torch.device, required_gb: float = 16.0) -> Tuple[bool, str]: """ Check if the device has enough memory for model execution """ if device.type == "cuda": total_memory = torch.cuda.get_device_properties(device).total_memory / 1e9 if total_memory < required_gb: return False, f"GPU has {total_memory:.1f}GB memory, but {required_gb}GB is recommended" else: return True, f"GPU has {total_memory:.1f}GB memory (sufficient)" return True, "Using non-CUDA device (memory check not applicable)" def format_duration(seconds: float) -> str: """ Format duration in seconds to human readable format """ seconds = float(seconds) if seconds < 60: return f"{seconds:.1f}s" minutes, seconds = divmod(seconds, 60) if minutes < 60: return f"{int(minutes)}m {seconds:.1f}s" hours, minutes = divmod(minutes, 60) return f"{int(hours)}h {int(minutes)}m" # --- NEW MEMORY-SAFE AND UNIFIED FUNCTIONS --- def extract_video_path(video): """Extract a file path from various potential video input types. Moved here to be a standalone utility.""" 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) and isinstance(getattr(video, attr), str): return getattr(video, attr) if isinstance(video, str): return video if isinstance(video, dict) and 'path' in video: return video['path'] return None # In utils.py, replace the existing function @torch.inference_mode() def _encode_visual_features_safely(frames_uint8, model_dict, fps_hint, batch_size=4, sync_batch_size=1, enable_profiling=False): """ Memory-safe visual feature extraction using multithreaded preprocessing. """ dev = model_dict.device logger.info(f"--- Starting Safe Visual Feature Extraction ---") logger.info(f"Input frames: {len(frames_uint8)}, SigLIP Batch: {batch_size}, Sync Batch: {sync_batch_size}") mm.unload_all_models(); import gc; gc.collect() if dev.type == 'cuda': torch.cuda.empty_cache() visual_features = {} pil_list = [Image.fromarray(f).convert("RGB") for f in frames_uint8] # --- STAGE 1: SigLIP2 (with Multithreaded Preprocessing) --- all_siglip_feats = [] try: logger.info("[SigLIP2] Moving model to device...") model_dict.siglip2_model.to(dev) from concurrent.futures import ThreadPoolExecutor # Use half the CPU cores for safety, just like before num_workers = max(1, os.cpu_count() // 2) logger.info(f"[SigLIP2] Using {num_workers} CPU threads for parallel preprocessing.") with SimpleProfiler("SigLIP2 Batching", enabled=enable_profiling): for i in tqdm(range(0, len(pil_list), batch_size), desc="Processing SigLIP2 batches"): batch_pils = pil_list[i:i+batch_size] # --- MULTITHREADED PREPROCESSING --- with ThreadPoolExecutor(max_workers=num_workers) as executor: siglip_list = list(executor.map(model_dict.siglip2_preprocess, batch_pils)) # --- END OF MULTITHREADED --- clip_frames = torch.stack(siglip_list, dim=0).to(dev) batch_feat = _encode_video_with_siglip2_safely(clip_frames, model_dict) all_siglip_feats.append(batch_feat.cpu()) del batch_pils, siglip_list, clip_frames, batch_feat if all_siglip_feats: visual_features['siglip2_feat'] = torch.cat(all_siglip_feats, dim=0).unsqueeze(0).to(dev) finally: logger.info("[SigLIP2] Offloading model from device."); model_dict.siglip2_model.to("cpu") del all_siglip_feats if dev.type == 'cuda': torch.cuda.empty_cache() # --- END STAGE 1 --- # --- STAGE 2: Syncformer (remains the same) --- try: logger.info("[Syncformer] Moving model to device...") model_dict.syncformer_model.to(dev) with SimpleProfiler("Syncformer Processing", enabled=enable_profiling): images = torch.from_numpy(np.array(frames_uint8)).permute(0, 3, 1, 2) sync_frames = model_dict.syncformer_preprocess(images).unsqueeze(0).to(dev) visual_features['syncformer_feat'] = encode_video_with_sync(sync_frames, model_dict, batch_size=sync_batch_size) finally: logger.info("[Syncformer] Offloading model from device."); model_dict.syncformer_model.to("cpu") if dev.type == 'cuda': torch.cuda.empty_cache() # --- END STAGE 2 --- audio_len_in_s = len(frames_uint8) / fps_hint return AttributeDict(visual_features), audio_len_in_s def feature_process_unified(video_input, image_input, prompt, model_dict, cfg, negative_prompt="", fps_hint=24.0, batch_size=4, sync_batch_size=1, enable_profiling=False): """ Unified, memory-safe feature processing for either a video path or an image tensor. """ frames_uint8 = None; fps = fps_hint if image_input is not None: logger.info("Processing features from IMAGE tensor input.") frames_uint8 = (image_input.cpu().clip(0, 1) * 255).byte().numpy() if image_input.dtype != torch.uint8 else image_input.cpu().numpy() elif video_input is not None: logger.info("Processing features from VIDEO path input.") video_path = extract_video_path(video_input) if not video_path: raise ValueError("Invalid video path provided.") video_reader = decord.VideoReader(video_path) total_frames = len(video_reader) fps = video_reader.get_avg_fps() frame_indices = np.linspace(0, total_frames - 1, num=total_frames, dtype=int) frames_uint8 = video_reader.get_batch(frame_indices).asnumpy() if frames_uint8 is None: raise ValueError("No valid video or image frames to process.") visual_feats, audio_len_in_s = _encode_visual_features_safely(frames_uint8, model_dict, fps, batch_size, sync_batch_size, enable_profiling) # Use the provided negative prompt, or fall back to a default. neg_prompt = negative_prompt if negative_prompt and negative_prompt.strip() else "noisy, harsh" prompts = [neg_prompt, prompt] clap_model = model_dict.clap_model try: clap_model.to(model_dict.device) text_feat_res, _ = encode_text_feat(prompts, model_dict) finally: clap_model.to("cpu") text_feat = text_feat_res[1:]; uncond_text_feat = text_feat_res[:1] if cfg.model_config.model_kwargs.text_length < text_feat.shape[1]: text_seq_length = cfg.model_config.model_kwargs.text_length text_feat = text_feat[:, :text_seq_length] uncond_text_feat = uncond_text_feat[:, :text_seq_length] text_feats = AttributeDict({'text_feat': text_feat, 'uncond_text_feat': uncond_text_feat}) return visual_feats, text_feats, audio_len_in_s # In utils.py, add this entire block at the end of the file from diffusers.utils.torch_utils import randn_tensor from hunyuanvideo_foley.utils.schedulers.scheduling_flow_match_discrete import FlowMatchDiscreteScheduler def _retrieve_timesteps(scheduler, num_inference_steps, device, **kwargs): """ Helper function forked from the library's pipeline.py """ scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) timesteps = scheduler.timesteps return timesteps, num_inference_steps def _prepare_latents(scheduler, batch_size, num_channels_latents, length, dtype, device): """ Helper function forked from the library's pipeline.py """ shape = (batch_size, num_channels_latents, int(length)) latents = randn_tensor(shape, device=device, dtype=dtype) if hasattr(scheduler, "init_noise_sigma"): latents = latents * scheduler.init_noise_sigma return latents def denoise_process_safely(visual_feats, text_feats, audio_len_in_s, model_dict, cfg, **kwargs): """ Memory-safe fork of denoise_process that ensures all tensors are on the correct device. """ target_dtype = model_dict.foley_model.dtype autocast_enabled = target_dtype != torch.float32 device = model_dict.device # --- Device Correction Fix --- visual_feats.siglip2_feat = visual_feats.siglip2_feat.to(device) visual_feats.syncformer_feat = visual_feats.syncformer_feat.to(device) text_feats.text_feat = text_feats.text_feat.to(device) text_feats.uncond_text_feat = text_feats.uncond_text_feat.to(device) # --- End of Fix --- guidance_scale = kwargs.get('guidance_scale', 4.5) num_inference_steps = kwargs.get('num_inference_steps', 50) batch_size = kwargs.get('batch_size', 1) scheduler = FlowMatchDiscreteScheduler( shift=cfg.diffusion_config.sample_flow_shift, reverse=cfg.diffusion_config.flow_reverse, solver=cfg.diffusion_config.flow_solver, use_flux_shift=cfg.diffusion_config.sample_use_flux_shift, flux_base_shift=cfg.diffusion_config.flux_base_shift, flux_max_shift=cfg.diffusion_config.flux_max_shift, ) timesteps, _ = _retrieve_timesteps(scheduler, num_inference_steps, device) latents = _prepare_latents( scheduler, batch_size=batch_size, num_channels_latents=cfg.model_config.model_kwargs.audio_vae_latent_dim, length=audio_len_in_s * cfg.model_config.model_kwargs.audio_frame_rate, dtype=target_dtype, device=device, ) for i, t in tqdm(enumerate(timesteps), total=len(timesteps), desc="Denoising steps"): latent_input = torch.cat([latents] * 2) if guidance_scale > 1.0 else latents latent_input = scheduler.scale_model_input(latent_input, t) t_expand = t.repeat(latent_input.shape[0]) siglip2_feat = visual_feats.siglip2_feat.repeat(batch_size, 1, 1) uncond_siglip2_feat = model_dict.foley_model.get_empty_clip_sequence(bs=batch_size, len=siglip2_feat.shape[1]).to(device) siglip2_feat_input = torch.cat([uncond_siglip2_feat, siglip2_feat], dim=0) if guidance_scale > 1.0 else siglip2_feat syncformer_feat = visual_feats.syncformer_feat.repeat(batch_size, 1, 1) uncond_syncformer_feat = model_dict.foley_model.get_empty_sync_sequence(bs=batch_size, len=syncformer_feat.shape[1]).to(device) syncformer_feat_input = torch.cat([uncond_syncformer_feat, syncformer_feat], dim=0) if guidance_scale > 1.0 else syncformer_feat text_feat_repeated = text_feats.text_feat.repeat(batch_size, 1, 1) uncond_text_feat_repeated = text_feats.uncond_text_feat.repeat(batch_size, 1, 1) text_feat_input = torch.cat([uncond_text_feat_repeated, text_feat_repeated], dim=0) if guidance_scale > 1.0 else text_feat_repeated with torch.autocast(device_type=device.type, enabled=autocast_enabled, dtype=target_dtype): noise_pred = model_dict.foley_model( x=latent_input, t=t_expand, cond=text_feat_input, clip_feat=siglip2_feat_input, sync_feat=syncformer_feat_input, return_dict=True, )["x"] noise_pred = noise_pred.to(dtype=torch.float32) if guidance_scale > 1.0: noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) latents = scheduler.step(noise_pred, t, latents, return_dict=False)[0] with torch.no_grad(): audio = model_dict.dac_model.decode(latents) audio = audio.float().cpu() audio = audio[:, :int(audio_len_in_s * model_dict.dac_model.sample_rate)] return audio, model_dict.dac_model.sample_rate def create_node_exit_values(silent_audio, passthrough_video=None, passthrough_images=None, message="Process skipped or failed."): """ Creates a standardized tuple of return values for exiting a node early. Handles passthrough of video or image inputs. """ audio_output = {"waveform": torch.zeros((1, 1, 1)), "sample_rate": 48000} if silent_audio else None # Prioritize passing through the image tensor if it exists if passthrough_images is not None: frames_output = passthrough_images else: frames_output = torch.zeros((1, 1, 1, 3), dtype=torch.float32) # Pass through the video path if it exists video_path_output = extract_video_path(passthrough_video) if passthrough_video else "" return (video_path_output, frames_output, audio_output, message) def _encode_video_with_siglip2_safely(pixel_values, model_dict): """ A wrapper to handle different versions of the transformers library for SigLIP2. """ if hasattr(model_dict.siglip2_model, 'get_image_features'): # Older transformers versions return model_dict.siglip2_model.get_image_features(pixel_values=pixel_values) else: # Newer transformers versions return model_dict.siglip2_model(pixel_values=pixel_values).image_embeds def _batched_frames_to_tensor(pil_list, preprocess_func, batch_size, device): """ Preprocesses and moves frames to device in small batches to save VRAM. """ tensor_chunks = [] for i in range(0, len(pil_list), batch_size): batch_pil = pil_list[i : i + batch_size] batch_preprocessed = [preprocess_func(im) for im in batch_pil] batch_tensor = torch.stack(batch_preprocessed, dim=0).to(device) tensor_chunks.append(batch_tensor) # Immediately clear the list to save RAM del batch_pil, batch_preprocessed, batch_tensor # Once all chunks are on the device, concatenate them. # This is much more memory-efficient than creating one giant tensor at once. full_tensor = torch.cat(tensor_chunks, dim=0).unsqueeze(0) del tensor_chunks return full_tensor