444 lines
17 KiB
Python
444 lines
17 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.5",
|
|
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.5",
|
|
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.5",
|
|
filename="tokenizer.json",
|
|
local_dir=checkpoint_path,
|
|
token=hf_token,
|
|
)
|
|
hf_hub_download(
|
|
repo_id="fishaudio/fish-speech-1.5",
|
|
filename="tokenizer_config.json",
|
|
local_dir=checkpoint_path,
|
|
token=hf_token,
|
|
)
|
|
hf_hub_download(
|
|
repo_id="fishaudio/fish-speech-1.5",
|
|
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.5",
|
|
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.5",
|
|
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.5",
|
|
filename="tokenizer.json",
|
|
local_dir=checkpoint_path,
|
|
token=hf_token,
|
|
)
|
|
hf_hub_download(
|
|
repo_id="fishaudio/fish-speech-1.5",
|
|
filename="tokenizer_config.json",
|
|
local_dir=checkpoint_path,
|
|
token=hf_token,
|
|
)
|
|
hf_hub_download(
|
|
repo_id="fishaudio/fish-speech-1.5",
|
|
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,)
|