separated out BatchedTranscriptionEncoderNode

__init__.py

@@ -1,4 +1,4 @@
-from .custom_nodes.stt import AudioTranscriptionNode
+from .custom_nodes.stt import AudioTranscriptionNode, BatchedTranscriptionEncoderNode
 from .custom_nodes.uvr import UVR5Node
 from .custom_nodes.rvc import RVCNode
 from .custom_nodes.loaders import DownloadAudio, LoadAudio, LoadWhisperModelNode, LoadRVCModelNode, LoadHubertModel, LoadPitchExtractionParams
@@ -26,7 +26,8 @@ NODE_CLASS_MAPPINGS = {
     "MergeLatentBatches": MergeLatentBatches,
     "ImageRepeatInterleavedNode": ImageRepeatInterleavedNode,
     "LatentRepeatInterleavedNode": LatentRepeatInterleavedNode,
-    "DownloadAudio": DownloadAudio
+    "DownloadAudio": DownloadAudio,
+    "BatchedTranscriptionEncoderNode": BatchedTranscriptionEncoderNode
 }
 
 # A dictionary that contains the friendly/humanly readable titles for the nodes
@@ -46,5 +47,6 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "MergeImageBatches": "🌺Merge Image Batches",
     "MergeLatentBatches": "🌺Merge Latent Batches",
     "ImageRepeatInterleavedNode": "🌺Image Repeat Interleaved",
-    "LatentRepeatInterleavedNode": "🌺Latent Repeat Interleaved"
+    "LatentRepeatInterleavedNode": "🌺Latent Repeat Interleaved",
+    "BatchedTranscriptionEncoderNode": "🌺Batched CLIP Transcription Encode (Prompt)"
 }

custom_nodes/loaders.py

@@ -1,7 +1,5 @@
 from io import BytesIO
 import os
-import subprocess
-import sys
 from pytube import YouTube
 from .settings import PITCH_EXTRACTION_OPTIONS
 from ..lib import BASE_MODELS_DIR
