SayanoAI-Comfy-RVC/custom_nodes/loaders.py

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)