diff --git a/__init__.py b/__init__.py index 81425e0..1d8bebe 100644 --- a/__init__.py +++ b/__init__.py @@ -1,8 +1,7 @@ from .custom_nodes.audio_nodes import LoadAudio, PreviewAudio, MergeAudioNode, DownloadAudio, AudioBatchValueNode -from .custom_nodes.stt import AudioTranscriptionNode, BatchedTranscriptionEncoderNode +from .custom_nodes.stt import AudioTranscriptionNode, BatchedTranscriptionEncoderNode, LoadWhisperModelNode from .custom_nodes.uvr import UVR5Node -from .custom_nodes.rvc import RVCNode -from .custom_nodes.loaders import LoadWhisperModelNode, LoadRVCModelNode, LoadHubertModel, LoadPitchExtractionParams +from .custom_nodes.rvc import LoadHubertModel, LoadPitchExtractionParams, LoadRVCModelNode, RVCNode from .custom_nodes.utils import Any2ListNode, List2AnyNode, MergeImageBatches, MergeLatentBatches, ImageRepeatInterleavedNode, LatentRepeatInterleavedNode, SimpleMathNode, SliceNode, SortImagesNode, ZipImagesNode # Set the web directory, any .js file in that directory will be loaded by the frontend as a frontend extension diff --git a/custom_nodes/loaders.py b/custom_nodes/loaders.py deleted file mode 100644 index b97ac3f..0000000 --- a/custom_nodes/loaders.py +++ /dev/null @@ -1,212 +0,0 @@ -import os -from .settings import PITCH_EXTRACTION_OPTIONS -from ..lib import BASE_MODELS_DIR -from ..lib.model_utils import load_hubert -from ..lib.utils import 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_INDEX, RVC_MODELS, download_file -from ..config import config - -input_path = folder_paths.get_input_directory() -temp_path = folder_paths.get_temp_directory() -CATEGORY = "🌺RVC-Studio/loaders" - -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'] - -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 = CATEGORY - - FUNCTION = 'load_params' - - def load_params(self, **params): - if "rmvpe" in params.get("f0_method",""): model_downloader("rmvpe.pt") - return (params,) - - @classmethod - def IS_CHANGED(cls, **params): - return get_hash(**params) - -class LoadHubertModel: - @classmethod - def INPUT_TYPES(cls): - model_list = ["content-vec-best.safetensors"] + get_filenames(root=BASE_MODELS_DIR,folder=".",exts=["pt","safetensors"],format_func=os.path.basename) - model_list = list(set(model_list)) # dedupe - - return { - 'required': { - 'model': (model_list,{"default": "content-vec-best.safetensors"}), - }, - } - - RETURN_TYPES = ('HUBERT_MODEL', ) - RETURN_NAMES = ('hubert_model', ) - - CATEGORY = CATEGORY - - 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 - index_list = ["None"] + RVC_INDEX + get_filenames(root=BASE_MODELS_DIR,folder="RVC",exts=["index"],format_func=lambda x: f"RVC/.index/{os.path.basename(x)}") - index_list = list(set(index_list)) # dedupe - - return { - 'required': { - 'model': (model_list,{"default": model_list[0]}), - }, - "optional": { - "index": (index_list,{"default": "None"}), - } - } - - RETURN_TYPES = ('RVC_MODEL', 'STRING') - RETURN_NAMES = ('model', 'model_name') - - CATEGORY = CATEGORY - - FUNCTION = 'load_model' - - - def load_model(self, model, index="None"): - model_path = file_index = None - 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}" - if download_file((model_path, download_link)): print(f"successfully downloaded: {model_path}") - - if not index=="None": - file_index = os.path.join(BASE_MODELS_DIR,subfolder,".index",os.path.basename(index)) - if not os.path.isfile(file_index): - download_link = f"{RVC_DOWNLOAD_LINK}{index}" - if download_file((file_index, download_link)): print(f"successfully downloaded: {file_index}") - except Exception as e: - print(f"Error in {self.__class__.__name__}: {e}") - raise e - finally: return (lambda:get_vc(model_path, file_index),filename.split(".")[0]) - - @classmethod - def IS_CHANGED(cls, model, index): - return get_hash(model, index) - -class LoadWhisperModelNode: - - @classmethod - def INPUT_TYPES(cls): - return { - 'required': { - 'model_id': (model_ids,{"default": "openai/whisper-base.en"}), - }, - "optional": { - "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 = CATEGORY - - FUNCTION = 'load_model' - - - def load_model(self, model_id, 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 = {} - - def pipe(): 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, - ) - return ([pipe,model_id], ) - - @classmethod - def IS_CHANGED(cls, model_id, max_new_tokens, chunk_length_s, batch_size): - return get_hash(model_id, max_new_tokens, chunk_length_s, batch_size) - diff --git a/custom_nodes/rvc.py b/custom_nodes/rvc.py index be6f4ef..bbd2a8d 100644 --- a/custom_nodes/rvc.py +++ b/custom_nodes/rvc.py @@ -1,11 +1,18 @@ import os import shutil +from ..config import config + +from ..lib.model_utils import load_hubert + +from .utils import model_downloader +from .settings import PITCH_EXTRACTION_OPTIONS +from .settings.downloader import RVC_DOWNLOAD_LINK, RVC_INDEX, RVC_MODELS, download_file from ..lib.audio import SUPPORTED_AUDIO, audio_to_bytes, bytes_to_audio, load_input_audio, save_input_audio -from ..vc_infer_pipeline import vc_single +from ..vc_infer_pipeline import get_vc, vc_single import folder_paths -from ..lib.utils import get_hash, get_optimal_torch_device +from ..lib.utils import get_filenames, get_hash, get_optimal_torch_device from ..lib import BASE_CACHE_DIR, BASE_MODELS_DIR input_path = folder_paths.get_input_directory() @@ -16,8 +23,134 @@ cache_dir = os.path.join(BASE_CACHE_DIR,"rvc") 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 +CATEGORY = "🌺RVC-Studio/rvc" +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 = CATEGORY + + FUNCTION = 'load_params' + + def load_params(self, **params): + if "rmvpe" in params.get("f0_method",""): model_downloader("rmvpe.pt") + return (params,) + + @classmethod + def IS_CHANGED(cls, **params): + return get_hash(**params) + +class LoadHubertModel: + @classmethod + def INPUT_TYPES(cls): + model_list = ["content-vec-best.safetensors"] + get_filenames(root=BASE_MODELS_DIR,folder=".",exts=["pt","safetensors"],format_func=os.path.basename) + model_list = list(set(model_list)) # dedupe + + return { + 'required': { + 'model': (model_list,{"default": "content-vec-best.safetensors"}), + }, + } + + RETURN_TYPES = ('HUBERT_MODEL', ) + RETURN_NAMES = ('hubert_model', ) + + CATEGORY = CATEGORY + + 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 + index_list = ["None"] + RVC_INDEX + get_filenames(root=BASE_MODELS_DIR,folder="RVC",exts=["index"],format_func=lambda x: f"RVC/.index/{os.path.basename(x)}") + index_list = list(set(index_list)) # dedupe + + return { + 'required': { + 'model': (model_list,{"default": model_list[0]}), + }, + "optional": { + "index": (index_list,{"default": "None"}), + } + } + + RETURN_TYPES = ('RVC_MODEL', 'STRING') + RETURN_NAMES = ('model', 'model_name') + + CATEGORY = CATEGORY + + FUNCTION = 'load_model' + + + def load_model(self, model, index="None"): + model_path = file_index = None + 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}" + if download_file((model_path, download_link)): print(f"successfully downloaded: {model_path}") + + if not index=="None": + file_index = os.path.join(BASE_MODELS_DIR,subfolder,".index",os.path.basename(index)) + if not os.path.isfile(file_index): + download_link = f"{RVC_DOWNLOAD_LINK}{index}" + if download_file((file_index, download_link)): print(f"successfully downloaded: {file_index}") + except Exception as e: + print(f"Error in {self.__class__.__name__}: {e}") + raise e + finally: return (lambda:get_vc(model_path, file_index),filename.split(".")[0]) + + @classmethod + def IS_CHANGED(cls, model, index): + return get_hash(model, index) + class RVCNode: def __init__(self): @@ -53,7 +186,7 @@ class RVCNode: FUNCTION = "convert" - CATEGORY = "🌺RVC-Studio/rvc" + CATEGORY = CATEGORY def convert(self, audio, model, hubert_model, pitch_extraction_params, f0_up_key, format="flac", use_cache=True): diff --git a/custom_nodes/stt.py b/custom_nodes/stt.py index 2cd8c2b..f55e70f 100644 --- a/custom_nodes/stt.py +++ b/custom_nodes/stt.py @@ -3,7 +3,7 @@ import os from pprint import pprint import subprocess import sys - +from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline import numpy as np import torch @@ -11,7 +11,7 @@ from .settings import SUPPORTED_LANGUAGES from ..lib import BASE_CACHE_DIR, BASE_MODELS_DIR from .utils import increment_filename_no_overwrite -from ..lib.utils import get_hash +from ..lib.utils import get_hash, get_optimal_torch_device import folder_paths from ..lib.audio import bytes_to_audio, remix_audio import spacy @@ -102,6 +102,73 @@ def init_spacy_model(language="en",use_sentiment=False): return model + +class LoadWhisperModelNode: + + @classmethod + def INPUT_TYPES(cls): + 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', + ] + return { + 'required': { + 'model_id': (model_ids,{"default": "openai/whisper-base.en"}), + }, + "optional": { + "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 = CATEGORY + + FUNCTION = 'load_model' + + + def load_model(self, model_id, 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 = {} + + def pipe(): 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, + ) + return ([pipe,model_id], ) + + @classmethod + def IS_CHANGED(cls, model_id, max_new_tokens, chunk_length_s, batch_size): + return get_hash(model_id, max_new_tokens, chunk_length_s, batch_size) + class AudioTranscriptionNode: def __init__(self): pass