diff --git a/__init__.py b/__init__.py index bea6778..d5df604 100644 --- a/__init__.py +++ b/__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)" } \ No newline at end of file diff --git a/custom_nodes/loaders.py b/custom_nodes/loaders.py index 087f9a1..47775ba 100644 --- a/custom_nodes/loaders.py +++ b/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 diff --git a/custom_nodes/settings/__init__.py b/custom_nodes/settings/__init__.py index 3214723..6d4ef31 100644 --- a/custom_nodes/settings/__init__.py +++ b/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 \ No newline at end of file +# SERVERS = get_servers() +# RVC_INFERENCE_URL = SERVERS.RVC_INFERENCE_URL +# UVR_INFERENCE_URL = SERVERS.UVR_INFERENCE_URL \ No newline at end of file diff --git a/custom_nodes/stt.py b/custom_nodes/stt.py index fbcb4bd..f668ecc 100644 --- a/custom_nodes/stt.py +++ b/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): diff --git a/custom_nodes/utils.py b/custom_nodes/utils.py index 6cb525d..cfde510 100644 --- a/custom_nodes/utils.py +++ b/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):