367 lines
15 KiB
Python
367 lines
15 KiB
Python
|
|
import os
|
|
import shutil
|
|
import numpy as np
|
|
import yt_dlp
|
|
import torch
|
|
from .settings import MERGE_OPTIONS
|
|
from .helper import MultipleTypeProxy, increment_filename_no_overwrite
|
|
from ..lib.audio import MAX_INT16, SUPPORTED_AUDIO, AudioProcessor, audio_to_bytes, get_audio, load_input_audio, pad_audio, remix_audio, save_input_audio
|
|
from ..lib.helper import get_filenames, get_hash, get_merge_func
|
|
import folder_paths
|
|
|
|
CATEGORY = "🌺RVC-Studio/audio"
|
|
input_path = folder_paths.get_input_directory()
|
|
temp_path = folder_paths.get_temp_directory()
|
|
|
|
def to_audio_dict(audio: "np.ndarray", sr: int):
|
|
audio = np.atleast_2d(audio)
|
|
waveform = torch.from_numpy(audio.reshape((-1,audio.shape[0]))).unsqueeze(0)
|
|
return dict(waveform=waveform,sample_rate=sr)
|
|
|
|
class LoadAudio:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
input_dir = os.path.join(input_path,"audio")
|
|
os.makedirs(input_dir,exist_ok=True)
|
|
files = get_filenames(root=input_dir,exts=SUPPORTED_AUDIO,format_func=os.path.basename)
|
|
|
|
return {
|
|
"required": {
|
|
"audio": (files,),
|
|
"sr": (["None",16000,44100,48000],{"default": "None"}),
|
|
}}
|
|
|
|
CATEGORY = CATEGORY
|
|
|
|
RETURN_TYPES = ("STRING","VHS_AUDIO","AUDIO")
|
|
RETURN_NAMES = ("audio_name","vhs_audio","audio")
|
|
FUNCTION = "load_audio"
|
|
|
|
def load_audio(self, audio, sr):
|
|
audio_path = os.path.join(input_path,"audio",audio)
|
|
widgetId = get_hash(audio_path)
|
|
audio_name = os.path.basename(audio).split(".")[0]
|
|
sr = None if sr=="None" else int(sr)
|
|
audio = load_input_audio(audio_path,sr=sr)
|
|
|
|
return {"ui": {"preview": [{"filename": os.path.basename(audio_path), "type": "input", "subfolder": "audio", "widgetId": widgetId}]}, "result": (audio_name, lambda:audio_to_bytes(*audio), to_audio_dict(*audio))}
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, audio, sr):
|
|
return get_hash(audio,sr)
|
|
|
|
class AudioInfoNode:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"audio": (MultipleTypeProxy('AUDIO,VHS_AUDIO'),),
|
|
}}
|
|
|
|
CATEGORY = CATEGORY
|
|
|
|
RETURN_TYPES = ("VHS_AUDIO","AUDIO","FLOAT", "INT")
|
|
RETURN_NAMES = ("vhs_audio","audio","seconds","sr")
|
|
FUNCTION = "get_info"
|
|
|
|
def get_info(self, audio):
|
|
input_audio = get_audio(audio)
|
|
seconds = len(input_audio[0])/input_audio[1]
|
|
return (lambda:audio_to_bytes(*input_audio),to_audio_dict(*input_audio),float(seconds),int(input_audio[1]))
|
|
|
|
|
|
class DownloadAudio:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
|
|
return {
|
|
"required": {
|
|
"url": ("STRING", {"default": ""})
|
|
},
|
|
"optional": {
|
|
"sr": (["None",16000,44100,48000],{"default": "None"}),
|
|
"song_name": ("STRING",{"default": ""},),
|
|
"format": (SUPPORTED_AUDIO,{"default": "flac"})
|
|
}
|
|
}
|
|
|
|
CATEGORY = CATEGORY
|
|
|
|
RETURN_TYPES = ("STRING","VHS_AUDIO","AUDIO")
|
|
RETURN_NAMES = ("audio_name","vhs_audio","audio")
|
|
FUNCTION = "download_audio"
|
|
|
|
def download_audio(self, url, sr="None", song_name="", format="flac"):
|
|
|
|
assert "youtube" in url, "Please provide a valid youtube URL!"
|
|
|
|
widgetId = get_hash(url, sr, format)
|
|
sr = None if sr=="None" else int(sr)
|
|
audio_name = widgetId if song_name=="" else song_name
|
|
input_dir = os.path.join(input_path,"audio")
|
|
os.makedirs(input_dir,exist_ok=True)
|
|
audio_path = os.path.join(input_dir,f"{audio_name}.{format}")
|
|
|
|
if os.path.isfile(audio_path): input_audio = load_input_audio(audio_path,sr=sr)
|
|
else:
|
|
ydl_opts = {
|
|
'format': 'bestaudio/best', # Download the best audio quality available
|
|
'postprocessors': [{
|
|
'key': 'FFmpegExtractAudio',
|
|
'preferredcodec': format,
|
|
}],
|
|
'outtmpl': os.path.splitext(audio_path)[0], # Output file name and directory
|
|
}
|
|
|
|
with yt_dlp.YoutubeDL(ydl_opts) as ydl:
|
|
ydl.download([url])
|
|
|
|
input_audio = load_input_audio(audio_path,sr=sr)
|
|
|
|
return {"ui": {"preview": [{"filename": os.path.basename(audio_path), "type": "input", "subfolder": "audio", "widgetId": widgetId}]}, "result": (audio_name, lambda:audio_to_bytes(*input_audio),to_audio_dict(*input_audio))}
|
|
|
|
class MergeAudioNode:
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"audio1": (MultipleTypeProxy('AUDIO,VHS_AUDIO'),),
|
|
"audio2": (MultipleTypeProxy('AUDIO,VHS_AUDIO'),),
|
|
},
|
|
"optional": {
|
|
"sr": (["None",32000,40000,44100,48000],{
|
|
"default": "None"
|
|
}),
|
|
"merge_type": (MERGE_OPTIONS,{"default": "median"}),
|
|
"normalize": ("BOOLEAN",{"default": True}),
|
|
"audio3_opt": (MultipleTypeProxy('AUDIO,VHS_AUDIO'),{"default": None}),
|
|
"audio4_opt": (MultipleTypeProxy('AUDIO,VHS_AUDIO'),{"default": None}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("VHS_AUDIO","AUDIO")
|
|
RETURN_NAMES = ("vhs_audio","audio")
|
|
OUTPUT_NODE = True
|
|
|
|
FUNCTION = "merge"
|
|
|
|
CATEGORY = CATEGORY
|
|
|
|
def merge(self, audio1, audio2, sr="None", merge_type="median", normalize=False, audio3_opt=None, audio4_opt=None):
|
|
|
|
input_audios = [get_audio(audio) for audio in [audio1, audio2, audio3_opt, audio4_opt] if audio is not None]
|
|
widgetId = get_hash(*[audio_to_bytes(*audio) for audio in input_audios],sr,merge_type,normalize)
|
|
audio_path = os.path.join(temp_path,"preview",f"{widgetId}.flac")
|
|
|
|
if os.path.isfile(audio_path): merged_audio = load_input_audio(audio_path)
|
|
else:
|
|
# input_audios = [bytes_to_audio(audio) for audio in audios]
|
|
merged_sr = min([sr for (_,sr) in input_audios]) if sr=="None" else sr
|
|
input_audios = [remix_audio(audio,merged_sr,norm=normalize) for audio in input_audios]
|
|
merge_func = get_merge_func(merge_type)
|
|
merged_audio = merge_func(pad_audio(*[audio for (audio,_) in input_audios],axis=0),axis=0), merged_sr
|
|
print(save_input_audio(audio_path,merged_audio))
|
|
del input_audios
|
|
if os.path.isfile(audio_path): merged_audio = load_input_audio(audio_path)
|
|
|
|
audio_name = os.path.basename(audio_path)
|
|
return {"ui": {"preview": [{"filename": audio_name, "type": "temp", "subfolder": "preview", "widgetId": widgetId}]}, "result": (lambda: audio_to_bytes(*merged_audio),to_audio_dict(*merged_audio))}
|
|
|
|
class PreviewAudio:
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"audio": (MultipleTypeProxy('AUDIO,VHS_AUDIO'),),
|
|
"filename": ("STRING",{"default": "test"}),
|
|
"save_format": (SUPPORTED_AUDIO,{"default": "flac"},),
|
|
"save_channels": ([1,2],{"default": 1}),
|
|
"overwrite_existing": ("BOOLEAN", {"default": True}),
|
|
"autoplay": ("BOOLEAN", {"default": True}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING","VHS_AUDIO","AUDIO")
|
|
RETURN_NAMES = ("output_path","vhs_audio","audio")
|
|
|
|
OUTPUT_NODE = True
|
|
|
|
CATEGORY = CATEGORY
|
|
|
|
FUNCTION = "save_audio"
|
|
|
|
def save_audio(self, audio, filename, save_format, save_channels, overwrite_existing, autoplay):
|
|
|
|
filename = filename.strip()
|
|
assert filename, "Filename cannot be empty"
|
|
|
|
base_output_dir = folder_paths.get_output_directory()
|
|
assert os.path.exists(base_output_dir), f"Output directory {base_output_dir} does not exist"
|
|
output_dir = os.path.join(base_output_dir, 'audio')
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
output_path = os.path.join(output_dir, f"{filename}.{save_format}")
|
|
|
|
if os.path.isfile(output_path) and not overwrite_existing:
|
|
output_path = increment_filename_no_overwrite(output_path)
|
|
|
|
input_audio = get_audio(audio)
|
|
print(save_input_audio(output_path,input_audio,to_int16=save_format!="mp3",to_stereo=save_channels==2))
|
|
|
|
tempdir = os.path.join(temp_path,"preview")
|
|
os.makedirs(tempdir, exist_ok=True)
|
|
widgetId = get_hash(output_path,save_channels)
|
|
audio_name = f"{widgetId}.{save_format}"
|
|
preview_file = os.path.join(tempdir,audio_name)
|
|
shutil.copyfile(output_path,preview_file)
|
|
return {"ui": {"preview": [{
|
|
"filename": audio_name, "type": "temp", "subfolder": "preview", "widgetId": widgetId, "autoplay": autoplay
|
|
}]}, "result": (output_path, lambda:audio_to_bytes(*input_audio), to_audio_dict(*input_audio))}
|
|
|
|
|
|
class ProcessAudioNode:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"normalize": ("BOOLEAN",{"default": True}),
|
|
"threshold_silence": ("BOOLEAN",{"default": True}),
|
|
"dynamic_threshold": ("BOOLEAN",{"default": True})
|
|
},
|
|
"optional": {
|
|
"audio": (MultipleTypeProxy('AUDIO,VHS_AUDIO'),{"default": None}),
|
|
"dynamic_threshold_sample_size": ("INT",{"default": 4000, "min": 160, "max": 48000, "step": 160}),
|
|
"dynamic_threshold_multiplier": ("FLOAT",{"default": 2., "min": 1.5, "step": .1}),
|
|
"dynamic_threshold_fill_method": (["median","interpolation"],{"default": "median"}),
|
|
"dynamic_threshold_kernel_size": ("INT",{"default": 5, "min": 3, "step": 2}),
|
|
"silence_threshold_db": ("INT",{"default": -50, "min": -120, "max": 0}),
|
|
"normalize_threshold_db": ("INT",{"default": -1, "min": -10, "max": 0})
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("AUDIO_PROCESSOR","VHS_AUDIO","AUDIO")
|
|
RETURN_NAMES = ("audio_processor","vhs_audio","audio")
|
|
|
|
CATEGORY = CATEGORY
|
|
|
|
FUNCTION = "process_audio"
|
|
|
|
def process_audio(self, normalize, threshold_silence, dynamic_threshold, audio=None,
|
|
dynamic_threshold_sample_size=16000,
|
|
dynamic_threshold_multiplier=2.0,
|
|
dynamic_threshold_fill_method="median",
|
|
dynamic_threshold_kernel_size=5,
|
|
silence_threshold_db=-50,
|
|
normalize_threshold_db=-1):
|
|
|
|
audio_processor = AudioProcessor(
|
|
normalize=normalize,
|
|
threshold_silence=threshold_silence,
|
|
dynamic_threshold=dynamic_threshold,
|
|
sample_size=dynamic_threshold_sample_size,
|
|
multiplier=dynamic_threshold_multiplier,
|
|
fill_method=dynamic_threshold_fill_method,
|
|
kernel_size=dynamic_threshold_kernel_size,
|
|
silence_threshold_db=silence_threshold_db,
|
|
normalize_threshold_db=normalize_threshold_db)
|
|
|
|
if audio is None:
|
|
vhs_audio = comfy_audio = audio
|
|
else:
|
|
output_audio = audio_processor(audio)
|
|
vhs_audio = lambda:audio_to_bytes(*output_audio)
|
|
comfy_audio = to_audio_dict(*output_audio)
|
|
|
|
return ( audio_processor, vhs_audio, comfy_audio)
|
|
|
|
|
|
class AudioBatchValueNode:
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"audio": (MultipleTypeProxy('AUDIO,VHS_AUDIO'),),
|
|
"num_segments": ('INT', {"default": 2, "min": 2, "max": 256, "step": 1}),
|
|
"output_min": ('FLOAT', {'default': 0., "min": -1000., "max": 1000., "step": .01}),
|
|
"output_max": ('FLOAT', {'default': 1., "min": 0., "max": 1000., "step": .01}),
|
|
"norm": (["scale","tanh","sigmoid"], {"default": "scale"}),
|
|
},
|
|
"optional": {
|
|
"silence_threshold": ("INT", {"default": 1000, "min": 1, "max": MAX_INT16, "step": 1}),
|
|
"duration_list": ("INT", {"default": 0, "min": 0, "forceInput": True}),
|
|
"print_output": ("BOOLEAN", {"default": False}),
|
|
"inverse": ("BOOLEAN", {"default": False}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("FLOAT","INT","INT")
|
|
RETURN_NAMES = ("FLOAT","INT","num_values")
|
|
|
|
FUNCTION = "get_frame_weights"
|
|
|
|
CATEGORY = CATEGORY
|
|
|
|
@staticmethod
|
|
def get_rms(audio): # root mean squared of audio segment
|
|
return np.sqrt(np.nanmean(audio**2))
|
|
|
|
def get_frame_weights(self, audio, num_segments, output_min, output_max, norm, silence_threshold=1000, duration_list=0, print_output=False, inverse=False):
|
|
assert output_max>=output_min, f"{output_max=} must be greater or equal to {output_min=}!"
|
|
|
|
audio_data = get_audio(audio)
|
|
audio,_ = remix_audio(audio_data,norm=True,to_int16=True)
|
|
num_values = int(num_segments)
|
|
audio_rms = np.nan_to_num(list(map(self.get_rms,np.array_split(audio.flatten()/silence_threshold, num_values))),nan=0)
|
|
audio_zscore = (audio_rms-audio_rms.mean())/audio_rms.std()
|
|
output_range = output_max-output_min
|
|
|
|
if norm=="tanh":
|
|
x_norm = np.tanh(audio_zscore) #tanh=-1 to 1
|
|
if inverse: x_norm*=-1
|
|
x_norm = (x_norm * output_range + output_max + output_min)/2
|
|
elif norm=="sigmoid":
|
|
x_norm = 1. / (1. + np.exp(-audio_zscore)) #sigmoid=0 to 1
|
|
if inverse: x_norm=1-x_norm
|
|
x_norm = x_norm * output_range + output_min
|
|
else:
|
|
x_min = audio_zscore.min()
|
|
x_norm = (audio_zscore - x_min) / (audio_zscore.max() - x_min) #scale=0 to 1
|
|
if inverse: x_norm=1-x_norm
|
|
x_norm = x_norm * output_range + output_min
|
|
|
|
if print_output:
|
|
print(f"{audio_rms.min()=} {audio_rms.max()=} {audio_rms.mean()=} {len(audio_rms)=}")
|
|
print(f"{x_norm.min()=} {x_norm.max()=} {x_norm.mean()=} {len(x_norm)=}")
|
|
|
|
if isinstance(duration_list,list):
|
|
segments = np.cumsum(duration_list)
|
|
x_norm = np.array_split(x_norm,segments)
|
|
x_norm = map(list,x_norm)
|
|
x_norm_int = [list(map(int,norms)) for norms in x_norm]
|
|
else:
|
|
x_norm_int = map(int,x_norm)
|
|
|
|
return (list(x_norm),list(x_norm_int),num_values)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"RVC-Studio.LoadAudio": LoadAudio,
|
|
"RVC-Studio.PreviewAudio": PreviewAudio,
|
|
"MergeAudioNode": MergeAudioNode,
|
|
"AudioBatchValueNode": AudioBatchValueNode,
|
|
"DownloadAudio": DownloadAudio,
|
|
"ProcessAudioNode": ProcessAudioNode,
|
|
"AudioInfoNode": AudioInfoNode
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"RVC-Studio.LoadAudio": "🌺Load Audio",
|
|
"DownloadAudio": "🌺Youtube Downloader",
|
|
"RVC-Studio.PreviewAudio": "🌺Save Audio",
|
|
"MergeAudioNode": "🌺Merge Audio",
|
|
"AudioBatchValueNode": "🌺Audio RMS Batch Values",
|
|
"ProcessAudioNode": "🌺Audio Processor",
|
|
"AudioInfoNode": "🌺Audio Info"
|
|
} |