382 lines
16 KiB
Python
382 lines
16 KiB
Python
import os
|
|
import sys
|
|
import time
|
|
from subprocess import Popen
|
|
import folder_paths
|
|
import cuda_malloc
|
|
import audiotsm
|
|
import audiotsm.io.wav
|
|
from pydub import AudioSegment
|
|
from srt import parse as SrtPare
|
|
from huggingface_hub import hf_hub_download
|
|
|
|
input_path = folder_paths.get_input_directory()
|
|
output_path = folder_paths.get_output_directory()
|
|
device = "cuda" if cuda_malloc.cuda_malloc_supported() else "cpu"
|
|
fish_tmp_out = os.path.join(output_path, "fish_speech")
|
|
os.makedirs(fish_tmp_out, exist_ok=True)
|
|
parent_directory = os.path.dirname(os.path.abspath(__file__))
|
|
checkpoint_path = os.path.join(parent_directory,"checkpoints")
|
|
|
|
class FishSpeech_INFER_SRT:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required":{
|
|
"text":("SRT",),
|
|
"prompt_audio": ("AUDIO",),
|
|
"prompt_text":("SRT",),
|
|
"if_mutiple_speaker":("BOOLEAN",{
|
|
"default": False
|
|
}),
|
|
"text2semantic_type":(["medium","large"],{
|
|
"default": "medium"
|
|
}),
|
|
"hf_token":("STRING",{
|
|
"default": "your token to download weights"
|
|
}),
|
|
"num_samples":("INT", {
|
|
"default":1
|
|
}),
|
|
"max_new_tokens": ("INT", {
|
|
"default":0
|
|
}),
|
|
"top_p":("FLOAT",{
|
|
"default": 0.7
|
|
}),
|
|
"repetition_penalty":("FLOAT",{
|
|
"default": 1.5
|
|
}),
|
|
"temperature":("FLOAT",{
|
|
"default": 0.7
|
|
}),
|
|
"compile":("BOOLEAN",{
|
|
"default": False
|
|
}),
|
|
"seed":("INT",{
|
|
"default": 42
|
|
}),
|
|
"half":("BOOLEAN",{
|
|
"default": False
|
|
}),
|
|
"iterative_prompt":
|
|
("BOOLEAN",{
|
|
"default": True
|
|
}),
|
|
"max_length":("INT",{
|
|
"default": 2048
|
|
}),
|
|
"chunk_length":("INT",{
|
|
"default": 30
|
|
}),
|
|
}}
|
|
|
|
CATEGORY = "AIFSH_FishSpeech"
|
|
RETURN_TYPES = ('AUDIO',)
|
|
OUTPUT_NODE = False
|
|
|
|
FUNCTION = "get_tts_wav"
|
|
|
|
def get_tts_wav(self,text,prompt_audio,prompt_text,if_mutiple_speaker,
|
|
text2semantic_type,hf_token,
|
|
num_samples,max_new_tokens,top_p,repetition_penalty,
|
|
temperature,compile,seed,half,iterative_prompt,max_length,
|
|
chunk_length):
|
|
|
|
filename = f"text2semantic-sft-{text2semantic_type}-v1-4k.pth"
|
|
t2s_model_path = os.path.join(checkpoint_path, filename)
|
|
if not os.path.isfile(t2s_model_path):
|
|
hf_hub_download(repo_id="fishaudio/fish-speech-1",filename=filename,local_dir=checkpoint_path,token=hf_token)
|
|
|
|
vq_model_path = os.path.join(checkpoint_path, "vq-gan-group-fsq-2x1024.pth")
|
|
if not os.path.isfile(vq_model_path):
|
|
hf_hub_download(repo_id="fishaudio/fish-speech-1",filename="vq-gan-group-fsq-2x1024.pth",local_dir=checkpoint_path,token=hf_token)
|
|
|
|
tokenizer_path = os.path.join(checkpoint_path, "tokenizer.json")
|
|
|
|
if not os.path.isfile(tokenizer_path):
|
|
hf_hub_download(repo_id="fishaudio/fish-speech-1",filename="tokenizer.json",local_dir=checkpoint_path,token=hf_token)
|
|
hf_hub_download(repo_id="fishaudio/fish-speech-1",filename="tokenizer_config.json",local_dir=checkpoint_path,token=hf_token)
|
|
hf_hub_download(repo_id="fishaudio/fish-speech-1",filename="special_tokens_map.json",local_dir=checkpoint_path,token=hf_token)
|
|
|
|
with open(text, 'r', encoding="utf-8") as file:
|
|
text_file_content = file.read()
|
|
with open(prompt_text, 'r', encoding="utf-8") as file:
|
|
prompt_text_file_content = file.read()
|
|
|
|
audio_seg = AudioSegment.from_file(prompt_audio)
|
|
prompt_subtitles = list(SrtPare(prompt_text_file_content))
|
|
python_exec = sys.executable or "python"
|
|
spk_aduio_dict = {}
|
|
if if_mutiple_speaker:
|
|
for i,prompt_sub in enumerate(prompt_subtitles):
|
|
start_time = prompt_sub.start.total_seconds() * 1000
|
|
end_time = prompt_sub.end.total_seconds() * 1000
|
|
speaker = 'SPK'+prompt_sub.content[0]
|
|
if spk_aduio_dict[speaker] is not None:
|
|
spk_aduio_dict[speaker] += audio_seg[start_time:end_time]
|
|
else:
|
|
spk_aduio_dict[speaker] = audio_seg[start_time:end_time]
|
|
for speaker in spk_aduio_dict.keys():
|
|
speaker_audio_seg = spk_aduio_dict[speaker]
|
|
speaker_audio = os.path.join(input_path, f"{speaker}.wav")
|
|
speaker_audio_seg.export(speaker_audio,format='wav')
|
|
npy_path = os.path.join(fish_tmp_out, os.path.basename(speaker_audio))
|
|
step_1 = f"{python_exec} {parent_directory}/tools/vqgan/inference.py -i {prompt_audio} -o {npy_path} -ckpt {vq_model_path} -d {device}"
|
|
print("step 1 ",step_1)
|
|
p = Popen(step_1,shell=True)
|
|
p.wait()
|
|
else:
|
|
npy_path = os.path.join(fish_tmp_out, "SPK0.wav")
|
|
step_1 = f"{python_exec} {parent_directory}/tools/vqgan/inference.py -i {prompt_audio} -o {npy_path} -ckpt {vq_model_path} -d {device}"
|
|
print("step 1 ",step_1)
|
|
p = Popen(step_1,shell=True)
|
|
p.wait()
|
|
|
|
config_name = f"dual_ar_2_codebook_{text2semantic_type}"
|
|
new_audio_seg = AudioSegment.silent(0)
|
|
for i, (prompt_sub,text_sub) in enumerate(zip(prompt_subtitles, list(SrtPare(text_file_content)))):
|
|
start_time = prompt_sub.start.total_seconds() * 1000
|
|
end_time = prompt_sub.end.total_seconds() * 1000
|
|
if i == 0:
|
|
new_audio_seg += audio_seg[:start_time]
|
|
|
|
refer_wav_seg = audio_seg[start_time:end_time]
|
|
refer_wav = os.path.join(fish_tmp_out, f"{i}_fishspeech_refer.wav")
|
|
refer_wav_seg.export(refer_wav, format='wav')
|
|
|
|
new_text = text_sub.content
|
|
new_prompt_text = prompt_sub.content
|
|
if if_mutiple_speaker:
|
|
new_text = new_text[1:]
|
|
new_prompt_text = new_prompt_text[1:]
|
|
npy_path = os.path.join(fish_tmp_out, f"SPK{new_text[0]}.npy")
|
|
out_put_path = os.path.join(fish_tmp_out, f"SPK{new_text[0]}")
|
|
else:
|
|
npy_path = os.path.join(fish_tmp_out, "SPK0.npy")
|
|
out_put_path = os.path.join(fish_tmp_out, "SPK0")
|
|
os.makedirs(out_put_path, exist_ok=True)
|
|
step_2 = f'{python_exec} {parent_directory}/tools/llama/generate.py --text "{new_text}" --prompt-text "{new_prompt_text}" \
|
|
--prompt-tokens {npy_path} --config-name {config_name} --num-samples {num_samples} --max-new-tokens {max_new_tokens} \
|
|
--top-p {top_p} --repetition-penalty {repetition_penalty} --temperature {temperature} \
|
|
--checkpoint-path {t2s_model_path} --tokenizer {checkpoint_path} {"--compile" if compile else "--no-compile"}\
|
|
--seed {seed} {"--half" if half else "--no-half"} {"--iterative-prompt" if iterative_prompt else "--no-iterative-prompt"} \
|
|
--max-length {max_length} --chunk-length {chunk_length} --output-path {out_put_path}'
|
|
print("step 2 ",step_2)
|
|
p2 = Popen(step_2,shell=True)
|
|
p2.wait()
|
|
|
|
step_2_npy = os.path.join(out_put_path,"codes_0.npy")
|
|
out_wav_path = os.path.join(out_put_path,f"{i}_fish_speech.wav")
|
|
step_3 = f"{python_exec} {parent_directory}/tools/vqgan/inference.py -i {step_2_npy} -o {out_wav_path} -ckpt {vq_model_path} -d {device}"
|
|
print("step 3 ",step_3)
|
|
p3 = Popen(step_3,shell=True)
|
|
p3.wait()
|
|
|
|
text_audio = AudioSegment.from_file(out_wav_path)
|
|
text_audio_dur_time = text_audio.duration_seconds * 1000
|
|
|
|
if i < len(prompt_subtitles) - 1:
|
|
nxt_start = prompt_subtitles[i+1].start.total_seconds() * 1000
|
|
dur_time = nxt_start - start_time
|
|
else:
|
|
org_dur_time = audio_seg.duration_seconds * 1000
|
|
dur_time = org_dur_time - start_time
|
|
|
|
ratio = text_audio_dur_time / dur_time
|
|
|
|
if text_audio_dur_time > dur_time:
|
|
tmp_audio = self.map_vocal(audio=text_audio,ratio=ratio,dur_time=dur_time,
|
|
wav_name=f"map_{i}_refer.wav",temp_folder=out_put_path)
|
|
tmp_audio += AudioSegment.silent(dur_time - tmp_audio.duration_seconds*1000)
|
|
else:
|
|
tmp_audio = text_audio + AudioSegment.silent(dur_time - text_audio_dur_time)
|
|
|
|
new_audio_seg += tmp_audio
|
|
|
|
infer_audio = os.path.join(output_path, f"{time.time()}_fishspeech_refer.wav")
|
|
new_audio_seg.export(infer_audio, format="wav")
|
|
return (infer_audio, )
|
|
|
|
def map_vocal(self,audio:AudioSegment,ratio:float,dur_time:float,wav_name:str,temp_folder:str):
|
|
tmp_path = f"{temp_folder}/map_{wav_name}"
|
|
audio.export(tmp_path, format="wav")
|
|
|
|
clone_path = f"{temp_folder}/cloned_{wav_name}"
|
|
reader = audiotsm.io.wav.WavReader(tmp_path)
|
|
|
|
writer = audiotsm.io.wav.WavWriter(clone_path,channels=reader.channels,
|
|
samplerate=reader.samplerate)
|
|
wsloa = audiotsm.wsola(channels=reader.channels,speed=ratio)
|
|
wsloa.run(reader=reader,writer=writer)
|
|
audio_extended = AudioSegment.from_file(clone_path)
|
|
return audio_extended[:dur_time]
|
|
|
|
|
|
class FishSpeech_INFER:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required":{
|
|
"prompt_audio": ("AUDIO",),
|
|
"text":("STRING",{
|
|
"multiline": True,
|
|
"default": "你好啊，世界！"
|
|
}),
|
|
"prompt_text_by_srt":("SRT",{
|
|
"multiline": True,
|
|
"default": "a man voice"
|
|
}),
|
|
"text2semantic_type":(["medium","large"],{
|
|
"default": "medium"
|
|
}),
|
|
"hf_token":("STRING",{
|
|
"default": "your token"
|
|
}),
|
|
"num_samples":("INT", {
|
|
"default":1
|
|
}),
|
|
"max_new_tokens": ("INT", {
|
|
"default":0
|
|
}),
|
|
"top_p":("FLOAT",{
|
|
"default": 0.7
|
|
}),
|
|
"repetition_penalty":("FLOAT",{
|
|
"default": 1.5
|
|
}),
|
|
"temperature":("FLOAT",{
|
|
"default": 0.7
|
|
}),
|
|
"compile":("BOOLEAN",{
|
|
"default": False
|
|
}),
|
|
"seed":("INT",{
|
|
"default": 42
|
|
}),
|
|
"half":("BOOLEAN",{
|
|
"default": False
|
|
}),
|
|
"iterative_prompt":
|
|
("BOOLEAN",{
|
|
"default": True
|
|
}),
|
|
"max_length":("INT",{
|
|
"default": 2048
|
|
}),
|
|
"chunk_length":("INT",{
|
|
"default": 30
|
|
}),
|
|
}}
|
|
|
|
CATEGORY = "AIFSH_FishSpeech"
|
|
RETURN_TYPES = ('AUDIO',)
|
|
OUTPUT_NODE = False
|
|
|
|
FUNCTION = "get_tts_wav"
|
|
|
|
def get_tts_wav(self,prompt_audio,text,prompt_text_by_srt,text2semantic_type,hf_token,
|
|
num_samples,max_new_tokens,top_p,repetition_penalty,
|
|
temperature,compile,seed,half,iterative_prompt,max_length,
|
|
chunk_length):
|
|
with open(prompt_text_by_srt, 'r', encoding="utf-8") as file:
|
|
file_content = file.read()
|
|
prompt_text = ' '.join([sub.content for sub in list(SrtPare(file_content))])
|
|
filename = f"text2semantic-sft-{text2semantic_type}-v1-4k.pth"
|
|
t2s_model_path = os.path.join(checkpoint_path, filename)
|
|
if not os.path.isfile(t2s_model_path):
|
|
hf_hub_download(repo_id="fishaudio/fish-speech-1",filename=filename,local_dir=checkpoint_path,token=hf_token)
|
|
|
|
vq_model_path = os.path.join(checkpoint_path, "vq-gan-group-fsq-2x1024.pth")
|
|
if not os.path.isfile(vq_model_path):
|
|
hf_hub_download(repo_id="fishaudio/fish-speech-1",filename="vq-gan-group-fsq-2x1024.pth",local_dir=checkpoint_path,token=hf_token)
|
|
|
|
tokenizer_path = os.path.join(checkpoint_path, "tokenizer.json")
|
|
|
|
if not os.path.isfile(tokenizer_path):
|
|
hf_hub_download(repo_id="fishaudio/fish-speech-1",filename="tokenizer.json",local_dir=checkpoint_path,token=hf_token)
|
|
hf_hub_download(repo_id="fishaudio/fish-speech-1",filename="tokenizer_config.json",local_dir=checkpoint_path,token=hf_token)
|
|
hf_hub_download(repo_id="fishaudio/fish-speech-1",filename="special_tokens_map.json",local_dir=checkpoint_path,token=hf_token)
|
|
|
|
python_exec = sys.executable or "python"
|
|
|
|
npy_path = os.path.join(fish_tmp_out, os.path.basename(prompt_audio))
|
|
step_1 = f"{python_exec} {parent_directory}/tools/vqgan/inference.py -i {prompt_audio} -o {npy_path} -ckpt {vq_model_path} -d {device}"
|
|
print("step 1 ",step_1)
|
|
p = Popen(step_1,shell=True)
|
|
p.wait()
|
|
|
|
|
|
config_name = f"dual_ar_2_codebook_{text2semantic_type}"
|
|
npy_path = os.path.join(fish_tmp_out, os.path.basename(prompt_audio)[:-4]+".npy")
|
|
step_2 = f'{python_exec} {parent_directory}/tools/llama/generate.py --text "{text}" --prompt-text "{prompt_text}" \
|
|
--prompt-tokens {npy_path} --config-name {config_name} --num-samples {num_samples} --max-new-tokens {max_new_tokens} \
|
|
--top-p {top_p} --repetition-penalty {repetition_penalty} --temperature {temperature} \
|
|
--checkpoint-path {t2s_model_path} --tokenizer {checkpoint_path} {"--compile" if compile else "--no-compile"}\
|
|
--seed {seed} {"--half" if half else "--no-half"} {"--iterative-prompt" if iterative_prompt else "--no-iterative-prompt"} \
|
|
--max-length {max_length} --chunk-length {chunk_length} --output-path {fish_tmp_out}'
|
|
print("step 2 ",step_2)
|
|
p2 = Popen(step_2,shell=True)
|
|
p2.wait()
|
|
|
|
step_2_npy = os.path.join(fish_tmp_out,"codes_0.npy")
|
|
out_wav_path = os.path.join(output_path,f"{time.time()}_fish_speech.wav")
|
|
step_3 = f"{python_exec} {parent_directory}/tools/vqgan/inference.py -i {step_2_npy} -o {out_wav_path} -ckpt {vq_model_path} -d {device}"
|
|
print("step 3 ",step_3)
|
|
p3 = Popen(step_3,shell=True)
|
|
p3.wait()
|
|
return (out_wav_path, )
|
|
|
|
class PreViewAudio:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required":
|
|
{"audio": ("AUDIO",),}
|
|
}
|
|
|
|
CATEGORY = "AIFSH_FishSpeech"
|
|
DESCRIPTION = "hello world!"
|
|
|
|
RETURN_TYPES = ()
|
|
|
|
OUTPUT_NODE = True
|
|
|
|
FUNCTION = "load_audio"
|
|
|
|
def load_audio(self, audio):
|
|
audio_name = os.path.basename(audio)
|
|
tmp_path = os.path.dirname(audio)
|
|
audio_root = os.path.basename(tmp_path)
|
|
return {"ui": {"audio":[audio_name,audio_root]}}
|
|
|
|
class LoadAudio:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
files = [f for f in os.listdir(input_path) if os.path.isfile(os.path.join(input_path, f)) and f.split('.')[-1] in ["wav", "mp3","WAV","flac","m4a"]]
|
|
return {"required":
|
|
{"audio": (sorted(files),)},
|
|
}
|
|
|
|
CATEGORY = "AIFSH_FishSpeech"
|
|
|
|
RETURN_TYPES = ("AUDIO",)
|
|
FUNCTION = "load_audio"
|
|
|
|
def load_audio(self, audio):
|
|
audio_path = folder_paths.get_annotated_filepath(audio)
|
|
return (audio_path,)
|
|
|
|
class LoadSRT:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
files = [f for f in os.listdir(input_path) if os.path.isfile(os.path.join(input_path, f)) and f.split('.')[-1] in ["srt", "txt"]]
|
|
return {"required":
|
|
{"srt": (sorted(files),)},
|
|
}
|
|
|
|
CATEGORY = "AIFSH_FishSpeech"
|
|
|
|
RETURN_TYPES = ("SRT",)
|
|
FUNCTION = "load_srt"
|
|
|
|
def load_srt(self, srt):
|
|
srt_path = folder_paths.get_annotated_filepath(srt)
|
|
return (srt_path,) |