add SenseVoiceNode

README.md

@@ -3,7 +3,7 @@
  * @Date: 2024-10-04 11:28:28
  * @version: 
  * @LastEditors: SpenserCai
- * @LastEditTime: 2024-10-04 21:14:03
+ * @LastEditTime: 2024-10-04 22:28:03
  * @Description: file content
 -->
 # ComfyUI-FunAudioLLM
@@ -22,6 +22,10 @@ Comfyui custom node for [FunAudioLLM](https://funaudiollm.github.io/) include [C
       <img src="./assets/LoadSpeakerModel.png" alt="zh-CN" />
     </details>
 
+### SenseVoice
+  - SenseVoice Version: 2024-10-04
+  - Support SenseVoice-Small
+
 ## How use
 ```bash
 apt update
@@ -40,4 +44,5 @@ If your network is unstable, you can pre-download the model from the following s
 - [CosyVoice-300M-SFT](https://modelscope.cn/models/iic/CosyVoice-300M-SFT) -> `ComfyUI/models/CosyVoice/CosyVoice-300M-SFT`
 - [CosyVoice-300M-SFT-25Hz](https://modelscope.cn/models/MachineS/CosyVoice-300M-SFT-25Hz) -> `ComfyUI/models/CosyVoice/CosyVoice-300M-SFT-25Hz`
 - [CosyVoice-300M-Instruct](https://modelscope.cn/models/iic/CosyVoice-300M-Instruct) -> `ComfyUI/models/CosyVoice/CosyVoice-300M-Instruct`
+- [SenseVoiceSmall](https://modelscope.cn/models/iic/SenseVoiceSmall) -> `ComfyUI/models/SenseVoice/SenseVoiceSmall`
      

__init__.py

@@ -3,7 +3,7 @@ Author: SpenserCai
 Date: 2024-10-04 12:14:22
 version: 
 LastEditors: SpenserCai
-LastEditTime: 2024-10-04 21:12:47
+LastEditTime: 2024-10-04 22:29:33
 Description: file content
 '''
 from .nodes.cosyvoice_nodes import *
@@ -38,6 +38,10 @@ NODE_CONFIG = {
         "class": CosyVoiceLoadSpeakerModelFromUrlNode,
         "name": "CosyVoice 从URL加载说话人模型"
     },
+    "SenseVoiceNode": {
+        "class": SenseVoiceNode,
+        "name": "SenseVoice 语音识别"
+    }
 }
 
 def generate_node_mappings(node_config):

funaudio_utils/download_models.py

@@ -3,7 +3,7 @@ Author: SpenserCai
 Date: 2024-10-04 13:54:01
 version: 
 LastEditors: SpenserCai
-LastEditTime: 2024-10-04 21:10:55
+LastEditTime: 2024-10-04 22:26:22
 Description: file content
 '''
 import modelscope
@@ -13,6 +13,7 @@ from modelscope import snapshot_download
 
 # Download the model
 base_cosyvoice_model_path = os.path.join(folder_paths.models_dir, "CosyVoice")
+base_sensevoice_model_path = os.path.join(folder_paths.models_dir, "SenseVoice")
 
 def download_cosyvoice_300m(is_25hz=False):
     model_name = "CosyVoice-300M"
@@ -34,6 +35,13 @@ def download_cosyvoice_300m_sft(is_25hz=False):
     snapshot_download(model_id=model_id, local_dir=model_dir)
     return model_name, model_dir
 
+def download_sensevoice_small():
+    model_name = "SenseVoiceSmall"
+    model_id = "iic/SenseVoiceSmall"
+    model_dir = os.path.join(base_sensevoice_model_path, model_name)
+    snapshot_download(model_id=model_id, local_dir=model_dir)
+    return model_name, model_dir
+
 def download_cosyvoice_300m_instruct():
     model_name = "CosyVoice-300M-Instruct"
     model_id = "iic/CosyVoice-300M-Instruct"

nodes/cosyvoice_nodes.py

@@ -3,12 +3,11 @@ Author: SpenserCai
 Date: 2024-10-04 12:13:28
 version: 
 LastEditors: SpenserCai
-LastEditTime: 2024-10-04 21:11:53
+LastEditTime: 2024-10-04 21:44:43
 Description: file content
 '''
 import os
 import folder_paths
-import random
 import numpy as np
 import torch
 from funaudio_utils.pre import FunAudioLLMTool

nodes/sensevoice_nodes.py

@@ -0,0 +1,63 @@
+import folder_paths
+import os
+import numpy as np
+from funasr import AutoModel
+from funaudio_utils.pre import FunAudioLLMTool
+from funaudio_utils.download_models import download_sensevoice_small
+from funasr.utils.postprocess_utils import rich_transcription_postprocess
+
+fAudioTool = FunAudioLLMTool()
+
+CATEGORY_NAME = "FunAudioLLM - SenseVoice"
+
+folder_paths.add_model_folder_path("SenseVoice", os.path.join(folder_paths.models_dir, "SenseVoice"))
+
+class SenseVoiceNode:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required":{
+                "audio":("AUDIO",),
+                "less_than_30s":("BOOLEAN",{
+                    "default": True
+                }),
+            }
+        }
+    
+    CATEGORY = CATEGORY_NAME
+    RETURN_TYPES = ("STRING",)
+
+    FUNCTION="generate"
+
+    def generate(self,audio, less_than_30s):
+        sensevoice_code_path = os.path.join(folder_paths.folder_names_and_paths["custom_nodes"],"ComfyUI-FunAudioLLM/sensevoice/model.py")
+        speech = audio["waveform"]
+        source_sr = audio["sample_rate"]
+        speech = fAudioTool.audio_resample(speech, source_sr)
+        speech = fAudioTool.postprocess(speech)
+        _, model_dir = download_sensevoice_small()
+        model_arg = {
+                "input":speech,
+                "cache":{},
+                "language":"auto",
+                "batch_size_s":60,
+        }
+        if not less_than_30s:
+            model = AutoModel(model=model_dir,
+                              trust_remote_code=True,
+                              remote_code=sensevoice_code_path,
+                              vad_model="fsmn-vad",
+                              vad_kwargs={"max_single_segment_time": 30000},
+                              device="cuda:0",
+                    )
+            model_arg["merge_vad"] = True
+            model_arg["merge_length_s"] = 15
+            
+        else:
+            model = AutoModel(model=model_dir,
+                              trust_remote_code=True,
+                              remote_code=sensevoice_code_path,
+                              device="cuda:0",
+                    )
+        output = model.generate(**model_arg)
+        return (rich_transcription_postprocess(output[0]["text"]),)