Files
jax-explorer-ComfyUI-VideoB…/nodes.py
T
2025-04-01 07:05:28 +03:00

1072 lines
47 KiB
Python

import os
import tempfile
import torchaudio
import uuid
import sys
import shutil
from collections.abc import Mapping
# Function to find ComfyUI directories
def get_comfyui_temp_dir():
"""Dynamically find the ComfyUI temp directory"""
# First check using folder_paths if available
try:
import folder_paths
comfy_dir = os.path.dirname(os.path.dirname(os.path.abspath(folder_paths.__file__)))
temp_dir = os.path.join(comfy_dir, "temp")
return temp_dir
except:
pass
# Try to locate based on current script location
try:
# This script is likely in a ComfyUI custom nodes directory
current_dir = os.path.dirname(os.path.abspath(__file__))
# Go up until we find the ComfyUI directory
potential_dir = current_dir
for _ in range(5): # Limit to 5 levels up
if os.path.exists(os.path.join(potential_dir, "comfy.py")):
return os.path.join(potential_dir, "temp")
potential_dir = os.path.dirname(potential_dir)
except:
pass
# Return None if we can't find it
return None
# Function to clean up any ComfyUI temp directories
def cleanup_comfyui_temp_directories():
"""Find and clean up any ComfyUI temp directories"""
comfyui_temp = get_comfyui_temp_dir()
if not comfyui_temp:
print("Could not locate ComfyUI temp directory")
return
comfyui_base = os.path.dirname(comfyui_temp)
# Check for the main temp directory
if os.path.exists(comfyui_temp):
try:
shutil.rmtree(comfyui_temp)
print(f"Removed ComfyUI temp directory: {comfyui_temp}")
except Exception as e:
print(f"Could not remove {comfyui_temp}: {str(e)}")
# If we can't remove it, try to rename it
try:
backup_name = f"{comfyui_temp}_backup_{uuid.uuid4().hex[:8]}"
os.rename(comfyui_temp, backup_name)
print(f"Renamed {comfyui_temp} to {backup_name}")
except:
pass
# Find and clean up any backup temp directories
try:
all_directories = [d for d in os.listdir(comfyui_base) if os.path.isdir(os.path.join(comfyui_base, d))]
for dirname in all_directories:
if dirname.startswith("temp_backup_"):
backup_path = os.path.join(comfyui_base, dirname)
try:
shutil.rmtree(backup_path)
print(f"Removed backup temp directory: {backup_path}")
except Exception as e:
print(f"Could not remove backup dir {backup_path}: {str(e)}")
except Exception as e:
print(f"Error cleaning up temp directories: {str(e)}")
# Create a module-level function to set up system-wide temp directory
def init_temp_directories():
"""Initialize global temporary directory settings"""
# First clean up any existing temp directories
cleanup_comfyui_temp_directories()
# Generate a unique base directory for this module
system_temp = tempfile.gettempdir()
unique_id = str(uuid.uuid4())[:8]
temp_base_path = os.path.join(system_temp, f"latentsync_{unique_id}")
os.makedirs(temp_base_path, exist_ok=True)
# Override environment variables that control temp directories
os.environ['TMPDIR'] = temp_base_path
os.environ['TEMP'] = temp_base_path
os.environ['TMP'] = temp_base_path
# Force Python's tempfile module to use our directory
tempfile.tempdir = temp_base_path
# Final check for ComfyUI temp directory
comfyui_temp = get_comfyui_temp_dir()
if comfyui_temp and os.path.exists(comfyui_temp):
try:
shutil.rmtree(comfyui_temp)
print(f"Removed ComfyUI temp directory: {comfyui_temp}")
except Exception as e:
print(f"Could not remove {comfyui_temp}, trying to rename: {str(e)}")
try:
backup_name = f"{comfyui_temp}_backup_{unique_id}"
os.rename(comfyui_temp, backup_name)
print(f"Renamed {comfyui_temp} to {backup_name}")
# Try to remove the renamed directory as well
try:
shutil.rmtree(backup_name)
print(f"Removed renamed temp directory: {backup_name}")
except:
pass
except:
print(f"Failed to rename {comfyui_temp}")
print(f"Set up system temp directory: {temp_base_path}")
return temp_base_path
# Function to clean up everything when the module exits
def module_cleanup():
"""Clean up all resources when the module is unloaded"""
global MODULE_TEMP_DIR
# Clean up our module temp directory
if MODULE_TEMP_DIR and os.path.exists(MODULE_TEMP_DIR):
try:
shutil.rmtree(MODULE_TEMP_DIR, ignore_errors=True)
print(f"Cleaned up module temp directory: {MODULE_TEMP_DIR}")
except:
pass
# Do a final sweep for any ComfyUI temp directories
cleanup_comfyui_temp_directories()
# Call this before anything else
MODULE_TEMP_DIR = init_temp_directories()
# Register the cleanup handler to run when Python exits
import atexit
atexit.register(module_cleanup)
# Now import regular dependencies
import math
import torch
import random
import torchaudio
import folder_paths
import numpy as np
import platform
import subprocess
import importlib.util
import importlib.machinery
import argparse
from omegaconf import OmegaConf
from PIL import Image
from decimal import Decimal, ROUND_UP
import requests
# Modify folder_paths module to use our temp directory
if hasattr(folder_paths, "get_temp_directory"):
original_get_temp = folder_paths.get_temp_directory
folder_paths.get_temp_directory = lambda: MODULE_TEMP_DIR
else:
# Add the function if it doesn't exist
setattr(folder_paths, 'get_temp_directory', lambda: MODULE_TEMP_DIR)
def import_inference_script(script_path):
"""Import a Python file as a module using its file path."""
if not os.path.exists(script_path):
raise ImportError(f"Script not found: {script_path}")
module_name = "latentsync_inference"
spec = importlib.util.spec_from_file_location(module_name, script_path)
if spec is None:
raise ImportError(f"Failed to create module spec for {script_path}")
module = importlib.util.module_from_spec(spec)
sys.modules[module_name] = module
try:
spec.loader.exec_module(module)
except Exception as e:
del sys.modules[module_name]
raise ImportError(f"Failed to execute module: {str(e)}")
return module
def check_ffmpeg():
try:
if platform.system() == "Windows":
# Check if ffmpeg exists in PATH
ffmpeg_path = shutil.which("ffmpeg.exe")
if ffmpeg_path is None:
# Look for ffmpeg in common locations
possible_paths = [
os.path.join(os.environ.get("ProgramFiles", "C:\\Program Files"), "ffmpeg", "bin"),
os.path.join(os.environ.get("ProgramFiles(x86)", "C:\\Program Files (x86)"), "ffmpeg", "bin"),
os.path.join(os.path.dirname(os.path.abspath(__file__)), "ffmpeg", "bin"),
]
for path in possible_paths:
if os.path.exists(os.path.join(path, "ffmpeg.exe")):
# Add to PATH
os.environ["PATH"] = path + os.pathsep + os.environ.get("PATH", "")
return True
print("FFmpeg not found. Please install FFmpeg and add it to PATH")
return False
return True
else:
subprocess.run(["ffmpeg", "-version"], capture_output=True, check=True)
return True
except (subprocess.CalledProcessError, FileNotFoundError):
print("FFmpeg not found. Please install FFmpeg")
return False
def check_and_install_dependencies():
if not check_ffmpeg():
raise RuntimeError("FFmpeg is required but not found")
required_packages = [
'omegaconf',
'pytorch_lightning',
'transformers',
'accelerate',
'huggingface_hub',
'einops',
'diffusers',
'ffmpeg-python'
]
def is_package_installed(package_name):
return importlib.util.find_spec(package_name) is not None
def install_package(package):
python_exe = sys.executable
try:
subprocess.check_call([python_exe, '-m', 'pip', 'install', package],
stdout=subprocess.PIPE,
stderr=subprocess.PIPE)
print(f"Successfully installed {package}")
except subprocess.CalledProcessError as e:
print(f"Error installing {package}: {str(e)}")
raise RuntimeError(f"Failed to install required package: {package}")
for package in required_packages:
if not is_package_installed(package):
print(f"Installing required package: {package}")
try:
install_package(package)
except Exception as e:
print(f"Warning: Failed to install {package}: {str(e)}")
raise
def normalize_path(path):
"""Normalize path to handle spaces and special characters"""
return os.path.normpath(path).replace('\\', '/')
def get_ext_dir(subpath=None, mkdir=False):
"""Get extension directory path, optionally with a subpath"""
# Get the directory containing this script
dir = os.path.dirname(os.path.abspath(__file__))
# Special case for temp directories
if subpath and ("temp" in subpath.lower() or "tmp" in subpath.lower()):
# Use our global temp directory instead
global MODULE_TEMP_DIR
sub_temp = os.path.join(MODULE_TEMP_DIR, subpath)
if mkdir and not os.path.exists(sub_temp):
os.makedirs(sub_temp, exist_ok=True)
return sub_temp
if subpath is not None:
dir = os.path.join(dir, subpath)
if mkdir and not os.path.exists(dir):
os.makedirs(dir, exist_ok=True)
return dir
def download_model(url, save_path):
"""Download a model from a URL and save it to the specified path."""
os.makedirs(os.path.dirname(save_path), exist_ok=True)
response = requests.get(url, stream=True)
with open(save_path, "wb") as f:
for chunk in response.iter_content(chunk_size=8192):
f.write(chunk)
def pre_download_models():
"""Pre-download all required models."""
models = {
"s3fd-e19a316812.pth": "https://www.adrianbulat.com/downloads/python-fan/s3fd-e19a316812.pth",
# Add other models here
}
cache_dir = os.path.join(MODULE_TEMP_DIR, "model_cache")
os.makedirs(cache_dir, exist_ok=True)
for model_name, url in models.items():
save_path = os.path.join(cache_dir, model_name)
if not os.path.exists(save_path):
print(f"Downloading {model_name}...")
download_model(url, save_path)
else:
print(f"{model_name} already exists in cache.")
def setup_models():
"""Setup and pre-download all required models."""
# Use our global temp directory
global MODULE_TEMP_DIR
# Pre-download additional models
pre_download_models()
# Existing setup logic for LatentSync models
cur_dir = get_ext_dir()
ckpt_dir = os.path.join(cur_dir, "checkpoints")
whisper_dir = os.path.join(ckpt_dir, "whisper")
os.makedirs(ckpt_dir, exist_ok=True)
os.makedirs(whisper_dir, exist_ok=True)
# Create a temp_downloads directory in our system temp
temp_downloads = os.path.join(MODULE_TEMP_DIR, "downloads")
os.makedirs(temp_downloads, exist_ok=True)
unet_path = os.path.join(ckpt_dir, "latentsync_unet.pt")
whisper_path = os.path.join(whisper_dir, "tiny.pt")
if not (os.path.exists(unet_path) and os.path.exists(whisper_path)):
print("Downloading required model checkpoints... This may take a while.")
try:
from huggingface_hub import snapshot_download
snapshot_download(repo_id="ByteDance/LatentSync-1.5",
allow_patterns=["latentsync_unet.pt", "whisper/tiny.pt"],
local_dir=ckpt_dir,
local_dir_use_symlinks=False,
cache_dir=temp_downloads)
print("Model checkpoints downloaded successfully!")
except Exception as e:
print(f"Error downloading models: {str(e)}")
print("\nPlease download models manually:")
print("1. Visit: https://huggingface.co/chunyu-li/LatentSync")
print("2. Download: latentsync_unet.pt and whisper/tiny.pt")
print(f"3. Place them in: {ckpt_dir}")
print(f" with whisper/tiny.pt in: {whisper_dir}")
raise RuntimeError("Model download failed. See instructions above.")
class LatentSyncNode:
def __init__(self):
# Make sure our temp directory is the current one
global MODULE_TEMP_DIR
if not os.path.exists(MODULE_TEMP_DIR):
os.makedirs(MODULE_TEMP_DIR, exist_ok=True)
# Ensure ComfyUI temp doesn't exist
comfyui_temp = "D:\\ComfyUI_windows\\temp"
if os.path.exists(comfyui_temp):
backup_name = f"{comfyui_temp}_backup_{uuid.uuid4().hex[:8]}"
try:
os.rename(comfyui_temp, backup_name)
except:
pass
check_and_install_dependencies()
setup_models()
@classmethod
def INPUT_TYPES(s):
return {"required": {
"images": ("IMAGE",),
"audio": ("AUDIO", ),
"seed": ("INT", {"default": 1247}),
"lips_expression": ("FLOAT", {"default": 1.5, "min": 1.0, "max": 3.0, "step": 0.1}),
"inference_steps": ("INT", {"default": 20, "min": 1, "max": 999, "step": 1}),
},}
CATEGORY = "LatentSyncNode"
RETURN_TYPES = ("IMAGE", "AUDIO")
RETURN_NAMES = ("images", "audio")
FUNCTION = "inference"
def process_batch(self, batch, use_mixed_precision=False):
with torch.cuda.amp.autocast(enabled=use_mixed_precision):
processed_batch = batch.float() / 255.0
if len(processed_batch.shape) == 3:
processed_batch = processed_batch.unsqueeze(0)
if processed_batch.shape[0] == 3:
processed_batch = processed_batch.permute(1, 2, 0)
if processed_batch.shape[-1] == 4:
processed_batch = processed_batch[..., :3]
return processed_batch
def inference(self, images, audio, seed, lips_expression=1.5, inference_steps=20):
# Use our module temp directory
global MODULE_TEMP_DIR
# Get GPU capabilities and memory
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
BATCH_SIZE = 4
use_mixed_precision = False
if torch.cuda.is_available():
gpu_mem = torch.cuda.get_device_properties(0).total_memory
# Convert to GB
gpu_mem_gb = gpu_mem / (1024 ** 3)
# Dynamically adjust batch size based on GPU memory
if gpu_mem_gb > 20: # High-end GPUs
BATCH_SIZE = 32
enable_tf32 = True
use_mixed_precision = True
elif gpu_mem_gb > 8: # Mid-range GPUs
BATCH_SIZE = 16
enable_tf32 = False
use_mixed_precision = True
else: # Lower-end GPUs
BATCH_SIZE = 8
enable_tf32 = False
use_mixed_precision = False
# Set performance options based on GPU capability
torch.backends.cudnn.benchmark = True
if enable_tf32:
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
# Clear GPU cache before processing
torch.cuda.empty_cache()
torch.cuda.set_per_process_memory_fraction(0.8)
# Create a run-specific subdirectory in our temp directory
run_id = ''.join(random.choice("abcdefghijklmnopqrstuvwxyz") for _ in range(5))
temp_dir = os.path.join(MODULE_TEMP_DIR, f"run_{run_id}")
os.makedirs(temp_dir, exist_ok=True)
# Ensure ComfyUI temp doesn't exist again (in case something recreated it)
comfyui_temp = "D:\\ComfyUI_windows\\temp"
if os.path.exists(comfyui_temp):
backup_name = f"{comfyui_temp}_backup_{uuid.uuid4().hex[:8]}"
try:
os.rename(comfyui_temp, backup_name)
except:
pass
temp_video_path = None
output_video_path = None
audio_path = None
try:
# Create temporary file paths in our system temp directory
temp_video_path = os.path.join(temp_dir, f"temp_{run_id}.mp4")
output_video_path = os.path.join(temp_dir, f"latentsync_{run_id}_out.mp4")
audio_path = os.path.join(temp_dir, f"latentsync_{run_id}_audio.wav")
# Get the extension directory
cur_dir = os.path.dirname(os.path.abspath(__file__))
# Process input frames
if isinstance(images, list):
frames = torch.stack(images).to(device)
else:
frames = images.to(device)
frames = (frames * 255).byte()
if len(frames.shape) == 3:
frames = frames.unsqueeze(0)
# Process audio with device awareness
waveform = audio["waveform"].to(device)
sample_rate = audio["sample_rate"]
if waveform.dim() == 3:
waveform = waveform.squeeze(0)
if sample_rate != 16000:
new_sample_rate = 16000
resampler = torchaudio.transforms.Resample(
orig_freq=sample_rate,
new_freq=new_sample_rate
).to(device)
waveform_16k = resampler(waveform)
waveform, sample_rate = waveform_16k, new_sample_rate
# Package resampled audio
resampled_audio = {
"waveform": waveform.unsqueeze(0),
"sample_rate": sample_rate
}
# Move waveform to CPU for saving
waveform_cpu = waveform.cpu()
torchaudio.save(audio_path, waveform_cpu, sample_rate)
# Move frames to CPU for saving to video
frames_cpu = frames.cpu()
try:
import torchvision.io as io
io.write_video(temp_video_path, frames_cpu, fps=25, video_codec='h264')
except TypeError:
import av
container = av.open(temp_video_path, mode='w')
stream = container.add_stream('h264', rate=25)
stream.width = frames_cpu.shape[2]
stream.height = frames_cpu.shape[1]
for frame in frames_cpu:
frame = av.VideoFrame.from_ndarray(frame.numpy(), format='rgb24')
packet = stream.encode(frame)
container.mux(packet)
packet = stream.encode(None)
container.mux(packet)
container.close()
# Define paths to required files and configs
inference_script_path = os.path.join(cur_dir, "scripts", "inference.py")
config_path = os.path.join(cur_dir, "configs", "unet", "stage2.yaml")
scheduler_config_path = os.path.join(cur_dir, "configs")
ckpt_path = os.path.join(cur_dir, "checkpoints", "latentsync_unet.pt")
whisper_ckpt_path = os.path.join(cur_dir, "checkpoints", "whisper", "tiny.pt")
# Create config and args
config = OmegaConf.load(config_path)
# Set the correct mask image path
mask_image_path = os.path.join(cur_dir, "latentsync", "utils", "mask.png")
# Make sure the mask image exists
if not os.path.exists(mask_image_path):
# Try to find it in the utils directory directly
alt_mask_path = os.path.join(cur_dir, "utils", "mask.png")
if os.path.exists(alt_mask_path):
mask_image_path = alt_mask_path
else:
print(f"Warning: Could not find mask image at expected locations")
# Set mask path in config
if hasattr(config, "data") and hasattr(config.data, "mask_image_path"):
config.data.mask_image_path = mask_image_path
args = argparse.Namespace(
unet_config_path=config_path,
inference_ckpt_path=ckpt_path,
video_path=temp_video_path,
audio_path=audio_path,
video_out_path=output_video_path,
seed=seed,
inference_steps=inference_steps,
guidance_scale=lips_expression, # Using lips_expression for the guidance_scale
scheduler_config_path=scheduler_config_path,
whisper_ckpt_path=whisper_ckpt_path,
device=device,
batch_size=BATCH_SIZE,
use_mixed_precision=use_mixed_precision,
temp_dir=temp_dir,
mask_image_path=mask_image_path
)
# Set PYTHONPATH to include our directories
package_root = os.path.dirname(cur_dir)
if package_root not in sys.path:
sys.path.insert(0, package_root)
if cur_dir not in sys.path:
sys.path.insert(0, cur_dir)
# Clean GPU cache before inference
if torch.cuda.is_available():
torch.cuda.empty_cache()
# Check and prevent ComfyUI temp creation again
if os.path.exists(comfyui_temp):
try:
os.rename(comfyui_temp, f"{comfyui_temp}_backup_{uuid.uuid4().hex[:8]}")
except:
pass
# Import the inference module
inference_module = import_inference_script(inference_script_path)
# Monkey patch any temp directory functions in the inference module
if hasattr(inference_module, 'get_temp_dir'):
inference_module.get_temp_dir = lambda *args, **kwargs: temp_dir
# Create subdirectories that the inference module might expect
inference_temp = os.path.join(temp_dir, "temp")
os.makedirs(inference_temp, exist_ok=True)
# Run inference
inference_module.main(config, args)
# Clean GPU cache after inference
if torch.cuda.is_available():
torch.cuda.empty_cache()
# Verify output file exists
if not os.path.exists(output_video_path):
raise FileNotFoundError(f"Output video not found at: {output_video_path}")
# Read the processed video - ensure it's loaded as CPU tensor
processed_frames = io.read_video(output_video_path, pts_unit='sec')[0]
processed_frames = processed_frames.float() / 255.0
# Ensure audio is on CPU before returning
if torch.cuda.is_available():
if hasattr(resampled_audio["waveform"], 'device') and resampled_audio["waveform"].device.type == 'cuda':
resampled_audio["waveform"] = resampled_audio["waveform"].cpu()
if hasattr(processed_frames, 'device') and processed_frames.device.type == 'cuda':
processed_frames = processed_frames.cpu()
return (processed_frames, resampled_audio)
except Exception as e:
print(f"Error during inference: {str(e)}")
import traceback
traceback.print_exc()
raise
finally:
# Clean up temporary files individually
for path in [temp_video_path, output_video_path, audio_path]:
if path and os.path.exists(path):
try:
os.remove(path)
print(f"Removed temporary file: {path}")
except Exception as e:
print(f"Failed to remove {path}: {str(e)}")
# Remove temporary run directory
if temp_dir and os.path.exists(temp_dir):
try:
shutil.rmtree(temp_dir, ignore_errors=True)
print(f"Removed run temporary directory: {temp_dir}")
except Exception as e:
print(f"Failed to remove temp run directory: {str(e)}")
# Clean up any ComfyUI temp directories again (in case they were created during execution)
cleanup_comfyui_temp_directories()
# Final GPU cache cleanup
if torch.cuda.is_available():
torch.cuda.empty_cache()
class VideoLengthAdjuster:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"images": ("IMAGE",),
"audio": ("AUDIO",),
"mode": (["normal", "pingpong", "loop_to_audio"], {"default": "normal"}),
"fps": ("FLOAT", {"default": 25.0, "min": 1.0, "max": 120.0}),
"silent_padding_sec": ("FLOAT", {"default": 0.5, "min": 0.1, "max": 3.0, "step": 0.1}),
}
}
CATEGORY = "LatentSyncNode"
RETURN_TYPES = ("IMAGE", "AUDIO")
RETURN_NAMES = ("images", "audio")
FUNCTION = "adjust"
def adjust(self, images, audio, mode, fps=25.0, silent_padding_sec=0.5):
waveform = audio["waveform"].squeeze(0)
sample_rate = int(audio["sample_rate"])
original_frames = [images[i] for i in range(images.shape[0])] if isinstance(images, torch.Tensor) else images.copy()
if mode == "normal":
# Add silent padding to the audio and then trim video to match
audio_duration = waveform.shape[1] / sample_rate
# Add silent padding to the audio
silence_samples = math.ceil(silent_padding_sec * sample_rate)
silence = torch.zeros((waveform.shape[0], silence_samples), dtype=waveform.dtype)
padded_audio = torch.cat([waveform, silence], dim=1)
# Calculate required frames based on the padded audio
padded_audio_duration = (waveform.shape[1] + silence_samples) / sample_rate
required_frames = int(padded_audio_duration * fps)
if len(original_frames) > required_frames:
# Trim video frames to match padded audio duration
adjusted_frames = original_frames[:required_frames]
else:
# If video is shorter than padded audio, keep all video frames
# and trim the audio accordingly
adjusted_frames = original_frames
required_samples = int(len(original_frames) / fps * sample_rate)
padded_audio = padded_audio[:, :required_samples]
return (
torch.stack(adjusted_frames),
{"waveform": padded_audio.unsqueeze(0), "sample_rate": sample_rate}
)
# This return statement is no longer needed as it's handled in the updated code
elif mode == "pingpong":
video_duration = len(original_frames) / fps
audio_duration = waveform.shape[1] / sample_rate
if audio_duration <= video_duration:
required_samples = int(video_duration * sample_rate)
silence = torch.zeros((waveform.shape[0], required_samples - waveform.shape[1]), dtype=waveform.dtype)
adjusted_audio = torch.cat([waveform, silence], dim=1)
return (
torch.stack(original_frames),
{"waveform": adjusted_audio.unsqueeze(0), "sample_rate": sample_rate}
)
else:
silence_samples = math.ceil(silent_padding_sec * sample_rate)
silence = torch.zeros((waveform.shape[0], silence_samples), dtype=waveform.dtype)
padded_audio = torch.cat([waveform, silence], dim=1)
total_duration = (waveform.shape[1] + silence_samples) / sample_rate
target_frames = math.ceil(total_duration * fps)
reversed_frames = original_frames[::-1][1:-1] # Remove endpoints
frames = original_frames + reversed_frames
while len(frames) < target_frames:
frames += frames[:target_frames - len(frames)]
return (
torch.stack(frames[:target_frames]),
{"waveform": padded_audio.unsqueeze(0), "sample_rate": sample_rate}
)
elif mode == "loop_to_audio":
# Add silent padding then simple loop
silence_samples = math.ceil(silent_padding_sec * sample_rate)
silence = torch.zeros((waveform.shape[0], silence_samples), dtype=waveform.dtype)
padded_audio = torch.cat([waveform, silence], dim=1)
total_duration = (waveform.shape[1] + silence_samples) / sample_rate
target_frames = math.ceil(total_duration * fps)
frames = original_frames.copy()
while len(frames) < target_frames:
frames += original_frames[:target_frames - len(frames)]
return (
torch.stack(frames[:target_frames]),
{"waveform": padded_audio.unsqueeze(0), "sample_rate": sample_rate}
)
class DG_VideoAudioMixer:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"images1": ("IMAGE", ),
"video_info1": ("VHS_VIDEOINFO", ),
"images2": ("IMAGE", ),
"video_info2": ("VHS_VIDEOINFO", ),
"bgm": ("AUDIO", ), # Add BGM as required input first
"bgm_volume": ("FLOAT", {"default": 0.3, "min": 0.0, "max": 1.0, "step": 0.05}),
},
"optional": {
"audio1": ("AUDIO", ),
"audio2": ("AUDIO", ),
"fade_in_sec": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.1}),
}
}
CATEGORY = "LatentSyncNode"
FUNCTION = "VideoAudioMixer"
TITLE = "DG_VideoAudioMixer"
RETURN_NAMES = ("images_output", "audio_output", "video_info_output")
RETURN_TYPES = ("IMAGE", "AUDIO", "VHS_VIDEOINFO")
def VideoAudioMixer(self, images1, video_info1, images2, video_info2, bgm=None, bgm_volume=0.3, audio1=None, audio2=None, fade_in_sec=1.0):
print(f"DEBUG: bgm={bgm is not None}, bgm_volume={bgm_volume}, fade_in_sec={fade_in_sec}")
# Verify frames
if not isinstance(images1, torch.Tensor) or not isinstance(images2, torch.Tensor):
raise ValueError("images1 and images2 must be frame tensors")
# Handle frame dimensions (assume [frames, h, w, c])
if images1.shape[1:] != images2.shape[1:]:
raise ValueError(f"Incompatible resolutions: images1 {images1.shape}, images2 {images2.shape}")
# Concatenate frames
concatenated_frames = torch.cat([images1, images2], dim=0) # [frames_total, h, w, c]
# Extract FPS from video_info
fps1 = video_info1["loaded_fps"]
fps2 = video_info2["loaded_fps"]
if fps1 != fps2:
print(f"Warning: Different FPS (video1: {fps1}, video2: {fps2}), using {fps1}")
output_fps = fps1
# Handle audio
audio_output = None
sample_rate = None
# Function to extract waveform and sample_rate from audio
def get_audio_data(audio_input, label=""):
print(f"{label}: Type of audio_input = {type(audio_input)}, value = {audio_input}")
if audio_input is None:
print(f"{label}: No audio provided")
return None, None
if isinstance(audio_input, Mapping):
try:
waveform = audio_input["waveform"].squeeze(0)
sample_rate = audio_input["sample_rate"]
print(f"{label}: Audio extracted from Mapping, waveform shape={waveform.shape}, sample_rate={sample_rate}")
return waveform, sample_rate
except KeyError as e:
print(f"{label}: Error - Missing key in Mapping: {e}")
return None, None
except Exception as e:
print(f"{label}: Error extracting from Mapping: {e}")
return None, None
elif callable(audio_input):
try:
audio_data = audio_input()
if isinstance(audio_data, dict) and "waveform" in audio_data:
waveform = audio_data["waveform"].squeeze(0)
print(f"{label}: Audio extracted from function, waveform shape={waveform.shape}, sample_rate={audio_data['sample_rate']}")
return waveform, audio_data["sample_rate"]
else:
print(f"{label}: Invalid function result: {audio_data}")
return None, None
except Exception as e:
print(f"{label}: Error evaluating function: {e}")
return None, None
elif isinstance(audio_input, dict) and "waveform" in audio_input:
waveform = audio_input["waveform"].squeeze(0)
print(f"{label}: Audio extracted from dictionary, waveform shape={waveform.shape}, sample_rate={audio_input['sample_rate']}")
return waveform, audio_input["sample_rate"]
else:
print(f"{label}: Audio format not recognized: {type(audio_input)}")
return None, None
# Extract audio data from all audio inputs
audio_waveform1, sample_rate1 = get_audio_data(audio1, "Audio1")
audio_waveform2, sample_rate2 = get_audio_data(audio2, "Audio2")
bgm_waveform, bgm_sample_rate = get_audio_data(bgm, "BGM")
# Determine target sample rate (prefer audio1, then audio2, then BGM, fallback to 44100)
sample_rate = sample_rate1 or sample_rate2 or bgm_sample_rate or 44100
print(f"Using sample rate: {sample_rate}Hz")
# Import torchaudio for resampling
import torchaudio
# Resample all audio to the target sample rate
if audio_waveform1 is not None and sample_rate1 != sample_rate:
resampler = torchaudio.transforms.Resample(
orig_freq=sample_rate1,
new_freq=sample_rate
)
audio_waveform1 = resampler(audio_waveform1)
print(f"Resampled audio1 from {sample_rate1}Hz to {sample_rate}Hz")
if audio_waveform2 is not None and sample_rate2 != sample_rate:
resampler = torchaudio.transforms.Resample(
orig_freq=sample_rate2,
new_freq=sample_rate
)
audio_waveform2 = resampler(audio_waveform2)
print(f"Resampled audio2 from {sample_rate2}Hz to {sample_rate}Hz")
if bgm_waveform is not None and bgm_sample_rate != sample_rate:
resampler = torchaudio.transforms.Resample(
orig_freq=bgm_sample_rate,
new_freq=sample_rate
)
bgm_waveform = resampler(bgm_waveform)
print(f"Resampled BGM from {bgm_sample_rate}Hz to {sample_rate}Hz")
# Calculate durations
duration1 = images1.shape[0] / fps1
duration2 = images2.shape[0] / fps2
total_duration = duration1 + duration2
total_samples = int(total_duration * sample_rate)
print(f"Video durations: Video1={duration1:.2f}s, Video2={duration2:.2f}s, Total={total_duration:.2f}s")
# Generate silent audio for missing primary audio inputs
if audio_waveform1 is None:
audio_waveform1 = torch.zeros((1, int(duration1 * sample_rate)))
print(f"Audio1: Silence generated, shape={audio_waveform1.shape}")
if audio_waveform2 is None:
audio_waveform2 = torch.zeros((1, int(duration2 * sample_rate)))
print(f"Audio2: Silence generated, shape={audio_waveform2.shape}")
# Match channel counts between primary audio streams
if audio_waveform1.shape[0] != audio_waveform2.shape[0]:
print(f"Channel count mismatch: audio1 has {audio_waveform1.shape[0]} channels, audio2 has {audio_waveform2.shape[0]} channels")
# If audio1 is mono and audio2 is stereo
if audio_waveform1.shape[0] == 1 and audio_waveform2.shape[0] == 2:
# Convert audio1 to stereo by duplicating the channel
audio_waveform1 = audio_waveform1.repeat(2, 1)
print(f"Converted audio1 to stereo: new shape={audio_waveform1.shape}")
# If audio1 is stereo and audio2 is mono
elif audio_waveform1.shape[0] == 2 and audio_waveform2.shape[0] == 1:
# Convert audio2 to stereo by duplicating the channel
audio_waveform2 = audio_waveform2.repeat(2, 1)
print(f"Converted audio2 to stereo: new shape={audio_waveform2.shape}")
# Concatenate the primary audio streams
primary_audio = torch.cat([audio_waveform1, audio_waveform2], dim=1)
print(f"Concatenated primary audio: shape={primary_audio.shape}")
# Check if primary audio has actual content (not just silence)
has_actual_audio = primary_audio.abs().max() > 0.01
print(f"Primary audio has actual content: {has_actual_audio}")
# Process background music if provided
if bgm_waveform is not None:
print(f"Processing BGM: shape={bgm_waveform.shape}")
# Match channel count with primary audio
primary_channels = primary_audio.shape[0]
if bgm_waveform.shape[0] != primary_channels:
if bgm_waveform.shape[0] == 1 and primary_channels == 2:
# Convert mono BGM to stereo
bgm_waveform = bgm_waveform.repeat(2, 1)
print(f"Converted mono BGM to stereo: new shape={bgm_waveform.shape}")
elif bgm_waveform.shape[0] == 2 and primary_channels == 1:
# Convert stereo BGM to mono
bgm_waveform = bgm_waveform.mean(dim=0, keepdim=True)
print(f"Converted stereo BGM to mono: new shape={bgm_waveform.shape}")
# Loop or trim BGM to match total audio length
if bgm_waveform.shape[1] < total_samples:
# BGM is shorter than needed, loop it
repeats_needed = (total_samples + bgm_waveform.shape[1] - 1) // bgm_waveform.shape[1]
bgm_repeated = bgm_waveform.repeat(1, repeats_needed)
bgm_waveform = bgm_repeated[:, :total_samples]
print(f"Looped BGM {repeats_needed} times to match duration, new shape={bgm_waveform.shape}")
elif bgm_waveform.shape[1] > total_samples:
# BGM is longer than needed, trim it
bgm_waveform = bgm_waveform[:, :total_samples]
print(f"Trimmed BGM to match duration, new shape={bgm_waveform.shape}")
# Apply fade-in effect
if fade_in_sec > 0:
fade_samples = int(fade_in_sec * sample_rate)
if fade_samples > 0 and fade_samples < bgm_waveform.shape[1]:
fade_curve = torch.linspace(0, 1, fade_samples)
for c in range(bgm_waveform.shape[0]):
bgm_waveform[c, :fade_samples] *= fade_curve
print(f"Applied {fade_in_sec}s fade-in to BGM")
# Mix BGM with primary audio
if has_actual_audio:
# For speech audio, we need a smoother approach than instant volume changes
# Use a sliding window average to detect audio presence, then smooth the volume control
# Step 1: Calculate audio energy over time with a sliding window
window_size = int(0.3 * sample_rate) # 300ms window, good for speech
primary_energy = torch.zeros(primary_audio.shape[1])
# Calculate energy profile
for i in range(primary_audio.shape[1]):
start = max(0, i - window_size//2)
end = min(primary_audio.shape[1], i + window_size//2)
window_data = primary_audio[:, start:end]
primary_energy[i] = window_data.abs().mean()
# Step 2: Apply smoothing to the energy profile
smoothing_window = int(0.5 * sample_rate) # 500ms smoothing window
smoothed_energy = torch.zeros_like(primary_energy)
for i in range(len(primary_energy)):
start = max(0, i - smoothing_window//2)
end = min(len(primary_energy), i + smoothing_window//2)
smoothed_energy[i] = primary_energy[start:end].mean()
# Step 3: Convert energy to volume level
# Set threshold - below this energy level, BGM will be at full volume
threshold = 0.005
# Set range - how quickly it transitions from min to max volume
range_factor = 0.01
# Create the volume mask
bgm_volume_mask = torch.ones(bgm_waveform.shape[1])
for i in range(min(len(smoothed_energy), bgm_volume_mask.shape[0])):
# Map energy to volume: higher energy = lower BGM volume
energy = smoothed_energy[i]
if energy > threshold:
# Linear mapping from energy to volume
volume_factor = max(bgm_volume, 1.0 - (energy - threshold) / range_factor)
bgm_volume_mask[i] = volume_factor
else:
# Below threshold, full volume
bgm_volume_mask[i] = 1.0
# Apply fade-in at the beginning of BGM
fade_samples = int(fade_in_sec * sample_rate)
if fade_samples > 0 and fade_samples < bgm_volume_mask.shape[0]:
fade_curve = torch.linspace(0, 1, fade_samples)
for i in range(fade_samples):
bgm_volume_mask[i] *= fade_curve[i]
# Apply the volume mask to all BGM channels
volume_adjusted_bgm = bgm_waveform.clone()
for c in range(bgm_waveform.shape[0]):
volume_adjusted_bgm[c, :min(bgm_waveform.shape[1], len(bgm_volume_mask))] *= bgm_volume_mask[:min(bgm_waveform.shape[1], len(bgm_volume_mask))]
print(f"Applied smooth BGM volume control with speech detection")
# Ensure primary_audio and BGM are the same length
if primary_audio.shape[1] < volume_adjusted_bgm.shape[1]:
# Pad primary audio with zeros
padding = torch.zeros(primary_audio.shape[0], volume_adjusted_bgm.shape[1] - primary_audio.shape[1])
primary_audio = torch.cat([primary_audio, padding], dim=1)
elif primary_audio.shape[1] > volume_adjusted_bgm.shape[1]:
# Pad BGM with zeros
padding = torch.zeros(volume_adjusted_bgm.shape[0], primary_audio.shape[1] - volume_adjusted_bgm.shape[1])
volume_adjusted_bgm = torch.cat([volume_adjusted_bgm, padding], dim=1)
audio_output = primary_audio + volume_adjusted_bgm
print(f"Mixed BGM with smooth volume control for speech audio")
else:
# If there's no actual audio content, use BGM at full volume
audio_output = bgm_waveform
print(f"Using BGM at full volume (no primary audio content)")
else:
# No BGM provided, just use primary audio
audio_output = primary_audio
print("No BGM provided, using only primary audio")
# Ensure audio_output has correct batch dimension
audio_output = audio_output.unsqueeze(0)
# Update video_info for output
video_info_output = {
"loaded_fps": output_fps,
"loaded_frame_count": concatenated_frames.shape[0],
"loaded_duration": concatenated_frames.shape[0] / output_fps,
"loaded_width": images1.shape[2], # Width
"loaded_height": images1.shape[1], # Height
"source_fps": video_info1["source_fps"],
"source_frame_count": video_info1["source_frame_count"] + video_info2["source_frame_count"],
"source_duration": video_info1["source_duration"] + video_info2["source_duration"],
"source_width": video_info1["source_width"],
"source_height": video_info1["source_height"],
}
# Prepare audio output
audio_output_dict = None
if audio_output is not None:
audio_output_dict = {
"waveform": audio_output,
"sample_rate": sample_rate
}
print(f"Final audio prepared: waveform shape={audio_output_dict['waveform'].shape}, sample_rate={sample_rate}")
else:
print("No final audio generated")
print(f"Output: frames={concatenated_frames.shape}, audio={audio_output.shape if audio_output is not None else 'none'}")
return (concatenated_frames, audio_output_dict, video_info_output)
# Node Mappings for ComfyUI
NODE_CLASS_MAPPINGS = {
"LatentSyncNode": LatentSyncNode,
"VideoLengthAdjuster": VideoLengthAdjuster,
"DG_VideoAudioMixer": DG_VideoAudioMixer,
}
# Display Names for ComfyUI
NODE_DISPLAY_NAME_MAPPINGS = {
"LatentSyncNode": "LatentSync1.5 Node",
"VideoLengthAdjuster": "Video Length Adjuster",
"DG_VideoAudioMixer": "DG Video Audio Mixer",
}