init
This commit is contained in:
aven
@@ -0,0 +1,220 @@
|
||||
import folder_paths
|
||||
import os
|
||||
import comfy.model_management as mm
|
||||
import time
|
||||
import torchaudio
|
||||
import torchvision.utils as vutils
|
||||
import torch
|
||||
import json
|
||||
|
||||
from modelscope.pipelines import pipeline
|
||||
from modelscope.utils.constant import Tasks
|
||||
from funasr import AutoModel
|
||||
from .format import Format2Subtitle
|
||||
|
||||
|
||||
|
||||
name_maps_ms = {
|
||||
"paraformer": "iic/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-pytorch",
|
||||
"paraformer-zh": "iic/speech_seaco_paraformer_large_asr_nat-zh-cn-16k-common-vocab8404-pytorch",
|
||||
"paraformer-en": "iic/speech_paraformer-large-vad-punc_asr_nat-en-16k-common-vocab10020",
|
||||
"paraformer-en-spk": "iic/speech_paraformer-large-vad-punc_asr_nat-en-16k-common-vocab10020",
|
||||
"paraformer-zh-streaming": "iic/speech_paraformer-large_asr_nat-zh-cn-16k-common-vocab8404-online",
|
||||
"fsmn-vad": "iic/speech_fsmn_vad_zh-cn-16k-common-pytorch",
|
||||
"ct-punc": "iic/punc_ct-transformer_cn-en-common-vocab471067-large",
|
||||
"ct-punc-c": "iic/punc_ct-transformer_zh-cn-common-vocab272727-pytorch",
|
||||
"fa-zh": "iic/speech_timestamp_prediction-v1-16k-offline",
|
||||
"cam++": "iic/speech_campplus_sv_zh-cn_16k-common",
|
||||
"Whisper-large-v2": "iic/speech_whisper-large_asr_multilingual",
|
||||
"Whisper-large-v3": "iic/Whisper-large-v3",
|
||||
"Qwen-Audio": "Qwen/Qwen-Audio",
|
||||
"emotion2vec_plus_large": "iic/emotion2vec_plus_large",
|
||||
"emotion2vec_plus_base": "iic/emotion2vec_plus_base",
|
||||
"emotion2vec_plus_seed": "iic/emotion2vec_plus_seed",
|
||||
"Whisper-large-v3-turbo": "iic/Whisper-large-v3-turbo",
|
||||
}
|
||||
|
||||
class AVSpeechTimestamp:
|
||||
infer_ins_cache = None
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"audio": ("AUDIO",),
|
||||
"text": ("STRING",),
|
||||
"unload_model": ("BOOLEAN", {"default": False}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING", "STRING", "ASRRESULT")
|
||||
RETURN_NAMES = ("text", "json_result", "asr_result")
|
||||
FUNCTION = "infer"
|
||||
CATEGORY = "Aven/AV-FunASR"
|
||||
DESCRIPTION = "get speech timestamp"
|
||||
|
||||
def infer(self, audio, text, unload_model):
|
||||
temp_dir = folder_paths.get_temp_directory()
|
||||
os.makedirs(temp_dir, exist_ok=True)
|
||||
|
||||
if AVSpeechTimestamp.infer_ins_cache is None:
|
||||
model_root = os.path.join(folder_paths.models_dir, "ASR/FunASR")
|
||||
model_dir = os.path.join(model_root, name_maps_ms["fa-zh"])
|
||||
|
||||
os.makedirs(model_dir, exist_ok=True)
|
||||
AVSpeechTimestamp.infer_ins_cache = pipeline(
|
||||
task=Tasks.speech_timestamp,
|
||||
model=model_dir,
|
||||
model_revision="v2.0.4",
|
||||
output_dir=temp_dir
|
||||
)
|
||||
# save audio
|
||||
audio_save_path = os.path.join(temp_dir, f"{int(time.time())}.wav")
|
||||
torchaudio.save(audio_save_path, audio['waveform'].squeeze(0), audio["sample_rate"])
|
||||
|
||||
rec_result = AVSpeechTimestamp.infer_ins_cache(input=(audio_save_path,text), data_type=("sound", "text"))
|
||||
if rec_result:
|
||||
rec_result = rec_result[0]
|
||||
|
||||
# infer
|
||||
if unload_model:
|
||||
import gc
|
||||
if AVSpeechTimestamp.infer_ins_cache is not None:
|
||||
AVSpeechTimestamp.infer_ins_cache = None
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
print("AVSpeechTimestamp memory cleanup successful")
|
||||
# jr = json.dumps(rec_result, indent=4)
|
||||
text = rec_result.get("text")
|
||||
jr = json.dumps(rec_result, ensure_ascii=False)
|
||||
print(text, jr)
|
||||
return (text, jr, rec_result)
|
||||
|
||||
|
||||
class AVASRTimestamp:
|
||||
infer_ins_cache = None
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"audio": ("AUDIO",),
|
||||
"batch_size_s": ("INT", {"default": 300, "min": 30, "max": 300, "step": 1}),
|
||||
"unload_model": ("BOOLEAN", {"default": False}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING", "STRING", "ASRRESULT")
|
||||
RETURN_NAMES = ("text", "json_result", "asr_result")
|
||||
FUNCTION = "infer"
|
||||
CATEGORY = "Aven/AV-FunASR"
|
||||
DESCRIPTION = "get speech timestamp"
|
||||
|
||||
def infer(self, audio, batch_size_s, unload_model):
|
||||
temp_dir = folder_paths.get_temp_directory()
|
||||
os.makedirs(temp_dir, exist_ok=True)
|
||||
|
||||
if AVASRTimestamp.infer_ins_cache is None:
|
||||
model_root = os.path.join(folder_paths.models_dir, "ASR/FunASR")
|
||||
model_dir = os.path.join(model_root, name_maps_ms["paraformer-zh"])
|
||||
vad_model = os.path.join(model_root, name_maps_ms["fsmn-vad"])
|
||||
os.makedirs(model_dir, exist_ok=True)
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
AVASRTimestamp.infer_ins_cache = AutoModel(
|
||||
model=model_dir,
|
||||
vad_model=vad_model,
|
||||
punc_model=None, #"ct-punc",
|
||||
device=device # GPU加速
|
||||
)
|
||||
# save audio
|
||||
audio_save_path = os.path.join(temp_dir, f"{int(time.time())}.wav")
|
||||
torchaudio.save(audio_save_path, audio['waveform'].squeeze(0), audio["sample_rate"])
|
||||
|
||||
rec_result = AVASRTimestamp.infer_ins_cache.generate(input=audio_save_path, batch_size_s=batch_size_s)
|
||||
if rec_result:
|
||||
rec_result = rec_result[0]
|
||||
|
||||
# infer
|
||||
if unload_model:
|
||||
import gc
|
||||
if AVSpeechTimestamp.infer_ins_cache is not None:
|
||||
AVSpeechTimestamp.infer_ins_cache = None
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
print("AVSpeechTimestamp memory cleanup successful")
|
||||
# jr = json.dumps(rec_result, indent=4)
|
||||
text = rec_result.get("text")
|
||||
jr = json.dumps(rec_result, ensure_ascii=False)
|
||||
print((text, jr, rec_result))
|
||||
return (text, jr, rec_result)
|
||||
|
||||
|
||||
class AVFormat2Subtitle:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"asr_result": ("ASRRESULT",),
|
||||
"text": ("STRING", {"default": None}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING", )
|
||||
RETURN_NAMES = ("subtitle", )
|
||||
FUNCTION = "format_subtitle"
|
||||
CATEGORY = "Aven/AV-FunASR"
|
||||
DESCRIPTION = "format asr result to subtitle"
|
||||
|
||||
def format_subtitle(self, asr_result, text=None):
|
||||
f = Format2Subtitle(asr_result, ori_text=text)
|
||||
content = f.pipeline()
|
||||
|
||||
return (content, )
|
||||
|
||||
|
||||
class AVSaveSubtitles:
|
||||
def __init__(self):
|
||||
self.output_dir = folder_paths.get_output_directory()
|
||||
self.type = "output"
|
||||
self.prefix_append = ""
|
||||
self.compress_level = 4
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"subtitles": ("STRING", {"tooltip": "The subtitles to save."}),
|
||||
"filename_prefix": ("STRING", {"default": "subtitles", "tooltip": "The prefix for the file to save. "})
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ()
|
||||
FUNCTION = "save_subtitles"
|
||||
|
||||
OUTPUT_NODE = True
|
||||
|
||||
CATEGORY = "Aven/AV-FunASR"
|
||||
DESCRIPTION = "Saves the subtitles to a file."
|
||||
|
||||
def save_subtitles(self, subtitles, filename_prefix="subtitles"):
|
||||
filename = f"{filename_prefix}-{time.time()}.srt"
|
||||
save_path = os.path.join(self.output_dir, filename)
|
||||
with open(save_path, "w") as f:
|
||||
f.write(subtitles)
|
||||
|
||||
return { "ui": { "text": save_path } }
|
||||
|
||||
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"AVSpeechTimestamp": AVSpeechTimestamp,
|
||||
"AVASRTimestamp": AVASRTimestamp,
|
||||
"AVFormat2Subtitle": AVFormat2Subtitle,
|
||||
"AVSaveSubtitles": AVSaveSubtitles,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"AVSpeechTimestamp": "AV Speech Timestamp",
|
||||
"AVASRTimestamp": "AV ASR Timestamp",
|
||||
"AVFormat2Subtitle": "AV Format to Subtitle",
|
||||
"AVSaveSubtitles": "AV Save Subtitles",
|
||||
}
|
||||
Reference in New Issue
Block a user