From 5e6f270308c58e4ed2df92c1429dc5451c75e6a4 Mon Sep 17 00:00:00 2001 From: SayanoAI <28544690+SayanoAI@users.noreply.github.com> Date: Tue, 11 Jun 2024 22:53:19 -0400 Subject: [PATCH] added hubert safetensors --- custom_nodes/loaders.py | 10 +- examples/complete-ai-voice-changer.json | 890 +++++++++--------------- lib/audio.py | 7 +- lib/infer_pack/loaders.py | 41 ++ lib/model_utils.py | 26 +- vc_infer_pipeline.py | 11 +- 6 files changed, 401 insertions(+), 584 deletions(-) create mode 100644 lib/infer_pack/loaders.py diff --git a/custom_nodes/loaders.py b/custom_nodes/loaders.py index 3e414ef..087f9a1 100644 --- a/custom_nodes/loaders.py +++ b/custom_nodes/loaders.py @@ -88,10 +88,10 @@ class LoadPitchExtractionParams: 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) + model_list = ["content-vec-best.safetensors"] + get_filenames(root=BASE_MODELS_DIR,folder="*",exts=["pt","safetensors"],format_func=os.path.basename) return { 'required': { - 'model': (model_list,{"default": "hubert_base.pt"}), + 'model': (model_list,{"default": "content-vec-best.safetensors"}), }, } @@ -253,7 +253,7 @@ class LoadAudio: FUNCTION = "load_audio" def load_audio(self, audio, sr): - audio_path = folder_paths.get_annotated_filepath(audio) + audio_path = os.path.join(input_path,audio) #folder_paths.get_annotated_filepath(audio) widgetId = get_hash(audio_path) audio_name = os.path.basename(audio).split(".")[0] sr = None if sr=="None" else int(sr) @@ -262,15 +262,13 @@ class LoadAudio: @classmethod def IS_CHANGED(cls, audio): - audio_path = folder_paths.get_annotated_filepath(audio) + audio_path = os.path.join(input_path,audio) print(f"{audio_path=}") return get_file_hash(audio_path) class DownloadAudio: @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": { diff --git a/examples/complete-ai-voice-changer.json b/examples/complete-ai-voice-changer.json index e11185c..0a4007c 100644 --- a/examples/complete-ai-voice-changer.json +++ b/examples/complete-ai-voice-changer.json @@ -1,61 +1,7 @@ { - "last_node_id": 50, - "last_link_id": 159, + "last_node_id": 54, + "last_link_id": 168, "nodes": [ - { - "id": 31, - "type": "JoinStrings", - "pos": [ - 1394, - -786 - ], - "size": { - "0": 315, - "1": 106 - }, - "flags": { - "collapsed": true - }, - "order": 6, - "mode": 0, - "inputs": [ - { - "name": "string1", - "type": "STRING", - "link": 66, - "widget": { - "name": "string1" - } - }, - { - "name": "string2", - "type": "STRING", - "link": 138, - "widget": { - "name": "string2" - } - } - ], - "outputs": [ - { - "name": "STRING", - "type": "STRING", - "links": [ - 64 - ], - "shape": 3, - "slot_index": 0 - } - ], - "properties": { - "Node name for S&R": "JoinStrings" - }, - "widgets_values": [ - "", - "", - "\\" - ] - }, { "id": 37, "type": "RVCNode", @@ -120,31 +66,6 @@ "paused": false, "params": {} }, - { - "hidden": false, - "paused": false, - "params": {} - }, - { - "hidden": false, - "paused": false, - "params": {} - }, - { - "hidden": false, - "paused": false, - "params": {} - }, - { - "hidden": false, - "paused": false, - "params": {} - }, - { - "hidden": false, - "paused": false, - "params": {} - }, { "hidden": false, "paused": false, @@ -152,128 +73,6 @@ } ] }, - { - "id": 49, - "type": "DownloadAudio", - "pos": [ - -538, - -794 - ], - "size": [ - 315, - 126 - ], - "flags": {}, - "order": 0, - "mode": 0, - "outputs": [ - { - "name": "audio_name", - "type": "STRING", - "links": [ - 138 - ], - "shape": 3, - "slot_index": 0 - }, - { - "name": "vhs_audio", - "type": "VHS_AUDIO", - "links": [ - 136, - 137 - ], - "shape": 3, - "slot_index": 1 - } - ], - "properties": { - "Node name for S&R": "DownloadAudio" - }, - "widgets_values": [ - "https://www.youtube.com/watch?v=DSnJ76vPi_8", - "None", - "Monologue of Three", - { - "hidden": false, - "paused": false, - "params": {} - }, - { - "hidden": false, - "paused": false, - "params": {} - }, - { - "hidden": false, - "paused": false, - "params": {} - }, - { - "hidden": false, - "paused": false, - "params": {} - }, - { - "hidden": false, - "paused": false, - "params": {} - }, - { - "hidden": false, - "paused": false, - "params": {} - }, - { - "hidden": false, - "paused": false, - "params": {} - }, - { - "hidden": false, - "paused": false, - "params": {} - }, - { - "hidden": false, - "paused": false, - "params": {} - } - ] - }, - { - "id": 5, - "type": "LoadHubertModel", - "pos": [ - 1010, - -145 - ], - "size": { - "0": 315, - "1": 58 - }, - "flags": { - "collapsed": true - }, - "order": 1, - "mode": 0, - "outputs": [ - { - "name": "hubert_model", - "type": "HUBERT_MODEL", - "links": [ - 58 - ], - "shape": 3 - } - ], - "properties": { - "Node name for S&R": "LoadHubertModel" - }, - "widgets_values": [ - "hubert_base.pt" - ] - }, { "id": 6, "type": "LoadPitchExtractionParams", @@ -286,7 +85,7 @@ "1": 178 }, "flags": {}, - "order": 2, + "order": 0, "mode": 0, "outputs": [ { @@ -310,46 +109,6 @@ 0.25 ] }, - { - "id": 4, - "type": "LoadRVCModelNode", - "pos": [ - 953, - -97 - ], - "size": { - "0": 315, - "1": 78 - }, - "flags": {}, - "order": 3, - "mode": 0, - "outputs": [ - { - "name": "model", - "type": "RVC_MODEL", - "links": [ - 57 - ], - "shape": 3 - }, - { - "name": "model_name", - "type": "STRING", - "links": [ - 66 - ], - "shape": 3, - "slot_index": 1 - } - ], - "properties": { - "Node name for S&R": "LoadRVCModelNode" - }, - "widgets_values": [ - "RVC/Sayano.pth" - ] - }, { "id": 33, "type": "MergeAudioNode", @@ -409,31 +168,6 @@ "paused": false, "params": {} }, - { - "hidden": false, - "paused": false, - "params": {} - }, - { - "hidden": false, - "paused": false, - "params": {} - }, - { - "hidden": false, - "paused": false, - "params": {} - }, - { - "hidden": false, - "paused": false, - "params": {} - }, - { - "hidden": false, - "paused": false, - "params": {} - }, { "hidden": false, "paused": false, @@ -499,31 +233,6 @@ "paused": false, "params": {} }, - { - "hidden": false, - "paused": false, - "params": {} - }, - { - "hidden": false, - "paused": false, - "params": {} - }, - { - "hidden": false, - "paused": false, - "params": {} - }, - { - "hidden": false, - "paused": false, - "params": {} - }, - { - "hidden": false, - "paused": false, - "params": {} - }, { "hidden": false, "paused": false, @@ -543,13 +252,13 @@ "1": 170 }, "flags": {}, - "order": 5, + "order": 6, "mode": 0, "inputs": [ { "name": "audio", "type": "VHS_AUDIO", - "link": 137 + "link": 163 } ], "outputs": [ @@ -603,7 +312,7 @@ { "name": "audio1", "type": "VHS_AUDIO", - "link": 158 + "link": 167 }, { "name": "audio2", @@ -645,31 +354,6 @@ "paused": false, "params": {} }, - { - "hidden": false, - "paused": false, - "params": {} - }, - { - "hidden": false, - "paused": false, - "params": {} - }, - { - "hidden": false, - "paused": false, - "params": {} - }, - { - "hidden": false, - "paused": false, - "params": {} - }, - { - "hidden": false, - "paused": false, - "params": {} - }, { "hidden": false, "paused": false, @@ -677,63 +361,6 @@ } ] }, - { - "id": 14, - "type": "UVR5Node", - "pos": [ - 179, - -579 - ], - "size": { - "0": 315, - "1": 170 - }, - "flags": {}, - "order": 9, - "mode": 0, - "inputs": [ - { - "name": "audio", - "type": "VHS_AUDIO", - "link": 99 - } - ], - "outputs": [ - { - "name": "primary_stem", - "type": "VHS_AUDIO", - "links": [ - 156 - ], - "shape": 3, - "slot_index": 0 - }, - { - "name": "secondary_stem", - "type": "VHS_AUDIO", - "links": [ - 157 - ], - "shape": 3, - "slot_index": 1 - }, - { - "name": "audio_passthrough", - "type": "VHS_AUDIO", - "links": null, - "shape": 3 - } - ], - "properties": { - "Node name for S&R": "UVR5Node" - }, - "widgets_values": [ - "karafan/MDX23C-8KFFT-InstVoc_HQ.ckpt", - 10, - 10, - "flac" - ] - }, { "id": 47, "type": "UVR5Node", @@ -791,63 +418,6 @@ "flac" ] }, - { - "id": 7, - "type": "UVR5Node", - "pos": [ - 173, - -800 - ], - "size": { - "0": 315, - "1": 170 - }, - "flags": {}, - "order": 8, - "mode": 0, - "inputs": [ - { - "name": "audio", - "type": "VHS_AUDIO", - "link": 98 - } - ], - "outputs": [ - { - "name": "primary_stem", - "type": "VHS_AUDIO", - "links": [ - 158 - ], - "shape": 3, - "slot_index": 0 - }, - { - "name": "secondary_stem", - "type": "VHS_AUDIO", - "links": [ - 159 - ], - "shape": 3, - "slot_index": 1 - }, - { - "name": "audio_passthrough", - "type": "VHS_AUDIO", - "links": null, - "shape": 3 - } - ], - "properties": { - "Node name for S&R": "UVR5Node" - }, - "widgets_values": [ - "UVR/model_bs_roformer_ep_317_sdr_12.9755.ckpt", - 10, - 10, - "flac" - ] - }, { "id": 43, "type": "MergeAudioNode", @@ -914,36 +484,6 @@ "paused": false, "params": {} }, - { - "hidden": false, - "paused": false, - "params": {} - }, - { - "hidden": false, - "paused": false, - "params": {} - }, - { - "hidden": false, - "paused": false, - "params": {} - }, - { - "hidden": false, - "paused": false, - "params": {} - }, - { - "hidden": false, - "paused": false, - "params": {} - }, - { - "hidden": false, - "paused": false, - "params": {} - }, { "hidden": false, "paused": false, @@ -963,13 +503,13 @@ "1": 170 }, "flags": {}, - "order": 4, + "order": 5, "mode": 0, "inputs": [ { "name": "audio", "type": "VHS_AUDIO", - "link": 136 + "link": 162 } ], "outputs": [ @@ -1024,7 +564,7 @@ { "name": "audio1", "type": "VHS_AUDIO", - "link": 159 + "link": 166 }, { "name": "audio2", @@ -1071,36 +611,6 @@ "paused": false, "params": {} }, - { - "hidden": false, - "paused": false, - "params": {} - }, - { - "hidden": false, - "paused": false, - "params": {} - }, - { - "hidden": false, - "paused": false, - "params": {} - }, - { - "hidden": false, - "paused": false, - "params": {} - }, - { - "hidden": false, - "paused": false, - "params": {} - }, - { - "hidden": false, - "paused": false, - "params": {} - }, { "hidden": false, "paused": false, @@ -1173,31 +683,6 @@ "paused": false, "params": {} }, - { - "hidden": false, - "paused": false, - "params": {} - }, - { - "hidden": false, - "paused": false, - "params": {} - }, - { - "hidden": false, - "paused": false, - "params": {} - }, - { - "hidden": false, - "paused": false, - "params": {} - }, - { - "hidden": false, - "paused": false, - "params": {} - }, { "hidden": false, "paused": false, @@ -1370,6 +855,297 @@ 10, "flac" ] + }, + { + "id": 4, + "type": "LoadRVCModelNode", + "pos": [ + 956, + -118 + ], + "size": { + "0": 315, + "1": 102 + }, + "flags": {}, + "order": 1, + "mode": 0, + "outputs": [ + { + "name": "model", + "type": "RVC_MODEL", + "links": [ + 57 + ], + "shape": 3 + }, + { + "name": "model_name", + "type": "STRING", + "links": [ + 66 + ], + "shape": 3, + "slot_index": 1 + } + ], + "properties": { + "Node name for S&R": "LoadRVCModelNode" + }, + "widgets_values": [ + "RVC/Sayano.pth", + "RVC/.index/added_IVF1063_Flat_nprobe_1_Sayano_v2.index" + ] + }, + { + "id": 5, + "type": "LoadHubertModel", + "pos": [ + 1011, + -156 + ], + "size": { + "0": 315, + "1": 58 + }, + "flags": { + "collapsed": true + }, + "order": 2, + "mode": 0, + "outputs": [ + { + "name": "hubert_model", + "type": "HUBERT_MODEL", + "links": [ + 58 + ], + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "LoadHubertModel" + }, + "widgets_values": [ + "content-vec-best.safetensors" + ] + }, + { + "id": 14, + "type": "UVR5Node", + "pos": [ + 179, + -579 + ], + "size": { + "0": 315, + "1": 170 + }, + "flags": {}, + "order": 9, + "mode": 0, + "inputs": [ + { + "name": "audio", + "type": "VHS_AUDIO", + "link": 99 + } + ], + "outputs": [ + { + "name": "primary_stem", + "type": "VHS_AUDIO", + "links": [ + 156 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "secondary_stem", + "type": "VHS_AUDIO", + "links": [ + 157 + ], + "shape": 3, + "slot_index": 1 + }, + { + "name": "audio_passthrough", + "type": "VHS_AUDIO", + "links": null, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "UVR5Node" + }, + "widgets_values": [ + "UVR/model_bs_roformer_ep_317_sdr_12.9755.ckpt", + 10, + 10, + "flac" + ] + }, + { + "id": 7, + "type": "UVR5Node", + "pos": [ + 173, + -800 + ], + "size": { + "0": 315, + "1": 170 + }, + "flags": {}, + "order": 8, + "mode": 0, + "inputs": [ + { + "name": "audio", + "type": "VHS_AUDIO", + "link": 98 + } + ], + "outputs": [ + { + "name": "primary_stem", + "type": "VHS_AUDIO", + "links": [ + 166 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "secondary_stem", + "type": "VHS_AUDIO", + "links": [ + 167 + ], + "shape": 3, + "slot_index": 1 + }, + { + "name": "audio_passthrough", + "type": "VHS_AUDIO", + "links": null, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "UVR5Node" + }, + "widgets_values": [ + "karafan/MDX23C-8KFFT-InstVoc_HQ.ckpt", + 10, + 10, + "flac" + ] + }, + { + "id": 53, + "type": "LoadAudio", + "pos": [ + -501, + -799 + ], + "size": [ + 210, + 126 + ], + "flags": {}, + "order": 3, + "mode": 0, + "outputs": [ + { + "name": "audio_name", + "type": "STRING", + "links": [ + 168 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "vhs_audio", + "type": "VHS_AUDIO", + "links": [ + 162, + 163 + ], + "shape": 3, + "slot_index": 1 + } + ], + "properties": { + "Node name for S&R": "LoadAudio" + }, + "widgets_values": [ + "Monologue of Three.mp3", + "None", + "image", + { + "hidden": false, + "paused": false, + "params": {} + } + ] + }, + { + "id": 31, + "type": "JoinStrings", + "pos": [ + 1379, + -786 + ], + "size": { + "0": 315, + "1": 106 + }, + "flags": { + "collapsed": true + }, + "order": 4, + "mode": 0, + "inputs": [ + { + "name": "string1", + "type": "STRING", + "link": 66, + "widget": { + "name": "string1" + } + }, + { + "name": "string2", + "type": "STRING", + "link": 168, + "widget": { + "name": "string2" + } + } + ], + "outputs": [ + { + "name": "STRING", + "type": "STRING", + "links": [ + 64 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "JoinStrings" + }, + "widgets_values": [ + "", + "", + "\\" + ] } ], "links": [ @@ -1485,30 +1261,6 @@ 0, "VHS_AUDIO" ], - [ - 136, - 49, - 1, - 11, - 0, - "VHS_AUDIO" - ], - [ - 137, - 49, - 1, - 44, - 0, - "VHS_AUDIO" - ], - [ - 138, - 49, - 0, - 31, - 1, - "STRING" - ], [ 141, 47, @@ -1582,20 +1334,44 @@ "VHS_AUDIO" ], [ - 158, + 162, + 53, + 1, + 11, + 0, + "VHS_AUDIO" + ], + [ + 163, + 53, + 1, + 44, + 0, + "VHS_AUDIO" + ], + [ + 166, 7, 0, + 50, + 0, + "VHS_AUDIO" + ], + [ + 167, + 7, + 1, 36, 0, "VHS_AUDIO" ], [ - 159, - 7, - 1, - 50, + 168, + 53, 0, - "VHS_AUDIO" + 31, + 1, + "STRING" ] ], "groups": [ @@ -1647,10 +1423,10 @@ "config": {}, "extra": { "ds": { - "scale": 0.6830134553650706, + "scale": 0.9090909090909091, "offset": [ - -663.7036134395787, - 1119.9858379127336 + -959.9219208282357, + 981.1146284125558 ] } }, diff --git a/lib/audio.py b/lib/audio.py index c84aad8..ddbaf1e 100644 --- a/lib/audio.py +++ b/lib/audio.py @@ -1,9 +1,8 @@ import base64 import io import os -from typing import Union import zlib -from .utils import get_hash, get_merge_func +from .utils import get_merge_func import numpy as np import librosa import soundfile as sf @@ -40,8 +39,8 @@ def load_audio(file, sr, **kwargs): .output("-", format="f32le", acodec="pcm_f32le", ac=1, ar=sr) .run(cmd=["ffmpeg", "-nostdin"], capture_stdout=True, capture_stderr=True) ) - except Exception as e: - raise RuntimeError(f"Failed to load audio: {e}") + except Exception as error: + raise RuntimeError(f"Failed to load audio: {error=} {file=}") return remix_audio((np.frombuffer(out, np.float32).flatten(), sr),**kwargs) diff --git a/lib/infer_pack/loaders.py b/lib/infer_pack/loaders.py new file mode 100644 index 0000000..f282a83 --- /dev/null +++ b/lib/infer_pack/loaders.py @@ -0,0 +1,41 @@ +import torch +from transformers import HubertModel, HubertConfig +import safetensors.torch as st +from safetensors import safe_open +import json + +class HubertModelWithFinalProj(HubertModel): + def __init__(self, config): + super().__init__(config) + + # The final projection layer is only used for backward compatibility. + # Following https://github.com/auspicious3000/contentvec/issues/6 + # Remove this layer is necessary to achieve the desired outcome. + self.final_proj = torch.nn.Linear(config.hidden_size, config.classifier_proj_size) + + @staticmethod + def load_safetensors(path: str, device="cpu"): + assert path.endswith(".safetensors"), f"{path} must end with '.safetensors'" + + with safe_open(path, framework="pt", device="cpu") as f: + metadata = f.metadata() + state_dict = {} + for key in f.keys(): + state_dict[key] = f.get_tensor(key) + model = HubertModelWithFinalProj(HubertConfig.from_dict(json.loads(metadata["config"]))) + model.load_state_dict(state_dict=state_dict) + return model.to(device) + + def save_safetensors(self, path: str): + assert path.endswith(".safetensors"), f"{path} must end with '.safetensors'" + + with open(path,"wb") as f: + state_dict = self.state_dict() + f.write(st.save(state_dict,dict(config=json.dumps(self.config.to_dict())))) + + def extract_features(self, source: torch.Tensor, version="v2", **kwargs): + with torch.no_grad(): + output_layer = 9 if version == "v1" else 12 + output = self(source.to(self.config.torch_dtype), output_hidden_states=True)["hidden_states"][output_layer] + features = self.final_proj(output) if version == "v1" else output + return features.to(source.dtype) \ No newline at end of file diff --git a/lib/model_utils.py b/lib/model_utils.py index 3a059f0..babb14f 100644 --- a/lib/model_utils.py +++ b/lib/model_utils.py @@ -1,10 +1,11 @@ import hashlib -import os import torch.nn.functional as F import librosa import torch +from .infer_pack.loaders import HubertModelWithFinalProj + def get_hash(model_path): try: with open(model_path, 'rb') as f: @@ -15,18 +16,21 @@ def get_hash(model_path): return model_hash -def load_hubert(model_path,config): +def load_hubert(model_path: str, config): try: - from fairseq import checkpoint_utils - models, _, _ = checkpoint_utils.load_model_ensemble_and_task([model_path],suffix="",) - hubert_model = models[0] - hubert_model = hubert_model.to(config.device) - if config.is_half: - hubert_model = hubert_model.half() + if model_path.endswith(".safetensors"): + return HubertModelWithFinalProj.load_safetensors(model_path, device=config.device) else: - hubert_model = hubert_model.float() - hubert_model.eval() - return hubert_model + from fairseq import checkpoint_utils + models, _, _ = checkpoint_utils.load_model_ensemble_and_task([model_path],suffix="",) + hubert_model = models[0] + hubert_model = hubert_model.to(config.device) + if config.is_half: + hubert_model = hubert_model.half() + else: + hubert_model = hubert_model.float() + hubert_model.eval() + return hubert_model except Exception as e: print(e) return None diff --git a/vc_infer_pipeline.py b/vc_infer_pipeline.py index 7284870..4d8dcd3 100644 --- a/vc_infer_pipeline.py +++ b/vc_infer_pipeline.py @@ -11,9 +11,8 @@ from .lib.model_utils import load_hubert, change_rms from .pitch_extraction import FeatureExtractor from .lib.audio import MAX_INT16, load_input_audio, remix_audio -from .lib import BASE_MODELS_DIR from .config import config -from .lib.utils import gc_collect, get_filenames +from .lib.utils import gc_collect # torchcrepe = lazyload("torchcrepe") # Fork Feature. Crepe algo for training and preprocess # torch = lazyload("torch") @@ -54,10 +53,10 @@ class VC(FeatureExtractor): "padding_mask": padding_mask, "output_layer": 9 if version == "v1" else 12, } - - with torch.no_grad(): - logits = model.extract_features(**inputs) - feats = model.final_proj(logits[0]) if version == "v1" else logits[0] + feats = model.extract_features(version=version,**inputs) + # with torch.no_grad(): + # logits = model.extract_features(**inputs) + # feats = model.final_proj(logits[0]) if version == "v1" else logits[0] if protect < 0.5 and pitch is not None and pitchf is not None: feats0 = feats.clone() if index is not None and big_npy is not None and index_rate > 0: