diff --git a/custom_nodes/stt.py b/custom_nodes/stt.py index ecf5081..726d1fc 100644 --- a/custom_nodes/stt.py +++ b/custom_nodes/stt.py @@ -20,7 +20,7 @@ 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, use_sentiment=False, prefix="", suffix="", **kwargs): +def extract_keywords(text: str, max_words: int, spacy_model, use_sentiment=False, prefix="", suffix="", weights=1., **kwargs): text = text.strip() sentiment = "" if use_sentiment: @@ -42,9 +42,11 @@ def extract_keywords(text: str, max_words: int, spacy_model, use_sentiment=False except Exception as error: print(f"{text=} {error=}") - return ", ".join(filter(None,[prefix,*tags,sentiment,suffix])).strip() + tags = ", ".join(tags) + if len(tags) and weights!=1.: tags = f"({tags}:{weights:.3f})" + return ", ".join(filter(None,[prefix,tags,sentiment,suffix])).strip() -def limit_sentence(text: str, max_words: int, spacy_model, use_sentiment=False, prefix="", suffix="", **kwargs): +def limit_sentence(text: str, max_words: int, spacy_model, use_sentiment=False, prefix="", suffix="", weights=1., **kwargs): text = text.strip() sentiment = "" if use_sentiment: @@ -57,6 +59,7 @@ def limit_sentence(text: str, max_words: int, spacy_model, use_sentiment=False, elif polarity>.2: sentiment="slight smile" topn = int(max_words) if max_words>0 else len(text) if topn>0: text = " ".join(text.split()[:topn]) + if len(text) and weights!=1.: text = f"({text}:{weights:.3f})" return ", ".join(filter(None,[prefix,text,sentiment,suffix])).strip() @@ -160,7 +163,7 @@ class BatchedTranscriptionEncoderNode: "prefix": ("STRING", {"default": "masterpiece, best quality", "multiline": True, "forceInput": True}), "suffix": ("STRING", {"default": "", "multiline": True, "forceInput": True}), "print_output": ('BOOLEAN', {'default': True}), - "weights": ("FLOAT",{"default": 1.}) + "weights": ("FLOAT",{"default": 1., "step": .01}) } } @@ -173,13 +176,12 @@ class BatchedTranscriptionEncoderNode: CATEGORY = CATEGORY @staticmethod - def process_text_chunks(ichunk,frame_interpolation,clip,*,use_tags,weights,**kwargs): + def process_text_chunks(ichunk,frame_interpolation,clip,*,use_tags,**kwargs): i,chunk=ichunk frame_interpolation=int(frame_interpolation) index = i*frame_interpolation if frame_interpolation!=0 else i - text_processor = extract_keywords if use_tags else limit_sentence - text = chunk["text"] if weights==1 else f"({chunk['text']}:{weights})" - text = text_processor(text,**kwargs) + textProcessor = extract_keywords if use_tags else limit_sentence + text = textProcessor(chunk["text"],**kwargs) timestamp = np.nan_to_num(np.array(chunk["timestamp"],dtype=float),nan=i) duration = max(abs(int(np.round(timestamp[1]-timestamp[0]))),1) # at least 1 sec duration if frame_interpolation>1: duration*=frame_interpolation @@ -217,7 +219,7 @@ class BatchedTranscriptionEncoderNode: if not max_chunks: max_chunks = len(transcription['chunks']) total_chunks = transcription['chunks'][:max_chunks] - max_frames = max(max_chunks, *filter(None,np.array([chunk["timestamp"] for chunk in transcription['chunks']]).flatten())) + max_frames = max(max_chunks, *filter(None,np.array([chunk["timestamp"] for chunk in total_chunks]).flatten())) last_chunk = None