oval
6851b26923
FastAPI + Celery + Svelte app for English→Danish video translation with voice cloning. Includes diarization, STT, translation, TTS, and audio mixing pipeline.
85 lines
2.3 KiB
Python
85 lines
2.3 KiB
Python
import logging
|
|
import torch
|
|
from typing import Optional
|
|
from pyannote.audio import Pipeline
|
|
from ..config import settings
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_diarization_pipeline: Optional[Pipeline] = None
|
|
|
|
|
|
def get_diarization_pipeline() -> Pipeline:
|
|
global _diarization_pipeline
|
|
if _diarization_pipeline is None:
|
|
import os
|
|
|
|
hf_token = settings.hf_token or os.environ.get("HUGGINGFACE_TOKEN")
|
|
logger.info("Loading diarization model on %s", settings.diarization_device)
|
|
_diarization_pipeline = Pipeline.from_pretrained(
|
|
settings.diarization_model,
|
|
use_auth_token=hf_token,
|
|
)
|
|
_diarization_pipeline.to(torch.device(settings.diarization_device))
|
|
return _diarization_pipeline
|
|
|
|
|
|
def diarize(audio_path: str) -> list[dict]:
|
|
pipeline = get_diarization_pipeline()
|
|
diarization = pipeline(audio_path)
|
|
|
|
segments = []
|
|
for turn, _, speaker in diarization.itertracks(yield_label=True):
|
|
segments.append(
|
|
{
|
|
"speaker_label": speaker,
|
|
"start_time": turn.start,
|
|
"end_time": turn.end,
|
|
}
|
|
)
|
|
|
|
segments.sort(key=lambda s: s["start_time"])
|
|
segments = _merge_adjacent_same_speaker(segments)
|
|
segments = _relabel_speakers(segments)
|
|
|
|
return segments
|
|
|
|
|
|
def _merge_adjacent_same_speaker(
|
|
segments: list[dict], max_gap: float = 0.5
|
|
) -> list[dict]:
|
|
if not segments:
|
|
return []
|
|
|
|
merged = [segments[0]]
|
|
for seg in segments[1:]:
|
|
prev = merged[-1]
|
|
if (
|
|
prev["speaker_label"] == seg["speaker_label"]
|
|
and (seg["start_time"] - prev["end_time"]) <= max_gap
|
|
):
|
|
prev["end_time"] = seg["end_time"]
|
|
else:
|
|
merged.append(seg)
|
|
return merged
|
|
|
|
|
|
def _relabel_speakers(segments: list[dict]) -> list[dict]:
|
|
label_map = {}
|
|
counter = 0
|
|
for seg in segments:
|
|
spk = seg["speaker_label"]
|
|
if spk not in label_map:
|
|
label = chr(65 + counter)
|
|
label_map[spk] = f"Speaker {label}"
|
|
counter += 1
|
|
seg["speaker_label"] = label_map[spk]
|
|
return segments
|
|
|
|
|
|
def unload_model():
|
|
global _diarization_pipeline
|
|
_diarization_pipeline = None
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|