feat(backend): setup wizard router, translation engines, export options, client-disconnect handling
- Add setup router (backend/api/routers/setup.py) for first-run wizard: system checks, engine probes, model downloads with progress - Add translation engines service with pluggable backends - Add utils/hf_progress for HuggingFace download progress streaming - Add PyInstaller runtime hooks (numpy compat, torch compiler disable) - Global exception handler short-circuits h11 LocalProtocolError and Starlette ClientDisconnect with HTTP 499 to silence noisy stack traces when users scrub or cancel video mid-stream - /dub/download-mp3 accepts bitrate query param (clamped 64–320kbps) - Refactor ASR/TTS backends, dub pipeline, engine management - Update backend.spec for PyInstaller packaging - Bump pyproject version to 0.2.0; refresh uv.lock Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
debpalash
co-authored by
Claude Opus 4.7
Claude Opus 4.7
parent
d1fd0e5fcb
commit
994c6cf065
+89
-9
@@ -6,8 +6,16 @@
|
||||
# on the heavy ML deps because PyInstaller's static analysis misses their
|
||||
# runtime-imported submodules, C extensions, and data files.
|
||||
#
|
||||
# Cross-platform: targets mac-ARM, mac-Intel, Linux x64, Windows x64. mlx
|
||||
# deps are gated on mac-ARM (sys.platform=='darwin' + machine=='arm64') so
|
||||
# PyInstaller on other hosts doesn't blow up trying to find mlx wheels.
|
||||
#
|
||||
# Run: uv run pyinstaller backend.spec --noconfirm --clean
|
||||
from PyInstaller.utils.hooks import collect_data_files, collect_all
|
||||
import platform
|
||||
import sys
|
||||
from PyInstaller.utils.hooks import collect_data_files, collect_all, collect_submodules
|
||||
|
||||
IS_MAC_ARM = sys.platform == "darwin" and platform.machine() == "arm64"
|
||||
|
||||
datas = []
|
||||
binaries = []
|
||||
@@ -27,22 +35,84 @@ hiddenimports = [
|
||||
'torch', 'torchaudio', 'soundfile', 'scipy', 'numpy',
|
||||
'numpy.random._pickle',
|
||||
|
||||
# Cross-platform primary ASR — WhisperX (faster-whisper + wav2vec2
|
||||
# alignment) is the default on every platform. faster-whisper is the
|
||||
# transcription engine; WhisperX adds forced alignment for ±10-30 ms
|
||||
# word timing, which directly improves dub lip-sync. Both backends are
|
||||
# registered in asr_backend.py; the user can switch via Settings.
|
||||
'whisperx', 'whisperx.alignment', 'whisperx.asr', 'whisperx.diarize',
|
||||
'whisperx.vad', 'whisperx.audio', 'whisperx.utils',
|
||||
'faster_whisper', 'faster_whisper.transcribe', 'faster_whisper.audio',
|
||||
'faster_whisper.utils', 'faster_whisper.tokenizer', 'faster_whisper.vad',
|
||||
'ctranslate2',
|
||||
|
||||
# Lightweight English TTS tier — ONNX-based, cross-platform. The ONNX
|
||||
# Runtime wheels ship platform-specific .so/.dll/.dylib which collect_all
|
||||
# picks up; the kittentts Python package is pure Python but has a couple
|
||||
# of asset files the bundler needs to include.
|
||||
'kittentts', 'onnxruntime',
|
||||
|
||||
# Pipeline
|
||||
'yt_dlp', 'demucs', 'demucs.separate',
|
||||
|
||||
# OmniVoice's own package
|
||||
'omnivoice', 'omnivoice.models', 'omnivoice.models.omnivoice',
|
||||
|
||||
# MLX Whisper on Apple Silicon (primary ASR path)
|
||||
'mlx', 'mlx_whisper',
|
||||
]
|
||||
|
||||
if IS_MAC_ARM:
|
||||
# MLX Whisper on Apple Silicon (optional speedup path). mlx's pure-Python
|
||||
# submodules (nn, utils, …) are imported lazily by mlx_whisper at
|
||||
# transcribe time and the plain dep tracer misses them. We deliberately
|
||||
# do NOT collect_all() mlx because that double-registers mlx.core with
|
||||
# nanobind and the binary aborts on the first mlx.core touch.
|
||||
hiddenimports.append('mlx_whisper')
|
||||
# mlx-audio engine multiplexer — Kokoro / CSM / Dia / Qwen3-TTS /
|
||||
# Chatterbox / MeloTTS / OuteTTS / … — gives mac-ARM users a rich
|
||||
# engine picker. Like mlx_whisper it's mac-ARM-only; also like
|
||||
# mlx_whisper we list it here but avoid collect_all() because it
|
||||
# depends on the same nanobind-registered mlx.core.
|
||||
hiddenimports += [
|
||||
'mlx_audio', 'mlx_audio.tts', 'mlx_audio.tts.utils',
|
||||
'mlx_audio.tts.models', 'mlx_audio.tts.generate',
|
||||
'mlx_audio.stt', 'mlx_audio.codec',
|
||||
]
|
||||
|
||||
# Note: we deliberately DON'T enumerate mlx submodules here. Any variant of
|
||||
# `collect_submodules('mlx')` or `collect_all('mlx')` — even filtered to
|
||||
# exclude mlx.core — reliably re-triggers the nanobind duplicate-key error
|
||||
# the first time anything imports mlx.core ("refusing to add duplicate key
|
||||
# 'cpu' to enumeration mlx.core.DeviceType"). Shipping without mlx in the
|
||||
# frozen bundle leaves mlx-whisper unavailable; asr_backend falls back to
|
||||
# pytorch-whisper (slower but functional on Apple Silicon). Revisit once
|
||||
# we have a minimal repro or a PyInstaller hook specifically for mlx.
|
||||
|
||||
# The nuclear option on heavy ML libs — pull every submodule, C ext, and
|
||||
# data file. Cost: bigger bundle. Benefit: we don't ship a binary that
|
||||
# ImportErrors the first time a user hits a code path.
|
||||
for pkg in ('torch', 'torchaudio', 'soundfile', 'scipy', 'numpy',
|
||||
'omnivoice', 'mlx', 'mlx_whisper', 'demucs', 'yt_dlp',
|
||||
'fastapi', 'uvicorn'):
|
||||
# Note: 'mlx' is intentionally NOT in this list. Calling collect_all('mlx')
|
||||
# alongside collect_all('mlx_whisper') causes the nanobind binding init to
|
||||
# run twice in the frozen bundle, crashing with
|
||||
# "Critical nanobind error: refusing to add duplicate key 'cpu'
|
||||
# to enumeration 'mlx.core.DeviceType'!"
|
||||
# the first time anything imports mlx.core. mlx_whisper already depends on
|
||||
# mlx and PyInstaller's dep tracer pulls the needed mlx submodules + the .so.
|
||||
_collect_pkgs = [
|
||||
'torch', 'torchaudio', 'soundfile', 'scipy', 'numpy',
|
||||
'omnivoice', 'demucs', 'yt_dlp', 'fastapi', 'uvicorn',
|
||||
# Primary cross-platform ASR. collect_all pulls CTranslate2's bundled
|
||||
# .so/.dylib/.dll plus its compiled kernel data. WhisperX ships its own
|
||||
# pure-Python code + some asset files (e.g. language metadata).
|
||||
'whisperx', 'faster_whisper', 'ctranslate2',
|
||||
# ONNX-based lightweight TTS. onnxruntime's collect_all pulls the
|
||||
# platform-appropriate .so/.dll/.dylib + CUDA providers when present.
|
||||
'kittentts', 'onnxruntime',
|
||||
]
|
||||
if IS_MAC_ARM:
|
||||
# Only attempt mlx_whisper collection on mac-ARM — no wheels exist for
|
||||
# Linux/Windows/mac-Intel, so collect_all would fail on CI for those.
|
||||
_collect_pkgs.append('mlx_whisper')
|
||||
|
||||
for pkg in _collect_pkgs:
|
||||
try:
|
||||
tmp_datas, tmp_binaries, tmp_hidden = collect_all(pkg)
|
||||
datas += tmp_datas
|
||||
@@ -69,10 +139,20 @@ a = Analysis(
|
||||
hiddenimports=hiddenimports,
|
||||
hookspath=[],
|
||||
hooksconfig={},
|
||||
runtime_hooks=[],
|
||||
runtime_hooks=[
|
||||
'backend/hooks/pyi_rth_numpy_compat.py',
|
||||
'backend/hooks/pyi_rth_torch_compiler_disable.py',
|
||||
],
|
||||
excludes=[
|
||||
# Desktop-only bloat we don't need inside the frozen backend.
|
||||
# Desktop-only bloat the frozen backend never uses.
|
||||
'tkinter', 'matplotlib', 'PIL.ImageQt', 'PyQt5', 'PyQt6',
|
||||
# CUDA / NVIDIA wheels on Apple Silicon — saves ~2 GB.
|
||||
# When we add a Windows/Linux CUDA build, remove these per-target.
|
||||
'nvidia', 'nvidia.cublas', 'nvidia.cudnn', 'nvidia.cuda_runtime',
|
||||
'nvidia.cuda_nvrtc', 'nvidia.nccl', 'nvidia.nvtx',
|
||||
'nvidia.curand', 'nvidia.cusolver', 'nvidia.cusparse',
|
||||
'nvidia.cufft', 'nvidia.cuda_cupti',
|
||||
'triton', 'flash_attn',
|
||||
],
|
||||
noarchive=False,
|
||||
optimize=0,
|
||||
|
||||
@@ -247,10 +247,16 @@ async def dub_ingest_url(req: DubIngestUrlRequest):
|
||||
os.makedirs(job_dir, exist_ok=True)
|
||||
|
||||
task_id = f"prep_{job_id}"
|
||||
source = {
|
||||
"kind": "url",
|
||||
"url": url,
|
||||
"fetch_subs": bool(req.fetch_subs),
|
||||
"sub_langs": req.sub_langs or None,
|
||||
}
|
||||
await task_manager.add_task(
|
||||
task_id, "prep",
|
||||
_ingest_gen, job_id, job_dir,
|
||||
{"kind": "url", "url": url}, None,
|
||||
source, None,
|
||||
)
|
||||
return JSONResponse(
|
||||
status_code=202,
|
||||
@@ -284,8 +290,13 @@ async def dub_transcribe_stream(job_id: str):
|
||||
if not asr_audio_target or not os.path.exists(asr_audio_target):
|
||||
raise HTTPException(status_code=404, detail="No audio available for transcription")
|
||||
|
||||
use_mlx = torch.backends.mps.is_available()
|
||||
asr_model = os.environ.get("ASR_MODEL", "mlx-community/whisper-large-v3-mlx")
|
||||
# ASR routing now goes through services.asr_backend so the active engine
|
||||
# (WhisperX by default, with faster-whisper / mlx / pytorch fallbacks) is
|
||||
# used consistently across platforms. dub pipelines specifically benefit
|
||||
# from WhisperX's wav2vec2 alignment (±10-30 ms word timing vs Whisper's
|
||||
# ±100-300 ms) — critical for lip-sync quality.
|
||||
from services.asr_backend import get_active_asr_backend
|
||||
_asr_backend = get_active_asr_backend(asr_pipe=getattr(_model, "_asr_pipe", None))
|
||||
scene_cuts = job.get("scene_cuts") or []
|
||||
|
||||
async def gen():
|
||||
@@ -312,6 +323,7 @@ async def dub_transcribe_stream(job_id: str):
|
||||
all_segments: list[dict] = []
|
||||
detected_lang = None
|
||||
next_seg_id = 0
|
||||
chunk_errors: list[str] = []
|
||||
|
||||
for i in range(chunks_n):
|
||||
if job.get("aborted"):
|
||||
@@ -326,44 +338,31 @@ async def dub_transcribe_stream(job_id: str):
|
||||
continue
|
||||
|
||||
def _transcribe_chunk(arr=chunk_arr, offset=t0, local_sr=sr):
|
||||
# Route through the active backend (WhisperX by default).
|
||||
# Backends all take a file path, so write the chunk first.
|
||||
try:
|
||||
if use_mlx:
|
||||
import mlx_whisper
|
||||
tmp = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
|
||||
tmp.close()
|
||||
try:
|
||||
sf.write(tmp.name, arr, local_sr)
|
||||
r = mlx_whisper.transcribe(
|
||||
tmp.name, path_or_hf_repo=asr_model, word_timestamps=True,
|
||||
)
|
||||
r = _asr_backend.transcribe(tmp.name, word_timestamps=True)
|
||||
finally:
|
||||
try: os.remove(tmp.name)
|
||||
except OSError: pass
|
||||
shifted = []
|
||||
for seg in r.get("segments", []) or []:
|
||||
shifted.append({
|
||||
"text": seg.get("text", ""),
|
||||
"timestamp": (float(seg.get("start", 0.0)) + offset,
|
||||
float(seg.get("end", 0.0)) + offset),
|
||||
})
|
||||
return {"chunks": shifted, "language": r.get("language")}
|
||||
else:
|
||||
r = _model._asr_pipe(
|
||||
{"array": arr, "sampling_rate": local_sr},
|
||||
return_timestamps=True, chunk_length_s=15, batch_size=1,
|
||||
)
|
||||
shifted = []
|
||||
for c in (r.get("chunks", []) if isinstance(r, dict) else []):
|
||||
for c in r.get("chunks", []) or []:
|
||||
ts = c.get("timestamp", (0.0, 0.0)) or (0.0, 0.0)
|
||||
a0 = (ts[0] if ts[0] is not None else 0.0) + offset
|
||||
a1 = (ts[1] if ts[1] is not None else 0.0) + offset
|
||||
shifted.append({"text": c.get("text", ""), "timestamp": (a0, a1)})
|
||||
return {"chunks": shifted, "language": r.get("language") if isinstance(r, dict) else None}
|
||||
return {"chunks": shifted, "language": r.get("language")}
|
||||
except Exception as e:
|
||||
logger.exception("chunk transcribe failed")
|
||||
logger.exception("chunk transcribe failed (backend=%s)", _asr_backend.id)
|
||||
return {"chunks": [], "language": None, "error": str(e)}
|
||||
|
||||
part = await loop.run_in_executor(_gpu_pool, _transcribe_chunk)
|
||||
if part.get("error"):
|
||||
chunk_errors.append(part["error"])
|
||||
if detected_lang is None and part.get("language"):
|
||||
detected_lang = part["language"]
|
||||
chunk_segs = segment_transcript(part, duration=t1, scene_cuts=scene_cuts)
|
||||
@@ -393,6 +392,34 @@ async def dub_transcribe_stream(job_id: str):
|
||||
yield _sse_event("aborted", {})
|
||||
return
|
||||
|
||||
# Empty-transcription guard: if every chunk came back with zero
|
||||
# segments we can't proceed to diarization/clone extraction. Emit an
|
||||
# actionable error so the UI can surface a Retry instead of silently
|
||||
# landing in an empty editor. Commonly caused by a first-run model
|
||||
# download failure, a PyTorch 2.6 weights_only regression inside
|
||||
# whisperx's VAD load, or an unsupported audio format.
|
||||
if not all_segments:
|
||||
# Deduplicate while preserving order so one root cause doesn't
|
||||
# repeat N times in the UI toast.
|
||||
seen = set()
|
||||
uniq: list[str] = []
|
||||
for msg in chunk_errors:
|
||||
if msg and msg not in seen:
|
||||
seen.add(msg)
|
||||
uniq.append(msg)
|
||||
if uniq:
|
||||
detail = "Transcription produced no segments. " + " | ".join(uniq[:3])
|
||||
else:
|
||||
detail = (
|
||||
"Transcription produced no segments. The audio may be silent, "
|
||||
"too short, or in an unsupported format. Try re-uploading or "
|
||||
"check that the source has an audible speech track."
|
||||
)
|
||||
logger.error("transcribe yielded 0 segments (job=%s): %s", job_id, detail)
|
||||
yield _sse_event("error", {"detail": detail, "retryable": True})
|
||||
yield _sse_event("done", {})
|
||||
return
|
||||
|
||||
def _diarize():
|
||||
diar_pipe = get_diarization_pipeline()
|
||||
try:
|
||||
@@ -481,37 +508,28 @@ async def dub_transcribe(job_id: str):
|
||||
|
||||
detected_lang = None
|
||||
|
||||
if torch.backends.mps.is_available():
|
||||
# Route through services.asr_backend — picks WhisperX / faster-whisper
|
||||
# / mlx / pytorch based on what's installed + user preference. Works
|
||||
# identically on all platforms; the older mlx-vs-pytorch branching
|
||||
# here duplicated the logic in asr_backend.py and skipped WhisperX.
|
||||
from services.asr_backend import get_active_asr_backend
|
||||
_asr = get_active_asr_backend(asr_pipe=getattr(_model, "_asr_pipe", None))
|
||||
try:
|
||||
import mlx_whisper
|
||||
asr_model = os.environ.get("ASR_MODEL", "mlx-community/whisper-large-v3-mlx")
|
||||
logger.info(f"Transcribing via MLX CoreML Engine ({asr_model})...")
|
||||
result = mlx_whisper.transcribe(
|
||||
asr_audio_target,
|
||||
path_or_hf_repo=asr_model,
|
||||
word_timestamps=True
|
||||
)
|
||||
logger.info("Transcribing full audio via %s ...", _asr.id)
|
||||
result = _asr.transcribe(asr_audio_target, word_timestamps=True)
|
||||
detected_lang = result.get("language")
|
||||
|
||||
if "segments" in result:
|
||||
result["chunks"] = []
|
||||
for seg in result["segments"]:
|
||||
result["chunks"].append({
|
||||
"text": seg["text"],
|
||||
"timestamp": (seg["start"], seg["end"])
|
||||
})
|
||||
except Exception as e:
|
||||
logger.error(f"MLX Whisper failed, falling back to PyTorch: {e}")
|
||||
audio_np, sr = sf.read(asr_audio_target, dtype="float32")
|
||||
if audio_np.ndim > 1: audio_np = audio_np.mean(axis=1)
|
||||
bs = 16 if torch.cuda.is_available() else 2
|
||||
result = _model._asr_pipe({"array": audio_np, "sampling_rate": sr}, return_timestamps=True, chunk_length_s=15, batch_size=bs)
|
||||
detected_lang = (result.get("language") if isinstance(result, dict) else None)
|
||||
else:
|
||||
logger.error("ASR backend %s failed: %s", _asr.id, e)
|
||||
# Last-resort fallback — in-memory pytorch whisper via the TTS
|
||||
# model's pipeline. Guaranteed present since the TTS model is
|
||||
# already loaded to reach this code path.
|
||||
audio_np, sr = sf.read(asr_audio_target, dtype="float32")
|
||||
if audio_np.ndim > 1: audio_np = audio_np.mean(axis=1)
|
||||
bs = 16 if torch.cuda.is_available() else 1
|
||||
result = _model._asr_pipe({"array": audio_np, "sampling_rate": sr}, return_timestamps=True, chunk_length_s=15, batch_size=bs)
|
||||
result = _model._asr_pipe(
|
||||
{"array": audio_np, "sampling_rate": sr},
|
||||
return_timestamps=True, chunk_length_s=15, batch_size=bs,
|
||||
)
|
||||
detected_lang = (result.get("language") if isinstance(result, dict) else None)
|
||||
|
||||
job["source_lang"] = (detected_lang or "en").split("_")[0][:2].lower()
|
||||
|
||||
@@ -158,6 +158,36 @@ async def dub_list_tracks(job_id: str):
|
||||
return {"tracks": job.get("dubbed_tracks", {})}
|
||||
|
||||
|
||||
def _write_burn_srt(job: dict, exports_dir: str, stamp: str, dual: bool) -> str | None:
|
||||
"""Build a temp SRT from job segments for use with ffmpeg's subtitles filter.
|
||||
|
||||
Returned path is already ffmpeg-filter-safe (plain ASCII basename under exports_dir).
|
||||
Returns None if there are no segments to render.
|
||||
"""
|
||||
segments = job.get("segments", [])
|
||||
if not segments:
|
||||
return None
|
||||
lines = []
|
||||
for i, seg in enumerate(segments):
|
||||
lines.append(str(i + 1))
|
||||
lines.append(f"{_format_srt_time(seg['start'])} --> {_format_srt_time(seg['end'])}")
|
||||
lines.append(_pick_subtitle_text(seg, dual))
|
||||
lines.append("")
|
||||
sub_path = os.path.join(exports_dir, f"burn_subs_{stamp}.srt")
|
||||
with open(sub_path, "w", encoding="utf-8") as f:
|
||||
f.write("\n".join(lines))
|
||||
return sub_path
|
||||
|
||||
|
||||
def _ffmpeg_filter_escape(path: str) -> str:
|
||||
"""Escape a path for use inside an ffmpeg filter value (subtitles=...).
|
||||
|
||||
ffmpeg's filter parser treats `:` as an option separator and `\\`, `'` specially.
|
||||
Backslashes first, then colons, then single quotes.
|
||||
"""
|
||||
return path.replace("\\", "\\\\").replace(":", "\\:").replace("'", "\\'")
|
||||
|
||||
|
||||
@router.get("/dub/download/{job_id}")
|
||||
@router.get("/dub/download/{job_id}/{filename}")
|
||||
async def dub_download(
|
||||
@@ -166,6 +196,8 @@ async def dub_download(
|
||||
default_track: str = Query("original"),
|
||||
include_tracks: str = Query("", description="Comma-separated list of tracks to include (e.g. 'original,de,es'). Empty = include all."),
|
||||
save_path: str = Query("", description="Absolute destination path. If set, mux output is copied there and JSON returned instead of FileResponse."),
|
||||
burn_subs: bool = Query(False, description="Burn subtitles into the video stream (forces re-encode). Uses dual-subtitle layout when dual=1."),
|
||||
dual: bool = Query(False, description="When burn_subs=1, render translated on top of italicised original."),
|
||||
):
|
||||
job = _get_job(job_id)
|
||||
if not job:
|
||||
@@ -193,6 +225,8 @@ async def dub_download(
|
||||
output_path = os.path.join(exports_dir, f"dubbed_video_{stamp}.mp4")
|
||||
ffmpeg = find_ffmpeg()
|
||||
|
||||
sub_path = _write_burn_srt(job, exports_dir, stamp, dual) if burn_subs else None
|
||||
|
||||
cmd = [ffmpeg, "-i", video_path]
|
||||
input_idx = 1
|
||||
|
||||
@@ -209,24 +243,37 @@ async def dub_download(
|
||||
tracks_to_process.append({"lang_code": lang_code, "idx": input_idx, "info": track_info})
|
||||
input_idx += 1
|
||||
|
||||
cmd += ["-map", "0:v:0"]
|
||||
filter_parts: list[str] = []
|
||||
video_map = "0:v:0"
|
||||
if sub_path:
|
||||
esc = _ffmpeg_filter_escape(sub_path)
|
||||
filter_parts.append(f"[0:v]subtitles='{esc}'[vout]")
|
||||
video_map = "[vout]"
|
||||
|
||||
cmd += ["-map", video_map]
|
||||
if include_original:
|
||||
cmd += ["-map", "0:a:0"]
|
||||
|
||||
if bg_idx is not None:
|
||||
filters = []
|
||||
for i, t in enumerate(tracks_to_process):
|
||||
out_label = f"[aout{i}]"
|
||||
filters.append(f"[{bg_idx}:a][{t['idx']}:a]amix=inputs=2:duration=longest:dropout_transition=2:weights=0.8 1.2{out_label}")
|
||||
filter_parts.append(f"[{bg_idx}:a][{t['idx']}:a]amix=inputs=2:duration=longest:dropout_transition=2:weights=0.8 1.2{out_label}")
|
||||
t["out_label"] = out_label
|
||||
cmd += ["-filter_complex", ";".join(filters)]
|
||||
for t in tracks_to_process:
|
||||
cmd += ["-map", t["out_label"]]
|
||||
else:
|
||||
for t in tracks_to_process:
|
||||
cmd += ["-map", f"{t['idx']}:a:0"]
|
||||
|
||||
cmd += ["-c:v", "copy", "-c:a", "aac", "-b:a", "192k"]
|
||||
if filter_parts:
|
||||
cmd += ["-filter_complex", ";".join(filter_parts)]
|
||||
|
||||
# Burning subs forces a video re-encode; stream-copy otherwise to keep mux cheap.
|
||||
if sub_path:
|
||||
cmd += ["-c:v", "libx264", "-preset", "medium", "-crf", "20", "-pix_fmt", "yuv420p"]
|
||||
else:
|
||||
cmd += ["-c:v", "copy"]
|
||||
cmd += ["-c:a", "aac", "-b:a", "192k"]
|
||||
|
||||
audio_stream_idx = 0
|
||||
if include_original:
|
||||
@@ -593,7 +640,7 @@ async def dub_export_segments_zip(job_id: str):
|
||||
|
||||
@router.get("/dub/download-mp3/{job_id}")
|
||||
@router.get("/dub/download-mp3/{job_id}/{filename}")
|
||||
async def dub_download_mp3(job_id: str, lang: str = Query(None), preserve_bg: bool = Query(True), save_path: str = Query("")):
|
||||
async def dub_download_mp3(job_id: str, lang: str = Query(None), preserve_bg: bool = Query(True), save_path: str = Query(""), bitrate: str = Query("192k")):
|
||||
job = _get_job(job_id)
|
||||
if not job:
|
||||
raise HTTPException(status_code=404, detail="Job not found")
|
||||
@@ -632,7 +679,15 @@ async def dub_download_mp3(job_id: str, lang: str = Query(None), preserve_bg: bo
|
||||
logger.error(f"Failed to mix audio for MP3: {e}")
|
||||
|
||||
mp3_path = os.path.join(exports_dir, f"dubbed_{lang_label}_{stamp}.mp3")
|
||||
cmd = [ffmpeg, "-i", source_path, "-codec:a", "libmp3lame", "-b:a", "192k", "-y", mp3_path]
|
||||
# Accept '128', '192k' etc. — normalize to ffmpeg's 'Nk' form and clamp
|
||||
# to a sensible range so a malformed value can't stall encoding.
|
||||
_br = str(bitrate or "192k").lower().rstrip("k") or "192"
|
||||
try:
|
||||
_br_int = max(64, min(int(_br), 320))
|
||||
except ValueError:
|
||||
_br_int = 192
|
||||
br_arg = f"{_br_int}k"
|
||||
cmd = [ffmpeg, "-i", source_path, "-codec:a", "libmp3lame", "-b:a", br_arg, "-y", mp3_path]
|
||||
try:
|
||||
rc, _, stderr = await run_ffmpeg(cmd, timeout=600.0)
|
||||
if rc != 0:
|
||||
|
||||
@@ -219,7 +219,23 @@ async def dub_translate(req: TranslateRequest):
|
||||
translated = await loop.run_in_executor(_cpu_pool, _translate_argos)
|
||||
return {"translated": translated, "target_lang": req.target_lang, "source_lang": src_lang}
|
||||
|
||||
# Legacy / API Deep_Translator logic
|
||||
# Legacy / API Deep_Translator logic.
|
||||
# Preflight the optional `deep_translator` dep once so we fail with a
|
||||
# single actionable error instead of N identical per-segment
|
||||
# ModuleNotFoundErrors that flood the UI's error badge.
|
||||
try:
|
||||
import deep_translator # noqa: F401
|
||||
except ImportError:
|
||||
friendly = (
|
||||
f"The '{provider}' translation engine needs the optional "
|
||||
f"`deep_translator` Python package, which isn't installed in "
|
||||
f"this backend. Install it with `uv pip install deep_translator` "
|
||||
f"(or `pip install deep_translator`) and restart the server, or "
|
||||
f"switch the Engine dropdown to Argos (local, bundled), NLLB "
|
||||
f"(local, heavier), or OpenAI (LLM)."
|
||||
)
|
||||
return JSONResponse(status_code=400, content={"error": friendly})
|
||||
|
||||
src_arg = TRANSLATE_CODES.get(src_lang, src_lang) or "auto"
|
||||
|
||||
def _build_translator(src, tgt):
|
||||
|
||||
@@ -16,7 +16,7 @@ from fastapi import APIRouter, HTTPException
|
||||
from pydantic import BaseModel
|
||||
|
||||
from core import prefs
|
||||
from services import tts_backend, asr_backend, llm_backend
|
||||
from services import tts_backend, asr_backend, llm_backend, translation_engines
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
@@ -60,6 +60,81 @@ def list_llm_backends():
|
||||
return {"active": llm_backend.active_backend_id(), "backends": llm_backend.list_backends()}
|
||||
|
||||
|
||||
@router.get("/engines/translation")
|
||||
def list_translation_engines():
|
||||
"""Translation engines with per-engine pip-package availability.
|
||||
|
||||
Separate from the tts/asr/llm "family" endpoints because these are
|
||||
pip-installable on demand rather than select-from-what's-available.
|
||||
The UI uses this to show a one-click Install chip when the user picks
|
||||
an engine whose Python dependency isn't importable yet.
|
||||
"""
|
||||
return {
|
||||
"engines": translation_engines.list_engines(),
|
||||
"sandboxed": translation_engines.is_frozen(),
|
||||
}
|
||||
|
||||
|
||||
@router.post("/engines/translation/{engine_id}/install")
|
||||
async def install_translation_engine(engine_id: str):
|
||||
entry = translation_engines.get_engine(engine_id)
|
||||
if not entry:
|
||||
raise HTTPException(status_code=404, detail=f"Unknown translation engine: {engine_id!r}")
|
||||
if translation_engines.is_frozen():
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=(
|
||||
"Engine install is disabled in the packaged build — the "
|
||||
"bundled Python environment is read-only and signed. Run the "
|
||||
"source/dev install (`uv sync`) if you need to add an engine."
|
||||
),
|
||||
)
|
||||
pkg = entry.get("pip_package")
|
||||
if not pkg:
|
||||
return {"status": "already_installed", "engine": engine_id, "reason": "no pip package required"}
|
||||
if translation_engines.is_installed(engine_id):
|
||||
return {"status": "already_installed", "engine": engine_id}
|
||||
rc, out = await translation_engines.run_pip(["install", pkg])
|
||||
if rc != 0:
|
||||
raise HTTPException(status_code=500, detail=f"pip install {pkg} failed ({rc}): {out[-1000:]}")
|
||||
# Probe again so the response reflects post-install reality; site-packages
|
||||
# is visible immediately but importlib may have cached a failure.
|
||||
import importlib
|
||||
importlib.invalidate_caches()
|
||||
ok = translation_engines.is_installed(engine_id)
|
||||
return {
|
||||
"status": "installed" if ok else "installed_but_probe_failed",
|
||||
"engine": engine_id,
|
||||
"package": pkg,
|
||||
"log_tail": out[-800:],
|
||||
"restart_required": not ok,
|
||||
}
|
||||
|
||||
|
||||
@router.delete("/engines/translation/{engine_id}")
|
||||
async def uninstall_translation_engine(engine_id: str):
|
||||
entry = translation_engines.get_engine(engine_id)
|
||||
if not entry:
|
||||
raise HTTPException(status_code=404, detail=f"Unknown translation engine: {engine_id!r}")
|
||||
if entry.get("builtin"):
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=(
|
||||
f"{entry['display_name']} is built-in and cannot be uninstalled. "
|
||||
"It shares its Python dependency with core features."
|
||||
),
|
||||
)
|
||||
if translation_engines.is_frozen():
|
||||
raise HTTPException(status_code=400, detail="Engine uninstall is disabled in packaged builds.")
|
||||
pkg = entry.get("pip_package")
|
||||
if not pkg:
|
||||
return {"status": "no_op", "engine": engine_id}
|
||||
rc, out = await translation_engines.run_pip(["uninstall", "-y", pkg])
|
||||
if rc != 0:
|
||||
raise HTTPException(status_code=500, detail=f"pip uninstall {pkg} failed ({rc}): {out[-1000:]}")
|
||||
return {"status": "uninstalled", "engine": engine_id, "package": pkg, "log_tail": out[-800:]}
|
||||
|
||||
|
||||
class SelectEngineRequest(BaseModel):
|
||||
family: str # "tts" | "asr" | "llm"
|
||||
backend_id: str
|
||||
|
||||
@@ -0,0 +1,535 @@
|
||||
"""First-run setup endpoints — model presence + live download progress.
|
||||
|
||||
`GET /setup/status` reports whether the primary model weights are cached on
|
||||
disk + how much disk space remains. The frontend uses this on boot to decide
|
||||
whether to show a setup wizard or the main UI.
|
||||
|
||||
`GET /setup/download-stream` is SSE that forwards every tqdm update emitted
|
||||
by `huggingface_hub` through the monkey-patch in `utils/hf_progress`. The
|
||||
frontend subscribes once and renders per-file progress bars until the wizard
|
||||
completes.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import platform as _platform
|
||||
import shutil
|
||||
import sys
|
||||
from typing import Optional
|
||||
|
||||
from fastapi import APIRouter, HTTPException
|
||||
from fastapi.responses import StreamingResponse
|
||||
from pydantic import BaseModel
|
||||
|
||||
from utils import hf_progress
|
||||
|
||||
logger = logging.getLogger("omnivoice.setup")
|
||||
router = APIRouter()
|
||||
|
||||
# Minimum free disk space before we'd even attempt a full model download.
|
||||
# Rough budget: ~6 GB for OmniVoice + Whisper-large-v3 + scratch; leave 4 GB
|
||||
# of headroom so the machine isn't pinned on disk after install.
|
||||
MIN_FREE_GB = 10
|
||||
|
||||
# Where HuggingFace caches downloads by default. If the user has overridden
|
||||
# via HF_HOME or HUGGINGFACE_HUB_CACHE, we honour it — nothing to move.
|
||||
def _hf_cache_dir() -> str:
|
||||
return (
|
||||
os.environ.get("HF_HUB_CACHE")
|
||||
or os.environ.get("HUGGINGFACE_HUB_CACHE")
|
||||
or os.environ.get("HF_HOME")
|
||||
or os.path.expanduser("~/.cache/huggingface")
|
||||
)
|
||||
|
||||
|
||||
def _disk_free_gb(path: str) -> float:
|
||||
try:
|
||||
return shutil.disk_usage(path).free / (1024 ** 3)
|
||||
except Exception:
|
||||
return 0.0
|
||||
|
||||
|
||||
# Every model the app knows about. `required=True` means the app doesn't
|
||||
# function end-to-end without it (wizard blocks on these). `required=False`
|
||||
# models are optional — ship with them uninstalled, user opts in from
|
||||
# Settings > Models.
|
||||
KNOWN_MODELS = [
|
||||
{
|
||||
"repo_id": "k2-fsa/OmniVoice",
|
||||
"label": "OmniVoice TTS (600+ languages, zero-shot)",
|
||||
"role": "TTS",
|
||||
"size_gb": 2.4,
|
||||
"required": True,
|
||||
},
|
||||
{
|
||||
# Cross-platform default ASR. CTranslate2-converted whisper-large-v3,
|
||||
# loads via faster-whisper (asr_backend.py:FasterWhisperBackend).
|
||||
# Works on Linux/Windows/mac-Intel/mac-ARM with no mlx dependency.
|
||||
"repo_id": "Systran/faster-whisper-large-v3",
|
||||
"label": "Whisper large-v3 (faster-whisper — default, cross-platform)",
|
||||
"role": "ASR",
|
||||
"size_gb": 2.9,
|
||||
"required": True,
|
||||
},
|
||||
{
|
||||
"repo_id": "mlx-community/whisper-large-v3-mlx",
|
||||
"label": "Whisper large-v3 (MLX — optional mac-ARM speedup)",
|
||||
"role": "ASR",
|
||||
"size_gb": 3.0,
|
||||
# Optional everywhere — only loadable on mac-ARM dev installs. The
|
||||
# frozen .app can't load mlx reliably (nanobind duplicate-registration
|
||||
# aborts on first mlx.core touch), and mlx doesn't exist on
|
||||
# Linux/Windows/mac-Intel at all. Users on a mac-ARM dev install can
|
||||
# opt in from Settings → Models for ~10-20% lower latency vs faster-
|
||||
# whisper int8 on large-v3.
|
||||
"required": False,
|
||||
},
|
||||
{
|
||||
"repo_id": "openai/whisper-large-v3",
|
||||
"label": "Whisper large-v3 (PyTorch — last-resort fallback)",
|
||||
"role": "ASR",
|
||||
"size_gb": 3.1,
|
||||
# Optional fallback. The faster-whisper repo above is the primary
|
||||
# ASR; openai/whisper-large-v3 is only needed if the user explicitly
|
||||
# picks pytorch-whisper in Settings (CUDA-heavy workflows or when
|
||||
# faster-whisper breaks on a specific host).
|
||||
"required": False,
|
||||
},
|
||||
{
|
||||
"repo_id": "mlx-community/whisper-tiny-mlx",
|
||||
"label": "Whisper tiny (MLX ASR — fast fallback)",
|
||||
"role": "ASR",
|
||||
"size_gb": 0.08,
|
||||
"required": False,
|
||||
},
|
||||
{
|
||||
"repo_id": "pyannote/speaker-diarization-3.1",
|
||||
"label": "pyannote speaker diarisation (multi-speaker videos)",
|
||||
"role": "Diarisation",
|
||||
"size_gb": 0.8,
|
||||
"required": False,
|
||||
"note": "Needs an HF_TOKEN with license accepted.",
|
||||
},
|
||||
{
|
||||
"repo_id": "OpenMOSS-Team/MOSS-TTS-Nano",
|
||||
"label": "MOSS-TTS-Nano (20 langs, CPU-realtime)",
|
||||
"role": "TTS",
|
||||
"size_gb": 0.4,
|
||||
"required": False,
|
||||
},
|
||||
{
|
||||
# Lightweight English "Turbo" TTS. Optional — the wizard doesn't
|
||||
# auto-download this; users opt in from Settings → Models when they
|
||||
# want fast English narration without voice cloning.
|
||||
"repo_id": "KittenML/kitten-tts-mini-0.8",
|
||||
"label": "KittenTTS (English, 8 preset voices, CPU realtime)",
|
||||
"role": "TTS",
|
||||
"size_gb": 0.08,
|
||||
"required": False,
|
||||
},
|
||||
# ── mlx-audio engines (mac-ARM only; opt-in from Settings → Models) ──
|
||||
# These come through backend.services.tts_backend:MLXAudioBackend. The
|
||||
# backend is only available on Apple Silicon; non-mac users never see
|
||||
# these download buttons as active because the backend is unavailable.
|
||||
{
|
||||
"repo_id": "mlx-community/Kokoro-82M-bf16",
|
||||
"label": "Kokoro 82M (8 langs, small, mlx-audio default)",
|
||||
"role": "TTS",
|
||||
"size_gb": 0.15,
|
||||
"required": False,
|
||||
"note": "Apple Silicon only — via mlx-audio backend.",
|
||||
},
|
||||
{
|
||||
"repo_id": "mlx-community/csm-1b-8bit",
|
||||
"label": "CSM 1B (voice cloning, mlx-audio)",
|
||||
"role": "TTS",
|
||||
"size_gb": 1.1,
|
||||
"required": False,
|
||||
"note": "Apple Silicon only — via mlx-audio backend.",
|
||||
},
|
||||
{
|
||||
"repo_id": "mlx-community/Qwen3-TTS-1.7B-4bit",
|
||||
"label": "Qwen3-TTS 1.7B 4bit (voice design, mlx-audio)",
|
||||
"role": "TTS",
|
||||
"size_gb": 1.4,
|
||||
"required": False,
|
||||
"note": "Apple Silicon only — via mlx-audio backend.",
|
||||
},
|
||||
{
|
||||
"repo_id": "mlx-community/Dia-1.6B",
|
||||
"label": "Dia 1.6B (expressive, mlx-audio)",
|
||||
"role": "TTS",
|
||||
"size_gb": 3.2,
|
||||
"required": False,
|
||||
"note": "Apple Silicon only — via mlx-audio backend.",
|
||||
},
|
||||
{
|
||||
"repo_id": "mlx-community/OuteTTS-0.3-500M",
|
||||
"label": "OuteTTS 0.3 500M (voice clone, mlx-audio)",
|
||||
"role": "TTS",
|
||||
"size_gb": 1.0,
|
||||
"required": False,
|
||||
"note": "Apple Silicon only — via mlx-audio backend.",
|
||||
},
|
||||
]
|
||||
# Back-compat tuple view for code that expects (repo_id, label) pairs.
|
||||
REQUIRED_MODELS = [(m["repo_id"], m["label"]) for m in KNOWN_MODELS if m["required"]]
|
||||
|
||||
|
||||
def _is_cached(repo_id: str) -> bool:
|
||||
"""Best-effort check: does HF have this repo in its cache on disk?
|
||||
We don't validate the specific file set — presence of the repo dir is
|
||||
close enough for a first-run gate."""
|
||||
try:
|
||||
from huggingface_hub import scan_cache_dir
|
||||
info = scan_cache_dir()
|
||||
for entry in info.repos:
|
||||
if entry.repo_id == repo_id and entry.size_on_disk > 0:
|
||||
return True
|
||||
return False
|
||||
except Exception as e:
|
||||
logger.debug("scan_cache_dir failed: %s", e)
|
||||
# Pessimistic: if we can't tell, report missing so the wizard appears
|
||||
# and the user sees progress instead of a silent hang.
|
||||
return False
|
||||
|
||||
|
||||
@router.get("/setup/status")
|
||||
def setup_status():
|
||||
"""Snapshot the setup state so the client can pick its boot screen.
|
||||
|
||||
Returns everything the wizard needs to decide: missing model list, disk
|
||||
headroom, HF cache path (for the user's information + "clear cache" ops).
|
||||
"""
|
||||
missing = [
|
||||
{"repo_id": rid, "label": label}
|
||||
for (rid, label) in REQUIRED_MODELS
|
||||
if not _is_cached(rid)
|
||||
]
|
||||
cache = _hf_cache_dir()
|
||||
free_gb = _disk_free_gb(cache)
|
||||
return {
|
||||
"models_ready": len(missing) == 0,
|
||||
"missing": missing,
|
||||
"hf_cache_dir": cache,
|
||||
"disk_free_gb": round(free_gb, 2),
|
||||
"min_free_gb": MIN_FREE_GB,
|
||||
"enough_disk": free_gb >= MIN_FREE_GB,
|
||||
}
|
||||
|
||||
|
||||
@router.get("/setup/download-stream")
|
||||
async def setup_download_stream():
|
||||
"""SSE: forward every HuggingFace download tqdm update as a JSON event.
|
||||
|
||||
The client connects on mount, then kicks a separate `POST /setup/download`
|
||||
(or invokes a normal ASR/TTS call that triggers the download). This
|
||||
endpoint stays open until the client closes it.
|
||||
"""
|
||||
# Buffered queue so fast-emitting tqdm updates don't drop events on slow
|
||||
# clients. Bounded so a stuck consumer can't grow memory indefinitely.
|
||||
queue: asyncio.Queue = asyncio.Queue(maxsize=512)
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
def listener(event):
|
||||
# tqdm lives on a background thread (hf's downloader). We need to
|
||||
# marshal events onto the FastAPI event loop before enqueueing.
|
||||
try:
|
||||
loop.call_soon_threadsafe(_safe_put, queue, event)
|
||||
except RuntimeError:
|
||||
# Loop closed between events — client has gone away, safe to drop.
|
||||
pass
|
||||
|
||||
listener_id = hf_progress.register_listener(listener)
|
||||
|
||||
async def gen():
|
||||
try:
|
||||
while True:
|
||||
try:
|
||||
event = await asyncio.wait_for(queue.get(), timeout=30.0)
|
||||
except asyncio.TimeoutError:
|
||||
# Heartbeat every 30 s so intermediaries don't time out.
|
||||
yield ": keepalive\n\n"
|
||||
continue
|
||||
yield f"data: {json.dumps(event)}\n\n"
|
||||
finally:
|
||||
hf_progress.unregister_listener(listener_id)
|
||||
|
||||
return StreamingResponse(
|
||||
gen(),
|
||||
media_type="text/event-stream",
|
||||
headers={
|
||||
"Cache-Control": "no-cache, no-transform",
|
||||
"X-Accel-Buffering": "no",
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _safe_put(queue: asyncio.Queue, event) -> None:
|
||||
"""Non-blocking enqueue — drop oldest on overflow rather than block the
|
||||
tqdm thread."""
|
||||
try:
|
||||
queue.put_nowait(event)
|
||||
except asyncio.QueueFull:
|
||||
try:
|
||||
queue.get_nowait()
|
||||
queue.put_nowait(event)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
@router.get("/models")
|
||||
def list_models():
|
||||
"""Catalogue every known model + its on-disk install state.
|
||||
|
||||
The frontend Models tab reads this to draw install/delete buttons. We
|
||||
don't walk disk for every model — instead `scan_cache_dir()` returns
|
||||
*everything* HF has cached, and we look up each known repo in that map.
|
||||
One os-walk regardless of model count.
|
||||
"""
|
||||
cached_by_repo: dict[str, dict] = {}
|
||||
try:
|
||||
from huggingface_hub import scan_cache_dir
|
||||
info = scan_cache_dir()
|
||||
for entry in info.repos:
|
||||
cached_by_repo[entry.repo_id] = {
|
||||
"size_on_disk": entry.size_on_disk,
|
||||
"last_accessed": entry.last_accessed,
|
||||
"nb_files": entry.nb_files,
|
||||
}
|
||||
except Exception as e:
|
||||
logger.warning("scan_cache_dir failed: %s", e)
|
||||
|
||||
out = []
|
||||
for m in KNOWN_MODELS:
|
||||
cached = cached_by_repo.get(m["repo_id"])
|
||||
out.append({
|
||||
**m,
|
||||
"installed": cached is not None and cached["size_on_disk"] > 0,
|
||||
"size_on_disk_bytes": cached["size_on_disk"] if cached else 0,
|
||||
"nb_files": cached["nb_files"] if cached else 0,
|
||||
})
|
||||
return {
|
||||
"models": out,
|
||||
"total_installed_bytes": sum(m["size_on_disk_bytes"] for m in out),
|
||||
"hf_cache_dir": _hf_cache_dir(),
|
||||
}
|
||||
|
||||
|
||||
class InstallModelRequest(BaseModel):
|
||||
repo_id: str
|
||||
|
||||
|
||||
@router.post("/models/install")
|
||||
async def install_model(req: InstallModelRequest):
|
||||
"""Download one HF repo snapshot; progress goes through the shared
|
||||
`/setup/download-stream` SSE feed. Returns immediately so the UI can
|
||||
start listening to the stream.
|
||||
|
||||
Matching by repo_id only — no version pinning today. HF's default-branch
|
||||
"main" / "refs/heads/main" is what snapshot_download picks."""
|
||||
if req.repo_id not in [m["repo_id"] for m in KNOWN_MODELS]:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=(
|
||||
f"Unknown model: {req.repo_id!r}. Known: "
|
||||
+ ", ".join(m["repo_id"] for m in KNOWN_MODELS)
|
||||
),
|
||||
)
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
def _do():
|
||||
try:
|
||||
from huggingface_hub import snapshot_download
|
||||
logger.info("model install starting: %s", req.repo_id)
|
||||
snapshot_download(repo_id=req.repo_id)
|
||||
logger.info("model install done: %s", req.repo_id)
|
||||
except Exception as e:
|
||||
logger.warning("model install failed for %s: %s", req.repo_id, e)
|
||||
|
||||
# Non-blocking — client polls /models or listens on the SSE.
|
||||
loop.create_task(asyncio.to_thread(_do))
|
||||
return {"status": "install_started", "repo_id": req.repo_id}
|
||||
|
||||
|
||||
@router.delete("/models/{repo_id:path}")
|
||||
def delete_model(repo_id: str):
|
||||
"""Remove every cached revision of a repo from the HF cache. Frees disk
|
||||
+ lets the user re-install a fresh copy via POST /models/install."""
|
||||
try:
|
||||
from huggingface_hub import scan_cache_dir
|
||||
info = scan_cache_dir()
|
||||
commits = [
|
||||
rev.commit_hash
|
||||
for entry in info.repos if entry.repo_id == repo_id
|
||||
for rev in entry.revisions
|
||||
]
|
||||
if not commits:
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail=(
|
||||
f"Model {repo_id!r} isn't installed. Nothing to delete — "
|
||||
"run POST /models/install first if you want a fresh download."
|
||||
),
|
||||
)
|
||||
strategy = info.delete_revisions(*commits)
|
||||
strategy.execute()
|
||||
return {
|
||||
"deleted": True,
|
||||
"repo_id": repo_id,
|
||||
"freed_bytes": strategy.expected_freed_size,
|
||||
}
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail=(
|
||||
f"Could not delete {repo_id}: {e}. "
|
||||
"Close any process using the model (e.g. the app's main dub job) and retry."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@router.get("/setup/recommendations")
|
||||
def recommendations():
|
||||
"""Return a curated model preset for the caller's device + architecture.
|
||||
|
||||
The Settings / first-run Models tab uses this to render a prominent
|
||||
"Install recommended" card so users don't have to pick from 14 models.
|
||||
Logic mirrors the engine availability matrix:
|
||||
- mac-ARM gets the rich mlx-audio stack (Kokoro) + MLX-Whisper speedup
|
||||
- mac-Intel + Linux + Windows get the cross-platform subset
|
||||
- CUDA hosts optionally get the pytorch-whisper fallback baked in
|
||||
"""
|
||||
is_mac_arm = sys.platform == "darwin" and _platform.machine() == "arm64"
|
||||
is_mac_intel = sys.platform == "darwin" and _platform.machine() == "x86_64"
|
||||
is_linux = sys.platform.startswith("linux")
|
||||
is_windows = sys.platform == "win32"
|
||||
|
||||
has_cuda = False
|
||||
try:
|
||||
import torch
|
||||
has_cuda = bool(torch.cuda.is_available())
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Device label — used as the card title.
|
||||
if is_mac_arm:
|
||||
device_label = f"Apple Silicon ({_platform.machine()})"
|
||||
elif is_mac_intel:
|
||||
device_label = "macOS Intel (x86_64)"
|
||||
elif is_windows:
|
||||
device_label = "Windows x64" + (" + CUDA" if has_cuda else "")
|
||||
elif is_linux:
|
||||
device_label = "Linux x64" + (" + CUDA" if has_cuda else "")
|
||||
else:
|
||||
device_label = f"{sys.platform} / {_platform.machine()}"
|
||||
|
||||
# Pick the preset for this device.
|
||||
if is_mac_arm:
|
||||
recommended_ids = [
|
||||
"k2-fsa/OmniVoice", # required — 600+ lang zero-shot
|
||||
"Systran/faster-whisper-large-v3", # required — WhisperX ASR
|
||||
"mlx-community/whisper-large-v3-mlx", # optional mac speedup
|
||||
"mlx-community/Kokoro-82M-bf16", # mlx-audio fast TTS
|
||||
"KittenML/kitten-tts-mini-0.8", # English turbo tier
|
||||
]
|
||||
rationale = (
|
||||
"Apple Silicon gets the full stack: OmniVoice for multilingual clone + "
|
||||
"WhisperX (faster-whisper weights) for cross-platform ASR + MLX-Whisper "
|
||||
"for the Apple-optimised speedup + Kokoro (mlx-audio) for fast local "
|
||||
"English + KittenTTS as a CPU-realtime backup."
|
||||
)
|
||||
else:
|
||||
recommended_ids = [
|
||||
"k2-fsa/OmniVoice", # required
|
||||
"Systran/faster-whisper-large-v3", # required
|
||||
"KittenML/kitten-tts-mini-0.8", # English turbo — cross-platform
|
||||
]
|
||||
if has_cuda:
|
||||
# A CUDA box can actually run pytorch-whisper well; ship it as a
|
||||
# fallback so the user can pin it in Settings → Engines later.
|
||||
recommended_ids.append("openai/whisper-large-v3")
|
||||
rationale = (
|
||||
"Cross-platform stack + pytorch-whisper as a CUDA-accelerated "
|
||||
"ASR fallback. MLX / mlx-audio are Apple-Silicon-only and don't "
|
||||
"apply here."
|
||||
)
|
||||
else:
|
||||
rationale = (
|
||||
"Cross-platform stack: OmniVoice (multilingual clone) + WhisperX "
|
||||
"(faster-whisper ASR) + KittenTTS (English turbo, CPU-realtime). "
|
||||
"Clean install, every model runs on CPU."
|
||||
)
|
||||
|
||||
# Cross-reference against KNOWN_MODELS so we can attach size + label to
|
||||
# each recommended entry, and flag which ones are already installed.
|
||||
known_by_id = {m["repo_id"]: m for m in KNOWN_MODELS}
|
||||
cached_ids: set[str] = set()
|
||||
try:
|
||||
from huggingface_hub import scan_cache_dir
|
||||
info = scan_cache_dir()
|
||||
cached_ids = {
|
||||
entry.repo_id for entry in info.repos if entry.size_on_disk > 0
|
||||
}
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
entries = []
|
||||
for rid in recommended_ids:
|
||||
meta = known_by_id.get(rid, {})
|
||||
entries.append({
|
||||
"repo_id": rid,
|
||||
"label": meta.get("label", rid),
|
||||
"role": meta.get("role", ""),
|
||||
"size_gb": meta.get("size_gb", 0),
|
||||
"required": bool(meta.get("required", False)),
|
||||
"note": meta.get("note"),
|
||||
"installed": rid in cached_ids,
|
||||
})
|
||||
|
||||
# Headline number for the "Install recommended (~X GB)" CTA — only
|
||||
# count models not yet on disk so users with a warm cache see a low
|
||||
# remaining number instead of the full bundle size.
|
||||
to_download_gb = sum(e["size_gb"] for e in entries if not e["installed"])
|
||||
all_installed = all(e["installed"] for e in entries)
|
||||
|
||||
return {
|
||||
"device": {
|
||||
"os": sys.platform,
|
||||
"arch": _platform.machine(),
|
||||
"is_mac_arm": is_mac_arm,
|
||||
"is_mac_intel": is_mac_intel,
|
||||
"is_linux": is_linux,
|
||||
"is_windows": is_windows,
|
||||
"has_cuda": has_cuda,
|
||||
"label": device_label,
|
||||
},
|
||||
"rationale": rationale,
|
||||
"models": entries,
|
||||
"download_gb_remaining": round(to_download_gb, 2),
|
||||
"total_gb": round(sum(e["size_gb"] for e in entries), 2),
|
||||
"all_installed": all_installed,
|
||||
}
|
||||
|
||||
|
||||
@router.post("/setup/warmup")
|
||||
async def setup_warmup():
|
||||
"""Trigger a model load in the background so the first dub doesn't pay
|
||||
the cold-start tax. Progress flows through the SSE stream."""
|
||||
loop = asyncio.get_event_loop()
|
||||
|
||||
async def _do_warmup():
|
||||
try:
|
||||
from services.model_manager import get_model
|
||||
await get_model()
|
||||
except Exception as e:
|
||||
logger.warning("setup/warmup: model load failed: %s", e)
|
||||
|
||||
# Don't await — let it run in the background; client watches SSE.
|
||||
loop.create_task(_do_warmup())
|
||||
return {"status": "warmup_started"}
|
||||
@@ -37,7 +37,7 @@ def system_info():
|
||||
"crash_log_path": CRASH_LOG_PATH,
|
||||
"idle_timeout_seconds": IDLE_TIMEOUT_SECONDS,
|
||||
"model_checkpoint": os.environ.get("OMNIVOICE_MODEL", "k2-fsa/OmniVoice"),
|
||||
"asr_model": os.environ.get("ASR_MODEL", "mlx-community/whisper-large-v3-mlx"),
|
||||
"asr_model": os.environ.get("ASR_MODEL", "Systran/faster-whisper-large-v3"),
|
||||
"translate_provider": os.environ.get("TRANSLATE_PROVIDER", "google"),
|
||||
"has_hf_token": bool(os.environ.get("HF_TOKEN")),
|
||||
"device": get_best_device(),
|
||||
|
||||
@@ -0,0 +1,28 @@
|
||||
"""Numpy compatibility shim for PyInstaller.
|
||||
|
||||
Two things we have to guarantee in a frozen build:
|
||||
|
||||
1. `numpy` gets imported before anything that monkey-touches its C API
|
||||
(scipy, torch, torchaudio) so the extension loader runs in a clean
|
||||
state and we don't hit "module compiled against API version X but
|
||||
this version of numpy is Y" surprises.
|
||||
|
||||
2. NUMPY_EXPERIMENTAL_ARRAY_FUNCTION is left at its default. Some
|
||||
frozen builds end up with it unset because PyInstaller strips certain
|
||||
env-based defaults — explicitly mirror numpy's own default so
|
||||
downstream libs don't disable fast paths by accident.
|
||||
|
||||
This hook runs BEFORE the main script via `runtime_hooks=[...]` in the
|
||||
spec, so the import order is deterministic regardless of which library
|
||||
the user's code touches first.
|
||||
"""
|
||||
import os
|
||||
|
||||
os.environ.setdefault("NUMPY_EXPERIMENTAL_ARRAY_FUNCTION", "1")
|
||||
|
||||
try:
|
||||
import numpy # noqa: F401 (side-effect import: primes the C extension)
|
||||
except ImportError:
|
||||
# Let the main app fail loudly with its own error message rather
|
||||
# than crashing the runtime hook itself.
|
||||
pass
|
||||
@@ -0,0 +1,21 @@
|
||||
"""Disable torch.compile / dynamo / inductor in PyInstaller-frozen builds.
|
||||
|
||||
`torch.compile` uses a mix of C++ ABI magic and source-file introspection
|
||||
that breaks inside a frozen bundle:
|
||||
- Dynamo's FX graph builder walks `__file__` paths that no longer exist
|
||||
once the Python source is shipped as PYC inside the bundle.
|
||||
- TorchInductor tries to read and JIT-compile additional C++ kernels at
|
||||
runtime, relying on a compiler toolchain the end user's Mac won't have.
|
||||
- Frozen modules load in a different order than the source tree, which
|
||||
occasionally surfaces circular-import issues inside dynamo guards.
|
||||
|
||||
Setting these env vars before torch is imported anywhere else is enough
|
||||
to keep PyTorch on its eager path. Negligible perf hit for inference-only
|
||||
workloads — we're not training here.
|
||||
"""
|
||||
import os
|
||||
|
||||
os.environ.setdefault("TORCHDYNAMO_DISABLE", "1")
|
||||
os.environ.setdefault("TORCH_COMPILE_DISABLE", "1")
|
||||
os.environ.setdefault("TORCHINDUCTOR_DISABLE", "1")
|
||||
os.environ.setdefault("PYTORCH_DISABLE_PER_OP_PROFILING", "1")
|
||||
+26
-3
@@ -85,7 +85,7 @@ import time
|
||||
import threading
|
||||
from contextlib import asynccontextmanager
|
||||
from fastapi import FastAPI, Request
|
||||
from fastapi.responses import JSONResponse, RedirectResponse
|
||||
from fastapi.responses import JSONResponse, RedirectResponse, Response
|
||||
from fastapi.staticfiles import StaticFiles
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
import traceback
|
||||
@@ -98,7 +98,13 @@ from core.tasks import task_manager
|
||||
from core import job_store
|
||||
from services.model_manager import idle_worker
|
||||
|
||||
from api.routers import system, profiles, exports, generation, dub_core, dub_generate, dub_export, dub_translate, projects, glossary, engines, tools
|
||||
from api.routers import system, profiles, exports, generation, dub_core, dub_generate, dub_export, dub_translate, projects, glossary, engines, tools, setup
|
||||
from utils import hf_progress
|
||||
|
||||
# Install the HuggingFace tqdm patch early — every downstream library import
|
||||
# that triggers `hf_hub_download` (transformers, mlx_whisper, etc.) must see
|
||||
# the patched class, not the original.
|
||||
hf_progress.install()
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
@@ -122,6 +128,13 @@ app = FastAPI(title="OmniVoice Studio API", version="0.4.0", lifespan=lifespan)
|
||||
|
||||
@app.exception_handler(Exception)
|
||||
async def global_exception_handler(request: Request, exc: Exception):
|
||||
# Client disconnected mid-stream (browser canceled a <video>/range fetch).
|
||||
# The response is already partially sent — trying to wrap it in a 500 just
|
||||
# produces a second protocol error. Log a one-liner and bail.
|
||||
exc_name = type(exc).__name__
|
||||
if exc_name in ("LocalProtocolError", "ClientDisconnect") or "Content-Length" in str(exc):
|
||||
logger.info("Client disconnect during %s (%s)", request.url, exc_name)
|
||||
return Response(status_code=499)
|
||||
try:
|
||||
# Serialize writes so concurrent unhandled exceptions don't interleave frames.
|
||||
with _crash_log_lock, open(CRASH_LOG_PATH, "a") as f:
|
||||
@@ -131,7 +144,16 @@ async def global_exception_handler(request: Request, exc: Exception):
|
||||
except Exception:
|
||||
logger.exception("Failed to write crash log")
|
||||
logger.exception("Unhandled exception for %s", request.url)
|
||||
return JSONResponse({"detail": str(exc)}, status_code=500)
|
||||
# CORSMiddleware doesn't always get a shot at `exception_handler`-created
|
||||
# responses, which leaves the browser reporting every 500 as a bare CORS
|
||||
# error. Attach the headers manually so the real `detail` bubbles up.
|
||||
origin = request.headers.get("origin", "")
|
||||
headers: dict[str, str] = {}
|
||||
if origin and (origin in _allowed or "*" in _allowed):
|
||||
headers["Access-Control-Allow-Origin"] = origin
|
||||
headers["Access-Control-Allow-Credentials"] = "true"
|
||||
headers["Vary"] = "Origin"
|
||||
return JSONResponse({"detail": str(exc)}, status_code=500, headers=headers)
|
||||
|
||||
_allowed = os.environ.get(
|
||||
"OMNIVOICE_ALLOWED_ORIGINS",
|
||||
@@ -161,6 +183,7 @@ app.include_router(projects.router)
|
||||
app.include_router(glossary.router)
|
||||
app.include_router(engines.router)
|
||||
app.include_router(tools.router)
|
||||
app.include_router(setup.router)
|
||||
|
||||
frontend_path = os.path.join(os.path.dirname(__file__), "..", "frontend", "dist")
|
||||
if os.path.exists(frontend_path):
|
||||
|
||||
@@ -66,6 +66,13 @@ class TranslateRequest(BaseModel):
|
||||
class DubIngestUrlRequest(BaseModel):
|
||||
url: str
|
||||
job_id: Optional[str] = None
|
||||
# When true and the URL is a caption-bearing host (YouTube, Vimeo, TED…),
|
||||
# ask yt-dlp to also download the original-language + any additional
|
||||
# sub_langs as VTT. The UI uses this to seed a transcript without running
|
||||
# Whisper, and optionally to skip the Translate step for languages that
|
||||
# YouTube auto-translates for us.
|
||||
fetch_subs: Optional[bool] = False
|
||||
sub_langs: Optional[List[str]] = None
|
||||
|
||||
class ProjectSaveRequest(BaseModel):
|
||||
name: str
|
||||
|
||||
@@ -3,16 +3,23 @@ ASR adapter interface — Phase 3.3 (ROADMAP.md).
|
||||
|
||||
One protocol, multiple engines. Today we ship:
|
||||
|
||||
• MLXWhisperBackend — mlx-whisper on Apple Silicon (default when MPS is
|
||||
available). Keeps today's word-level timestamp
|
||||
output shape.
|
||||
• PyTorchWhisperBackend — fallback on CUDA / CPU using the existing
|
||||
• FasterWhisperBackend — CTranslate2-based (the engine WhisperX uses).
|
||||
Default on Linux, Windows, mac-Intel. Also fast
|
||||
on mac-ARM so we use it as the cross-platform
|
||||
baseline and only prefer MLX on mac-ARM when
|
||||
explicitly installed.
|
||||
• MLXWhisperBackend — mlx-whisper on Apple Silicon. Optional speedup,
|
||||
only available when mlx wheels install (mac-ARM).
|
||||
• PyTorchWhisperBackend — last-resort fallback using the existing
|
||||
`_asr_pipe` on the TTS model.
|
||||
|
||||
Both return the raw Whisper output dict so `services.segmentation.
|
||||
segment_transcript(...)` can keep working unchanged.
|
||||
segment_transcript(...)` can keep working unchanged — new backends normalise
|
||||
their output to the `{"chunks": [{"text", "timestamp": (start, end)}]}`
|
||||
shape the segmenter expects.
|
||||
|
||||
Selection via `OMNIVOICE_ASR_BACKEND` (default: auto-detect).
|
||||
Selection via `OMNIVOICE_ASR_BACKEND` (default: auto-detect, prefers
|
||||
faster-whisper because it's available on every platform we ship to).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -44,7 +51,296 @@ class ASRBackend(ABC):
|
||||
"""
|
||||
|
||||
|
||||
# ── MLX Whisper (Apple Silicon default) ─────────────────────────────────────
|
||||
# ── WhisperX (cross-platform default — forced-alignment word timing) ────────
|
||||
|
||||
|
||||
class WhisperXBackend(ASRBackend):
|
||||
id = "whisperx"
|
||||
display_name = "WhisperX (faster-whisper + wav2vec2 forced alignment)"
|
||||
|
||||
def __init__(self):
|
||||
self._model_name = os.environ.get("ASR_MODEL_WHISPERX", "large-v3")
|
||||
self._asr = None
|
||||
self._align_cache = {} # language_code → (align_model, metadata)
|
||||
self._device, self._compute_type = self._pick_device()
|
||||
|
||||
@staticmethod
|
||||
def _pick_device() -> tuple[str, str]:
|
||||
# CUDA fp16 when available; otherwise CPU int8 (fastest CPU path,
|
||||
# negligible WER regression vs fp32 for whisper-large-v3).
|
||||
try:
|
||||
import torch
|
||||
if torch.cuda.is_available():
|
||||
return "cuda", "float16"
|
||||
except Exception:
|
||||
pass
|
||||
return "cpu", "int8"
|
||||
|
||||
@classmethod
|
||||
def is_available(cls) -> tuple[bool, str]:
|
||||
try:
|
||||
import whisperx # noqa: F401
|
||||
return True, "ready"
|
||||
except ImportError as e:
|
||||
return False, f"whisperx not installed: {e}"
|
||||
|
||||
def _ensure_asr(self):
|
||||
if self._asr is not None:
|
||||
return
|
||||
import whisperx
|
||||
import torch
|
||||
logger.info(
|
||||
"whisperx loading ASR %s on %s (%s)",
|
||||
self._model_name, self._device, self._compute_type,
|
||||
)
|
||||
# PyTorch 2.6 flipped `torch.load(weights_only=True)` to default,
|
||||
# which breaks pyannote 3.x's VAD checkpoint (that whisperx ships):
|
||||
# each load surfaces a different missing global — `omegaconf.*`,
|
||||
# `typing.Any`, etc. The VAD file ships inside the whisperx wheel,
|
||||
# so it's as trusted as whisperx itself. Two-layer defence:
|
||||
# (a) allowlist the known pickle globals so the secure load path
|
||||
# actually succeeds, and
|
||||
# (b) monkey-patch `torch.load` to force `weights_only=False` as
|
||||
# a belt-and-braces fallback for anything we missed.
|
||||
self._allow_vad_pickle_globals()
|
||||
import torch.serialization as _ts
|
||||
_orig_top = torch.load
|
||||
_orig_inner = _ts.load
|
||||
def _patched(*args, **kwargs):
|
||||
# Force — Lightning explicitly passes weights_only=True, so a
|
||||
# setdefault wouldn't override it. The VAD pickle ships in the
|
||||
# whisperx wheel; trust is the same as trusting whisperx itself.
|
||||
kwargs["weights_only"] = False
|
||||
return _orig_inner(*args, **kwargs)
|
||||
torch.load = _patched
|
||||
_ts.load = _patched
|
||||
try:
|
||||
self._asr = whisperx.load_model(
|
||||
self._model_name,
|
||||
device=self._device,
|
||||
compute_type=self._compute_type,
|
||||
# vad_method="silero" is the default; keep it so short gaps
|
||||
# get cleaned up before transcription.
|
||||
)
|
||||
finally:
|
||||
torch.load = _orig_top
|
||||
_ts.load = _orig_inner
|
||||
|
||||
@staticmethod
|
||||
def _allow_vad_pickle_globals():
|
||||
"""Register the pickle classes that pyannote's VAD checkpoint contains.
|
||||
|
||||
Without this, PyTorch 2.6's secure unpickler refuses to load the file
|
||||
even if the call explicitly passes `weights_only=False` later — the
|
||||
allowlist is per-process and harmless to re-apply. Each class we add
|
||||
is one that has surfaced in the wild from pyannote/omegaconf/pytorch-
|
||||
lightning pickles; extending the list is safe.
|
||||
"""
|
||||
try:
|
||||
import torch.serialization as _ts
|
||||
except Exception:
|
||||
return
|
||||
add = getattr(_ts, "add_safe_globals", None)
|
||||
if add is None:
|
||||
return # older torch — secure unpickler didn't exist
|
||||
|
||||
allow = []
|
||||
# omegaconf config containers — the immediate cause of the error
|
||||
# pyannote's VAD emits (`GLOBAL omegaconf.listconfig.ListConfig`).
|
||||
try:
|
||||
from omegaconf.listconfig import ListConfig
|
||||
from omegaconf.dictconfig import DictConfig
|
||||
from omegaconf.base import ContainerMetadata, Metadata
|
||||
allow += [ListConfig, DictConfig, ContainerMetadata, Metadata]
|
||||
except Exception:
|
||||
pass
|
||||
# Python typing primitives that show up in config annotations.
|
||||
try:
|
||||
import typing
|
||||
allow += [typing.Any]
|
||||
except Exception:
|
||||
pass
|
||||
# pytorch-lightning's OrderedDict-backed state dict helpers.
|
||||
try:
|
||||
from collections import OrderedDict, defaultdict
|
||||
allow += [OrderedDict, defaultdict]
|
||||
except Exception:
|
||||
pass
|
||||
if allow:
|
||||
try:
|
||||
add(allow)
|
||||
except Exception as e:
|
||||
logger.debug("add_safe_globals failed (harmless): %s", e)
|
||||
|
||||
def _get_align(self, language_code: str):
|
||||
"""Lazy-load the wav2vec2 alignment model for this language. WhisperX
|
||||
bundles aligners for ~20 major languages; for the others we fall back
|
||||
to faster-whisper's native word timestamps (already in result)."""
|
||||
if language_code in self._align_cache:
|
||||
return self._align_cache[language_code]
|
||||
import whisperx
|
||||
try:
|
||||
model, metadata = whisperx.load_align_model(
|
||||
language_code=language_code, device=self._device,
|
||||
)
|
||||
self._align_cache[language_code] = (model, metadata)
|
||||
return model, metadata
|
||||
except Exception as e:
|
||||
logger.info(
|
||||
"whisperx: no alignment model for language=%r (%s); "
|
||||
"falling back to Whisper's native word timestamps",
|
||||
language_code, e,
|
||||
)
|
||||
self._align_cache[language_code] = None
|
||||
return None
|
||||
|
||||
def transcribe(self, audio_path: str, *, word_timestamps: bool = True) -> dict:
|
||||
import whisperx
|
||||
self._ensure_asr()
|
||||
logger.info("whisperx transcribing %s (word_timestamps=%s)", audio_path, word_timestamps)
|
||||
audio = whisperx.load_audio(audio_path)
|
||||
result = self._asr.transcribe(audio)
|
||||
lang = result.get("language", "en")
|
||||
|
||||
# Forced alignment when available — drastically improves word boundary
|
||||
# accuracy (±10-30 ms vs Whisper's ±100-300 ms). Skip for rare-language
|
||||
# audio where no wav2vec2 aligner exists.
|
||||
if word_timestamps:
|
||||
align = self._get_align(lang)
|
||||
if align is not None:
|
||||
model_a, metadata = align
|
||||
try:
|
||||
result = whisperx.align(
|
||||
result["segments"], model_a, metadata, audio,
|
||||
self._device, return_char_alignments=False,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning("whisperx alignment failed: %s — using raw timestamps", e)
|
||||
|
||||
# Normalise to the shape segment_transcript(...) expects: chunks +
|
||||
# segments + language metadata. whisperx's post-align result has
|
||||
# `segments` with `words: [{word, start, end, score}]`.
|
||||
segments = result.get("segments", [])
|
||||
chunks = [
|
||||
{"text": seg.get("text", ""),
|
||||
"timestamp": (seg.get("start"), seg.get("end"))}
|
||||
for seg in segments
|
||||
]
|
||||
return {
|
||||
"chunks": chunks,
|
||||
"segments": [
|
||||
{
|
||||
"text": seg.get("text", ""),
|
||||
"start": seg.get("start"),
|
||||
"end": seg.get("end"),
|
||||
"words": seg.get("words", []) if word_timestamps else [],
|
||||
}
|
||||
for seg in segments
|
||||
],
|
||||
"language": lang,
|
||||
}
|
||||
|
||||
|
||||
# ── Faster-Whisper (cross-platform fallback) ────────────────────────────────
|
||||
|
||||
|
||||
class FasterWhisperBackend(ASRBackend):
|
||||
id = "faster-whisper"
|
||||
display_name = "Faster-Whisper (CTranslate2 — Linux/Windows/macOS)"
|
||||
|
||||
def __init__(self):
|
||||
# Defaulting to the CTranslate2-converted large-v3 repo. Matches
|
||||
# KNOWN_MODELS in api/routers/setup.py so the first-run wizard
|
||||
# downloads what the backend will actually load.
|
||||
self._model_name = os.environ.get(
|
||||
"ASR_MODEL_FASTER", "Systran/faster-whisper-large-v3"
|
||||
)
|
||||
self._model = None # lazy — first transcribe() loads weights
|
||||
|
||||
@classmethod
|
||||
def is_available(cls) -> tuple[bool, str]:
|
||||
try:
|
||||
import faster_whisper # noqa: F401
|
||||
return True, "ready"
|
||||
except ImportError as e:
|
||||
return False, f"faster-whisper not installed: {e}"
|
||||
|
||||
def _ensure_model(self):
|
||||
if self._model is not None:
|
||||
return
|
||||
from faster_whisper import WhisperModel
|
||||
# Device / compute-type auto-pick:
|
||||
# - CUDA present → GPU fp16
|
||||
# - Apple Silicon / CPU → CPU int8 (fastest on CPU, negligible
|
||||
# WER regression vs fp32 for whisper-large-v3)
|
||||
device, compute_type = "cpu", "int8"
|
||||
try:
|
||||
import torch
|
||||
if torch.cuda.is_available():
|
||||
device, compute_type = "cuda", "float16"
|
||||
except Exception:
|
||||
pass
|
||||
logger.info(
|
||||
"faster-whisper loading %s on %s (%s)",
|
||||
self._model_name, device, compute_type,
|
||||
)
|
||||
self._model = WhisperModel(
|
||||
self._model_name, device=device, compute_type=compute_type
|
||||
)
|
||||
|
||||
def transcribe(self, audio_path: str, *, word_timestamps: bool = True) -> dict:
|
||||
self._ensure_model()
|
||||
logger.info(
|
||||
"faster-whisper transcribing %s (word_timestamps=%s)",
|
||||
audio_path, word_timestamps,
|
||||
)
|
||||
# faster-whisper returns a generator of Segment objects + an Info
|
||||
# struct. Materialise the generator so downstream consumers can
|
||||
# index / re-iterate.
|
||||
segments_iter, info = self._model.transcribe(
|
||||
audio_path,
|
||||
word_timestamps=word_timestamps,
|
||||
vad_filter=True, # built-in Silero VAD — cleaner segment starts
|
||||
)
|
||||
segments = list(segments_iter)
|
||||
# Normalise to the shape segment_transcript(...) expects: a dict with
|
||||
# `chunks` (for backwards compat with mlx output) AND `segments` +
|
||||
# `language` (so callers that peek at language metadata keep working).
|
||||
chunks = [
|
||||
{"text": seg.text, "timestamp": (seg.start, seg.end)}
|
||||
for seg in segments
|
||||
]
|
||||
out = {
|
||||
"chunks": chunks,
|
||||
"segments": [
|
||||
{
|
||||
"text": seg.text,
|
||||
"start": seg.start,
|
||||
"end": seg.end,
|
||||
"words": (
|
||||
[
|
||||
{
|
||||
"word": w.word,
|
||||
"start": w.start,
|
||||
"end": w.end,
|
||||
"probability": w.probability,
|
||||
}
|
||||
for w in (seg.words or [])
|
||||
]
|
||||
if word_timestamps
|
||||
else []
|
||||
),
|
||||
}
|
||||
for seg in segments
|
||||
],
|
||||
"language": info.language,
|
||||
"language_probability": info.language_probability,
|
||||
"duration": info.duration,
|
||||
}
|
||||
return out
|
||||
|
||||
|
||||
# ── MLX Whisper (Apple Silicon optional) ────────────────────────────────────
|
||||
|
||||
|
||||
class MLXWhisperBackend(ASRBackend):
|
||||
@@ -142,6 +438,8 @@ class PyTorchWhisperBackend(ASRBackend):
|
||||
|
||||
|
||||
_REGISTRY: dict[str, type[ASRBackend]] = {
|
||||
"whisperx": WhisperXBackend,
|
||||
"faster-whisper": FasterWhisperBackend,
|
||||
"mlx-whisper": MLXWhisperBackend,
|
||||
"pytorch-whisper": PyTorchWhisperBackend,
|
||||
}
|
||||
@@ -161,12 +459,37 @@ def list_backends() -> list[dict]:
|
||||
|
||||
|
||||
def _auto_detect() -> str:
|
||||
"""Pick the best available ASR engine for the current hardware."""
|
||||
"""Pick the best available ASR engine for the current hardware.
|
||||
|
||||
Preference order:
|
||||
1. whisperx — faster-whisper transcription + wav2vec2 forced
|
||||
alignment (±10-30 ms word timing). Best for the
|
||||
dub pipeline because lip-sync quality depends on
|
||||
word-boundary accuracy.
|
||||
2. faster-whisper — transcription only (no forced alignment). Slightly
|
||||
looser word boundaries but strictly faster; safe
|
||||
fallback when whisperx isn't installed.
|
||||
3. mlx-whisper — mac-ARM speedup if installed (~10-20% latency win
|
||||
vs faster-whisper int8 on Apple Silicon for
|
||||
large-v3). Optional; faster-whisper remains the
|
||||
baseline so we don't diverge mac-only behaviour.
|
||||
4. pytorch-whisper — last resort; requires the TTS model to be loaded
|
||||
so it can reuse `_asr_pipe`.
|
||||
"""
|
||||
ok, _ = WhisperXBackend.is_available()
|
||||
if ok:
|
||||
return "whisperx"
|
||||
ok, _ = FasterWhisperBackend.is_available()
|
||||
if ok:
|
||||
return "faster-whisper"
|
||||
try:
|
||||
import torch
|
||||
if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
|
||||
ok, _ = MLXWhisperBackend.is_available()
|
||||
if ok:
|
||||
return "mlx-whisper"
|
||||
except Exception:
|
||||
pass
|
||||
return "pytorch-whisper"
|
||||
|
||||
|
||||
@@ -187,6 +510,10 @@ def get_active_asr_backend(*, asr_pipe=None) -> ASRBackend:
|
||||
return PyTorchWhisperBackend(asr_pipe=asr_pipe)
|
||||
if bid == "mlx-whisper":
|
||||
return MLXWhisperBackend()
|
||||
if bid == "faster-whisper":
|
||||
return FasterWhisperBackend()
|
||||
if bid == "whisperx":
|
||||
return WhisperXBackend()
|
||||
if bid not in _REGISTRY:
|
||||
raise ValueError(f"Unknown ASR backend: {bid!r}. Known: {list(_REGISTRY)}")
|
||||
return _REGISTRY[bid]()
|
||||
|
||||
@@ -277,11 +277,29 @@ def run_proc_factory(job_id: str):
|
||||
return run_proc
|
||||
|
||||
|
||||
def yt_download_sync(url: str, job_dir: str) -> tuple[str, str]:
|
||||
"""Blocking yt-dlp download into `job_dir`. Returns (video_path, title)."""
|
||||
def yt_download_sync(
|
||||
url: str,
|
||||
job_dir: str,
|
||||
*,
|
||||
fetch_subs: bool = False,
|
||||
sub_langs: list[str] | None = None,
|
||||
) -> tuple[str, str, list[str]]:
|
||||
"""Blocking yt-dlp download into `job_dir`.
|
||||
|
||||
Returns (video_path, title, downloaded_sub_files).
|
||||
|
||||
When `fetch_subs` is True we also ask yt-dlp to download both
|
||||
manually-uploaded (`writesubtitles=True`) and auto-generated / auto-
|
||||
translated captions (`writeautomaticsub=True`). This is how we pull
|
||||
YouTube's free machine translations without needing a Google API key —
|
||||
yt-dlp talks to the same public endpoints the YouTube player does.
|
||||
`sub_langs` controls which language tracks to ask for; default `['all']`
|
||||
grabs whatever the uploader / auto-translator makes available.
|
||||
"""
|
||||
import glob
|
||||
import yt_dlp
|
||||
outtmpl = os.path.join(job_dir, "original.%(ext)s")
|
||||
ydl_opts = {
|
||||
ydl_opts: dict = {
|
||||
"outtmpl": outtmpl,
|
||||
"format": "bv*[ext=mp4]+ba[ext=m4a]/b[ext=mp4]/bv*+ba/b",
|
||||
"merge_output_format": "mp4",
|
||||
@@ -291,14 +309,86 @@ def yt_download_sync(url: str, job_dir: str) -> tuple[str, str]:
|
||||
"restrictfilenames": True,
|
||||
"socket_timeout": 30,
|
||||
}
|
||||
if fetch_subs:
|
||||
ydl_opts.update({
|
||||
"writesubtitles": True,
|
||||
"writeautomaticsub": True,
|
||||
"subtitleslangs": list(sub_langs) if sub_langs else ["all"],
|
||||
"subtitlesformat": "vtt",
|
||||
})
|
||||
with yt_dlp.YoutubeDL(ydl_opts) as ydl:
|
||||
info = ydl.extract_info(url, download=True)
|
||||
path = ydl.prepare_filename(info)
|
||||
root, _ = os.path.splitext(path)
|
||||
mp4 = root + ".mp4"
|
||||
if os.path.exists(mp4):
|
||||
return mp4, info.get("title") or os.path.basename(mp4)
|
||||
return path, info.get("title") or os.path.basename(path)
|
||||
video_path = mp4
|
||||
else:
|
||||
video_path = path
|
||||
sub_files: list[str] = []
|
||||
if fetch_subs:
|
||||
# yt-dlp names captions "<base>.<lang>.vtt"; scoop them all.
|
||||
base = os.path.splitext(video_path)[0]
|
||||
sub_files = sorted(glob.glob(base + ".*.vtt"))
|
||||
return video_path, info.get("title") or os.path.basename(video_path), sub_files
|
||||
|
||||
|
||||
def parse_vtt_segments(vtt_path: str) -> list[dict]:
|
||||
"""Very small WEBVTT parser → list of {start, end, text}.
|
||||
|
||||
Purpose: turn yt-dlp's downloaded caption track into the same segment
|
||||
shape the dub pipeline's transcript step produces, so the UI can seed
|
||||
its editor from YouTube captions instead of running Whisper. We don't
|
||||
care about positioning cues or styling — just the timed text.
|
||||
"""
|
||||
try:
|
||||
with open(vtt_path, "r", encoding="utf-8") as f:
|
||||
raw = f.read()
|
||||
except OSError:
|
||||
return []
|
||||
|
||||
def _ts(s: str) -> float:
|
||||
# Format: HH:MM:SS.mmm or MM:SS.mmm
|
||||
parts = s.strip().split(":")
|
||||
try:
|
||||
parts = [float(p.replace(",", ".")) for p in parts]
|
||||
except ValueError:
|
||||
return 0.0
|
||||
if len(parts) == 3:
|
||||
h, m, sec = parts
|
||||
elif len(parts) == 2:
|
||||
h, m, sec = 0.0, parts[0], parts[1]
|
||||
else:
|
||||
return 0.0
|
||||
return h * 3600.0 + m * 60.0 + sec
|
||||
|
||||
segments: list[dict] = []
|
||||
blocks = raw.replace("\r\n", "\n").split("\n\n")
|
||||
for block in blocks:
|
||||
lines = [ln for ln in block.split("\n") if ln.strip() and not ln.startswith("WEBVTT") and not ln.startswith("NOTE")]
|
||||
if not lines:
|
||||
continue
|
||||
# Skip numeric cue ID line if present
|
||||
if "-->" not in lines[0] and len(lines) > 1:
|
||||
lines = lines[1:]
|
||||
if not lines or "-->" not in lines[0]:
|
||||
continue
|
||||
ts_line = lines[0]
|
||||
try:
|
||||
left, right = ts_line.split("-->")
|
||||
# Drop any settings after the right timestamp ("00:01.000 align:left line:0%")
|
||||
right = right.strip().split(" ")[0]
|
||||
start = _ts(left)
|
||||
end = _ts(right)
|
||||
except Exception:
|
||||
continue
|
||||
text = " ".join(ln.strip() for ln in lines[1:]).strip()
|
||||
# Strip inline styling like <c.colorE5E5E5>foo</c> or <00:00:01.200>
|
||||
import re as _re
|
||||
text = _re.sub(r"<[^>]+>", "", text)
|
||||
if text:
|
||||
segments.append({"start": start, "end": end, "text": text})
|
||||
return segments
|
||||
|
||||
|
||||
async def ingest_pipeline(
|
||||
@@ -312,12 +402,18 @@ async def ingest_pipeline(
|
||||
Stages: download_start, download_done, extract_start, extract_done,
|
||||
demucs_start, demucs_done, scene_start, scene_done, ready, error, cancelled.
|
||||
"""
|
||||
youtube_subs_by_lang: dict[str, list[dict]] = {}
|
||||
try:
|
||||
if source.get("kind") == "url":
|
||||
url = source["url"]
|
||||
fetch_subs = bool(source.get("fetch_subs"))
|
||||
sub_langs = source.get("sub_langs") or None
|
||||
yield prep_event("download_start", url=url)
|
||||
try:
|
||||
video_path, title = await asyncio.to_thread(yt_download_sync, url, job_dir)
|
||||
video_path, title, sub_files = await asyncio.to_thread(
|
||||
yt_download_sync, url, job_dir,
|
||||
fetch_subs=fetch_subs, sub_langs=sub_langs,
|
||||
)
|
||||
except Exception as e:
|
||||
yield prep_event("error", stage="download", error=str(e)[:300])
|
||||
shutil.rmtree(job_dir, ignore_errors=True)
|
||||
@@ -327,7 +423,22 @@ async def ingest_pipeline(
|
||||
size = os.path.getsize(video_path)
|
||||
except OSError:
|
||||
size = 0
|
||||
yield prep_event("download_done", title=title, size=size, filename=filename)
|
||||
# Parse each downloaded .<lang>.vtt into a segment list we can
|
||||
# stash on the job. The router will merge this into the job dict
|
||||
# after ingest completes so /dub/transcribe-stream (or a future
|
||||
# "use-youtube-subs" toggle) can seed segments from it.
|
||||
for sf_path in sub_files:
|
||||
base = os.path.splitext(os.path.basename(sf_path))[0]
|
||||
# "original.en" → lang "en"; fallback to the trailing token.
|
||||
lang_tag = base.rsplit(".", 1)[-1] if "." in base else "und"
|
||||
segs = parse_vtt_segments(sf_path)
|
||||
if segs:
|
||||
youtube_subs_by_lang[lang_tag] = segs
|
||||
yield prep_event(
|
||||
"download_done",
|
||||
title=title, size=size, filename=filename,
|
||||
youtube_subs=sorted(youtube_subs_by_lang.keys()),
|
||||
)
|
||||
else:
|
||||
video_path = source["path"]
|
||||
filename = filename_hint or os.path.basename(video_path)
|
||||
@@ -391,6 +502,7 @@ async def ingest_pipeline(
|
||||
"segments": None,
|
||||
"dubbed_tracks": {},
|
||||
"scene_cuts": scene_cuts,
|
||||
"youtube_subs": youtube_subs_by_lang or None,
|
||||
}
|
||||
put_job(job_id, full_job)
|
||||
save_job(job_id, full_job, filename, dur, content_hash)
|
||||
@@ -413,6 +525,7 @@ async def ingest_pipeline(
|
||||
"segments": None,
|
||||
"dubbed_tracks": {},
|
||||
"scene_cuts": [],
|
||||
"youtube_subs": youtube_subs_by_lang or None,
|
||||
}
|
||||
put_job(job_id, partial)
|
||||
save_job(job_id, partial, filename, dur, content_hash)
|
||||
|
||||
@@ -0,0 +1,184 @@
|
||||
"""
|
||||
Translation engine registry + UI-driven install/uninstall.
|
||||
|
||||
This is the single source of truth for which translation providers we know
|
||||
about, what pip package they need, and whether that package is importable
|
||||
right now. The Engine dropdown in the Dub tab reads list_engines() to
|
||||
decide which options are ready-to-use vs. "needs install".
|
||||
|
||||
Why a registry rather than inline probes in dub_translate.py? The UI wants
|
||||
to render the availability table BEFORE the user clicks Translate, so we
|
||||
don't surface a cryptic ModuleNotFoundError for every segment. Having the
|
||||
registry live next to the dub_translate dispatch also means adding a new
|
||||
engine is one entry here + one branch in _build_translator.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import importlib
|
||||
import logging
|
||||
import os
|
||||
import shutil
|
||||
import sys
|
||||
|
||||
logger = logging.getLogger("omnivoice.translation_engines")
|
||||
|
||||
|
||||
# Engine ID → registry entry. Keyed by the `provider` string sent from the
|
||||
# frontend (must match the values of `translateProvider` in the store).
|
||||
REGISTRY: dict[str, dict] = {
|
||||
"argos": {
|
||||
"id": "argos",
|
||||
"display_name": "Argos (Local, Fast)",
|
||||
"pip_package": "argostranslate",
|
||||
"probe_module": "argostranslate",
|
||||
"category": "offline",
|
||||
"needs_key": False,
|
||||
"builtin": True,
|
||||
"notes": "Pure-CPU offline translator. Downloads a ~50MB language pack on first use per pair.",
|
||||
},
|
||||
"nllb": {
|
||||
"id": "nllb",
|
||||
"display_name": "NLLB-200 (Local, Heavy)",
|
||||
"pip_package": None, # uses HF transformers — already a core dep
|
||||
"probe_module": "transformers",
|
||||
"category": "offline",
|
||||
"needs_key": False,
|
||||
"builtin": True,
|
||||
"notes": "Meta's 200-language NMT model. Large download (~2.4GB), best offline quality.",
|
||||
},
|
||||
"google": {
|
||||
"id": "google",
|
||||
"display_name": "Google Translate (Online, Free)",
|
||||
"pip_package": "deep_translator",
|
||||
"probe_module": "deep_translator",
|
||||
"category": "online",
|
||||
"needs_key": False,
|
||||
"notes": "Free web endpoint via deep_translator. Rate-limited by Google; no API key required.",
|
||||
},
|
||||
"deepl": {
|
||||
"id": "deepl",
|
||||
"display_name": "DeepL (Online, Key)",
|
||||
"pip_package": "deep_translator",
|
||||
"probe_module": "deep_translator",
|
||||
"category": "online",
|
||||
"needs_key": True,
|
||||
"notes": "High-quality EU MT. Free tier: 500K chars/month. Set DEEPL_API_KEY.",
|
||||
},
|
||||
"microsoft": {
|
||||
"id": "microsoft",
|
||||
"display_name": "Microsoft Translator (Online, Key)",
|
||||
"pip_package": "deep_translator",
|
||||
"probe_module": "deep_translator",
|
||||
"category": "online",
|
||||
"needs_key": True,
|
||||
"notes": "Azure Cognitive Services. Free tier: 2M chars/month. Set MICROSOFT_API_KEY.",
|
||||
},
|
||||
"mymemory": {
|
||||
"id": "mymemory",
|
||||
"display_name": "MyMemory (Online, No Key)",
|
||||
"pip_package": "deep_translator",
|
||||
"probe_module": "deep_translator",
|
||||
"category": "online",
|
||||
"needs_key": False,
|
||||
"notes": "Crowdsourced MT. Free, 5K chars/day anonymous; more with an email param.",
|
||||
},
|
||||
"openai": {
|
||||
"id": "openai",
|
||||
"display_name": "LLM (OpenAI-compatible)",
|
||||
"pip_package": "openai",
|
||||
"probe_module": "openai",
|
||||
"category": "llm",
|
||||
"needs_key": True,
|
||||
"notes": (
|
||||
"Any OpenAI-compatible endpoint: GPT-4/5 (OpenAI), Claude (via OpenRouter), "
|
||||
"Gemini (OpenAI-compat mode), DeepSeek, Qwen, Ollama, LM Studio. "
|
||||
"Set TRANSLATE_BASE_URL + TRANSLATE_API_KEY + TRANSLATE_MODEL."
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def is_frozen() -> bool:
|
||||
"""True when running inside a packaged Tauri / PyInstaller bundle.
|
||||
|
||||
In that case the Python site-packages is read-only and signed, so we
|
||||
refuse install/uninstall requests instead of corrupting the bundle.
|
||||
"""
|
||||
return bool(getattr(sys, "frozen", False) or os.environ.get("OMNIVOICE_FROZEN"))
|
||||
|
||||
|
||||
def _probe(entry: dict) -> tuple[bool, str]:
|
||||
mod = entry.get("probe_module")
|
||||
if not mod:
|
||||
return True, "no module required"
|
||||
try:
|
||||
importlib.import_module(mod)
|
||||
return True, "ready"
|
||||
except ImportError as e:
|
||||
return False, f"import {mod!r} failed: {e}"
|
||||
|
||||
|
||||
def list_engines() -> list[dict]:
|
||||
"""Return a UI-ready list with per-engine availability stamped in."""
|
||||
out = []
|
||||
for e in REGISTRY.values():
|
||||
installed, reason = _probe(e)
|
||||
out.append({
|
||||
**e,
|
||||
"installed": installed,
|
||||
"availability_reason": reason,
|
||||
})
|
||||
return out
|
||||
|
||||
|
||||
def get_engine(engine_id: str) -> dict | None:
|
||||
return REGISTRY.get(engine_id)
|
||||
|
||||
|
||||
def is_installed(engine_id: str) -> bool:
|
||||
entry = REGISTRY.get(engine_id)
|
||||
if not entry:
|
||||
return False
|
||||
ok, _ = _probe(entry)
|
||||
return ok
|
||||
|
||||
|
||||
def _installer_cmd() -> list[str]:
|
||||
"""Prefer `uv pip` (the dev install's default), fall back to `python -m pip`.
|
||||
|
||||
Using `python -m pip` ensures we target the same interpreter the server
|
||||
is running under — avoids the classic "pip installed into the wrong venv"
|
||||
footgun.
|
||||
"""
|
||||
if shutil.which("uv"):
|
||||
return ["uv", "pip"]
|
||||
return [sys.executable, "-m", "pip"]
|
||||
|
||||
|
||||
async def run_pip(args: list[str], timeout: float = 600.0) -> tuple[int, str]:
|
||||
"""Run a pip command async and return (rc, combined_output).
|
||||
|
||||
Combines stdout + stderr so the UI can surface a useful tail on failure
|
||||
(pip's "ERROR: ..." lines go to stderr).
|
||||
"""
|
||||
cmd = _installer_cmd() + args
|
||||
logger.info("pip: %s", " ".join(cmd))
|
||||
try:
|
||||
proc = await asyncio.create_subprocess_exec(
|
||||
*cmd,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.STDOUT,
|
||||
)
|
||||
except FileNotFoundError as e:
|
||||
return 1, f"installer not found: {e}"
|
||||
try:
|
||||
stdout, _ = await asyncio.wait_for(proc.communicate(), timeout=timeout)
|
||||
except asyncio.TimeoutError:
|
||||
try:
|
||||
proc.kill()
|
||||
except ProcessLookupError:
|
||||
pass
|
||||
return 1, f"pip timed out after {timeout:.0f}s"
|
||||
out = stdout.decode(errors="replace") if stdout else ""
|
||||
return proc.returncode or 0, out
|
||||
@@ -322,11 +322,224 @@ class MossTTSNanoBackend(TTSBackend):
|
||||
return wav
|
||||
|
||||
|
||||
# ── KittenTTS (lightweight English "Turbo" tier) ────────────────────────────
|
||||
|
||||
|
||||
class KittenTTSBackend(TTSBackend):
|
||||
"""KittenML/KittenTTS — 25-80 MB ONNX model, 8 preset voices, English only.
|
||||
|
||||
Fills the ElevenLabs-Flash niche: when the caller just needs quick English
|
||||
narration (voiceover, demo reads, short phrases) with no reference sample.
|
||||
Runs CPU-realtime on any platform — no torch, no CUDA, no mlx. The
|
||||
trade-off vs OmniVoice is obvious:
|
||||
- No voice cloning (fixed preset voices)
|
||||
- English only
|
||||
- Much faster + much smaller install
|
||||
|
||||
Preset voice is chosen via `extras["voice"]` (defaults to "Jasper"). Any
|
||||
`ref_audio` / `instruct` / `language` arg is ignored with a log line so
|
||||
the common call-site doesn't need to know which engine it's talking to.
|
||||
"""
|
||||
|
||||
id = "kittentts"
|
||||
display_name = "KittenTTS (English, 8 preset voices, CPU realtime)"
|
||||
|
||||
PRESET_VOICES = [
|
||||
"expr-voice-2-m", "expr-voice-2-f",
|
||||
"expr-voice-3-m", "expr-voice-3-f",
|
||||
"expr-voice-4-m", "expr-voice-4-f",
|
||||
"expr-voice-5-m", "expr-voice-5-f",
|
||||
]
|
||||
DEFAULT_VOICE = "expr-voice-2-f"
|
||||
|
||||
def __init__(self):
|
||||
self._model = None
|
||||
|
||||
@classmethod
|
||||
def is_available(cls) -> tuple[bool, str]:
|
||||
try:
|
||||
import kittentts # noqa: F401
|
||||
return True, "ready"
|
||||
except ImportError as e:
|
||||
return False, f"kittentts not installed: {e}"
|
||||
|
||||
@property
|
||||
def sample_rate(self) -> int:
|
||||
# KittenTTS emits 24 kHz mono per its ONNX model config.
|
||||
return 24000
|
||||
|
||||
@property
|
||||
def supported_languages(self) -> list[str]:
|
||||
return ["en"]
|
||||
|
||||
def _ensure_loaded(self):
|
||||
if self._model is not None:
|
||||
return
|
||||
from kittentts import KittenTTS
|
||||
checkpoint = os.environ.get(
|
||||
"OMNIVOICE_KITTENTTS_MODEL", "KittenML/kitten-tts-mini-0.8"
|
||||
)
|
||||
logger.info("Loading KittenTTS from %s", checkpoint)
|
||||
self._model = KittenTTS(checkpoint)
|
||||
|
||||
def generate(self, text: str, **kw) -> torch.Tensor:
|
||||
import numpy as np
|
||||
self._ensure_loaded()
|
||||
|
||||
language = kw.get("language")
|
||||
if language and language.lower() not in {"en", "english", "auto"}:
|
||||
logger.info(
|
||||
"KittenTTS is English-only; ignoring language=%r — "
|
||||
"use OmniVoice for multilingual synthesis.",
|
||||
language,
|
||||
)
|
||||
|
||||
voice = kw.get("voice") or self.DEFAULT_VOICE
|
||||
if voice not in self.PRESET_VOICES:
|
||||
logger.info(
|
||||
"KittenTTS: unknown voice %r, falling back to %r. Valid: %s",
|
||||
voice, self.DEFAULT_VOICE, self.PRESET_VOICES,
|
||||
)
|
||||
voice = self.DEFAULT_VOICE
|
||||
|
||||
speed = float(kw.get("speed", 1.0))
|
||||
wav_np = self._model.generate(text, voice=voice, speed=speed)
|
||||
if not isinstance(wav_np, np.ndarray):
|
||||
wav_np = np.asarray(wav_np)
|
||||
wav = torch.from_numpy(wav_np).float()
|
||||
if wav.ndim == 1:
|
||||
wav = wav.unsqueeze(0)
|
||||
elif wav.ndim == 2 and wav.shape[0] > 1:
|
||||
wav = wav.mean(dim=0, keepdim=True)
|
||||
return wav
|
||||
|
||||
|
||||
# ── MLX-Audio (mac-ARM engine multiplexer) ──────────────────────────────────
|
||||
|
||||
|
||||
class MLXAudioBackend(TTSBackend):
|
||||
"""Blaizzy/mlx-audio — Apple-Silicon-only wrapper over 14+ TTS engines
|
||||
(Kokoro, CSM, Dia, Qwen3-TTS, Chatterbox, MeloTTS, OuteTTS, Spark,
|
||||
Higgs-Audio, Voxtral, LongCat-AudioDiT, KugelAudio, MingOmni, Soprano).
|
||||
|
||||
Exposed as a single backend with a `model_id` selector so the Settings
|
||||
UI can surface an engine picker within one adapter. The user switches
|
||||
models by setting `OMNIVOICE_MLX_AUDIO_MODEL` or picking from the UI —
|
||||
no code change per engine. Default is Kokoro (82M, multilingual, small).
|
||||
|
||||
Availability: requires mlx (Apple Silicon only). Skipped entirely on
|
||||
Linux/Windows/mac-Intel; the dep is platform-gated in pyproject.toml.
|
||||
"""
|
||||
|
||||
id = "mlx-audio"
|
||||
display_name = "MLX-Audio (mac-ARM, 14+ engines: Kokoro, CSM, Dia, Qwen3, …)"
|
||||
|
||||
# A curated subset surfaced by default — the full mlx-audio roster is
|
||||
# larger but these cover the useful tiers: small multilingual (Kokoro),
|
||||
# voice-clone (CSM), voice-design (Qwen3), European (Kugel), lightweight
|
||||
# VITS (MeloTTS). Users can point at any HF repo via OMNIVOICE_MLX_AUDIO_MODEL.
|
||||
CURATED_MODELS = {
|
||||
"kokoro": "mlx-community/Kokoro-82M-bf16",
|
||||
"csm": "mlx-community/csm-1b-8bit",
|
||||
"qwen3-tts": "mlx-community/Qwen3-TTS-1.7B-4bit",
|
||||
"dia": "mlx-community/Dia-1.6B",
|
||||
"chatterbox": "mlx-community/Chatterbox",
|
||||
"melotts": "mlx-community/MeloTTS",
|
||||
"outetts": "mlx-community/OuteTTS-0.3-500M",
|
||||
}
|
||||
DEFAULT_MODEL_KEY = "kokoro"
|
||||
|
||||
def __init__(self):
|
||||
self._model = None
|
||||
self._sr = 24000 # most mlx-audio engines emit 24 kHz mono
|
||||
key = os.environ.get("OMNIVOICE_MLX_AUDIO_MODEL", self.DEFAULT_MODEL_KEY)
|
||||
# Accept either a curated key ("kokoro") or a full HF repo id
|
||||
# ("mlx-community/Kokoro-82M-bf16") — flexibility for power users.
|
||||
self._model_id = self.CURATED_MODELS.get(key, key)
|
||||
|
||||
@classmethod
|
||||
def is_available(cls) -> tuple[bool, str]:
|
||||
try:
|
||||
import mlx_audio # noqa: F401
|
||||
return True, "ready"
|
||||
except ImportError as e:
|
||||
return False, (
|
||||
f"mlx-audio not installed: {e}. "
|
||||
"This backend is Apple Silicon only — available on mac-ARM dev "
|
||||
"installs; not shipped on Linux/Windows/mac-Intel."
|
||||
)
|
||||
|
||||
@property
|
||||
def sample_rate(self) -> int:
|
||||
return self._sr
|
||||
|
||||
@property
|
||||
def supported_languages(self) -> list[str]:
|
||||
# Per-model; Kokoro supports 8, Qwen3 ~4, Kugel 24. Return "multi"
|
||||
# so the language picker doesn't gate by engine — each engine
|
||||
# silently ignores languages it doesn't know.
|
||||
return ["multi"]
|
||||
|
||||
def _ensure_loaded(self):
|
||||
if self._model is not None:
|
||||
return
|
||||
from mlx_audio.tts.utils import load_model
|
||||
logger.info("Loading mlx-audio model %s", self._model_id)
|
||||
self._model = load_model(self._model_id)
|
||||
|
||||
def generate(self, text: str, **kw) -> torch.Tensor:
|
||||
import numpy as np
|
||||
self._ensure_loaded()
|
||||
|
||||
voice = kw.get("voice")
|
||||
ref_audio = kw.get("ref_audio")
|
||||
language = kw.get("language")
|
||||
speed = float(kw.get("speed", 1.0))
|
||||
|
||||
# mlx-audio's generate(...) returns an iterator of result objects,
|
||||
# each with a .audio attribute. Different engines accept different
|
||||
# kwargs (voice for Kokoro, ref_audio for CSM, instruct for Qwen3)
|
||||
# — we pass them all and let the engine ignore what it doesn't use.
|
||||
kwargs = {"text": text, "speed": speed}
|
||||
if voice: kwargs["voice"] = voice
|
||||
if ref_audio: kwargs["ref_audio"] = ref_audio
|
||||
if language: kwargs["lang_code"] = language[:2].lower()
|
||||
|
||||
pieces = []
|
||||
try:
|
||||
for result in self._model.generate(**kwargs):
|
||||
audio = getattr(result, "audio", result)
|
||||
if hasattr(audio, "numpy"):
|
||||
audio = audio.numpy()
|
||||
pieces.append(np.asarray(audio, dtype=np.float32))
|
||||
except TypeError:
|
||||
# Some engines don't accept lang_code / ref_audio. Retry with
|
||||
# only the universal kwargs.
|
||||
pieces = []
|
||||
for result in self._model.generate(text=text, speed=speed):
|
||||
audio = getattr(result, "audio", result)
|
||||
if hasattr(audio, "numpy"):
|
||||
audio = audio.numpy()
|
||||
pieces.append(np.asarray(audio, dtype=np.float32))
|
||||
|
||||
if not pieces:
|
||||
raise RuntimeError(f"mlx-audio ({self._model_id}) produced no audio")
|
||||
wav_np = np.concatenate(pieces, axis=-1)
|
||||
wav = torch.from_numpy(wav_np).float()
|
||||
if wav.ndim == 1:
|
||||
wav = wav.unsqueeze(0)
|
||||
elif wav.ndim == 2 and wav.shape[0] > 1:
|
||||
wav = wav.mean(dim=0, keepdim=True)
|
||||
return wav
|
||||
|
||||
|
||||
# ── Registry ────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
_REGISTRY: dict[str, type[TTSBackend]] = {
|
||||
"omnivoice": OmniVoiceBackend,
|
||||
"kittentts": KittenTTSBackend,
|
||||
"mlx-audio": MLXAudioBackend,
|
||||
"voxcpm2": VoxCPM2Backend,
|
||||
"moss-tts-nano": MossTTSNanoBackend,
|
||||
}
|
||||
|
||||
@@ -0,0 +1,146 @@
|
||||
"""HuggingFace download progress — one monkey-patch, every `hf_hub_download`
|
||||
reports bytes downloaded through a central callback.
|
||||
|
||||
Pattern lifted from jamiepine/voicebox's `backend/utils/hf_progress.py`:
|
||||
`huggingface_hub` uses tqdm for progress bars; we subclass it, intercept
|
||||
`update()` calls, and forward (filename, downloaded_bytes, total_bytes) to
|
||||
whatever callback is registered. No changes to calling sites across
|
||||
transformers / mlx_whisper / diffusers / accelerate — they all route through
|
||||
`hf_hub_download`, which uses the patched tqdm.
|
||||
|
||||
Usage:
|
||||
from utils.hf_progress import install, register_listener, unregister_listener
|
||||
|
||||
install() # once at app startup
|
||||
listener_id = register_listener(lambda ev: print(ev))
|
||||
# …models download, listener fires…
|
||||
unregister_listener(listener_id)
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import itertools
|
||||
import logging
|
||||
import threading
|
||||
from typing import Callable, Optional
|
||||
|
||||
logger = logging.getLogger("omnivoice.hf_progress")
|
||||
|
||||
# Event shape forwarded to listeners. Typed loosely on purpose — SSE encodes
|
||||
# it as JSON so consumers read the dict directly.
|
||||
# {
|
||||
# "filename": str, # desc on the tqdm bar, usually the HF file path
|
||||
# "downloaded": int, # bytes pulled so far
|
||||
# "total": int | None, # total bytes or None if unknown
|
||||
# "pct": float, # 0.0-1.0 (or 0.0 if total unknown)
|
||||
# "phase": "start"|"progress"|"done",
|
||||
# }
|
||||
ProgressEvent = dict
|
||||
Listener = Callable[[ProgressEvent], None]
|
||||
|
||||
_listeners: dict[int, Listener] = {}
|
||||
_listener_lock = threading.Lock()
|
||||
_listener_counter = itertools.count(1)
|
||||
_installed = False
|
||||
_install_lock = threading.Lock()
|
||||
|
||||
|
||||
def register_listener(cb: Listener) -> int:
|
||||
"""Register a callback that receives progress events. Returns an id that
|
||||
can be passed to `unregister_listener` when the listener is done."""
|
||||
with _listener_lock:
|
||||
lid = next(_listener_counter)
|
||||
_listeners[lid] = cb
|
||||
return lid
|
||||
|
||||
|
||||
def unregister_listener(lid: int) -> None:
|
||||
with _listener_lock:
|
||||
_listeners.pop(lid, None)
|
||||
|
||||
|
||||
def _emit(event: ProgressEvent) -> None:
|
||||
"""Fan out to all registered listeners. Never raise — a bad listener
|
||||
shouldn't break a download."""
|
||||
with _listener_lock:
|
||||
listeners = list(_listeners.values())
|
||||
for cb in listeners:
|
||||
try:
|
||||
cb(event)
|
||||
except Exception as e: # noqa: BLE001
|
||||
logger.debug("hf_progress listener raised: %s", e)
|
||||
|
||||
|
||||
def install() -> None:
|
||||
"""Monkey-patch `huggingface_hub`'s tqdm so every download reports to our
|
||||
listeners. Safe to call multiple times — second call is a no-op."""
|
||||
global _installed
|
||||
with _install_lock:
|
||||
if _installed:
|
||||
return
|
||||
# `huggingface_hub.utils.__init__` does `from .tqdm import tqdm`,
|
||||
# which shadows the `tqdm` SUBMODULE with the CLASS of the same name
|
||||
# when accessed via attribute lookup. Pull the real module out of
|
||||
# sys.modules after an explicit import so we patch the right thing.
|
||||
try:
|
||||
import sys
|
||||
import huggingface_hub.utils.tqdm # noqa: F401
|
||||
hf_tqdm_module = sys.modules.get("huggingface_hub.utils.tqdm")
|
||||
if hf_tqdm_module is None:
|
||||
raise ImportError("huggingface_hub.utils.tqdm not in sys.modules after import")
|
||||
except Exception as e: # noqa: BLE001
|
||||
logger.warning(
|
||||
"hf_progress.install: huggingface_hub.utils.tqdm missing (%s); "
|
||||
"progress tracking disabled.", e,
|
||||
)
|
||||
return
|
||||
|
||||
original = getattr(hf_tqdm_module, "tqdm", None)
|
||||
if original is None or not isinstance(original, type):
|
||||
logger.warning("hf_progress.install: no `tqdm` class on the module; aborting")
|
||||
return
|
||||
|
||||
class TrackedTqdm(original): # type: ignore[misc,valid-type]
|
||||
"""tqdm subclass that emits a progress event on every update."""
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
# Emit once on construction so the UI can show the file
|
||||
# before a single byte is read. Some tqdm variants don't
|
||||
# populate `desc` / `n` as attributes — use getattr so a
|
||||
# patched tqdm never crashes the whole model load.
|
||||
try:
|
||||
desc = getattr(self, "desc", None)
|
||||
total = int(getattr(self, "total", 0) or 0)
|
||||
_emit({
|
||||
"filename": str(desc or "download"),
|
||||
"downloaded": 0,
|
||||
"total": total,
|
||||
"pct": 0.0,
|
||||
"phase": "start",
|
||||
})
|
||||
except Exception:
|
||||
# Never let progress telemetry break a real download.
|
||||
pass
|
||||
|
||||
def update(self, n=1):
|
||||
super().update(n)
|
||||
try:
|
||||
desc = getattr(self, "desc", None)
|
||||
total = int(getattr(self, "total", 0) or 0)
|
||||
done = int(getattr(self, "n", 0) or 0)
|
||||
pct = (done / total) if total > 0 else 0.0
|
||||
_emit({
|
||||
"filename": str(desc or "download"),
|
||||
"downloaded": done,
|
||||
"total": total,
|
||||
"pct": pct,
|
||||
"phase": "done" if (total > 0 and done >= total) else "progress",
|
||||
})
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Stash the original for inspection / uninstall, then swap.
|
||||
hf_tqdm_module._omnivoice_original_tqdm = original # type: ignore[attr-defined]
|
||||
hf_tqdm_module.tqdm = TrackedTqdm # type: ignore[assignment]
|
||||
_installed = True
|
||||
logger.info("hf_progress: installed tqdm patch on huggingface_hub.utils.tqdm")
|
||||
+36
-3
@@ -8,7 +8,7 @@ version = "0.2.0"
|
||||
description = "OmniVoice: Towards Omnilingual Zero-Shot Text-to-Speech with Diffusion Language Models"
|
||||
readme = "README.md"
|
||||
license = "Apache-2.0"
|
||||
requires-python = ">=3.10"
|
||||
requires-python = ">=3.11"
|
||||
authors = [{name = "Han Zhu"}]
|
||||
keywords = [
|
||||
"tts",
|
||||
@@ -41,14 +41,41 @@ dependencies = [
|
||||
"numpy",
|
||||
"soundfile",
|
||||
"psutil>=7.2.2",
|
||||
"pyannote-audio>=4.0.4",
|
||||
# Pinned to 3.x — pyannote 4.x removed `use_auth_token` from `Inference`
|
||||
# which whisperx 3.4.2 still passes, blowing up `whisperx.load_model()`
|
||||
# with TypeError. whisperx tests against pyannote 3.3.2+, so we track
|
||||
# that range and revisit when whisperx releases a 4-compatible build.
|
||||
"pyannote-audio>=3.3.2,<4.0",
|
||||
"pyinstaller>=6.19.0",
|
||||
"imageio-ffmpeg>=0.6.0",
|
||||
"pedalboard>=0.9.14",
|
||||
"mlx-whisper>=0.2.1",
|
||||
# Primary ASR — cross-platform, CTranslate2-based under the hood. WhisperX
|
||||
# adds wav2vec2 forced alignment (±10-30 ms word timing vs Whisper's own
|
||||
# ±100-300 ms) which directly improves lip-sync on the dub pipeline.
|
||||
# Pulls `faster-whisper` transitively, so a WhisperX install also provides
|
||||
# the plain faster-whisper backend as a fallback for rare-language audio
|
||||
# where no wav2vec2 alignment model exists.
|
||||
"whisperx>=3.1.0",
|
||||
"faster-whisper>=1.0.0",
|
||||
# Apple Silicon-only speedup; skipped everywhere else so `uv sync` can
|
||||
# succeed on Linux/Windows/mac-Intel (no mlx wheels exist for those).
|
||||
"mlx-whisper>=0.2.1 ; sys_platform == 'darwin' and platform_machine == 'arm64'",
|
||||
# Apple Silicon-only rich TTS library — 14+ engines (Kokoro, CSM, Dia,
|
||||
# Qwen3-TTS, Chatterbox, MeloTTS, OuteTTS, Spark, Higgs-Audio, Voxtral,
|
||||
# …). Gives mac-ARM users a broad engine picker in Settings. Also gated
|
||||
# by platform markers because it depends on mlx (Apple Silicon only).
|
||||
"mlx-audio>=0.3.0 ; sys_platform == 'darwin' and platform_machine == 'arm64'",
|
||||
"demucs>=4.0.1",
|
||||
"yt-dlp>=2024.12.13",
|
||||
"alembic>=1.13",
|
||||
# Lightweight English TTS "Turbo" tier — 25-80 MB ONNX model, 8 preset
|
||||
# voices (Bella, Jasper, Luna, Bruno, Rosie, Hugo, Kiki, Leo), CPU
|
||||
# realtime on any platform. Complements OmniVoice's 2.4 GB multilingual
|
||||
# zero-shot clone: when the caller just needs fast English narration with
|
||||
# no reference sample, this is ~100× smaller + ~10× faster. Pinned to
|
||||
# the exact wheel because the project is in developer preview and the
|
||||
# 0.x API is explicitly unstable.
|
||||
"kittentts @ https://github.com/KittenML/KittenTTS/releases/download/0.8.1/kittentts-0.8.1-py3-none-any.whl",
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
@@ -109,6 +136,12 @@ constraint-dependencies = [
|
||||
"torchaudio==2.8.0",
|
||||
]
|
||||
|
||||
[tool.hatch.metadata]
|
||||
# Needed so the KittenTTS wheel-URL dep in `project.dependencies` is accepted
|
||||
# by hatchling's metadata validator. KittenTTS isn't on PyPI (dev preview),
|
||||
# so pulling it via GH Releases URL is the only option today.
|
||||
allow-direct-references = true
|
||||
|
||||
[tool.hatch.build.targets.sdist]
|
||||
include = ["omnivoice"]
|
||||
|
||||
|
||||
Reference in New Issue
Block a user