migrated rvc+stt nodes

2024-08-12 19:44:04 -04:00
parent be423fe38c
commit 0929b117e4
4 changed files with 208 additions and 221 deletions
@@ -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
@@ -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)
-
@@ -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):
        
@@ -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