diff --git a/__init__.py b/__init__.py index 4e9c0bd..064e33c 100644 --- a/__init__.py +++ b/__init__.py @@ -1,5 +1,30 @@ -from .videochat.nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS +# from .videochat.nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS +from .videochat.nodes import ModelLoader, VideoSummary, LoadVideoUpload +from .nodes.audio import ResampleAudio, ConcatAudio, SplitAudio, JoinAudio, AudioCombine, ConvertAudioChannels +NODE_CLASS_MAPPINGS = { + "VCW_ModelLoader": ModelLoader, + "VCW_VideoSummary": VideoSummary, + "VCW_LoadVideo": LoadVideoUpload, + "ResampleAudio": ResampleAudio, + "ConcatAudio": ConcatAudio, + "SplitAudio": SplitAudio, + "JoinAudio": JoinAudio, + "AudioCombine": AudioCombine, + "ConvertAudioChannels": ConvertAudioChannels, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "VCW_ModelLoader": "Load Model (VideoChat)", + "VCW_VideoSummary": "Video Summary (VideoChat)", + "VCW_LoadVideo": "Load Video (Upload) (VideoChat) ", + "ResampleAudio": "Resample Audio", + "ConcatAudio": "Concat Audio", + "SplitAudio": "Split Audio", + "JoinAudio": "Join Audio", + "AudioCombine": "Audio Combine", + "ConvertAudioChannels": "Convert Audio Channels", +} WEB_DIRECTORY = "./web" __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"] \ No newline at end of file diff --git a/nodes/audio.py b/nodes/audio.py new file mode 100644 index 0000000..1677655 --- /dev/null +++ b/nodes/audio.py @@ -0,0 +1,375 @@ + +import sys +import warnings +from typing import Any, Dict + +import torch +import torchaudio +from torchaudio.transforms import Resample + +from typing import TypedDict +from typing import Tuple +from torch import Tensor + +import comfy.model_management + +# pip install git+https://git@github.com/facebookresearch/audiocraft#egg=audiocraft +# from audiocraft.data.audio_utils import convert_audio + +# from .node_def import BASE_NODE_CATEGORY, AudioData + +BASE_NODE_CATEGORY = "jida" +AudioData = Dict[str, Any] + +NODE_CATEGORY = BASE_NODE_CATEGORY + "/audio" + +class AUDIO(TypedDict): + """ + Required Fields: + waveform (torch.Tensor): A tensor containing the audio data. Shape: [Batch, Channels, frames]. + sample_rate (int): The sample rate of the audio data. + """ + + waveform: Tensor + sample_rate: int + +def convert_audio_channels(wav: torch.Tensor, channels: int = 2) -> torch.Tensor: + """Convert audio to the given number of channels. + + Args: + wav (torch.Tensor): Audio wave of shape [B, C, T]. + channels (int): Expected number of channels as output. + Returns: + torch.Tensor: Downmixed or unchanged audio wave [B, C, T]. + """ + *shape, src_channels, length = wav.shape + if src_channels == channels: + pass + elif channels == 1: + # Case 1: + # The caller asked 1-channel audio, and the stream has multiple + # channels, downmix all channels. + wav = wav.mean(dim=-2, keepdim=True) + elif src_channels == 1: + # Case 2: + # The caller asked for multiple channels, but the input file has + # a single channel, replicate the audio over all channels. + wav = wav.expand(*shape, channels, length) + elif src_channels >= channels: + # Case 3: + # The caller asked for multiple channels, and the input file has + # more channels than requested. In that case return the first channels. + wav = wav[..., :channels, :] + else: + # Case 4: What is a reasonable choice here? + raise ValueError('The audio file has less channels than requested but is not mono.') + return wav + +# def convert_audio(wav: torch.Tensor, from_rate: float, +# to_rate: float, to_channels: int) -> torch.Tensor: +# """Convert audio to new sample rate and number of audio channels.""" +# wav = julius.resample_frac(wav, int(from_rate), int(to_rate)) +# wav = convert_audio_channels(wav, to_channels) +# return wav + +# https://github.com/christian-byrne/audio-separation-nodes-comfyui/blob/master/src/combine_video_with_audio.py +# https://github.com/kale4eat/ComfyUI-speech-dataset-toolkit/blob/main/edit_nodes.py +# https://github.com/eigenpunk/ComfyUI-audio/blob/main/util_nodes.py +# https://github.com/facebookresearch/audiocraft/blob/main/audiocraft/data/audio_utils.py +class ConvertAudioChannels: + """ + convert an AUDIO's number of channels + """ + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "audio": ("AUDIO",), + "to_channels": ("INT", {"default": 1, "min": 1, "max": 2, "step": 1}), + } + } + + RETURN_TYPES = ("AUDIO",) + RETURN_NAMES = ("audio",) + FUNCTION = "convert" + CATEGORY = NODE_CATEGORY + DESCRIPTION = "Convert an AUDIO's number of channels." + + def convert(self, audio, to_channels): + # converted = [] + clip = audio["waveform"] + expand_dim = len(clip.shape) == 2 + if expand_dim: + clip = clip.unsqueeze(0) + # conv_clip = convert_audio(clip, from_rate, to_rate, to_channels) + conv_clip = convert_audio_channels(clip, to_channels) + conv_clip = conv_clip.squeeze(0) if expand_dim else conv_clip + audio_data = {"waveform": conv_clip, "sample_rate": audio["sample_rate"]} + return (audio_data,) + + + +# combine two audio tracks(mono or stereo) by overlaying their waveforms +class AudioCombine: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "audio_1": ("AUDIO",), + "audio_2": ("AUDIO",), + }, + "optional": { + "method": ( + ["add", "mean", "subtract", "multiply", "divide"], + { + "default": "add", + "tooltip": "The method used to combine the audio waveforms.", + }, + ) + }, + } + + FUNCTION = "main" + RETURN_TYPES = ("AUDIO",) + RETURN_NAMES = ("audio",) + CATEGORY = NODE_CATEGORY + DESCRIPTION = "Combine two audio tracks by overlaying their waveforms." + + def main( + self, + audio_1: AUDIO, + audio_2: AUDIO, + method: str = "add", + ) -> Tuple[AUDIO]: + + waveform_1: torch.Tensor = audio_1["waveform"] + input_sample_rate_1: int = audio_1["sample_rate"] + + waveform_2: torch.Tensor = audio_2["waveform"] + input_sample_rate_2: int = audio_2["sample_rate"] + + # Resample the audio if the sample rates are different + if input_sample_rate_1 != input_sample_rate_2: + device: torch.device = comfy.model_management.get_torch_device() + if input_sample_rate_1 < input_sample_rate_2: + resample = Resample(input_sample_rate_1, input_sample_rate_2).to(device) + waveform_1: torch.Tensor = resample(waveform_1.to(device)) + waveform_1.to("cpu") + output_sample_rate = input_sample_rate_2 + else: + resample = Resample(input_sample_rate_2, input_sample_rate_1).to(device) + waveform_2: torch.Tensor = resample(waveform_2.to(device)) + waveform_2.to("cpu") + output_sample_rate = input_sample_rate_1 + else: + output_sample_rate = input_sample_rate_1 + + # Ensure the audio is the same length + min_length = min(waveform_1.shape[-1], waveform_2.shape[-1]) + if waveform_1.shape[-1] != min_length: + waveform_1 = waveform_1[..., :min_length] + if waveform_2.shape[-1] != min_length: + waveform_2 = waveform_2[..., :min_length] + + match method: + case "add": + waveform = waveform_1 + waveform_2 + case "subtract": + waveform = waveform_1 - waveform_2 + case "multiply": + waveform = waveform_1 * waveform_2 + case "divide": + waveform = waveform_1 / waveform_2 + case "mean": + waveform = (waveform_1 + waveform_2) / 2 + + return ( + { + "waveform": waveform, + "sample_rate": output_sample_rate, + }, + ) + +# join/combine multiple audio tracks into a single audio track +class JoinAudio: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "audios": ("AUDIO",), + "silent_interval": ("FLOAT", {"default": 0.0, "step": 0.001}), + } + } + + CATEGORY = NODE_CATEGORY + INPUT_IS_LIST = True + RETURN_TYPES = ("AUDIO",) + RETURN_NAMES = ("audio",) + FUNCTION = "join_audio" + DESCRIPTION = "Join/combine multiple audio tracks into a single audio track." + + def join_audio(self, audios: list[AudioData], silent_interval: float): + if len(audios) == 0: + raise ValueError("Audios size is zero.") + if not all( + [audio["sample_rate"] == audios[0]["sample_rate"] for audio in audios] + ): + raise ValueError("sample_rate must be the same.") + if not all( + [ + audio["waveform"].size(1) == audios[0]["waveform"].size(1) + for audio in audios + ] + ): + raise ValueError("All audio must be either stereo or monaural.") + + samples = int(silent_interval * audios[0]["sample_rate"]) + interval = torch.zeros((audios[0]["waveform"].size(1), samples)) + new_tensors = [] + for audio in audios: + new_tensors.append(audio) + new_tensors.append(interval) + + new_tensors.pop() + + new_waveforms = torch.concat(new_tensors, dim=-1) + new_audio = {"waveform": new_waveforms, "sample_rate": audios[0]["sample_rate"]} + return (new_audio,) + +# split an audio track into two audio tracks(left and right) +class SplitAudio: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "audio": ("AUDIO",), + "second": ("FLOAT", {"default": 0.0, "step": 0.001}), + } + } + + CATEGORY = NODE_CATEGORY + + RETURN_TYPES = ("AUDIO", "AUDIO") + RETURN_NAMES = ("audio1", "audio2") + FUNCTION = "split_audio" + DESCRIPTION = "Split an audio track into two audio tracks(left and right)." + + def split_audio(self, audio: AudioData, second: float): + sample = max(0, int(second * audio["sample_rate"]) - 1) + view1 = audio["waveform"][..., :sample] + view2 = audio["waveform"][..., sample:] + audio1 = { + "waveform": view1.detach().clone(), + "sample_rate": audio["sample_rate"], + } + audio2 = { + "waveform": view2.detach().clone(), + "sample_rate": audio["sample_rate"], + } + return (audio1, audio2) + +# concatenate two audio tracks +class ConcatAudio: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "audio1": ("AUDIO",), + "audio2": ("AUDIO",), + "silent_interval": ("FLOAT", {"default": 0.0, "step": 0.001}), + } + } + + CATEGORY = NODE_CATEGORY + DESCRIPTION = "Concatenate two audio tracks." + RETURN_TYPES = ("AUDIO",) + RETURN_NAMES = ("audio",) + FUNCTION = "concat_audio" + + def concat_audio( + self, audio1: AudioData, audio2: AudioData, silent_interval: float + ): + if audio1["sample_rate"] != audio2["sample_rate"]: + raise ValueError( + f"sample_rate is different.\naudio1: {0}\naudio2: {1}".format( + audio1["sample_rate"], audio2["sample_rate"] + ) + ) + + audio1_ch = "stereo" if audio1["waveform"].size(1) > 1 else "monaural" + audio2_ch = "stereo" if audio2["waveform"].size(1) > 1 else "monaural" + if not audio1_ch == audio2_ch: + raise ValueError(f"naudio1 is {audio1_ch} but audio2 is {audio2_ch}") + + samples = int(silent_interval * audio1["sample_rate"]) + interval = torch.zeros( + (audio1["waveform"].size(0), audio1["waveform"].size(1), samples) + ) + new_waveform = torch.concat( + [audio1["waveform"], interval, audio2["waveform"]], dim=-1 + ) + new_audio = {"waveform": new_waveform, "sample_rate": audio1["sample_rate"]} + return (new_audio,) + +# resample an audio track +class ResampleAudio: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "audio": ("AUDIO",), + "new_freq": ("INT", {"default": 32000, "min": 1, "max": 2 ** 32}), + "resampling_method": (["sinc_interp_hann", "sinc_interp_kaiser"],), + "lowpass_filter_width": ( + "INT", + {"default": 6, "min": 0, "max": 2**32}, + ), + "rolloff": ( + "FLOAT", + { + "default": 0.99, + "min": 0, + "max": sys.float_info.max, + }, + ), + }, + "optional": { + "beta": ( + "FLOAT", + { + "default": 14.769656459379492, + "min": 0.0, + "max": sys.float_info.max, + }, + ), + }, + } + + CATEGORY = NODE_CATEGORY + RETURN_TYPES = ("AUDIO",) + RETURN_NAMES = ("audio",) + FUNCTION = "resample" + DESCRIPTION = "Resample an audio track." + + def resample( + self, + audio: AudioData, + new_freq: int, + resampling_method: str, + lowpass_filter_width: int, + rolloff: float, + beta=None, + ): + transform = torchaudio.transforms.Resample( + orig_freq=audio["sample_rate"], + new_freq=new_freq, + resampling_method=resampling_method, + lowpass_filter_width=lowpass_filter_width, + rolloff=rolloff, + beta=beta, + ) + new_audio = {"waveform": transform(audio["waveform"]), "sample_rate": new_freq} + return (new_audio,) \ No newline at end of file diff --git a/nodes/utils.py b/nodes/utils.py new file mode 100644 index 0000000..7c8f87d --- /dev/null +++ b/nodes/utils.py @@ -0,0 +1,26 @@ +import wave + +def is_stereo(wav_file_path): + """ + Checks if a WAV file is stereo or mono. + + Args: + wav_file_path: Path to the WAV file. + + Returns: + True if the WAV file is stereo, False if it's mono. + """ + try: + with wave.open(wav_file_path, 'rb') as wf: + num_channels = wf.getnchannels() + return num_channels == 2 + except wave.Error as e: + print(f"Error processing {wav_file_path}: {e}") + return False + +# Example usage +file_path = 'path/to/your/audio.wav' +if is_stereo(file_path): + print(f"{file_path} is a stereo WAV file.") +else: + print(f"{file_path} is a mono WAV file.") \ No newline at end of file diff --git a/videochat/nodes.py b/videochat/nodes.py index 98f87c4..a10c0b4 100644 --- a/videochat/nodes.py +++ b/videochat/nodes.py @@ -173,14 +173,3 @@ class LoadVideoUpload: return True -NODE_CLASS_MAPPINGS = { - "VCW_ModelLoader": ModelLoader, - "VCW_VideoSummary": VideoSummary, - "VCW_LoadVideo": LoadVideoUpload, -} - -NODE_DISPLAY_NAME_MAPPINGS = { - "VCW_ModelLoader": "Load Model (VideoChat)", - "VCW_VideoSummary": "Video Summary (VideoChat)", - "VCW_LoadVideo": "Load Video (Upload) (VideoChat) ", -} \ No newline at end of file