initial refactor of rvc studio
This commit is contained in:
SayanoAI
@@ -0,0 +1,254 @@
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
from .settings import PITCH_EXTRACTION_OPTIONS
|
||||
from ..lib import BASE_MODELS_DIR
|
||||
from ..lib.model_utils import load_hubert
|
||||
from ..lib.utils import get_file_hash, get_filenames, get_hash, get_optimal_torch_device
|
||||
from .utils import model_downloader
|
||||
import torch
|
||||
import folder_paths
|
||||
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
|
||||
from ..vc_infer_pipeline import get_vc
|
||||
from .settings.downloader import RVC_DOWNLOAD_LINK, RVC_MODELS, download_file
|
||||
from ..lib.audio import SUPPORTED_AUDIO, audio_to_bytes, load_input_audio
|
||||
from ..config import config
|
||||
|
||||
input_path = folder_paths.get_input_directory()
|
||||
temp_path = folder_paths.get_temp_directory()
|
||||
|
||||
model_ids = [
|
||||
'openai/whisper-large-v3',
|
||||
'openai/whisper-large-v2',
|
||||
'openai/whisper-large',
|
||||
'openai/whisper-medium',
|
||||
'openai/whisper-small',
|
||||
'openai/whisper-base',
|
||||
'openai/whisper-tiny',
|
||||
'openai/whisper-medium.en',
|
||||
'openai/whisper-small.en',
|
||||
'openai/whisper-base.en',
|
||||
'openai/whisper-tiny.en',
|
||||
]
|
||||
|
||||
languages = ['en', 'fr', 'es']
|
||||
|
||||
spacy_model_map = {
|
||||
"en": "en_core_web_md",
|
||||
"fr": "fr_core_news_md",
|
||||
"es": "es_core_news_md"
|
||||
}
|
||||
|
||||
class LoadPitchExtractionParams:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
'required': {
|
||||
'f0_method': (PITCH_EXTRACTION_OPTIONS,{"default": "rmvpe"}),
|
||||
"f0_autotune": ("BOOLEAN",),
|
||||
"index_rate": ("FLOAT",{
|
||||
"default": .75,
|
||||
"min": 0., #Minimum value
|
||||
"max": 1., #Maximum value
|
||||
"step": .01, #Slider's step
|
||||
"display": "slider"
|
||||
}),
|
||||
"resample_sr": ([0,16000,32000,40000,44100,48000],{"default": 0}),
|
||||
"rms_mix_rate": ("FLOAT",{
|
||||
"default": 0.25,
|
||||
"min": 0., #Minimum value
|
||||
"max": 1., #Maximum value
|
||||
"step": .01, #Slider's step
|
||||
"display": "slider"
|
||||
}),
|
||||
"protect": ("FLOAT",{
|
||||
"default": 0.25,
|
||||
"min": 0., #Minimum value
|
||||
"max": .5, #Maximum value
|
||||
"step": .01, #Slider's step
|
||||
"display": "slider"
|
||||
})
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ('PITCH_EXTRACTION', )
|
||||
RETURN_NAMES = ('pitch_extraction_params', )
|
||||
|
||||
CATEGORY = "🌺RVC-Studio/loaders"
|
||||
|
||||
FUNCTION = 'load_params'
|
||||
|
||||
def load_params(self, **params):
|
||||
if "rmvpe" in params.get("f0_method",""): model_downloader("rmvpe.pt")
|
||||
return (params,)
|
||||
|
||||
class LoadHubertModel:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
model_list = ["hubert_base.pt"] + get_filenames(root=BASE_MODELS_DIR,folder="*",exts=["pt"],format_func=os.path.basename)
|
||||
return {
|
||||
'required': {
|
||||
'model': (model_list,{"default": "hubert_base.pt"}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ('HUBERT_MODEL', )
|
||||
RETURN_NAMES = ('hubert_model', )
|
||||
|
||||
CATEGORY = "🌺RVC-Studio/loaders"
|
||||
|
||||
FUNCTION = 'load_model'
|
||||
|
||||
def load_model(self, model):
|
||||
model_path = model_downloader(model)
|
||||
hubert_model = lambda:load_hubert(model_path,config=config)
|
||||
return (hubert_model,)
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, model):
|
||||
return get_hash(model)
|
||||
|
||||
class LoadRVCModelNode:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
model_list = RVC_MODELS + get_filenames(root=BASE_MODELS_DIR,folder="RVC",exts=["pth"],format_func=lambda x: f"RVC/{os.path.basename(x)}")
|
||||
model_list = list(set(model_list)) # dedupe
|
||||
return {
|
||||
'required': {
|
||||
'model': (model_list,{"default": "RVC/Sayano.pth"}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ('RVC_MODEL', 'STRING')
|
||||
RETURN_NAMES = ('model', 'model_name')
|
||||
|
||||
CATEGORY = "🌺RVC-Studio/loaders"
|
||||
|
||||
FUNCTION = 'load_model'
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, model):
|
||||
return get_hash(model)
|
||||
|
||||
def load_model(self, model):
|
||||
try:
|
||||
filename = os.path.basename(model)
|
||||
subfolder = os.path.dirname(model)
|
||||
model_path = os.path.join(BASE_MODELS_DIR,subfolder,filename)
|
||||
|
||||
if not os.path.isfile(model_path):
|
||||
download_link = f"{RVC_DOWNLOAD_LINK}{model}"
|
||||
params = model_path, download_link
|
||||
if download_file(params): print(f"successfully downloaded: {model_path}")
|
||||
model = lambda:get_vc(model_path)
|
||||
return (model,filename.split(".")[0])
|
||||
except Exception as e:
|
||||
print(f"Error in {self.__class__.__name__}: {e}")
|
||||
raise e
|
||||
|
||||
class LoadWhisperModelNode:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
'required': {
|
||||
'model_id': (model_ids,{"default": "openai/whisper-base.en"}),
|
||||
},
|
||||
"optional": {
|
||||
'language': (languages,{"default": "en"}),
|
||||
"max_new_tokens": ("INT", {"default": 128, "min": 16, "max": 1024, "display": "slider"}),
|
||||
"chunk_length_s": ("INT", {"default": 30, "min": 15, "max": 60, "display": "slider"}),
|
||||
"batch_size": ("INT", {"default": 16, "min": 1, "max": 128, "display": "slider"}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ('TRANSCRIPTION_MODEL', )
|
||||
RETURN_NAMES = ('model', )
|
||||
|
||||
CATEGORY = "🌺RVC-Studio/loaders"
|
||||
|
||||
FUNCTION = 'load_model'
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, model):
|
||||
return get_hash(model)
|
||||
|
||||
def load_model(self, model_id, language="en", max_new_tokens=128, chunk_length_s=12, batch_size=16):
|
||||
device = get_optimal_torch_device()
|
||||
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
|
||||
|
||||
model = AutoModelForSpeechSeq2Seq.from_pretrained(
|
||||
model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True
|
||||
)
|
||||
processor = AutoProcessor.from_pretrained(model_id)
|
||||
model.to(device)
|
||||
|
||||
generate_kwargs = {}
|
||||
if model_id.endswith('.en'):
|
||||
if language != 'en':
|
||||
raise ValueError(f'Model {model_id} only supports English language')
|
||||
else:
|
||||
generate_kwargs['language'] = language
|
||||
|
||||
def pipe():
|
||||
|
||||
def get_spacy_model():
|
||||
import spacy
|
||||
model_name = spacy_model_map[language]
|
||||
model_path = os.path.join(BASE_MODELS_DIR,model_name)
|
||||
|
||||
if not os.path.exists(model_path):
|
||||
subprocess.call([sys.executable, "-m", "spacy", "download", model_name])
|
||||
model = spacy.load(model_name)
|
||||
model.to_disk(model_path)
|
||||
|
||||
return spacy.load(model_path)
|
||||
|
||||
return pipeline(
|
||||
'automatic-speech-recognition',
|
||||
model=model,
|
||||
tokenizer=processor.tokenizer,
|
||||
feature_extractor=processor.feature_extractor,
|
||||
max_new_tokens=max_new_tokens,
|
||||
chunk_length_s=chunk_length_s,
|
||||
batch_size=batch_size,
|
||||
return_timestamps=True,
|
||||
torch_dtype=torch_dtype,
|
||||
device=device,
|
||||
generate_kwargs=generate_kwargs,
|
||||
),get_spacy_model
|
||||
return (pipe, )
|
||||
|
||||
class LoadAudio:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
input_dir = input_path
|
||||
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 = "🌺RVC-Studio/loaders"
|
||||
|
||||
RETURN_TYPES = ("STRING","VHS_AUDIO")
|
||||
RETURN_NAMES = ("audio_name","vhs_audio")
|
||||
FUNCTION = "load_audio"
|
||||
|
||||
def load_audio(self, audio, sr):
|
||||
audio_path = folder_paths.get_annotated_filepath(audio)
|
||||
print(f"{audio_path=} {audio=}")
|
||||
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 (audio_name,lambda:audio_to_bytes(*audio))
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, audio):
|
||||
audio_path = folder_paths.get_annotated_filepath(audio)
|
||||
print(f"{audio_path=}")
|
||||
return get_file_hash(audio_path)
|
||||
@@ -0,0 +1,54 @@
|
||||
import os
|
||||
|
||||
from ..lib.audio import SUPPORTED_AUDIO, audio_to_bytes, bytes_to_audio, load_input_audio, save_input_audio
|
||||
from ..lib.utils import get_hash
|
||||
from .utils import increment_filename_no_overwrite
|
||||
import folder_paths
|
||||
|
||||
class PreviewAudio:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"audio": ("VHS_AUDIO",),
|
||||
"filename": ("STRING",{"default": "test"}),
|
||||
"save_format": (SUPPORTED_AUDIO,{"default": "flac"},),
|
||||
"save_channels": ([1,2],{"default": 1}),
|
||||
"overwrite_existing": ("BOOLEAN", {"default": True}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING","VHS_AUDIO")
|
||||
RETURN_NAMES = ("output_path","vhs_audio")
|
||||
|
||||
OUTPUT_NODE = True
|
||||
|
||||
CATEGORY = "🌺RVC-Studio/output"
|
||||
|
||||
FUNCTION = "save_audio"
|
||||
|
||||
def save_audio(self, audio, filename, save_format, save_channels, overwrite_existing):
|
||||
|
||||
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 = bytes_to_audio(audio())
|
||||
print(save_input_audio(output_path,input_audio,to_int16=True,to_stereo=save_channels==2))
|
||||
|
||||
return output_path, lambda:audio_to_bytes(*input_audio)
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, audio):
|
||||
return get_hash(audio())
|
||||
@@ -0,0 +1,76 @@
|
||||
import os
|
||||
|
||||
from ..lib.audio import audio_to_bytes, bytes_to_audio, load_input_audio, save_input_audio
|
||||
|
||||
from ..vc_infer_pipeline import vc_single
|
||||
import folder_paths
|
||||
from ..lib.utils import get_hash, get_optimal_torch_device
|
||||
from ..lib import BASE_CACHE_DIR, BASE_MODELS_DIR
|
||||
|
||||
input_path = folder_paths.get_input_directory()
|
||||
temp_path = folder_paths.get_temp_directory()
|
||||
# output_path = folder_paths.get_output_directory()
|
||||
# base_path = os.path.dirname(input_path)
|
||||
node_path = os.path.join(BASE_MODELS_DIR,"custom_nodes/ComfyUI-UVR5")
|
||||
weights_path = os.path.join(BASE_MODELS_DIR, "uvr5")
|
||||
device = get_optimal_torch_device()
|
||||
is_half = True
|
||||
|
||||
class RVCNode:
|
||||
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
|
||||
|
||||
return {
|
||||
"required": {
|
||||
"audio": ("VHS_AUDIO",),
|
||||
"model": ("RVC_MODEL",),
|
||||
"hubert_model": ("HUBERT_MODEL",),
|
||||
"pitch_extraction_params": ("PITCH_EXTRACTION",),
|
||||
"f0_up_key": ("INT",{
|
||||
"default": 0,
|
||||
"min": -14, #Minimum value
|
||||
"max": 14, #Maximum value
|
||||
"step": 1, #Slider's step
|
||||
"display": "slider"
|
||||
}),
|
||||
"format":(["wav", "flac", "mp3"],{
|
||||
"default": "flac"
|
||||
}),
|
||||
},
|
||||
"optional": {
|
||||
"use_cache": ("BOOLEAN",{"default": True})
|
||||
}
|
||||
}
|
||||
|
||||
OUTPUT_NODE = True
|
||||
|
||||
RETURN_TYPES = ("VHS_AUDIO",)
|
||||
|
||||
FUNCTION = "convert"
|
||||
|
||||
CATEGORY = "🌺RVC-Studio/rvc"
|
||||
|
||||
def convert(self, audio, model, hubert_model, pitch_extraction_params, format, use_cache=True, **kwargs):
|
||||
|
||||
cache_name = os.path.join(BASE_CACHE_DIR,"rvc",f"{get_hash(audio(), model, *pitch_extraction_params.items())}.{format}")
|
||||
|
||||
if use_cache and os.path.isfile(cache_name): output_audio = load_input_audio(cache_name)
|
||||
else:
|
||||
input_audio = bytes_to_audio(audio())
|
||||
output_audio = vc_single(hubert_model=hubert_model(),input_audio=input_audio,**model(),**pitch_extraction_params,**kwargs)
|
||||
print(save_input_audio(cache_name, output_audio))
|
||||
|
||||
return (lambda:audio_to_bytes(*output_audio),)
|
||||
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, *args, **kwargs):
|
||||
print(f"{args=} {kwargs=}")
|
||||
return get_hash(*args, *kwargs.items())
|
||||
|
||||
|
||||
@@ -0,0 +1,30 @@
|
||||
from functools import lru_cache
|
||||
import os
|
||||
|
||||
from ...lib import BASE_CACHE_DIR, PersistedDict
|
||||
|
||||
MENU_ITEMS = {
|
||||
"Get help": "https://github.com/SayanoAI/RVC-Studio/discussions",
|
||||
"Report a Bug": "https://github.com/SayanoAI/RVC-Studio/issues",
|
||||
"About": """This project provides a comprehensive platform for training RVC models and generating AI voice covers.
|
||||
Check out this github for more info: https://github.com/SayanoAI/RVC-Studio
|
||||
"""
|
||||
}
|
||||
|
||||
DEVICE_OPTIONS = ["cpu","cuda"]
|
||||
PITCH_EXTRACTION_OPTIONS = ["rmvpe","rmvpe+"]
|
||||
MERGE_OPTIONS=["median","mean","min","max"]
|
||||
TTS_MODELS = ["edge","speecht5"]
|
||||
N_THREADS_OPTIONS=[1,2,4,8,12,16]
|
||||
SR_MAP = {"32k": 32000,"40k": 40000, "48k": 48000}
|
||||
|
||||
@lru_cache
|
||||
def get_servers():
|
||||
os.makedirs(BASE_CACHE_DIR,exist_ok=True)
|
||||
fname = os.path.join(BASE_CACHE_DIR,"servers.shelve")
|
||||
servers = PersistedDict(fname)
|
||||
return servers
|
||||
|
||||
SERVERS = get_servers()
|
||||
RVC_INFERENCE_URL = SERVERS.RVC_INFERENCE_URL
|
||||
UVR_INFERENCE_URL = SERVERS.UVR_INFERENCE_URL
|
||||
@@ -0,0 +1,121 @@
|
||||
import re
|
||||
from typing import IO, List, Tuple
|
||||
import unicodedata
|
||||
import requests
|
||||
import os
|
||||
import zipfile
|
||||
|
||||
from ...lib import BASE_CACHE_DIR, BASE_MODELS_DIR
|
||||
|
||||
RVC_DOWNLOAD_LINK = 'https://huggingface.co/datasets/SayanoAI/RVC-Studio/resolve/main/'
|
||||
MDX_MODELS = ["MDXNET/UVR-MDX-NET-vocal_FT.onnx"]
|
||||
KARAFAN_MODELS = ["karafan/MDX23C-8KFFT-InstVoc_HQ.ckpt"]
|
||||
VR_MODELS = ["UVR/UVR-DeEcho-DeReverb.pth","UVR/HP5-vocals+instrumentals.pth"]
|
||||
RVC_MODELS = [
|
||||
"RVC/Claire.pth",
|
||||
"RVC/Sayano.pth",
|
||||
"RVC/Mae_v2.pth",
|
||||
"RVC/Fuji.pth",
|
||||
"RVC/Monika.pth"]
|
||||
BASE_MODELS = ["hubert_base.pt", "rmvpe.pt"]
|
||||
VITS_MODELS = ["VITS/pretrained_ljs.pth"]
|
||||
PRETRAINED_MODELS = [
|
||||
"pretrained_v2/D48k.pth",
|
||||
"pretrained_v2/G48k.pth",
|
||||
"pretrained_v2/D32k.pth",
|
||||
"pretrained_v2/G32k.pth",
|
||||
"pretrained_v2/D40k.pth",
|
||||
"pretrained_v2/G40k.pth",
|
||||
"pretrained_v2/f0D48k.pth",
|
||||
"pretrained_v2/f0G48k.pth",
|
||||
"pretrained_v2/f0D40k.pth",
|
||||
"pretrained_v2/f0G40k.pth",
|
||||
"pretrained_v2/f0D32k.pth",
|
||||
"pretrained_v2/f0G32k.pth",
|
||||
|
||||
"pretrained_v2/f0_RIN_E3_40k_D.pth",
|
||||
"pretrained_v2/f0_RIN_E3_40k_G.pth",
|
||||
"pretrained_v2/f0Ov2Super32kD.pth",
|
||||
"pretrained_v2/f0Ov2Super32kG.pth",
|
||||
"pretrained_v2/f0Ov2Super40kD.pth",
|
||||
"pretrained_v2/f0Ov2Super40kG.pth",
|
||||
]
|
||||
LLM_MODELS = [
|
||||
"https://huggingface.co/TheBloke/Airoboros-L2-7B-2.1-GGUF/resolve/main/airoboros-l2-7b-2.1.Q4_K_M.gguf",
|
||||
"https://huggingface.co/TheBloke/Pygmalion-2-7B-GGUF/resolve/main/pygmalion-2-7b.Q4_K_M.gguf",
|
||||
"https://huggingface.co/TheBloke/Zarablend-MX-L2-7B-GGUF/resolve/main/zarablend-mx-l2-7b.Q4_K_M.gguf",
|
||||
"https://huggingface.co/TheBloke/MythoMax-L2-Kimiko-v2-13B-GGUF/resolve/main/mythomax-l2-kimiko-v2-13b.Q4_K_M.gguf"
|
||||
]
|
||||
STT_MODELS = [
|
||||
"https://alphacephei.com/vosk/models/vosk-model-en-us-0.22-lgraph.zip"
|
||||
]
|
||||
|
||||
def download_file(params: Tuple[str, str]):
|
||||
model_path, download_link = params
|
||||
if os.path.isfile(model_path): raise FileExistsError(f"{model_path} already exists!")
|
||||
os.makedirs(os.path.dirname(model_path), exist_ok=True)
|
||||
|
||||
with requests.get(download_link,stream=True) as r:
|
||||
r.raise_for_status()
|
||||
with open(model_path, 'wb') as f:
|
||||
for chunk in r.iter_content(chunk_size=8192):
|
||||
f.write(chunk)
|
||||
return True
|
||||
|
||||
def download_link_generator(download_link: str,model_list: List[str]):
|
||||
for model in model_list:
|
||||
model_path = os.path.join(BASE_MODELS_DIR,model)
|
||||
yield (model_path, f"{download_link}{model}")
|
||||
|
||||
def save_file(params: Tuple[str, any]):
|
||||
(data_path, datum) = params
|
||||
if "zip" in os.path.splitext(data_path)[-1]: return save_zipped_files(params) # unzip
|
||||
else:
|
||||
try:
|
||||
with open(data_path,"wb") as f:
|
||||
f.write(datum)
|
||||
return f"Successfully saved file to: {data_path}"
|
||||
except Exception as e:
|
||||
return f"Failed to save file: {e}"
|
||||
|
||||
def save_file_generator(save_dir: str, data: List[IO]):
|
||||
for datum in data:
|
||||
data_path = os.path.join(save_dir,datum.name)
|
||||
yield (data_path, datum.read())
|
||||
|
||||
def save_zipped_files(params: Tuple[str, any]):
|
||||
(data_path, datum) = params
|
||||
|
||||
try:
|
||||
print(f"saving zip file: {data_path}")
|
||||
temp_dir = os.path.join(BASE_CACHE_DIR,"zips")
|
||||
os.makedirs(temp_dir,exist_ok=True)
|
||||
name = os.path.basename(data_path)
|
||||
zip_path = os.path.join(temp_dir,name)
|
||||
|
||||
with open(zip_path,"wb") as f:
|
||||
f.write(datum)
|
||||
|
||||
print(f"extracting zip file: {zip_path}")
|
||||
with zipfile.ZipFile(zip_path, 'r') as zip_ref:
|
||||
zip_ref.extractall(os.path.dirname(data_path))
|
||||
print(f"finished extracting zip file")
|
||||
|
||||
os.remove(zip_path) # cleanup
|
||||
return f"Successfully saved files to: {data_path}"
|
||||
except Exception as e:
|
||||
return f"Failed to save files: {e}"
|
||||
|
||||
def slugify_filepath(filepath):
|
||||
# Split the path into directory and filename
|
||||
directory, filename = os.path.split(filepath)
|
||||
# Normalize the filename
|
||||
filename = unicodedata.normalize('NFKD', filename)
|
||||
# Encode the filename as ASCII and ignore errors
|
||||
filename = filename.encode('ascii', 'ignore').decode()
|
||||
# Convert the filename to lowercase
|
||||
filename = filename.lower()
|
||||
# Replace spaces and other unwanted characters with dashes
|
||||
filename = re.sub(r'[^a-z0-9.-]+', '-', filename)
|
||||
# Join the directory and the slugified filename
|
||||
return os.path.join(directory, filename)
|
||||
@@ -0,0 +1,204 @@
|
||||
import json
|
||||
import os
|
||||
from pprint import pprint
|
||||
|
||||
import numpy as np
|
||||
|
||||
from ..lib import BASE_CACHE_DIR
|
||||
from .utils import increment_filename_no_overwrite
|
||||
from ..lib.utils import get_hash
|
||||
import folder_paths
|
||||
from ..lib.audio import bytes_to_audio, remix_audio
|
||||
import textacy
|
||||
from textacy.extract import keyterms as kt
|
||||
|
||||
temp_path = folder_paths.get_temp_directory()
|
||||
|
||||
def extract_keywords(text: str, max_words: int, spacy_model, **kwargs):
|
||||
text = text.strip()
|
||||
doc = textacy.make_spacy_doc(text, lang=spacy_model)
|
||||
topn = int(max_words) if max_words>0 else len(text)
|
||||
include_pos = ["NOUN","ADJ","PROPN","VERB","NUM","ADP"]
|
||||
ngrams = [1,2]
|
||||
tags = []
|
||||
try:
|
||||
terms = kt.sgrank(doc, ngrams=ngrams, normalize="lower", topn=topn, include_pos=include_pos)
|
||||
tags = list(map(lambda v:v[0],sorted(terms,key=lambda v:v[1],reverse=True)[:topn]))
|
||||
except Exception as error:
|
||||
print(f"{text=} {error=}")
|
||||
|
||||
return ", ".join(tags).strip()
|
||||
|
||||
def limit_sentence(text: str, max_words: int, **kwargs):
|
||||
text = text.strip()
|
||||
topn = int(max_words) if max_words>0 else len(text)
|
||||
if topn>0: return " ".join(text.split()[:topn])
|
||||
else: return text
|
||||
|
||||
class AudioTranscriptionNode:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"pipeline": ('TRANSCRIPTION_MODEL',),
|
||||
"audio": ('VHS_AUDIO',),
|
||||
"format_newlines_on_punctuation": ('BOOLEAN', {'default': True}),
|
||||
"loop": ('BOOLEAN', {'default': True}),
|
||||
"use_tags": ('BOOLEAN', {'default': True}),
|
||||
"max_words": ('INT', {'default': 16, "min": 0, "max": 32, "display": "slider"}),
|
||||
"frame_interpolation": ("INT", {"default": 0, "min": 0, "max": 120})
|
||||
},
|
||||
"optional": {
|
||||
"save_transcription": ('BOOLEAN', {'default': False}),
|
||||
"save_chunks": ('BOOLEAN', {'default': False}),
|
||||
"save_filename": ('STRING', {'default': 'transcription'}),
|
||||
"overwrite_existing": ('BOOLEAN', {'default': True}),
|
||||
"print_output": ('BOOLEAN', {'default': True}),
|
||||
"use_cache": ("BOOLEAN",{"default": True})
|
||||
}
|
||||
}
|
||||
|
||||
OUTPUT_NODE = True
|
||||
RETURN_TYPES = ("STRING", "STRING", "STRING", "INT", "INT")
|
||||
RETURN_NAMES = ("transcription", "batch_prompt_text", "batch_values_text", "num_chunks","num_frames")
|
||||
|
||||
FUNCTION = "transcribe"
|
||||
|
||||
CATEGORY = "🌺RVC-Studio/stt"
|
||||
|
||||
@staticmethod
|
||||
def save_output(filename,text,overwrite_existing=False):
|
||||
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, 'transcriptions')
|
||||
|
||||
file_path = os.path.join(output_dir, filename)
|
||||
dir_name = os.path.dirname(file_path)
|
||||
os.makedirs(dir_name, exist_ok=True)
|
||||
if os.path.exists(file_path) and not overwrite_existing:
|
||||
file_path = increment_filename_no_overwrite(file_path)
|
||||
with open(file_path, 'w') as f:
|
||||
f.write(text)
|
||||
|
||||
@staticmethod
|
||||
def process_text_chunks(ichunk,text_processor,frame_interpolation,**kwargs):
|
||||
i,chunk=ichunk
|
||||
frame_interpolation=int(frame_interpolation)
|
||||
text = text_processor(chunk["text"],**kwargs)
|
||||
timestamp = np.nan_to_num(np.array(chunk["timestamp"],dtype=float),nan=i)
|
||||
index = int(timestamp[0]*frame_interpolation) if frame_interpolation>0 else i
|
||||
prompt = f'"{index}":"{text}"' if len(text) else f'"{index}":""'
|
||||
duration = max(abs(int(np.round(timestamp[1]-timestamp[0]))),1) # at least 1 sec duration
|
||||
if frame_interpolation>1: duration*=frame_interpolation
|
||||
value = f'"{index}":"({duration})"'
|
||||
return prompt,value
|
||||
|
||||
@staticmethod
|
||||
def load_cache(*args):
|
||||
results = None
|
||||
cache_name = os.path.join(BASE_CACHE_DIR,"stt",f"{get_hash(*args)}.json")
|
||||
os.makedirs(os.path.dirname(cache_name), exist_ok=True)
|
||||
if os.path.isfile(cache_name): results = json.load(cache_name)
|
||||
return results, cache_name
|
||||
|
||||
def transcribe(
|
||||
self, pipeline, audio,
|
||||
format_newlines_on_punctuation,loop,
|
||||
use_tags,max_words,frame_interpolation,
|
||||
save_transcription=False, save_chunks=False,
|
||||
save_filename="transcription", overwrite_existing=True,
|
||||
print_output=True,use_cache=True
|
||||
):
|
||||
print('Starting Transcription')
|
||||
|
||||
audio_data = bytes_to_audio(audio())
|
||||
if use_cache: result, cache_name = self.load_cache(pipeline,audio(),loop,use_tags,max_words,frame_interpolation)
|
||||
else: result = None
|
||||
|
||||
audio_frames = int(np.ceil(len(audio_data[0])/audio_data[1]))
|
||||
whisper_model,get_spacy_model=pipeline()
|
||||
|
||||
if result is None:
|
||||
audio,_ = remix_audio(audio_data,target_sr=16000,norm=True)
|
||||
result = whisper_model(audio)
|
||||
|
||||
if use_cache and result:
|
||||
with open(cache_name, "w") as f:
|
||||
json.dump(result, f)
|
||||
|
||||
print('Transcription Done')
|
||||
text = result["text"]
|
||||
total_chunks = result['chunks']
|
||||
if format_newlines_on_punctuation:
|
||||
punct = ['.', '?', '!']
|
||||
for p in punct:
|
||||
text = text.replace(f'{p} ', f'{p}\n')
|
||||
|
||||
if loop: # append first frame to chunk stack
|
||||
last_chunk = None
|
||||
for i in range(len(total_chunks)):
|
||||
last_chunk = total_chunks[-1-i]
|
||||
|
||||
if "timestamp" in last_chunk and len(last_chunk["timestamp"])>=1:
|
||||
last_chunk["timestamp"] = np.nan_to_num(np.array(last_chunk["timestamp"],dtype=float),nan=audio_frames)
|
||||
break
|
||||
if last_chunk is not None:
|
||||
timestamp = last_chunk["timestamp"]
|
||||
start_time = timestamp[-1]
|
||||
duration = max(1,audio_frames - start_time)
|
||||
total_chunks.append({
|
||||
"text": total_chunks[0]["text"],
|
||||
"timestamp": (start_time,start_time+duration)
|
||||
})
|
||||
else: # adds 1s frame to prevent audio from stopping early
|
||||
num_frames+=1
|
||||
|
||||
if use_tags:
|
||||
spacy_model,text_processor = get_spacy_model(),extract_keywords
|
||||
else:
|
||||
spacy_model,text_processor = None,limit_sentence
|
||||
|
||||
# split transcript into prompt based on timestamp
|
||||
# find length of each frame using timestamp
|
||||
prompts = []
|
||||
values = []
|
||||
for ichunk in enumerate(total_chunks):
|
||||
p, v = self.process_text_chunks(ichunk,text_processor,frame_interpolation,spacy_model=spacy_model,max_words=max_words)
|
||||
prompts.append(p)
|
||||
values.append(v)
|
||||
|
||||
num_chunks = len(total_chunks)
|
||||
batch_prompt = ',\n'.join(prompts)
|
||||
batch_values = ',\n'.join(values)
|
||||
num_frames = int(audio_frames*frame_interpolation if frame_interpolation>0 else audio_frames)
|
||||
|
||||
if print_output:
|
||||
pprint(total_chunks)
|
||||
print(batch_prompt)
|
||||
print(batch_values)
|
||||
print(f"{num_chunks=}, {audio_frames=}s, {num_frames=}")
|
||||
|
||||
if save_transcription:
|
||||
self.save_output(save_filename + '.txt', text, overwrite_existing=overwrite_existing)
|
||||
|
||||
if save_chunks:
|
||||
self.save_output(
|
||||
save_filename + '_prompt.json',
|
||||
json.dumps(json.loads("{"+batch_prompt+"}"),indent=2), # formats json with indent
|
||||
overwrite_existing=overwrite_existing)
|
||||
|
||||
self.save_output(
|
||||
save_filename + '_values.json',
|
||||
json.dumps(json.loads("{"+batch_values+"}"),indent=2), # formats json with indent
|
||||
overwrite_existing=overwrite_existing)
|
||||
|
||||
text = text_processor(text,max_words=max_words,spacy_model=spacy_model)
|
||||
return (text, batch_prompt, batch_values, num_chunks, num_frames)
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, *args, **kwargs):
|
||||
print(f"{args=} {kwargs=}")
|
||||
return get_hash(*args, *kwargs.items())
|
||||
@@ -0,0 +1,368 @@
|
||||
import itertools
|
||||
import os
|
||||
from pprint import pprint
|
||||
import re
|
||||
import tempfile
|
||||
import torch
|
||||
import folder_paths
|
||||
from .settings import MERGE_OPTIONS
|
||||
from ..lib.utils import gc_collect, get_hash
|
||||
from .settings.downloader import RVC_DOWNLOAD_LINK, download_file
|
||||
from ..lib import BASE_MODELS_DIR
|
||||
import numpy as np
|
||||
from ..lib.audio import MAX_INT16, audio_to_bytes, bytes_to_audio, merge_audio, remix_audio
|
||||
|
||||
CATEGORY = "🌺RVC-Studio/utils"
|
||||
temp_path = folder_paths.get_temp_directory()
|
||||
|
||||
def model_downloader(model):
|
||||
filename = os.path.basename(model)
|
||||
subfolder = os.path.dirname(model)
|
||||
model_path = os.path.join(BASE_MODELS_DIR,subfolder,filename)
|
||||
if not os.path.isfile(model_path):
|
||||
download_link = f"{RVC_DOWNLOAD_LINK}{model}"
|
||||
params = model_path, download_link
|
||||
if download_file(params): print(f"successfully downloaded: {model_path}")
|
||||
return model_path
|
||||
|
||||
def increment_filename_no_overwrite(proposed_path):
|
||||
output_dir, filename_ext = os.path.split(proposed_path)
|
||||
filename, file_format = os.splitext(filename_ext)
|
||||
files = os.listdir(output_dir)
|
||||
files = filter(lambda f: os.path.isfile(os.path.join(output_dir, f)), files)
|
||||
files = filter(lambda f: f.startswith(filename), files)
|
||||
files = filter(lambda f: f.endswith('.' + file_format), files)
|
||||
file_numbers = [re.search(r'_(\d+)\.', f) for f in files]
|
||||
file_numbers = [int(f.group(1)) for f in file_numbers if f]
|
||||
this_file_number = max(file_numbers) + 1 if file_numbers else 1
|
||||
output_path = os.path.join(output_dir, filename + f'_{this_file_number}.' + file_format)
|
||||
return output_path
|
||||
|
||||
class MergeImageBatches:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"images": ("IMAGE",),
|
||||
"iterate": ("BOOLEAN",{"default": False})
|
||||
},
|
||||
}
|
||||
|
||||
INPUT_IS_LIST = (True, False)
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "merge"
|
||||
|
||||
CATEGORY = CATEGORY
|
||||
|
||||
__archive__ = None
|
||||
|
||||
@staticmethod
|
||||
def iter_images(fname,images):
|
||||
size = len(images)
|
||||
if size==0: return images
|
||||
shape = images[0].shape
|
||||
if len(shape)>3: shape=shape[1:]
|
||||
|
||||
memmap = np.memmap(fname, mode='w+', dtype=np.float32, shape=(size,*shape))
|
||||
|
||||
for i,img in enumerate(itertools.chain(*images)):
|
||||
if i >= len(memmap): #need to expand memmap
|
||||
print(f"index {i} is larger than {len(memmap)}. expanding memmap to {size+i}...")
|
||||
memmap.flush()
|
||||
memmap = np.memmap(fname, dtype=np.float32, mode='r+',shape=(size+i,*shape), order='F')
|
||||
memmap[i]=img
|
||||
# del img
|
||||
# gc_collect()
|
||||
memmap.flush()
|
||||
|
||||
print(f"{fname=} size={i} shape={memmap.shape}")
|
||||
|
||||
# del images
|
||||
images = torch.from_numpy(memmap[:i])
|
||||
|
||||
return images
|
||||
|
||||
def __del__(self):
|
||||
try:
|
||||
if os.path.exists(self.__archive__): os.remove(self.__archive__)
|
||||
gc_collect()
|
||||
except Exception as error:
|
||||
print(f"Failed to delete {self.__archive__}: {error=}")
|
||||
|
||||
def merge(self, images, iterate):
|
||||
|
||||
if len(images) <= 1:
|
||||
return (images[0],)
|
||||
else:
|
||||
if hasattr(iterate,"pop"): iterate=iterate.pop()
|
||||
print(f"{iterate=}")
|
||||
|
||||
if iterate:
|
||||
if self.__archive__ is None:
|
||||
with tempfile.NamedTemporaryFile(delete=False) as ntf:
|
||||
self.__archive__ = ntf.name
|
||||
images = self.iter_images(self.__archive__,images)
|
||||
else:
|
||||
images = torch.cat(images)
|
||||
|
||||
gc_collect()
|
||||
print(f"Merged images: {len(images)=} {images.shape=}")
|
||||
|
||||
return (images,)
|
||||
|
||||
|
||||
class MergeLatentBatches:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": {
|
||||
"latents": ("LATENT",),
|
||||
}
|
||||
}
|
||||
|
||||
INPUT_IS_LIST = True
|
||||
|
||||
RETURN_TYPES = ("LATENT",)
|
||||
FUNCTION = "merge"
|
||||
|
||||
CATEGORY = CATEGORY
|
||||
|
||||
@staticmethod
|
||||
def merge_tensor(tensors):
|
||||
return torch.cat(tensors, dim=0)
|
||||
|
||||
def merge(self, latents):
|
||||
merged_latents = {}
|
||||
gc_collect()
|
||||
merged_latents["samples"] = self.merge_tensor([latent["samples"] for latent in latents])
|
||||
merged_latents["noise_mask"] = self.merge_tensor([latent["noise_mask"] for latent in latents])
|
||||
merged_latents["batch_index"] = self.merge_tensor([latent["batch_index"] for latent in latents])
|
||||
|
||||
print(f'Merged latents: {len(merged_latents["samples"])=} {latents["batch_index"]=}')
|
||||
return (merged_latents,)
|
||||
|
||||
class AudioBatchValueNode:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"audio": ('VHS_AUDIO',),
|
||||
"num_frames": ('INT', {"min": 1, "max": 128, "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, "display": "slider"}),
|
||||
"frame_multiplier": ("INT", {"default": 1, "min": 1, "max": 120, "step": 1}),
|
||||
"print_output": ("BOOLEAN", {"default": False}),
|
||||
"inverse": ("BOOLEAN", {"default": False}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING","INT")
|
||||
RETURN_NAMES = ("batch_value_text","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_frames, output_min, output_max, norm,
|
||||
silence_threshold=1000, frame_multiplier=1, print_output=False, inverse=False):
|
||||
assert output_max>=output_min, f"{output_max=} must be greater or equal to {output_min=}!"
|
||||
|
||||
audio_data = bytes_to_audio(audio())
|
||||
audio,_ = remix_audio(audio_data,norm=True,to_int16=True)
|
||||
num_values = int(num_frames*frame_multiplier)
|
||||
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
|
||||
|
||||
print(f"{audio_rms.min()=} {audio_rms.max()=} {audio_rms.mean()=}")
|
||||
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
|
||||
print(f"{x_norm.min()=} {x_norm.max()=} {x_norm.mean()=}")
|
||||
|
||||
batch_value = ",\n".join([f'{n}:({v})' for n,v in enumerate(x_norm)])
|
||||
if print_output: pprint(batch_value)
|
||||
return (batch_value,num_values)
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, *args, **kwargs):
|
||||
print(f"{args=} {kwargs=}")
|
||||
return get_hash(*args, *kwargs.items())
|
||||
|
||||
class ImageRepeatInterleavedNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"images": ("IMAGE",),
|
||||
"repeats": ("INT", {"default": 1, "min": 1, "max": 4096}),
|
||||
},
|
||||
"optional": {
|
||||
"fps": ("INT", {"default": 1, "min": 1})
|
||||
}
|
||||
}
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
INPUT_IS_LIST = True
|
||||
OUTPUT_IS_LIST = (True, )
|
||||
|
||||
FUNCTION = "rebatch"
|
||||
|
||||
CATEGORY = CATEGORY
|
||||
|
||||
def rebatch(self, images, repeats, fps=1):
|
||||
repeats = np.array(repeats).flatten()*fps
|
||||
print(f"{repeats=}")
|
||||
all_images = []
|
||||
for img in images:
|
||||
for i in range(img.shape[0]):
|
||||
all_images.append(img[i:i+1])
|
||||
|
||||
# expand repeat
|
||||
if len(repeats)==1:
|
||||
repeats = repeats.repeat(len(all_images),axis=0)
|
||||
elif len(repeats)<len(all_images):
|
||||
pad_width = len(all_images) - len(repeats)
|
||||
repeats = np.pad(repeats, (0, pad_width), mode='constant', constant_values=1)
|
||||
|
||||
for i,img in enumerate(all_images):
|
||||
if repeats[i]>1: all_images[i] = img.expand(repeats[i], *img.shape).flatten(0,1)
|
||||
else: all_images[i] = img
|
||||
print(f"{img.shape=} {all_images[i].shape=}")
|
||||
|
||||
print(f"{len(images)=} {len(all_images)=}")
|
||||
gc_collect()
|
||||
|
||||
return (all_images,)
|
||||
|
||||
class LatentRepeatInterleavedNode:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"latents": ("LATENT",),
|
||||
"repeats": ("INT", {"default": 1, "min": 1, "max": 4096}),
|
||||
},
|
||||
"optional": {
|
||||
"fps": ("INT", {"default": 1, "min": 1})
|
||||
}
|
||||
}
|
||||
RETURN_TYPES = ("LATENT",)
|
||||
INPUT_IS_LIST = True
|
||||
OUTPUT_IS_LIST = (True, )
|
||||
|
||||
FUNCTION = "merge"
|
||||
|
||||
CATEGORY = CATEGORY
|
||||
|
||||
@staticmethod
|
||||
def repeat(latents, repeat):
|
||||
if latents.ndim>3: #batched
|
||||
expanded_latents=(latents.expand(repeat, *latents.shape).flatten(0,1))
|
||||
else:
|
||||
expanded_latents=(latents.expand(repeat, *latents.shape))
|
||||
|
||||
print(f"{len(latents)}*{repeat}={len(expanded_latents)} ({latents.shape=} => {expanded_latents.shape=})")
|
||||
return expanded_latents
|
||||
|
||||
@staticmethod
|
||||
def rebatch(samples, repeats):
|
||||
offset = 0
|
||||
for i in range(len(samples)):
|
||||
|
||||
samples[i]["batch_index"] = [i+offset for i in range(repeats[i])]
|
||||
offset += repeats[i]
|
||||
|
||||
samples[i]["samples"] = LatentRepeatInterleavedNode.repeat(samples[i]["samples"],repeats[i])
|
||||
|
||||
if "noise_mask" in samples[i]:
|
||||
samples[i]["noise_mask"] = LatentRepeatInterleavedNode.repeat(samples[i]["noise_mask"],repeats[i])
|
||||
|
||||
return samples
|
||||
|
||||
def merge(self, latents, repeats, fps=1):
|
||||
repeats = np.array(repeats).flatten()*fps
|
||||
print(f"{repeats=}")
|
||||
|
||||
# combine all latents
|
||||
new_latents = []
|
||||
for latent in latents:
|
||||
if latent["samples"].ndim==3: #unbatched
|
||||
new_latents.append(latent)
|
||||
else:
|
||||
for i in range(len(latent["samples"])):
|
||||
new_latent = {
|
||||
"samples": latent["samples"][i],
|
||||
}
|
||||
if "noise_mask" in latent: new_latent["noise_mask"] = latent["noise_mask"][i]
|
||||
new_latents.append(new_latent)
|
||||
|
||||
# expand repeat
|
||||
if len(repeats)==1:
|
||||
repeats = repeats.repeat(len(new_latents),axis=0)
|
||||
elif len(repeats)<len(new_latents):
|
||||
pad_width = len(new_latents) - len(repeats)
|
||||
repeats = np.pad(repeats, (0, pad_width), mode='constant', constant_values=1)
|
||||
|
||||
new_latents = self.rebatch(new_latents,repeats.astype(int))
|
||||
gc_collect()
|
||||
|
||||
return (new_latents,)
|
||||
|
||||
class MergeAudioNode:
|
||||
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"audio1": ("VHS_AUDIO",),
|
||||
"audio2": ("VHS_AUDIO",),
|
||||
"sr": (["None",32000,40000,44100,48000],{
|
||||
"default": "None"
|
||||
}),
|
||||
"merge_type": (MERGE_OPTIONS,{"default": "mean"}),
|
||||
"normalize": ("BOOLEAN",{"default": True})
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("VHS_AUDIO",)
|
||||
RETURN_NAMES = ("vhs_audio",)
|
||||
|
||||
FUNCTION = "merge"
|
||||
|
||||
CATEGORY = CATEGORY
|
||||
|
||||
def merge(self, audio1, audio2, sr, merge_type, normalize):
|
||||
if sr=="None": sr=None
|
||||
input_audio1 = bytes_to_audio(audio1())
|
||||
input_audio2 = bytes_to_audio(audio2())
|
||||
merged_audio = merge_audio(input_audio1,input_audio2,sr=sr,to_int16=True,norm=normalize,merge_type=merge_type)
|
||||
return (lambda: audio_to_bytes(*merged_audio),)
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, *args, **kwargs):
|
||||
print(f"{args=} {kwargs=}")
|
||||
return get_hash(*args, *kwargs.items())
|
||||
@@ -0,0 +1,86 @@
|
||||
import os
|
||||
|
||||
from ..lib.audio import audio_to_bytes, bytes_to_audio, save_input_audio
|
||||
import folder_paths
|
||||
from ..lib.utils import get_filenames, get_hash, get_optimal_torch_device
|
||||
from ..lib import BASE_CACHE_DIR, BASE_MODELS_DIR, karafan
|
||||
from ..uvr5_cli import Separator
|
||||
from .settings.downloader import KARAFAN_MODELS, MDX_MODELS, RVC_DOWNLOAD_LINK, VR_MODELS, download_file
|
||||
|
||||
temp_path = folder_paths.get_temp_directory()
|
||||
cache_dir = os.path.join(BASE_CACHE_DIR,"uvr")
|
||||
device = get_optimal_torch_device()
|
||||
is_half = True
|
||||
|
||||
class UVR5Node:
|
||||
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
|
||||
model_list = MDX_MODELS + VR_MODELS + KARAFAN_MODELS + get_filenames(root=BASE_MODELS_DIR,format_func=os.path.basename)
|
||||
|
||||
return {
|
||||
"required": {
|
||||
"audio": ("VHS_AUDIO",),
|
||||
"model": (model_list,{
|
||||
"default": "UVR/HP5-vocals+instrumentals.pth"
|
||||
}),
|
||||
"agg":("INT",{
|
||||
"default": 10,
|
||||
"min": 0, #Minimum value
|
||||
"max": 20, #Maximum value
|
||||
"step": 1, #Slider's step
|
||||
"display": "slider"
|
||||
}),
|
||||
"format":(["wav", "flac", "mp3"],{
|
||||
"default": "flac"
|
||||
}),
|
||||
},
|
||||
"optional": {
|
||||
"use_cache": ("BOOLEAN",{"default": True})
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("VHS_AUDIO","VHS_AUDIO","VHS_AUDIO")
|
||||
RETURN_NAMES = ("primary_stem","secondary_stem","audio_passthrough")
|
||||
|
||||
FUNCTION = "split"
|
||||
|
||||
CATEGORY = "🌺RVC-Studio/uvr"
|
||||
|
||||
def split(self, audio, model, **kwargs):
|
||||
filename = os.path.basename(model)
|
||||
subfolder = os.path.dirname(model)
|
||||
model_path = os.path.join(BASE_MODELS_DIR,subfolder,filename)
|
||||
if not os.path.isfile(model_path):
|
||||
download_link = f"{RVC_DOWNLOAD_LINK}{model}"
|
||||
params = model_path, download_link
|
||||
if download_file(params): print(f"successfully downloaded: {model_path}")
|
||||
|
||||
input_audio = bytes_to_audio(audio())
|
||||
audio_path = os.path.join(temp_path,"uvr",f"{get_hash(audio(),model, *kwargs.items())}.wav")
|
||||
if not os.path.isfile(audio_path):
|
||||
os.makedirs(os.path.dirname(audio_path),exist_ok=True)
|
||||
print(save_input_audio(audio_path,input_audio))
|
||||
|
||||
if "karafan" in model_path:
|
||||
vocals, instrumental, input_audio = karafan.inference.Process(audio_path,cache_dir=cache_dir,**kwargs)
|
||||
else:
|
||||
model = Separator(
|
||||
model_path=model_path,
|
||||
device=device,
|
||||
is_half="cuda" in str(device),
|
||||
cache_dir=cache_dir,
|
||||
**kwargs
|
||||
)
|
||||
vocals, instrumental, input_audio = model.run_inference(audio_path,format=kwargs.get("format"))
|
||||
|
||||
return (lambda:audio_to_bytes(*vocals), lambda:audio_to_bytes(*instrumental), audio)
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls, *args, **kwargs):
|
||||
print(f"{args=} {kwargs=}")
|
||||
return get_hash(*args, *kwargs.items())
|
||||
Reference in New Issue
Block a user