@@ -36,12 +34,6 @@ model_ids = [
 
 languages = ['en', 'fr', 'es']
 
-spacy_model_map = {
-    "en": "en_core_web_md",
-    "fr": "fr_core_news_md",
-    "es": "es_core_news_md"
-}
-
 class LoadPitchExtractionParams:
     @classmethod
     def INPUT_TYPES(cls):
@@ -88,7 +80,9 @@ class LoadPitchExtractionParams:
 class LoadHubertModel:
     @classmethod
     def INPUT_TYPES(cls):
-        model_list = ["content-vec-best.safetensors"] + get_filenames(root=BASE_MODELS_DIR,folder="*",exts=["pt","safetensors"],format_func=os.path.basename)
+        model_list = ["content-vec-best.safetensors"] + get_filenames(root=BASE_MODELS_DIR,folder=".",exts=["pt","safetensors"],format_func=os.path.basename)
+        model_list = list(set(model_list)) # dedupe
+        
         return {
             'required': {
                 'model': (model_list,{"default": "content-vec-best.safetensors"}),
@@ -170,7 +164,6 @@ class LoadWhisperModelNode:
                 'model_id': (model_ids,{"default": "openai/whisper-base.en"}),
             },
             "optional": {
-                'language': (languages,{"default": "en"}),
                 "max_new_tokens": ("INT", {"default": 128, "min": 16, "max": 1024, "display": "slider"}),
                 "chunk_length_s": ("INT", {"default": 30, "min": 15, "max": 60, "display": "slider"}),
                 "batch_size": ("INT", {"default": 16, "min": 1, "max": 128, "display": "slider"}),
@@ -188,7 +181,7 @@ class LoadWhisperModelNode:
     def IS_CHANGED(cls, model):
         return get_hash(model)
 
-    def load_model(self, model_id, language="en", max_new_tokens=128, chunk_length_s=12, batch_size=16):
+    def load_model(self, model_id, max_new_tokens=128, chunk_length_s=12, batch_size=16):
         device = get_optimal_torch_device()
         torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 
@@ -198,28 +191,9 @@ class LoadWhisperModelNode:
         processor = AutoProcessor.from_pretrained(model_id)
         model.to(device)
 
-        generate_kwargs = {}
-        if model_id.endswith('.en'):
-            if language != 'en':
-                raise ValueError(f'Model {model_id} only supports English language')
-        else:
-            generate_kwargs['language'] = language
+        # generate_kwargs = {}
 
-        def pipe():
-            
-            def get_spacy_model():
-                import spacy
-                model_name = spacy_model_map[language]
-                model_path = os.path.join(BASE_MODELS_DIR,model_name)
-                
-                if not os.path.exists(model_path):
-                    subprocess.call([sys.executable, "-m", "spacy", "download", model_name])
-                    model = spacy.load(model_name)
-                    model.to_disk(model_path)
-
-                return spacy.load(model_path)
-                
-            return pipeline(
+        def pipe(): return pipeline(
                 'automatic-speech-recognition',
                 model=model,
                 tokenizer=processor.tokenizer,
@@ -230,9 +204,9 @@ class LoadWhisperModelNode:
                 return_timestamps=True,
                 torch_dtype=torch_dtype,
                 device=device,
-                generate_kwargs=generate_kwargs,
-            ),get_spacy_model
-        return (pipe, )
+                # generate_kwargs=generate_kwargs,
+            )
+        return ([pipe,model_id], )
 
 class LoadAudio:
     @classmethod

custom_nodes/settings/__init__.py

@@ -17,14 +17,14 @@ MERGE_OPTIONS=["median","mean","min","max"]
 TTS_MODELS = ["edge","speecht5"]
 N_THREADS_OPTIONS=[1,2,4,8,12,16]
 SR_MAP = {"32k": 32000,"40k": 40000, "48k": 48000}
+SUPPORTED_LANGUAGES = ['en', 'fr', 'es', "ja", "zh"]
+# @lru_cache
+# def get_servers():
+#     os.makedirs(BASE_CACHE_DIR,exist_ok=True)
+#     fname = os.path.join(BASE_CACHE_DIR,"servers.shelve")
+#     servers = PersistedDict(fname)
+#     return servers
 
-@lru_cache
-def get_servers():
-    os.makedirs(BASE_CACHE_DIR,exist_ok=True)
-    fname = os.path.join(BASE_CACHE_DIR,"servers.shelve")
-    servers = PersistedDict(fname)
-    return servers
-
-SERVERS = get_servers()
-RVC_INFERENCE_URL = SERVERS.RVC_INFERENCE_URL
-UVR_INFERENCE_URL = SERVERS.UVR_INFERENCE_URL
+# SERVERS = get_servers()
+# RVC_INFERENCE_URL = SERVERS.RVC_INFERENCE_URL
+# UVR_INFERENCE_URL = SERVERS.UVR_INFERENCE_URL

custom_nodes/stt.py

@@ -1,10 +1,15 @@
 import json
 import os
 from pprint import pprint
+import subprocess
+import sys
 
 import numpy as np
+import torch
 
-from ..lib import BASE_CACHE_DIR
+from .settings import SUPPORTED_LANGUAGES
+
+from ..lib import BASE_CACHE_DIR, BASE_MODELS_DIR
 from .utils import increment_filename_no_overwrite
 from ..lib.utils import get_hash
 import folder_paths
@@ -13,8 +18,9 @@ import textacy
 from textacy.extract import keyterms as kt
 
 temp_path = folder_paths.get_temp_directory()
+CATEGORY = "🌺RVC-Studio/stt"
 
-def extract_keywords(text: str, max_words: int, spacy_model, **kwargs):
+def extract_keywords(text: str, max_words: int, spacy_model, prefix="", suffix="", **kwargs):
     text = text.strip()
     doc = textacy.make_spacy_doc(text, lang=spacy_model)
     topn = int(max_words) if max_words>0 else len(text)
@@ -27,13 +33,14 @@ def extract_keywords(text: str, max_words: int, spacy_model, **kwargs):
     except Exception as error:
         print(f"{text=} {error=}")
 
-    return ", ".join(tags).strip()
+    return ", ".join(filter(None,[prefix,*tags,suffix])).strip()
 
-def limit_sentence(text: str, max_words: int, **kwargs):
+def limit_sentence(text: str, max_words: int, prefix="", suffix="", **kwargs):
     text = text.strip()
     topn = int(max_words) if max_words>0 else len(text)
-    if topn>0: return " ".join(text.split()[:topn])
-    else: return text
+    if topn>0: text = " ".join(text.split()[:topn])
+    
+    return ", ".join(filter(None,[prefix,text,suffix])).strip()
 
 class AudioTranscriptionNode:
     def __init__(self):
@@ -44,17 +51,10 @@ class AudioTranscriptionNode:
         return {
             "required": {
                 "pipeline": ('TRANSCRIPTION_MODEL',),
-                "audio": ('VHS_AUDIO',),
-                "format_newlines_on_punctuation": ('BOOLEAN', {'default': True}),
-                "loop": ('BOOLEAN', {'default': True}),
-                "use_tags": ('BOOLEAN', {'default': True}),
-                "max_words": ('INT', {'default': 16, "min": 0, "max": 32, "display": "slider"}),
-                "frame_interpolation": ("INT", {"default": 0, "min": 0, "max": 120})
+                "audio": ('VHS_AUDIO',)
             },
             "optional": {
-                "save_transcription": ('BOOLEAN', {'default': False}),
-                "save_chunks": ('BOOLEAN', {'default': False}),
-                "save_filename": ('STRING', {'default': 'transcription'}),
+                "save_filename": ('STRING', {'default': ''}),
                 "overwrite_existing": ('BOOLEAN', {'default': True}),
                 "print_output": ('BOOLEAN', {'default': True}),
                 "use_cache": ("BOOLEAN",{"default": True})
@@ -62,12 +62,12 @@ class AudioTranscriptionNode:
         }
 
     OUTPUT_NODE = True
-    RETURN_TYPES = ("STRING", "STRING", "STRING", "INT", "INT")
-    RETURN_NAMES = ("transcription", "batch_prompt_text", "batch_values_text", "num_chunks","num_frames")
+    RETURN_TYPES = ("TRANSCRIPTION", "INT")
+    RETURN_NAMES = ("transcription", "audio_frames")
 
     FUNCTION = "transcribe"
 
-    CATEGORY = "🌺RVC-Studio/stt"
+    CATEGORY = CATEGORY
 
     @staticmethod
     def save_output(filename,text,overwrite_existing=False):
@@ -82,60 +82,125 @@ class AudioTranscriptionNode:
             file_path = increment_filename_no_overwrite(file_path)
         with open(file_path, 'w') as f:
             f.write(text)
+    
+    @staticmethod
+    def load_cache(*args,save_filename=""):
+        results = None
+        if save_filename: save_filename=get_hash(*args)
+        cache_name = os.path.join(BASE_CACHE_DIR,"stt",f"{save_filename}.json")
+        
+        if os.path.isfile(cache_name):
+            with open(cache_name,"r") as f:
+                results = json.load(f)
+        return results, cache_name
+
+    def transcribe(self, pipeline, audio,save_filename="", overwrite_existing=True,print_output=True,use_cache=True):
+        print('Starting Transcription')
+        audio_data = bytes_to_audio(audio())
+        whisper_model,model_id=pipeline
+        if use_cache: transcription, cache_name = self.load_cache(model_id,audio(),save_filename=save_filename)
+        else: transcription = None
+
+        audio_frames = int(np.ceil(len(audio_data[0])/audio_data[1]))
+
+        if transcription is None:
+            audio,_ = remix_audio(audio_data,target_sr=16000,norm=True)
+            transcription = whisper_model()(audio)
+
+        if use_cache and transcription is not None:
+            self.save_output(cache_name, json.dumps(transcription, indent=2), overwrite_existing=overwrite_existing)
+
+        print('Transcription Done')
+        
+
+        if print_output:
+            pprint(transcription)
+            print(f"{audio_frames=}s")
+
+        return (transcription, audio_frames)
+    
+    @classmethod
+    def IS_CHANGED(cls, *args, **kwargs):
+        print(f"{args=} {kwargs=}")
+        return get_hash(*args, *kwargs.items())
+    
+
+class BatchedTranscriptionEncoderNode:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+
+        return {
+            "required": {
+                "transcription": ('TRANSCRIPTION',),
+                "num_frames": ('INT', {'forceInput': True,}),
+                "clip": ('CLIP',)
+            },
+            "optional": {
+                "loop": ('BOOLEAN', {'default': False}),
+                "use_tags": ('BOOLEAN', {'default': False}),
+                'language': (SUPPORTED_LANGUAGES,{"default": "en"}),
+                "max_words": ('INT', {'default': 16, "min": 0, "max": 32, "display": "slider"}),
+                "frame_interpolation": ("INT", {"default": 0, "min": 0, "max": 120}),
+                "prefix": ("STRING", {"default": "masterpiece, best quality", "multiline": True, "forceInput": True}),
+                "suffix": ("STRING", {"default": "", "multiline": True, "forceInput": True}),
+                "print_output": ('BOOLEAN', {'default': True}),
+            }
+        }
+
+    OUTPUT_NODE = True
+    RETURN_TYPES = ("CONDITIONING", "STRING", "INT", "INT", "INT")
+    RETURN_NAMES = ("conditioning", "batch_prompt_text", "batch_duration_int", "num_chunks", "num_frames")
+
+    FUNCTION = "get_prompt"
+
+    CATEGORY = CATEGORY
 
     @staticmethod
-    def process_text_chunks(ichunk,text_processor,frame_interpolation,**kwargs):
+    def process_text_chunks(ichunk,text_processor,frame_interpolation,clip,**kwargs):
         i,chunk=ichunk
         frame_interpolation=int(frame_interpolation)
         text = text_processor(chunk["text"],**kwargs)
         timestamp = np.nan_to_num(np.array(chunk["timestamp"],dtype=float),nan=i)
-        index = int(timestamp[0]*frame_interpolation) if frame_interpolation>0 else i
-        prompt = f'"{index}":"{text}"' if len(text) else f'"{index}":""'
+        # index = int(timestamp[0]*frame_interpolation) if frame_interpolation>0 else i
+        # prompt = f'"{index}":"{text}"' if len(text) else f'"{index}":""'
         duration = max(abs(int(np.round(timestamp[1]-timestamp[0]))),1) # at least 1 sec duration
         if frame_interpolation>1: duration*=frame_interpolation
-        value = f'"{index}":"({duration})"'
-        return prompt,value
+        # value = f'"{index}":"({duration})"'
+        tokens = clip.tokenize(text)
+        cond, pooled = clip.encode_from_tokens(tokens, return_pooled=True)
+        
+        return text, duration, cond, pooled
     
     @staticmethod
-    def load_cache(*args):
-        results = None
-        cache_name = os.path.join(BASE_CACHE_DIR,"stt",f"{get_hash(*args)}.json")
-        os.makedirs(os.path.dirname(cache_name), exist_ok=True)
-        if os.path.isfile(cache_name): results = json.load(cache_name)
-        return results, cache_name
+    def get_spacy_model(language="en"):
+        import spacy
+        spacy_model_map = {
+            "en": "en_core_web_md",
+            "fr": "fr_core_news_md",
+            "es": "es_core_news_md",
+            "ja": "ja_core_news_md",
+            "zh": "zh_core_web_md",
+        }
+        model_name = spacy_model_map[language]
+        model_path = os.path.join(BASE_MODELS_DIR,model_name)
+        
+        if not os.path.exists(model_path):
+            subprocess.check_call([sys.executable, "-m", "spacy", "download", model_name])
+            model = spacy.load(model_name)
+            model.to_disk(model_path)
 
-    def transcribe(
-        self, pipeline, audio,
-        format_newlines_on_punctuation,loop,
-        use_tags,max_words,frame_interpolation,
-        save_transcription=False, save_chunks=False,
-        save_filename="transcription", overwrite_existing=True,
-        print_output=True,use_cache=True
+        return spacy.load(model_path)
+
+    def get_prompt(
+        self, transcription, num_frames, clip, language="en", loop=False, use_tags=False, max_words=16,
+        frame_interpolation=0, print_output=True, prefix="", suffix=""
     ):
-        print('Starting Transcription')
-
-        audio_data = bytes_to_audio(audio())
-        if use_cache: result, cache_name = self.load_cache(pipeline,audio(),loop,use_tags,max_words,frame_interpolation)
-        else: result = None
-
-        audio_frames = int(np.ceil(len(audio_data[0])/audio_data[1]))
-        whisper_model,get_spacy_model=pipeline()
-
-        if result is None:
-            audio,_ = remix_audio(audio_data,target_sr=16000,norm=True)
-            result = whisper_model(audio)
-
-        if use_cache and result:
-            with open(cache_name, "w") as f:
-                json.dump(result, f)
 
         print('Transcription Done')
-        text = result["text"]
-        total_chunks = result['chunks']
-        if format_newlines_on_punctuation:
-            punct = ['.', '?', '!']
-            for p in punct:
-                text = text.replace(f'{p} ', f'{p}\n')
+        total_chunks = transcription['chunks']
 
         if loop: # append first frame to chunk stack
             last_chunk = None
@@ -143,12 +208,12 @@ class AudioTranscriptionNode:
                 last_chunk = total_chunks[-1-i]
 
                 if "timestamp" in last_chunk and len(last_chunk["timestamp"])>=1:
-                    last_chunk["timestamp"] = np.nan_to_num(np.array(last_chunk["timestamp"],dtype=float),nan=audio_frames)
+                    last_chunk["timestamp"] = np.nan_to_num(np.array(last_chunk["timestamp"],dtype=float),nan=num_frames)
                     break
             if last_chunk is not None:
                 timestamp = last_chunk["timestamp"]
                 start_time = timestamp[-1]
-                duration = max(1,audio_frames - start_time)
+                duration = max(1,num_frames - start_time)
                 total_chunks.append({
                     "text": total_chunks[0]["text"],
                     "timestamp": (start_time,start_time+duration)
@@ -157,46 +222,36 @@ class AudioTranscriptionNode:
             num_frames+=1
 
         if use_tags:
-            spacy_model,text_processor = get_spacy_model(),extract_keywords
+            spacy_model,text_processor = self.get_spacy_model(language),extract_keywords
         else:
             spacy_model,text_processor = None,limit_sentence
 
         # split transcript into prompt based on timestamp
         # find length of each frame using timestamp
-        prompts = []
-        values = []
+        batch_prompt_text = []
+        batch_values_int = []
+        cond = []
+        pooled = []
         for ichunk in enumerate(total_chunks):
-            p, v = self.process_text_chunks(ichunk,text_processor,frame_interpolation,spacy_model=spacy_model,max_words=max_words)
-            prompts.append(p)
-            values.append(v)
+            t, d, c, pc = self.process_text_chunks(ichunk,text_processor,frame_interpolation,clip,prefix=prefix,suffix=suffix,spacy_model=spacy_model,max_words=max_words)
+            batch_prompt_text.append(t)
+            batch_values_int.append(d)
+            cond.append(c)
+            pooled.append(pc)
         
         num_chunks = len(total_chunks)
-        batch_prompt = ',\n'.join(prompts)
-        batch_values = ',\n'.join(values)
-        num_frames = int(audio_frames*frame_interpolation if frame_interpolation>0 else audio_frames)
+        final_pooled_output = torch.cat(pooled, dim=0)
+        final_conditioning = torch.cat(cond, dim=0)
+        conditioning = [[final_conditioning,{"pooled_output": final_pooled_output}]]
+        del pooled, cond
 
         if print_output:
-            pprint(total_chunks)
-            print(batch_prompt)
-            print(batch_values)
-            print(f"{num_chunks=}, {audio_frames=}s, {num_frames=}")
-            
-        if save_transcription:
-            self.save_output(save_filename + '.txt', text, overwrite_existing=overwrite_existing)
+            pprint(conditioning)
+            print(batch_prompt_text)
+            print(batch_values_int)
+            print(f"{num_chunks=}, {num_frames=}")
 
-        if save_chunks:
-            self.save_output(
-                save_filename + '_prompt.json',
-                json.dumps(json.loads("{"+batch_prompt+"}"),indent=2), # formats json with indent
-                overwrite_existing=overwrite_existing)
-            
-            self.save_output(
-                save_filename + '_values.json',
-                json.dumps(json.loads("{"+batch_values+"}"),indent=2), # formats json with indent
-                overwrite_existing=overwrite_existing)
-
-        text = text_processor(text,max_words=max_words,spacy_model=spacy_model)
-        return (text, batch_prompt, batch_values, num_chunks, num_frames)
+        return (conditioning, batch_prompt_text, batch_values_int, num_chunks, num_frames)
     
     @classmethod
     def IS_CHANGED(cls, *args, **kwargs):

custom_nodes/utils.py

@@ -150,7 +150,7 @@ class AudioBatchValueNode:
         return {
             "required": {
                 "audio": ('VHS_AUDIO',),
-                "num_frames": ('INT', {"min": 1, "max": 128, "step": 1}),
+                "num_segments": ('INT', {"min": 1, "max": 128, "step": 1, "forceInput": True}),
                 "output_min": ('FLOAT', {'default': 0., "min": -1000., "max": 1000., "step": .01}),
                 "output_max": ('FLOAT', {'default': 1., "min": 0., "max": 1000., "step": .01}),
                 "norm": (["scale","tanh","sigmoid"], {"default": "scale"}),
@@ -163,8 +163,8 @@ class AudioBatchValueNode:
             }
         }
 
-    RETURN_TYPES = ("STRING","INT")
-    RETURN_NAMES = ("batch_value_text","num_values")
+    RETURN_TYPES = ("FLOAT","INT","INT")
+    RETURN_NAMES = ("FLOAT","INT","num_values")
 
     FUNCTION = "get_frame_weights"
 
@@ -174,13 +174,13 @@ class AudioBatchValueNode:
     def get_rms(audio): # root mean squared of audio segment
         return np.sqrt(np.nanmean(audio**2))
 
-    def get_frame_weights(self, audio, num_frames, output_min, output_max, norm,
+    def get_frame_weights(self, audio, num_segments, output_min, output_max, norm,
                           silence_threshold=1000, frame_multiplier=1, print_output=False, inverse=False):
         assert output_max>=output_min, f"{output_max=} must be greater or equal to {output_min=}!"
 
         audio_data = bytes_to_audio(audio())
         audio,_ = remix_audio(audio_data,norm=True,to_int16=True)
-        num_values = int(num_frames*frame_multiplier)
+        num_values = int(num_segments*frame_multiplier)
         audio_rms = np.nan_to_num(list(map(self.get_rms,np.array_split(audio.flatten()/silence_threshold, num_values))),nan=0)
         audio_zscore = (audio_rms-audio_rms.mean())/audio_rms.std()
         output_range = output_max-output_min
@@ -203,7 +203,7 @@ class AudioBatchValueNode:
 
         batch_value = ",\n".join([f'{n}:({v})' for n,v in enumerate(x_norm)])
         if print_output: pprint(batch_value)
-        return (batch_value,num_values)
+        return (list(x_norm),list(map(int,x_norm)),num_values)
     
     @classmethod
     def IS_CHANGED(cls, *args, **kwargs):