commit 6840820f8d34ffb273e8776b8ba73a8fb0e89d17 Author: aven Date: Sun Jun 8 16:07:40 2025 +0800 init diff --git a/.github/workflows/publish_action.yml b/.github/workflows/publish_action.yml new file mode 100644 index 0000000..be84400 --- /dev/null +++ b/.github/workflows/publish_action.yml @@ -0,0 +1,26 @@ +name: Publish to Comfy registry +on: + workflow_dispatch: + push: + branches: + - master + - main + paths: + - "pyproject.toml" + +permissions: + issues: write + +jobs: + publish-node: + name: Publish Custom Node to registry + runs-on: ubuntu-latest + if: ${{ github.repository_owner == 'avenstack' }} + steps: + - name: Check out code + uses: actions/checkout@v4 + - name: Publish Custom Node + uses: Comfy-Org/publish-node-action@v1 + with: + ## Add your own personal access token to your Github Repository secrets and reference it here. + personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }} diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..dc98536 --- /dev/null +++ b/.gitignore @@ -0,0 +1,3 @@ +*.pyc +__pycache__/ +.hisotry/ \ No newline at end of file diff --git a/README.md b/README.md new file mode 100644 index 0000000..afd2575 --- /dev/null +++ b/README.md @@ -0,0 +1,63 @@ +## ComfyUI FunAsr Nodes + +使用FunAsr进行语音识别，并转为文本或字幕文件。 + +![workflow](example_workflow/workflow.webp) + +## 功能 + +- [x] 语音识别 +- [x] 语音时间戳预测 +- [x] 语音转字幕 + +## 安装 + +``` +cd ComfyUI/custom_nodes +git clone https://github.com/avenstack/ComfyUI-AV-FunASR.git +cd ComfyUI-AV-FunASR +pip install -r requirements.txt +``` + +## 模型下载 + +1. [语音识别](https://modelscope.cn/models/iic/speech_seaco_paraformer_large_asr_nat-zh-cn-16k-common-vocab8404-pytorch/files) +2. [语音端点检测](modelscope.cn/models/iic/speech_fsmn_vad_zh-cn-16k-common-pytorch/files) +3. [语音时间戳预测](modelscope.cn/models/iic/speech_timestamp_prediction-v1-16k-offline) + +模型存放目录：`models/ASR/FunASR/iic` +``` + iic + ├── speech_fsmn_vad_zh-cn-16k-common-pytorch + │   ├── README.md + │   ├── am.mvn + │   ├── config.yaml + │   ├── configuration.json + │   ├── example + │   ├── fig + │   └── model.pt + ├── speech_seaco_paraformer_large_asr_nat-zh-cn-16k-common-vocab8404-pytorch + │   ├── README.md + │   ├── am.mvn + │   ├── asr_example_hotword.wav + │   ├── config.yaml + │   ├── configuration.json + │   ├── example + │   ├── fig + │   ├── model.pt + │   ├── seg_dict + │   └── tokens.json + └── speech_timestamp_prediction-v1-16k-offline + ├── README.md + ├── am.mvn + ├── config.yaml + ├── configuration.json + ├── example + ├── model.pt + ├── seg_dict + └── tokens.json +``` + +## 鸣谢 + +- [ComfyUI](https://github.com/comfyanonymous/ComfyUI) \ No newline at end of file diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..d721463 --- /dev/null +++ b/__init__.py @@ -0,0 +1,3 @@ +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] \ No newline at end of file diff --git a/example_workflow/av-audio2subtitle-with-prompt.json b/example_workflow/av-audio2subtitle-with-prompt.json new file mode 100644 index 0000000..01a43a2 --- /dev/null +++ b/example_workflow/av-audio2subtitle-with-prompt.json @@ -0,0 +1,311 @@ +{ + "id": "377e0192-2ad2-408a-871c-6a86cf058f73", + "revision": 0, + "last_node_id": 33, + "last_link_id": 50, + "nodes": [ + { + "id": 2, + "type": "LoadAudio", + "pos": [ + 182.18014526367188, + 389.2699890136719 + ], + "size": [ + 270, + 136 + ], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "AUDIO", + "type": "AUDIO", + "links": [ + 24 + ] + } + ], + "properties": { + "cnr_id": "comfy-core", + "ver": "0.3.32", + "Node name for S&R": "LoadAudio" + }, + "widgets_values": [ + "4_jidaohuman-megatts3_0.flac", + null, + null + ] + }, + { + "id": 25, + "type": "PrimitiveStringMultiline", + "pos": [ + 186.9199676513672, + 591.9599609375 + ], + "size": [ + 260, + 163 + ], + "flags": { + "collapsed": false + }, + "order": 1, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "STRING", + "type": "STRING", + "links": [ + 47, + 48 + ] + } + ], + "properties": { + "cnr_id": "comfy-core", + "ver": "0.3.32", + "Node name for S&R": "PrimitiveStringMultiline" + }, + "widgets_values": [ + "直观来说，测试时试错的步骤只有当能带来信息增益的情况下才有益，然而人们并没有在RL训练中告诉模型试错和反思带来的信息增益。" + ] + }, + { + "id": 18, + "type": "AVSpeechTimestamp", + "pos": [ + 489.56988525390625, + 374.6600341796875 + ], + "size": [ + 210, + 128 + ], + "flags": { + "collapsed": false + }, + "order": 2, + "mode": 0, + "inputs": [ + { + "name": "audio", + "type": "AUDIO", + "link": 24 + }, + { + "name": "text", + "type": "STRING", + "widget": { + "name": "text" + }, + "link": 47 + } + ], + "outputs": [ + { + "name": "text", + "type": "STRING", + "links": [] + }, + { + "name": "json_result", + "type": "STRING", + "links": [] + }, + { + "name": "asr_result", + "type": "ASRRESULT", + "links": [ + 43 + ] + } + ], + "properties": { + "Node name for S&R": "AVSpeechTimestamp" + }, + "widgets_values": [ + "", + false + ] + }, + { + "id": 30, + "type": "AVFormat2Subtitle", + "pos": [ + 518.5700073242188, + 561.3599853515625 + ], + "size": [ + 210, + 88 + ], + "flags": { + "collapsed": false + }, + "order": 3, + "mode": 0, + "inputs": [ + { + "name": "asr_result", + "type": "ASRRESULT", + "link": 43 + }, + { + "name": "text", + "type": "STRING", + "widget": { + "name": "text" + }, + "link": 48 + } + ], + "outputs": [ + { + "name": "subtitle", + "type": "STRING", + "links": [ + 45, + 50 + ] + } + ], + "properties": { + "Node name for S&R": "AVFormat2Subtitle" + }, + "widgets_values": [ + "" + ] + }, + { + "id": 31, + "type": "AVSaveSubtitles", + "pos": [ + 778.79541015625, + 357.02178955078125 + ], + "size": [ + 270, + 82 + ], + "flags": { + "collapsed": true + }, + "order": 4, + "mode": 0, + "inputs": [ + { + "name": "subtitles", + "type": "STRING", + "widget": { + "name": "subtitles" + }, + "link": 45 + } + ], + "outputs": [], + "properties": { + "Node name for S&R": "AVSaveSubtitles" + }, + "widgets_values": [ + "", + "subtitles" + ] + }, + { + "id": 33, + "type": "Bjornulf_ShowText", + "pos": [ + 814.5744018554688, + 427.3533935546875 + ], + "size": [ + 650.6199951171875, + 437.7900085449219 + ], + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [ + { + "name": "text_value", + "type": "STRING", + "link": 50 + } + ], + "outputs": [], + "properties": { + "cnr_id": "bjornulf_custom_nodes", + "ver": "1.0.0", + "Node name for S&R": "Bjornulf_ShowText" + }, + "widgets_values": [ + "1\n00:00:00,170 --> 00:00:02,330\n直观来说 测试时试错的步骤\n\n2\n00:00:02,330 --> 00:00:04,590\n只有当能带来信息增益的情况下才有益\n\n3\n00:00:04,590 --> 00:00:06,870\n然而人们并没有在RL训练\n\n4\n00:00:06,870 --> 00:00:09,330\n中告诉模型试错和反思带来的信息增益\n\n5\n00:00:09,330 --> 00:00:09,470\n\n\n" + ] + } + ], + "links": [ + [ + 24, + 2, + 0, + 18, + 0, + "AUDIO" + ], + [ + 43, + 18, + 2, + 30, + 0, + "ASRRESULT" + ], + [ + 45, + 30, + 0, + 31, + 0, + "STRING" + ], + [ + 47, + 25, + 0, + 18, + 1, + "STRING" + ], + [ + 48, + 25, + 0, + 30, + 1, + "STRING" + ], + [ + 50, + 30, + 0, + 33, + 0, + "STRING" + ] + ], + "groups": [], + "config": {}, + "extra": { + "frontendVersion": "1.18.9", + "VHS_latentpreview": false, + "VHS_latentpreviewrate": 0, + "VHS_MetadataImage": true, + "VHS_KeepIntermediate": true + }, + "version": 0.4 +} \ No newline at end of file diff --git a/example_workflow/av-audio2subtitle-without-prompt.json b/example_workflow/av-audio2subtitle-without-prompt.json new file mode 100644 index 0000000..ba40555 --- /dev/null +++ b/example_workflow/av-audio2subtitle-without-prompt.json @@ -0,0 +1,241 @@ +{ + "id": "377e0192-2ad2-408a-871c-6a86cf058f73", + "revision": 0, + "last_node_id": 35, + "last_link_id": 53, + "nodes": [ + { + "id": 2, + "type": "LoadAudio", + "pos": [ + 187.18014526367188, + 368.2699890136719 + ], + "size": [ + 270, + 136 + ], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "AUDIO", + "type": "AUDIO", + "links": [ + 51 + ] + } + ], + "properties": { + "cnr_id": "comfy-core", + "ver": "0.3.32", + "Node name for S&R": "LoadAudio" + }, + "widgets_values": [ + "4_jidaohuman-megatts3_0.flac", + null, + null + ] + }, + { + "id": 31, + "type": "AVSaveSubtitles", + "pos": [ + 819.79541015625, + 213.02178955078125 + ], + "size": [ + 270, + 82 + ], + "flags": { + "collapsed": true + }, + "order": 3, + "mode": 0, + "inputs": [ + { + "name": "subtitles", + "type": "STRING", + "widget": { + "name": "subtitles" + }, + "link": 45 + } + ], + "outputs": [], + "properties": { + "Node name for S&R": "AVSaveSubtitles" + }, + "widgets_values": [ + "", + "subtitles" + ] + }, + { + "id": 34, + "type": "AVASRTimestamp", + "pos": [ + 467.7745361328125, + 206.14337158203125 + ], + "size": [ + 270, + 122 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [ + { + "name": "audio", + "type": "AUDIO", + "link": 51 + } + ], + "outputs": [ + { + "name": "text", + "type": "STRING", + "links": null + }, + { + "name": "json_result", + "type": "STRING", + "links": null + }, + { + "name": "asr_result", + "type": "ASRRESULT", + "links": [ + 52 + ] + } + ], + "properties": { + "Node name for S&R": "AVASRTimestamp" + }, + "widgets_values": [ + 300, + false + ] + }, + { + "id": 30, + "type": "AVFormat2Subtitle", + "pos": [ + 485.57000732421875, + 415.3599853515625 + ], + "size": [ + 210, + 88 + ], + "flags": { + "collapsed": false + }, + "order": 2, + "mode": 0, + "inputs": [ + { + "name": "asr_result", + "type": "ASRRESULT", + "link": 52 + } + ], + "outputs": [ + { + "name": "subtitle", + "type": "STRING", + "links": [ + 45, + 53 + ] + } + ], + "properties": { + "Node name for S&R": "AVFormat2Subtitle" + }, + "widgets_values": [ + "" + ] + }, + { + "id": 35, + "type": "Bjornulf_ShowText", + "pos": [ + 804.8345336914062, + 269.0433654785156 + ], + "size": [ + 398, + 387 + ], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [ + { + "name": "text_value", + "type": "STRING", + "link": 53 + } + ], + "outputs": [], + "properties": { + "cnr_id": "bjornulf_custom_nodes", + "ver": "1.0.0", + "Node name for S&R": "Bjornulf_ShowText" + }, + "widgets_values": [ + "1\n00:00:00,150 --> 00:00:02,610\n直观来说测试时试错的步骤只有\n\n2\n00:00:02,610 --> 00:00:05,110\n当能带来信息增益的情况下才有益\n\n3\n00:00:05,410 --> 00:00:07,130\n然而人们并没有在all训练中\n\n4\n00:00:07,130 --> 00:00:10,225\n告诉模型试错和反思带来的信息增益\n\n" + ] + } + ], + "links": [ + [ + 45, + 30, + 0, + 31, + 0, + "STRING" + ], + [ + 51, + 2, + 0, + 34, + 0, + "AUDIO" + ], + [ + 52, + 34, + 2, + 30, + 0, + "ASRRESULT" + ], + [ + 53, + 30, + 0, + 35, + 0, + "STRING" + ] + ], + "groups": [], + "config": {}, + "extra": { + "frontendVersion": "1.18.9", + "VHS_latentpreview": false, + "VHS_latentpreviewrate": 0, + "VHS_MetadataImage": true, + "VHS_KeepIntermediate": true + }, + "version": 0.4 +} \ No newline at end of file diff --git a/example_workflow/workflow.webp b/example_workflow/workflow.webp new file mode 100644 index 0000000..5055df7 Binary files /dev/null and b/example_workflow/workflow.webp differ diff --git a/format.py b/format.py new file mode 100644 index 0000000..8f1b228 --- /dev/null +++ b/format.py @@ -0,0 +1,136 @@ +import jieba +import string +import json +import Levenshtein +import re + +class Format2Subtitle: + def __init__(self, asr_result, ori_text): + self.asr_result = asr_result + self.ori_text = ori_text + def ms_to_srt_time(self, ms): + """ + 将毫秒转换为SRT时间格式 (HH:MM:SS,mmm) + 例如：150毫秒 -> 00:00:00,150 + """ + hours = ms // 3600000 # 1小时 = 3,600,000毫秒 + ms %= 3600000 + minutes = ms // 60000 # 1分钟 = 60,000毫秒 + ms %= 60000 + seconds = ms // 1000 # 1秒 = 1,000毫秒 + milliseconds = ms % 1000 + return f"{hours:02d}:{minutes:02d}:{seconds:02d},{milliseconds:03d}" + # 查询标点符号位置 + def find_punctuation_indices(self, words): + punctuation = """，。；！？!,;.?:""" + for i, word in enumerate(words): + if word in punctuation: + return i + return -1 + + # 清除句子中的断句符号 + def remove_sentence_punctuation(self, sentence): + """去除句子中的断句符号""" + punctuation = """，。；！？!,;.?:()（）""" + for p in punctuation: + sentence = sentence.replace(p, ' ') + sentence = sentence.strip() + return sentence + + def align_char_timestamps(self, original_sentence, recognized_text, word_timestamps, text_length=8): + """ + 将词语级时间戳对齐到字符级时间戳 + + 参数: + original_sentence (str): 带标点的原始句子 + recognized_text (str): 空格分隔的无标点识别文本 + word_timestamps (list): 词语级时间戳列表 [[start1, end1], [start2, end2], ...] + + 返回: + list: 原始字符级时间戳列表 [[char_start, char_end], ...] + """ + if not original_sentence: + original_sentence = "".join(recognized_text.split()) + + words = jieba.lcut(original_sentence) + + # print(len(original_sentence), len(word_timestamps), len(recognized_text)) + recognized_words = recognized_text.split() + + # 4. 为每个字符分配时间戳 + char_timestamps = [] + + while True: + if len(words) < 10: + break + # 查询列表中的标点符号的位置 + punctuation_indices = self.find_punctuation_indices(words) + # print("标点符号位置：",punctuation_indices) + # 如果位置小于10，怎使用这个位置分割 + if punctuation_indices > 10 or punctuation_indices <= 2: + _words = words[:text_length] + words = words[text_length:] + else: + _words = words[:punctuation_indices] + words = words[punctuation_indices:] + + text = ''.join(_words) + # 计算词语持续时间 + best_index = -1 + best_score = 9999 + for index in range(text_length*4): + _recognized_words = recognized_words[:index] + _recognized_text = ''.join(_recognized_words) + if text == _recognized_text: + break + # 计算编辑距离 + distance = Levenshtein.distance(text, _recognized_text) + if distance < best_score: + best_index = index + best_score = distance + _recognized_words = recognized_words[:best_index] + recognized_words = recognized_words[best_index:] + + timestamp = word_timestamps[:best_index] + word_timestamps = word_timestamps[best_index:] + word_start = timestamp[0][0] + word_end = timestamp[-1][1] + + char_timestamps.append([word_start, word_end, self.remove_sentence_punctuation(text)]) + # 处理最后一段 + if len(word_timestamps) > 0: + text = ''.join(words) + timestamp = word_timestamps + word_start = timestamp[0][0] + word_end = timestamp[-1][1] + char_timestamps.append([word_start, word_end, self.remove_sentence_punctuation(text)]) + return char_timestamps + + def format_subtitle(self, char_timestamps): + """ + [230, 2210, '一个能够从音频中生成面部'] + [230, 5430, '运动系数的工具'] + 转为字幕文件 + """ + srt_content = "" + counter = 1 + for item in char_timestamps: + start_time = self.ms_to_srt_time(item[0]) + end_time = self.ms_to_srt_time(item[1]) + text = item[2] + srt_content += f"{counter}\n{start_time} --> {end_time}\n{text}\n\n" + # print(srt_content) + counter += 1 + return srt_content + + def pipeline(self, output=None): + # print(rec_result) + recognized_text = self.asr_result['text'] + word_timestamps = self.asr_result['timestamp'] + original_sentence = self.ori_text + char_timestamps = self.align_char_timestamps(original_sentence, recognized_text, word_timestamps) + subtitle_text = self.format_subtitle(char_timestamps) + if output is not None: + with open(output, 'w', encoding='utf-8') as f: + f.write(subtitle_text) + return subtitle_text \ No newline at end of file diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..66244d6 --- /dev/null +++ b/nodes.py @@ -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", +} \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..3b168e1 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,15 @@ +[project] +name = "ComfyUI-AV-FunASR" +description = "FunASR wrapper for ComfyUI" +version = "1.0.0" +license = {file = "LICENSE"} +dependencies = ["modelscope", "sortedcontainers", "simplejson", "datasets==2.16.0", "Levenshtein", "soundfile"] + +[project.urls] +Repository = "https://github.com/avenstack/ComfyUI-AV-FunASR" +# Used by Comfy Registry https://comfyregistry.org + +[tool.comfy] +PublisherId = "av" +DisplayName = "ComfyUI-AV-FunASR" +Icon = "" diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..afdee3f --- /dev/null +++ b/requirements.txt @@ -0,0 +1,5 @@ +sortedcontainers +simplejson +datasets==2.16.0 +Levenshtein +modelscope \ No newline at end of file