From ef53ae847b48b641b7800122e78a5ccc5a6a5d4e Mon Sep 17 00:00:00 2001
From: AIFSH <1509359472@qq.com>
Date: Mon, 15 Jul 2024 11:44:22 +0800
Subject: [PATCH] first commit

---
 .gitignore       |  2 ++
 __init__.py      | 82 ++++++++++++++++++++++++++++++++++++++++++++++++
 js/innerText.js  | 18 +++++++++++
 requirements.txt |  5 +++
 4 files changed, 107 insertions(+)
 create mode 100644 .gitignore
 create mode 100644 __init__.py
 create mode 100644 js/innerText.js
 create mode 100644 requirements.txt

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..851d849
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,2 @@
+__pycache__
+pretrianed_models
\ No newline at end of file
diff --git a/__init__.py b/__init__.py
new file mode 100644
index 0000000..929bfc6
--- /dev/null
+++ b/__init__.py
@@ -0,0 +1,82 @@
+import os
+now_dir = os.path.dirname(os.path.abspath(__file__))
+
+import torchaudio
+from funasr import AutoModel
+# from funasr.utils.postprocess_utils import rich_transcription_postprocess
+from modelscope import snapshot_download
+
+pre_model_dir = os.path.join(now_dir,"pretrianed_models","SenseVoiceSmall")
+snapshot_download(model_id="iic/SenseVoiceSmall",local_dir=pre_model_dir)
+
+class SenseVoiceNode:
+
+    def __init__(self) -> None:
+        self.model = None
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required":{
+                "audio":("AUDIO",),
+                "batch_size_s":("INT",{
+                    "default": 0
+                })
+            }
+        }
+    
+    RETURN_TYPES = ("TEXT",)
+    #RETURN_NAMES = ("image_output_name",)
+
+    FUNCTION = "generate"
+
+    CATEGORY = "AIFSH_SenseVoice"
+
+    def generate(self,audio,batch_size_s):
+        if self.model is None:
+            self.model = AutoModel(model=pre_model_dir,
+                    vad_model="fsmn-vad",
+                    vad_kwargs={"max_single_segment_time": 30000},
+                    trust_remote_code=True, device="cuda:0")
+        audio_data = audio["waveform"].squeeze(0)
+        audio_rate = audio['sample_rate']
+        if audio_rate != 16000:
+            audio_data = torchaudio.transforms.Resample(audio_rate,16000)(audio_data)
+        
+        # audio_data = torchaudio.compliance.kaldi.fbank(audio_data)
+        audio_data = audio_data.numpy()
+        # print(audio_data.shape)
+        res = self.model.generate(
+            input=[audio_data],
+            cache={},
+            language="auto", # "zn", "en", "yue", "ja", "ko", "nospeech"
+            use_itn=False,
+            batch_size_s=batch_size_s, 
+        )
+        print(res[0]["text"])
+        return (res[0]["text"],)
+
+class TextNode:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "sense_voice_output":("TEXT",),
+                "text": ("STRING", {"multiline": True, "dynamicPrompts": True}),
+            }
+        }
+    RETURN_TYPES = ("TEXT",)
+    FUNCTION = "encode"
+    OUTPUT_NODE = True
+    CATEGORY = "AIFSH_SenseVoice"
+
+    def encode(self,sense_voice_output,text):
+        return {"ui":{"text":[sense_voice_output]}}
+
+
+WEB_DIRECTORY = "./js"
+
+NODE_CLASS_MAPPINGS = {
+    "TextNode":TextNode,
+    "SenseVoiceNode": SenseVoiceNode
+}
diff --git a/js/innerText.js b/js/innerText.js
new file mode 100644
index 0000000..0de5642
--- /dev/null
+++ b/js/innerText.js
@@ -0,0 +1,18 @@
+import {app} from "../../../scripts/app.js"
+
+function innerText(node,text){
+    const textWidget = node.widgets.find((w)=>w.name === "text")
+    textWidget.value = text
+}
+
+app.registerExtension({
+    name: "SenseVoice.innerText",
+    async beforeRegisterNodeDef(nodeType,nodeData,app){
+        if (nodeData?.name == "TextNode"){
+            nodeType.prototype.onExecuted = function(data){
+                console.log(data.text[0]);
+                innerText(this,data.text[0])
+            }
+        }
+    }
+})
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..0b2da35
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,5 @@
+modelscope
+huggingface
+huggingface_hub
+funasr>=1.1.0
+numpy<=1.26.4
\ No newline at end of file