import folder_paths import os import comfy.model_management as mm import time import torchaudio import torchvision.utils as vutils import torch import json import uuid from comfy.comfy_types import FileLocator 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"]) # 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" AVSpeechTimestamp.infer_ins_cache = AutoModel( model=model_dir, model_revision="v2.0.4", device=device, # GPU加速 disable_update=True ) # save audio uuidv4 = str(uuid.uuid4()) audio_save_path = os.path.join(temp_dir, f"{uuidv4}.wav") waveform = audio['waveform'] sr = audio["sample_rate"] waveform = torchaudio.functional.resample(waveform, sr, 16000) torchaudio.save(audio_save_path, waveform.squeeze(0), 16000) rec_result = AVSpeechTimestamp.infer_ins_cache.generate( input=(audio_save_path, text), data_type=("sound", "text"), ) # print(rec_result) 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加速 disable_update=True ) # save uuidv4 = str(uuid.uuid4()) audio_save_path = os.path.join(temp_dir, f"{uuidv4}.wav") # 重新采样为16k waveform = audio['waveform'] sr = audio["sample_rate"] waveform = torchaudio.functional.resample(waveform, sr, 16000) torchaudio.save(audio_save_path, waveform.squeeze(0), 16000) rec_result = AVASRTimestamp.infer_ins_cache.generate(input=audio_save_path, batch_size_s=batch_size_s) # print(rec_result) if rec_result: rec_result = rec_result[0] # infer if unload_model: import gc if AVASRTimestamp.infer_ins_cache is not None: AVASRTimestamp.infer_ins_cache = None gc.collect() torch.cuda.empty_cache() print("AVASRTimestamp 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_prefix += self.prefix_append full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir) results: list[FileLocator] = [] file = f"{filename}_{counter:05}_.srt" with open(os.path.join(full_output_folder, file), 'w', encoding='utf-8') as f: f.write(subtitles) results.append({ "filename": file, "subfolder": subfolder, "type": self.type }) counter += 1 return { "ui": { "subtitles": results } } 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", }