From 460f64dfed9991d0f5f2f8a8910121f31b10ddb4 Mon Sep 17 00:00:00 2001
From: SayanoAI <28544690+SayanoAI@users.noreply.github.com>
Date: Fri, 23 Aug 2024 18:26:13 -0400
Subject: [PATCH] refactored node mappings

---
 __init__.py         | 22 ++++++----------------
 custom_nodes/rvc.py | 16 +++++++++++++++-
 custom_nodes/stt.py | 17 ++++++++++++++++-
 3 files changed, 37 insertions(+), 18 deletions(-)

diff --git a/__init__.py b/__init__.py
index 1beab6a..45ac1ea 100644
--- a/__init__.py
+++ b/__init__.py
@@ -1,7 +1,7 @@
 from .custom_nodes.audio_nodes import NODE_CLASS_MAPPINGS as audio_nodes, NODE_DISPLAY_NAME_MAPPINGS as audio_nodes_name
-from .custom_nodes.stt import AudioTranscriptionNode, BatchedTranscriptionEncoderNode, LoadWhisperModelNode
+from .custom_nodes.stt import NODE_CLASS_MAPPINGS as stt_nodes, NODE_DISPLAY_NAME_MAPPINGS as stt_nodes_name
 from .custom_nodes.uvr import UVR5Node
-from .custom_nodes.rvc import LoadHubertModel, LoadPitchExtractionParams, LoadRVCModelNode, RVCNode
+from .custom_nodes.rvc import NODE_CLASS_MAPPINGS as rvc_nodes, NODE_DISPLAY_NAME_MAPPINGS as rvc_nodes_name
 from .custom_nodes.utility_nodes import NODE_CLASS_MAPPINGS as utility_nodes, NODE_DISPLAY_NAME_MAPPINGS as utility_nodes_name
 from .custom_nodes.musetalk_nodes import NODE_CLASS_MAPPINGS as musetalk_nodes, NODE_DISPLAY_NAME_MAPPINGS as musetalk_nodes_name
 # Set the web directory, any .js file in that directory will be loaded by the frontend as a frontend extension
@@ -11,13 +11,8 @@ WEB_DIRECTORY = "./web"
 # NOTE: names should be globally unique
 NODE_CLASS_MAPPINGS = {
     "UVR5Node": UVR5Node,
-    "AudioTranscriptionNode": AudioTranscriptionNode,
-    "LoadWhisperModelNode": LoadWhisperModelNode,
-    "LoadRVCModelNode": LoadRVCModelNode,
-    "RVCNode": RVCNode,
-    "LoadHubertModel": LoadHubertModel,
-    "LoadPitchExtractionParams": LoadPitchExtractionParams,
-    "BatchedTranscriptionEncoderNode": BatchedTranscriptionEncoderNode,
+    **rvc_nodes,
+    **stt_nodes,
     **audio_nodes,
     **musetalk_nodes,
     **utility_nodes
@@ -26,13 +21,8 @@ NODE_CLASS_MAPPINGS = {
 # A dictionary that contains the friendly/humanly readable titles for the nodes
 NODE_DISPLAY_NAME_MAPPINGS = {
     "UVR5Node": "🌺Vocal Removal",
-    "AudioTranscriptionNode": "🌺Transcribe Audio",
-    "LoadWhisperModelNode": "🌺Load Whisper Model",
-    "LoadRVCModelNode": "🌺Load RVC Model",
-    "RVCNode": "🌺Voice Changer",
-    "LoadHubertModel": "🌺Load Hubert Model",
-    "LoadPitchExtractionParams": "🌺Load Pitch Extraction Params",
-    "BatchedTranscriptionEncoderNode": "🌺Batched CLIP Transcription Encode (Prompt)",
+    **rvc_nodes_name,
+    **stt_nodes_name,
     **audio_nodes_name,
     **musetalk_nodes_name,
     **utility_nodes_name
diff --git a/custom_nodes/rvc.py b/custom_nodes/rvc.py
index bbd2a8d..b701afe 100644
--- a/custom_nodes/rvc.py
+++ b/custom_nodes/rvc.py
@@ -208,6 +208,20 @@ class RVCNode:
         preview_file = os.path.join(tempdir,audio_name)
         if not os.path.isfile(preview_file): shutil.copyfile(cache_name,preview_file)
         return {"ui": {"preview": [{"filename": audio_name, "type": "temp", "subfolder": "preview", "widgetId": widgetId}]}, "result": (lambda:audio_to_bytes(*output_audio),)}
-
     
+# A dictionary that contains all nodes you want to export with their names
+# NOTE: names should be globally unique
+NODE_CLASS_MAPPINGS = {
+    "LoadRVCModelNode": LoadRVCModelNode,
+    "RVCNode": RVCNode,
+    "LoadHubertModel": LoadHubertModel,
+    "LoadPitchExtractionParams": LoadPitchExtractionParams,
+}
 
+# A dictionary that contains the friendly/humanly readable titles for the nodes
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "LoadRVCModelNode": "🌺Load RVC Model",
+    "RVCNode": "🌺Voice Changer",
+    "LoadHubertModel": "🌺Load Hubert Model",
+    "LoadPitchExtractionParams": "🌺Load Pitch Extraction Params",
+}
\ No newline at end of file
diff --git a/custom_nodes/stt.py b/custom_nodes/stt.py
index f55e70f..35e40c9 100644
--- a/custom_nodes/stt.py
+++ b/custom_nodes/stt.py
@@ -365,4 +365,19 @@ class BatchedTranscriptionEncoderNode:
             print(f"{duration_list=}")
             print(f"{num_chunks=}, {max_chunks=}, {num_frames=}")
 
-        return (conditioning, batch_prompt_text, list(map(int,duration_list)), num_chunks, num_frames, text_list)
\ No newline at end of file
+        return (conditioning, batch_prompt_text, list(map(int,duration_list)), num_chunks, num_frames, text_list)
+    
+# A dictionary that contains all nodes you want to export with their names
+# NOTE: names should be globally unique
+NODE_CLASS_MAPPINGS = {
+    "AudioTranscriptionNode": AudioTranscriptionNode,
+    "LoadWhisperModelNode": LoadWhisperModelNode,
+    "BatchedTranscriptionEncoderNode": BatchedTranscriptionEncoderNode
+}
+
+# A dictionary that contains the friendly/humanly readable titles for the nodes
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "AudioTranscriptionNode": "🌺Transcribe Audio",
+    "LoadWhisperModelNode": "🌺Load Whisper Model",
+    "BatchedTranscriptionEncoderNode": "🌺Batched CLIP Transcription Encode (Prompt)",
+}
\ No newline at end of file