diff --git a/nodes.py b/nodes.py index e468031..026758f 100644 --- a/nodes.py +++ b/nodes.py @@ -1,8 +1,10 @@ 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(): @@ -733,14 +735,338 @@ class VideoLengthAdjuster: {"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", } \ No newline at end of file