204 lines
8.0 KiB
Python
204 lines
8.0 KiB
Python
import json
|
|
import os
|
|
from pprint import pprint
|
|
|
|
import numpy as np
|
|
|
|
from ..lib import BASE_CACHE_DIR
|
|
from .utils import increment_filename_no_overwrite
|
|
from ..lib.utils import get_hash
|
|
import folder_paths
|
|
from ..lib.audio import bytes_to_audio, remix_audio
|
|
import textacy
|
|
from textacy.extract import keyterms as kt
|
|
|
|
temp_path = folder_paths.get_temp_directory()
|
|
|
|
def extract_keywords(text: str, max_words: int, spacy_model, **kwargs):
|
|
text = text.strip()
|
|
doc = textacy.make_spacy_doc(text, lang=spacy_model)
|
|
topn = int(max_words) if max_words>0 else len(text)
|
|
include_pos = ["NOUN","ADJ","PROPN","VERB","NUM","ADP"]
|
|
ngrams = [1,2]
|
|
tags = []
|
|
try:
|
|
terms = kt.sgrank(doc, ngrams=ngrams, normalize="lower", topn=topn, include_pos=include_pos)
|
|
tags = list(map(lambda v:v[0],sorted(terms,key=lambda v:v[1],reverse=True)[:topn]))
|
|
except Exception as error:
|
|
print(f"{text=} {error=}")
|
|
|
|
return ", ".join(tags).strip()
|
|
|
|
def limit_sentence(text: str, max_words: int, **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
|
|
|
|
class AudioTranscriptionNode:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
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})
|
|
},
|
|
"optional": {
|
|
"save_transcription": ('BOOLEAN', {'default': False}),
|
|
"save_chunks": ('BOOLEAN', {'default': False}),
|
|
"save_filename": ('STRING', {'default': 'transcription'}),
|
|
"overwrite_existing": ('BOOLEAN', {'default': True}),
|
|
"print_output": ('BOOLEAN', {'default': True}),
|
|
"use_cache": ("BOOLEAN",{"default": True})
|
|
}
|
|
}
|
|
|
|
OUTPUT_NODE = True
|
|
RETURN_TYPES = ("STRING", "STRING", "STRING", "INT", "INT")
|
|
RETURN_NAMES = ("transcription", "batch_prompt_text", "batch_values_text", "num_chunks","num_frames")
|
|
|
|
FUNCTION = "transcribe"
|
|
|
|
CATEGORY = "🌺RVC-Studio/stt"
|
|
|
|
@staticmethod
|
|
def save_output(filename,text,overwrite_existing=False):
|
|
base_output_dir = folder_paths.get_output_directory()
|
|
assert os.path.exists(base_output_dir), f"Output directory {base_output_dir} does not exist"
|
|
output_dir = os.path.join(base_output_dir, 'transcriptions')
|
|
|
|
file_path = os.path.join(output_dir, filename)
|
|
dir_name = os.path.dirname(file_path)
|
|
os.makedirs(dir_name, exist_ok=True)
|
|
if os.path.exists(file_path) and not overwrite_existing:
|
|
file_path = increment_filename_no_overwrite(file_path)
|
|
with open(file_path, 'w') as f:
|
|
f.write(text)
|
|
|
|
@staticmethod
|
|
def process_text_chunks(ichunk,text_processor,frame_interpolation,**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}":""'
|
|
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
|
|
|
|
@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 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
|
|
):
|
|
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')
|
|
|
|
if loop: # append first frame to chunk stack
|
|
last_chunk = None
|
|
for i in range(len(total_chunks)):
|
|
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)
|
|
break
|
|
if last_chunk is not None:
|
|
timestamp = last_chunk["timestamp"]
|
|
start_time = timestamp[-1]
|
|
duration = max(1,audio_frames - start_time)
|
|
total_chunks.append({
|
|
"text": total_chunks[0]["text"],
|
|
"timestamp": (start_time,start_time+duration)
|
|
})
|
|
else: # adds 1s frame to prevent audio from stopping early
|
|
num_frames+=1
|
|
|
|
if use_tags:
|
|
spacy_model,text_processor = get_spacy_model(),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 = []
|
|
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)
|
|
|
|
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)
|
|
|
|
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)
|
|
|
|
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)
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, *args, **kwargs):
|
|
print(f"{args=} {kwargs=}")
|
|
return get_hash(*args, *kwargs.items()) |