Files
AIFSH-CosyVoice-ComfyUI/nodes.py
T
2024-07-11 16:57:04 +08:00

140 lines
5.1 KiB
Python

import torch
import random
import librosa
import zipfile
import torchaudio
import numpy as np
import os,sys
now_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.append(now_dir)
pretrained_models = os.path.join(now_dir,"pretrained_models")
from modelscope import snapshot_download
from cosyvoice.cli.cosyvoice import CosyVoice
from cosyvoice.utils.file_utils import load_wav
sft_spk_list = ['中文女', '中文男', '日语男', '粤语女', '英文女', '英文男', '韩语女']
inference_mode_list = ['预训练音色', '3s极速复刻', '跨语种复刻', '自然语言控制']
def set_all_random_seed(seed):
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
max_val = 0.8
prompt_sr, target_sr = 16000, 22050
def postprocess(speech, top_db=60, hop_length=220, win_length=440):
speech, _ = librosa.effects.trim(
speech, top_db=top_db,
frame_length=win_length,
hop_length=hop_length
)
if speech.abs().max() > max_val:
speech = speech / speech.abs().max() * max_val
speech = torch.concat([speech, torch.zeros(1, int(target_sr * 0.2))], dim=1)
return speech
class TextNode:
@classmethod
def INPUT_TYPES(s):
return {"required": {"text": ("STRING", {"multiline": True, "dynamicPrompts": True})}}
RETURN_TYPES = ("TEXT",)
FUNCTION = "encode"
CATEGORY = "AIFSH_CosyVoice"
def encode(self,text):
return (text, )
class CosyVoiceNode:
def __init__(self):
self.model_dir = None
self.cosyvoice = None
@classmethod
def INPUT_TYPES(s):
return {
"required":{
"tts_text":("TEXT",),
"inference_mode":(inference_mode_list,{
"default": "预训练音色"
}),
"sft_dropdown":(sft_spk_list,{
"default":"中文女"
}),
"seed":("INT",{
"default": 42
})
},
"optional":{
"prompt_text":("TEXT",),
"prompt_wav": ("AUDIO",),
"instruct_text":("TEXT",),
}
}
RETURN_TYPES = ("AUDIO",)
#RETURN_NAMES = ("image_output_name",)
FUNCTION = "generate"
#OUTPUT_NODE = False
CATEGORY = "AIFSH_CosyVoice"
def generate(self,tts_text,inference_mode,sft_dropdown,seed,
prompt_text=None,prompt_wav=None,instruct_text=None):
if inference_mode == '自然语言控制':
model_dir = os.path.join(pretrained_models,"CosyVoice-300M-Instruct")
snapshot_download(model_id="iic/CosyVoice-300M-Instruct",local_dir=model_dir)
assert instruct_text is not None, "in 自然语言控制 mode, instruct_text can't be none"
if inference_mode in ["跨语种复刻",'3s极速复刻']:
model_dir = os.path.join(pretrained_models,"CosyVoice-300M")
snapshot_download(model_id="iic/CosyVoice-300M",local_dir=model_dir)
assert prompt_wav is not None, "in 跨语种复刻 or 3s极速复刻 mode, prompt_wav can't be none"
if inference_mode == "3s极速复刻":
assert len(prompt_text) > 0, "prompt文本为空,您是否忘记输入prompt文本?"
if inference_mode == "预训练音色":
model_dir = os.path.join(pretrained_models,"CosyVoice-300M-SFT")
snapshot_download(model_id="iic/CosyVoice-300M-Instruct",local_dir=model_dir)
if self.model_dir != model_dir:
self.model_dir = model_dir
self.cosyvoice = CosyVoice(model_dir)
if prompt_wav:
waveform = prompt_wav['waveform'].squeeze(0)
source_sr = prompt_wav['sample_rate']
speech = waveform.mean(dim=0,keepdim=True)
if source_sr != prompt_sr:
speech = torchaudio.transforms.Resample(orig_freq=source_sr, new_freq=prompt_sr)(speech)
if inference_mode == '预训练音色':
print('get sft inference request')
print(self.model_dir)
set_all_random_seed(seed)
output = self.cosyvoice.inference_sft(tts_text, sft_dropdown)
elif inference_mode == '3s极速复刻':
print('get zero_shot inference request')
print(self.model_dir)
prompt_speech_16k = postprocess(speech)
set_all_random_seed(seed)
output = self.cosyvoice.inference_zero_shot(tts_text, prompt_text, prompt_speech_16k)
elif inference_mode == '跨语种复刻':
print('get cross_lingual inference request')
print(self.model_dir)
prompt_speech_16k = postprocess(speech)
set_all_random_seed(seed)
output = self.cosyvoice.inference_cross_lingual(tts_text, prompt_speech_16k)
else:
print('get instruct inference request')
set_all_random_seed(seed)
print(self.model_dir)
output = self.cosyvoice.inference_instruct(tts_text, sft_dropdown, instruct_text)
audio = {"waveform": [output['tts_speech']],"sample_rate":target_sr}
return (audio,)