80 lines
2.4 KiB
Python
80 lines
2.4 KiB
Python
import logging
|
|
from pathlib import Path
|
|
import runpod
|
|
from runpod.serverless import os
|
|
import loaders
|
|
import file_operations
|
|
|
|
import process
|
|
|
|
logging.basicConfig(level=logging.DEBUG)
|
|
|
|
input_path: Path = Path(os.environ.get("VERBANOTE_INPUT_PATH", "/in"))
|
|
output_path: Path = Path(os.environ.get("VERBANOTE_OUTPUT_PATH", "/out"))
|
|
|
|
access_token: str = os.environ.get("VERBANOTE_HF_TOKEN", "")
|
|
|
|
loaders.prep()
|
|
diarize_pipeline = loaders.diarization(access_token)
|
|
# TODO download model into Docker image beforehand
|
|
whisper_model = loaders.whispermodel()
|
|
|
|
|
|
def handler(job):
|
|
input: dict = job["input"]
|
|
url: str | None = input.get("url")
|
|
lang: str | None = input.get("lang")
|
|
word_timestamps: str | None = input.get("word_timestamps")
|
|
|
|
if not url:
|
|
return {"error": "no file link provided"}
|
|
|
|
try:
|
|
audiofile = loaders.audiofile(url, input_path=input_path)
|
|
except Exception:
|
|
return {"error": "audiofile import failed"}
|
|
|
|
diarized = process.diarize(audiofile, diarize_pipeline, output_path)
|
|
diarized_groups = process.save_diarized_audio_files(
|
|
diarized, audiofile, output_path
|
|
)
|
|
uploaded_diarization: str = file_operations.upload_to_oxo(file=diarized, expires=1)
|
|
process.transcribe(
|
|
model=whisper_model,
|
|
diarized_groups=diarized_groups,
|
|
files_path=output_path,
|
|
lang=lang or "en",
|
|
word_timestamps=word_timestamps or True,
|
|
)
|
|
transcription: process.TxtTranscription = process.output_txt(
|
|
diarized_groups, output_path
|
|
)
|
|
uploaded_transcription: str = file_operations.upload_to_oxo(
|
|
file=transcription.file, expires=1
|
|
)
|
|
|
|
return {
|
|
"audiofile": str(audiofile),
|
|
"diarization_url": uploaded_diarization,
|
|
"diarization": diarized_groups,
|
|
"transcription_url": uploaded_transcription,
|
|
"transcription_text": transcription.text,
|
|
}
|
|
|
|
|
|
# speakers = {
|
|
# # speaker, textboxcolor, speaker color
|
|
# "SPEAKER_00": ("SPEAKER00", "white", "darkgreen"),
|
|
# "SPEAKER_01": ("SPEAKER01", "white", "darkorange"),
|
|
# "SPEAKER_02": ("SPEAKER02", "white", "darkred"),
|
|
# "SPEAKER_03": ("SPEAKER03", "white", "darkblue"),
|
|
# "SPEAKER_04": ("SPEAKER04", "white", "darkyellow"),
|
|
# "SPEAKER_05": ("SPEAKER05", "white", "lightgreen"),
|
|
# "SPEAKER_06": ("SPEAKER06", "white", "lightred"),
|
|
# "SPEAKER_07": ("SPEAKER07", "white", "lightblue"),
|
|
# }
|
|
|
|
|
|
if __name__ == "__main__":
|
|
runpod.serverless.start({"handler": handler})
|