subtatix 0.1.0__tar.gz

subtatix-0.1.0/PKG-INFO

@@ -0,0 +1,129 @@
+Metadata-Version: 2.4
+Name: subtatix
+Version: 0.1.0
+Summary: CLI for generating and translating SRT subtitles with WhisperX
+Keywords: cli,subtitles,srt,whisperx,transcription,translation
+Author: Chris Paganon
+Author-email: Chris Paganon <info@chrispaganon.com>
+License-Expression: BSD-2-Clause
+Classifier: Development Status :: 3 - Alpha
+Classifier: Environment :: Console
+Classifier: Intended Audience :: End Users/Desktop
+Classifier: License :: OSI Approved :: BSD License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Multimedia :: Sound/Audio :: Speech
+Classifier: Topic :: Text Processing
+Classifier: Topic :: Utilities
+Requires-Dist: accelerate>=1.13.0
+Requires-Dist: sentencepiece>=0.2.1
+Requires-Dist: torch>=2.8.0
+Requires-Dist: transformers>=4.57.6
+Requires-Dist: typer>=0.25.1
+Requires-Dist: whisperx>=3.8.5
+Requires-Python: >=3.12
+Project-URL: Homepage, https://github.com/chris-paganon/subtatix
+Project-URL: Repository, https://github.com/chris-paganon/subtatix
+Project-URL: Issues, https://github.com/chris-paganon/subtatix/issues
+Description-Content-Type: text/markdown
+
+# Subtatix
+
+`Subtatix` is a small CLI for generating `.srt` subtitles from audio or video files with [WhisperX](https://github.com/m-bain/whisperX), with optional subtitle translation.
+
+It transcribes the input, aligns subtitle timings with WhisperX, and can then translate the resulting subtitle lines into another language.
+
+## Requirements
+
+- Python 3.12+
+- `ffmpeg` installed separately and available on your `PATH`
+- Enough disk space for model downloads and caching
+
+The first run will be slower because WhisperX and translation models need to be downloaded. Subsequent runs reuse the cached models and do not need to download them again unless the cache is cleared.
+
+`ffmpeg` is an external system dependency. It is not installed by `pip`, `uvx`, or `uv tool install`.
+
+## Installation
+
+Run without installing:
+
+```bash
+uvx subtatix --help
+```
+
+Install as a tool with `uv`:
+
+```bash
+uv tool install subtatix
+```
+
+Install with `pip`:
+
+```bash
+pip install subtatix
+```
+
+## Usage
+
+Run the CLI:
+
+```bash
+subtatix input.mp4
+```
+
+Transcribe to a specific output path:
+
+```bash
+subtatix input.mp4 --output some-path/some-file-name
+```
+
+`--output` is a base path, not a full `.srt` filename. This writes `some-path/some-file-name.srt`. If you also translate to Spanish, it writes `some-path/some-file-name.es.srt`.
+
+Set the source language explicitly:
+
+```bash
+subtatix input.mp4 --source-language en
+```
+
+Translate after transcription:
+
+```bash
+subtatix input.mp4 --to es
+```
+
+This writes both the original transcription SRT and the translated SRT by default.
+
+If CUDA runs out of memory on larger files, reduce the batch size or force CPU mode:
+
+```bash
+subtatix input.mp4 --batch-size 4
+subtatix input.mp4 --device cpu
+```
+
+To discard the original transcription and only keep the translated output:
+
+```bash
+subtatix input.mp4 --to es --discard-transcription
+```
+
+Passing an `--output` value that ends in `.srt` is rejected. Use a base path such as `--output subtitles` instead.
+
+List supported language codes:
+
+```bash
+subtatix --list-languages
+subtatix --list-target-languages
+```
+
+## Models
+
+By default, transcription uses WhisperX with the Whisper model `large-v2`. This is a good general default when you want higher transcription quality and aligned subtitle timings, but it is heavier and slower than smaller Whisper models.
+
+Translation uses `facebook/nllb-200-1.3B`. The CLI accepts simple target codes such as `en`, `es`, `fr`, `de`, `pt`, `ja`, `ko`, `zh`, and also raw NLLB codes such as `spa_Latn`.
+
+Other model options can also be used:
+
+- For transcription, you can pass another Whisper model with `--model`, such as `small`, `medium`, or `large-v3`, depending on your speed and accuracy needs.
+- For translation, the code currently defaults to the NLLB model above, but the translation layer is built around Hugging Face seq2seq models and could be adapted to use a different multilingual translation model if needed.

subtatix-0.1.0/README.md

@@ -0,0 +1,98 @@
+# Subtatix
+
+`Subtatix` is a small CLI for generating `.srt` subtitles from audio or video files with [WhisperX](https://github.com/m-bain/whisperX), with optional subtitle translation.
+
+It transcribes the input, aligns subtitle timings with WhisperX, and can then translate the resulting subtitle lines into another language.
+
+## Requirements
+
+- Python 3.12+
+- `ffmpeg` installed separately and available on your `PATH`
+- Enough disk space for model downloads and caching
+
+The first run will be slower because WhisperX and translation models need to be downloaded. Subsequent runs reuse the cached models and do not need to download them again unless the cache is cleared.
+
+`ffmpeg` is an external system dependency. It is not installed by `pip`, `uvx`, or `uv tool install`.
+
+## Installation
+
+Run without installing:
+
+```bash
+uvx subtatix --help
+```
+
+Install as a tool with `uv`:
+
+```bash
+uv tool install subtatix
+```
+
+Install with `pip`:
+
+```bash
+pip install subtatix
+```
+
+## Usage
+
+Run the CLI:
+
+```bash
+subtatix input.mp4
+```
+
+Transcribe to a specific output path:
+
+```bash
+subtatix input.mp4 --output some-path/some-file-name
+```
+
+`--output` is a base path, not a full `.srt` filename. This writes `some-path/some-file-name.srt`. If you also translate to Spanish, it writes `some-path/some-file-name.es.srt`.
+
+Set the source language explicitly:
+
+```bash
+subtatix input.mp4 --source-language en
+```
+
+Translate after transcription:
+
+```bash
+subtatix input.mp4 --to es
+```
+
+This writes both the original transcription SRT and the translated SRT by default.
+
+If CUDA runs out of memory on larger files, reduce the batch size or force CPU mode:
+
+```bash
+subtatix input.mp4 --batch-size 4
+subtatix input.mp4 --device cpu
+```
+
+To discard the original transcription and only keep the translated output:
+
+```bash
+subtatix input.mp4 --to es --discard-transcription
+```
+
+Passing an `--output` value that ends in `.srt` is rejected. Use a base path such as `--output subtitles` instead.
+
+List supported language codes:
+
+```bash
+subtatix --list-languages
+subtatix --list-target-languages
+```
+
+## Models
+
+By default, transcription uses WhisperX with the Whisper model `large-v2`. This is a good general default when you want higher transcription quality and aligned subtitle timings, but it is heavier and slower than smaller Whisper models.
+
+Translation uses `facebook/nllb-200-1.3B`. The CLI accepts simple target codes such as `en`, `es`, `fr`, `de`, `pt`, `ja`, `ko`, `zh`, and also raw NLLB codes such as `spa_Latn`.
+
+Other model options can also be used:
+
+- For transcription, you can pass another Whisper model with `--model`, such as `small`, `medium`, or `large-v3`, depending on your speed and accuracy needs.
+- For translation, the code currently defaults to the NLLB model above, but the translation layer is built around Hugging Face seq2seq models and could be adapted to use a different multilingual translation model if needed.

subtatix-0.1.0/pyproject.toml

@@ -0,0 +1,44 @@
+[project]
+name = "subtatix"
+version = "0.1.0"
+description = "CLI for generating and translating SRT subtitles with WhisperX"
+readme = "README.md"
+license = "BSD-2-Clause"
+authors = [
+    { name = "Chris Paganon", email = "info@chrispaganon.com" }
+]
+requires-python = ">=3.12"
+keywords = ["cli", "subtitles", "srt", "whisperx", "transcription", "translation"]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Environment :: Console",
+    "Intended Audience :: End Users/Desktop",
+    "License :: OSI Approved :: BSD License",
+    "Operating System :: OS Independent",
+    "Programming Language :: Python",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.12",
+    "Topic :: Multimedia :: Sound/Audio :: Speech",
+    "Topic :: Text Processing",
+    "Topic :: Utilities",
+]
+dependencies = [
+    "accelerate>=1.13.0",
+    "sentencepiece>=0.2.1",
+    "torch>=2.8.0",
+    "transformers>=4.57.6",
+    "typer>=0.25.1",
+    "whisperx>=3.8.5",
+]
+
+[project.urls]
+Homepage = "https://github.com/chris-paganon/subtatix"
+Repository = "https://github.com/chris-paganon/subtatix"
+Issues = "https://github.com/chris-paganon/subtatix/issues"
+
+[project.scripts]
+subtatix = "subtatix.__main__:main"
+
+[build-system]
+requires = ["uv_build>=0.11.13,<0.12.0"]
+build-backend = "uv_build"

subtatix-0.1.0/src/subtatix/__init__.py

@@ -0,0 +1 @@
+__all__ = []

subtatix-0.1.0/src/subtatix/__main__.py

@@ -0,0 +1,5 @@
+from subtatix.cli import main
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

subtatix-0.1.0/src/subtatix/cli.py

@@ -0,0 +1,273 @@
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+from typing import Annotated
+
+import typer
+from click.exceptions import ClickException
+
+from subtatix.errors import SubtatixError
+
+from subtatix.subtitles import (
+    DEFAULT_MODEL,
+    SUPPORTED_SOURCE_LANGUAGE_CODES,
+    require_ffmpeg,
+    transcribe_to_srt,
+)
+from subtatix.runtime import configure_runtime_noise
+from subtatix.translation import (
+    DEFAULT_TRANSLATION_BATCH_SIZE,
+    SUPPORTED_TARGET_LANGUAGE_CODES,
+    get_available_nllb_languages,
+    translate_subtitles,
+)
+
+app = typer.Typer(
+    add_completion=False,
+    context_settings={"help_option_names": ["-h", "--help"]},
+    help=(
+        "Transcribe an audio or video file to SRT with WhisperX. "
+        "Without --to, the tool only transcribes. Passing --to also translates the "
+        "subtitles and keeps the original transcribed SRT unless "
+        "--discard-transcription is used. Output names are generated as "
+        "'.srt' and translated variants like '.es.srt'. "
+        "Use the language listing flags to inspect supported source, mapped target, "
+        "and raw NLLB codes."
+    ),
+)
+
+class ProgressBar:
+    def __init__(self, label: str, total: int, width: int = 28) -> None:
+        self._label = label
+        self._total = max(1, total)
+        self._width = width
+        self._current = 0
+        self._visible = False
+
+    def __enter__(self) -> "ProgressBar":
+        return self
+
+    def __exit__(self, exc_type: object, exc: object, tb: object) -> None:
+        if exc_type is None and self._current < self._total:
+            self._current = self._total
+            self._render()
+        if self._visible:
+            sys.stderr.write("\n")
+            sys.stderr.flush()
+
+    def update_to(self, value: int) -> None:
+        bounded = min(self._total, max(0, value))
+        if bounded != self._current:
+            self._current = bounded
+            self._render()
+
+    def reset(self) -> None:
+        if self._current != 0:
+            self._clear_line()
+            self._current = 0
+            self._visible = False
+
+    def write_message(self, message: str) -> None:
+        if self._visible:
+            sys.stderr.write("\n")
+        sys.stderr.write(f"{message}\n")
+        sys.stderr.flush()
+        if self._visible and self._current < self._total:
+            self._render()
+
+    def _render(self) -> None:
+        filled = round((self._current / self._total) * self._width)
+        empty = self._width - filled
+        percent = round((self._current / self._total) * 100)
+        self._visible = True
+        sys.stderr.write(
+            f"\r{self._label} [{'#' * filled}{'.' * empty}] {percent:>3}%"
+        )
+        sys.stderr.flush()
+
+    def _clear_line(self) -> None:
+        sys.stderr.write("\r" + (" " * (self._width + len(self._label) + 10)) + "\r")
+        sys.stderr.flush()
+
+
+@app.command()
+def run(
+    input_file: Annotated[
+        Path | None,
+        typer.Argument(help="Path to the input audio or video file."),
+    ] = None,
+    target_language: Annotated[
+        str | None,
+        typer.Option(
+            "--to",
+            "--target-language",
+            "-t",
+            help=(
+                "Translate to this language. Use one of the mapped Whisper target "
+                "codes or a raw NLLB code like 'spa_Latn'. If omitted, the tool only "
+                "transcribes. Use --list-target-languages to inspect supported values."
+            ),
+        ),
+    ] = None,
+    source_language: Annotated[
+        str | None,
+        typer.Option(
+            "--source-language",
+            "-s",
+            help=(
+                "Optional Whisper source language code to skip language detection, "
+                "for example 'en', 'es', or 'fr'."
+            ),
+        ),
+    ] = None,
+    output: Annotated[
+        Path | None,
+        typer.Option(
+            "--output",
+            "-o",
+            help=(
+                "Base output path without a .srt suffix. Subtatix writes "
+                "'.srt' and translated variants like '.es.srt'. If a directory is "
+                "provided, the default filename based on the input file is used inside it."
+            ),
+        ),
+    ] = None,
+    model: Annotated[
+        str,
+        typer.Option(help=f"Whisper model name to use. Default: {DEFAULT_MODEL}."),
+    ] = DEFAULT_MODEL,
+    batch_size: Annotated[
+        int,
+        typer.Option(
+            "--batch-size",
+            help="Batch size for Whisper inference. Reduce this if you run out of GPU memory.",
+        ),
+    ] = 8,
+    device: Annotated[
+        str,
+        typer.Option(
+            "--device",
+            help=(
+                "Execution device for WhisperX: 'auto', 'cuda', or 'cpu'. "
+                "Default: auto."
+            ),
+        ),
+    ] = "auto",
+    discard_transcription: Annotated[
+        bool,
+        typer.Option(
+            "--discard-transcription",
+            help=(
+                "When used with --to, do not save the original untranslated SRT file. "
+                "Without --to, the transcribed SRT is still written."
+            ),
+            is_flag=True,
+        ),
+    ] = False,
+    list_languages: Annotated[
+        bool,
+        typer.Option(
+            "--list-languages",
+            help=(
+                "List source Whisper codes and the convenience target codes. "
+                "Use --list-target-languages for the full raw NLLB target list."
+            ),
+            is_flag=True,
+        ),
+    ] = False,
+    list_source_languages: Annotated[
+        bool,
+        typer.Option(
+            "--list-source-languages",
+            help="List supported Whisper source language codes.",
+            is_flag=True,
+        ),
+    ] = False,
+    list_target_languages: Annotated[
+        bool,
+        typer.Option(
+            "--list-target-languages",
+            help=(
+                "List the convenience target codes for --to, followed by the full raw "
+                "NLLB target language codes accepted by --to."
+            ),
+            is_flag=True,
+        ),
+    ] = False,
+) -> None:
+    if list_languages or list_source_languages or list_target_languages:
+        if input_file is not None:
+            raise typer.BadParameter(
+                "INPUT_FILE cannot be used with language listing options."
+            )
+        if list_languages or list_source_languages:
+            typer.echo("Source languages (Whisper codes):")
+            typer.echo(", ".join(SUPPORTED_SOURCE_LANGUAGE_CODES))
+            typer.echo()
+        if list_languages or list_target_languages:
+            typer.echo(
+                "Convenience target languages (--to Whisper-style codes mapped to NLLB):"
+            )
+            typer.echo(", ".join(SUPPORTED_TARGET_LANGUAGE_CODES))
+            typer.echo()
+        if list_languages:
+            typer.echo(
+                "Use --list-target-languages to also show the full raw NLLB target list."
+            )
+        if list_target_languages:
+            typer.echo("Raw NLLB target languages (--to NLLB codes):")
+            typer.echo(", ".join(get_available_nllb_languages()))
+        return
+
+    if input_file is None:
+        raise typer.BadParameter("Missing argument 'INPUT_FILE'.")
+
+    configure_runtime_noise()
+    require_ffmpeg()
+    write_original_srt = target_language is None or not discard_transcription
+    with ProgressBar("Transcription", 100) as transcription_progress:
+        document = transcribe_to_srt(
+            input_file=input_file,
+            model_name=model,
+            batch_size=batch_size,
+            output_file=output,
+            write_output=write_original_srt,
+            source_language=source_language,
+            device_preference=device,
+            log=transcription_progress.write_message,
+            progress_callback=lambda percent: transcription_progress.update_to(
+                round(percent)
+            ),
+            progress_reset=transcription_progress.reset,
+        )
+    if write_original_srt:
+        typer.echo(document.subtitle_path)
+
+    if target_language:
+        total_batches = max(
+            1,
+            (len(document.cues) + DEFAULT_TRANSLATION_BATCH_SIZE - 1)
+            // DEFAULT_TRANSLATION_BATCH_SIZE,
+        )
+        with ProgressBar("Translation", total_batches) as translation_progress:
+            translated_path = translate_subtitles(
+                document=document,
+                target_language=target_language,
+                progress_callback=lambda batch_index, total: translation_progress.update_to(
+                    batch_index
+                ),
+            )
+        typer.echo(translated_path)
+
+
+def main() -> int:
+    try:
+        app(standalone_mode=False)
+    except SubtatixError as error:
+        ClickException(str(error)).show()
+        return 1
+    except ClickException as error:
+        error.show()
+        return error.exit_code
+    return 0

subtatix-0.1.0/src/subtatix/errors.py

@@ -0,0 +1,5 @@
+from __future__ import annotations
+
+
+class SubtatixError(Exception):
+    """Expected user-facing error that should be rendered without a traceback."""

subtatix-0.1.0/src/subtatix/runtime.py

@@ -0,0 +1,65 @@
+from __future__ import annotations
+
+import gc
+import logging
+import warnings
+
+import torch
+
+from subtatix.errors import SubtatixError
+
+
+def get_device(preferred: str = "auto") -> str:
+    normalized = preferred.strip().lower()
+    if normalized not in {"auto", "cpu", "cuda"}:
+        raise SubtatixError(
+            f"Unsupported device '{preferred}'. Use 'auto', 'cpu', or 'cuda'."
+        )
+    if normalized == "cpu":
+        return "cpu"
+    if normalized == "cuda":
+        if not torch.cuda.is_available():
+            raise SubtatixError(
+                "CUDA was requested but is not available on this system."
+            )
+        return "cuda"
+    if torch.cuda.is_available():
+        return "cuda"
+    return "cpu"
+
+
+def get_whisperx_runtime(preferred_device: str = "auto") -> tuple[str, str]:
+    device = get_device(preferred_device)
+    if device == "cuda":
+        return device, "float16"
+    return device, "float32"
+
+
+def configure_runtime_noise() -> None:
+    for logger_name in (
+        "whisperx",
+        "whisperx.asr",
+        "whisperx.vads",
+        "whisperx.vads.pyannote",
+        "lightning",
+        "lightning.pytorch",
+        "lightning.pytorch.utilities.migration",
+        "lightning.pytorch.utilities.migration.utils",
+        "pytorch_lightning",
+        "pytorch_lightning.utilities.migration",
+        "pytorch_lightning.utilities.migration.utils",
+    ):
+        logging.getLogger(logger_name).setLevel(logging.ERROR)
+
+    warnings.filterwarnings(
+        "ignore",
+        message=r"TensorFloat-32 \(TF32\) has been disabled.*",
+        module=r"pyannote\.audio\.utils\.reproducibility",
+    )
+
+
+def release_memory() -> None:
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.ipc_collect()

subtatix-0.1.0/src/subtatix/subtitles.py

@@ -0,0 +1,256 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+import shutil
+from pathlib import Path
+from typing import Callable
+
+from faster_whisper.tokenizer import _LANGUAGE_CODES
+import whisperx
+
+from subtatix.errors import SubtatixError
+from subtatix.runtime import get_whisperx_runtime, release_memory
+
+
+DEFAULT_MODEL = "large-v2"
+SUPPORTED_SOURCE_LANGUAGE_CODES = tuple(sorted(_LANGUAGE_CODES))
+
+
+@dataclass(frozen=True)
+class SubtitleCue:
+    start: float
+    end: float
+    text: str
+
+
+@dataclass(frozen=True)
+class SubtitleDocument:
+    source_language: str
+    subtitle_path: Path
+    cues: list[SubtitleCue]
+
+
+def require_ffmpeg() -> None:
+    if shutil.which("ffmpeg") is None:
+        raise SubtatixError(
+            "ffmpeg is required but was not found on PATH. Install ffmpeg and try again."
+        )
+
+
+def normalize_source_language(language: str) -> str:
+    language_code = language.strip().lower()
+    if language_code in SUPPORTED_SOURCE_LANGUAGE_CODES:
+        return language_code
+
+    raise SubtatixError(
+        f"Unsupported source language '{language}'. Use a Whisper language code like "
+        "'en', 'es', or 'fr'."
+    )
+
+
+def resolve_output_path(input_file: Path, output_file: Path | None) -> Path:
+    default_output_path = input_file.with_suffix(".srt")
+    if output_file is None:
+        return default_output_path
+
+    output_file = output_file.expanduser()
+    if output_file.exists() and output_file.is_dir():
+        return (output_file / default_output_path.name).resolve()
+    if output_file.suffix == ".srt":
+        raise SubtatixError(
+            "Output paths must not end in '.srt'. Pass a base output path like "
+            "'some-path/some-file-name' so Subtatix can write '.srt' and translated variants like "
+            "'.es.srt'."
+        )
+    return output_file.with_suffix(".srt").resolve()
+
+
+def format_srt_timestamp(seconds: float) -> str:
+    total_milliseconds = max(0, round(seconds * 1000))
+    hours, remainder = divmod(total_milliseconds, 3_600_000)
+    minutes, remainder = divmod(remainder, 60_000)
+    secs, milliseconds = divmod(remainder, 1000)
+    return f"{hours:02}:{minutes:02}:{secs:02},{milliseconds:03}"
+
+
+def build_cues(aligned_transcription: dict) -> list[SubtitleCue]:
+    cues: list[SubtitleCue] = []
+    for segment in aligned_transcription["segments"]:
+        text = segment["text"].strip()
+        if not text:
+            continue
+        cues.append(
+            SubtitleCue(
+                start=float(segment["start"]),
+                end=float(segment["end"]),
+                text=text,
+            )
+        )
+    return cues
+
+
+def write_srt(subtitle_path: Path, cues: list[SubtitleCue]) -> None:
+    subtitle_path.parent.mkdir(parents=True, exist_ok=True)
+    blocks = []
+    for index, cue in enumerate(cues, start=1):
+        blocks.append(
+            "\n".join(
+                [
+                    str(index),
+                    (
+                        f"{format_srt_timestamp(cue.start)} --> "
+                        f"{format_srt_timestamp(cue.end)}"
+                    ),
+                    cue.text,
+                ]
+            )
+        )
+    subtitle_path.write_text("\n\n".join(blocks) + "\n", encoding="utf-8")
+
+
+def is_cuda_oom(error: RuntimeError) -> bool:
+    message = str(error).lower()
+    return "cuda" in message and "out of memory" in message
+
+
+def iter_retry_batch_sizes(batch_size: int) -> list[int]:
+    sizes: list[int] = []
+    current = max(1, batch_size)
+    while current not in sizes:
+        sizes.append(current)
+        if current == 1:
+            break
+        current = max(1, current // 2)
+    return sizes
+
+
+def transcribe_with_backoff(
+    whisper_model: object,
+    audio: object,
+    batch_size: int,
+    source_language: str | None,
+    device: str,
+    log: Callable[[str], None] | None = None,
+    progress_callback: Callable[[float], None] | None = None,
+    progress_reset: Callable[[], None] | None = None,
+) -> tuple[dict, int]:
+    attempts = iter_retry_batch_sizes(batch_size)
+    last_error: RuntimeError | None = None
+    for attempt_batch_size in attempts:
+        if log is not None:
+            log(f"Starting transcription with batch size {attempt_batch_size}.")
+        try:
+            transcription = whisper_model.transcribe(
+                audio,
+                batch_size=attempt_batch_size,
+                language=source_language,
+                progress_callback=progress_callback,
+            )
+            return transcription, attempt_batch_size
+        except RuntimeError as error:
+            if device != "cuda" or not is_cuda_oom(error):
+                raise
+            last_error = error
+            next_attempt_index = attempts.index(attempt_batch_size) + 1
+            if log is not None and next_attempt_index < len(attempts):
+                log(
+                    "CUDA out of memory at batch size "
+                    f"{attempt_batch_size}; retrying with batch size "
+                    f"{attempts[next_attempt_index]}."
+                )
+            if progress_reset is not None and next_attempt_index < len(attempts):
+                progress_reset()
+            release_memory()
+    assert last_error is not None
+    raise SubtatixError(
+        "CUDA ran out of memory during transcription even after retrying with "
+        f"smaller batch sizes {attempts}. Retry with --device cpu, a smaller "
+        "--model, or a lower --batch-size."
+    ) from last_error
+
+
+def transcribe_to_srt(
+    input_file: Path,
+    model_name: str = DEFAULT_MODEL,
+    batch_size: int = 8,
+    output_file: Path | None = None,
+    write_output: bool = True,
+    source_language: str | None = None,
+    device_preference: str = "auto",
+    log: Callable[[str], None] | None = None,
+    progress_callback: Callable[[float], None] | None = None,
+    progress_reset: Callable[[], None] | None = None,
+) -> SubtitleDocument:
+    input_file = input_file.expanduser().resolve()
+    if not input_file.is_file():
+        raise SubtatixError(f"Input file not found: {input_file}")
+    output_path = resolve_output_path(input_file, output_file)
+    normalized_source_language = (
+        normalize_source_language(source_language)
+        if source_language is not None
+        else None
+    )
+
+    device, compute_type = get_whisperx_runtime(device_preference)
+    if log is not None:
+        log(
+            f"Using device: {device} (compute_type={compute_type}, "
+            f"requested={device_preference}, initial_batch_size={batch_size})."
+        )
+    whisper_model = None
+    align_model = None
+    audio = None
+    transcription = None
+    aligned_transcription = None
+
+    try:
+        whisper_model = whisperx.load_model(
+            model_name,
+            device,
+            compute_type=compute_type,
+            language=normalized_source_language,
+        )
+
+        audio = whisperx.load_audio(str(input_file))
+        transcription, _ = transcribe_with_backoff(
+            whisper_model=whisper_model,
+            audio=audio,
+            batch_size=batch_size,
+            source_language=normalized_source_language,
+            device=device,
+            log=log,
+            progress_callback=progress_callback,
+            progress_reset=progress_reset,
+        )
+        language = transcription["language"]
+        if normalized_source_language is None and log is not None:
+            log(f"Auto-detected source language: {language}.")
+
+        align_model, align_metadata = whisperx.load_align_model(
+            language_code=language,
+            device=device,
+        )
+        aligned_transcription = whisperx.align(
+            transcription["segments"],
+            align_model,
+            align_metadata,
+            audio,
+            device,
+            return_char_alignments=False,
+        )
+        aligned_transcription["language"] = language
+        cues = build_cues(aligned_transcription)
+        if write_output:
+            write_srt(output_path, cues)
+        return SubtitleDocument(
+            source_language=language,
+            subtitle_path=output_path,
+            cues=cues,
+        )
+    finally:
+        del whisper_model
+        del align_model
+        del audio
+        del transcription
+        del aligned_transcription
+        release_memory()

subtatix-0.1.0/src/subtatix/translation.py

@@ -0,0 +1,151 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+import math
+from pathlib import Path
+from typing import Callable
+
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+
+from subtatix.errors import SubtatixError
+from subtatix.runtime import get_device, release_memory
+from subtatix.subtitles import SubtitleCue, SubtitleDocument, write_srt
+
+
+DEFAULT_TRANSLATION_MODEL = "facebook/nllb-200-1.3B"
+DEFAULT_TRANSLATION_BATCH_SIZE = 16
+
+
+@dataclass(frozen=True)
+class LanguageSpec:
+    nllb_code: str
+    suffix: str
+
+
+LANGUAGE_SPECS = {
+    "ca": LanguageSpec("cat_Latn", "ca"),
+    "de": LanguageSpec("deu_Latn", "de"),
+    "en": LanguageSpec("eng_Latn", "en"),
+    "es": LanguageSpec("spa_Latn", "es"),
+    "fr": LanguageSpec("fra_Latn", "fr"),
+    "it": LanguageSpec("ita_Latn", "it"),
+    "ja": LanguageSpec("jpn_Jpan", "ja"),
+    "ko": LanguageSpec("kor_Hang", "ko"),
+    "nl": LanguageSpec("nld_Latn", "nl"),
+    "pt": LanguageSpec("por_Latn", "pt"),
+    "ru": LanguageSpec("rus_Cyrl", "ru"),
+    "zh": LanguageSpec("zho_Hans", "zh"),
+}
+SUPPORTED_TARGET_LANGUAGE_CODES = tuple(sorted(LANGUAGE_SPECS))
+
+_TRANSLATION_MODELS: dict[str, tuple[object, object, str]] = {}
+_NLLB_LANGUAGE_CODES: tuple[str, ...] | None = None
+
+
+def resolve_language(language: str) -> LanguageSpec:
+    raw_language = language.strip()
+    key = raw_language.lower()
+    if key in LANGUAGE_SPECS:
+        return LANGUAGE_SPECS[key]
+
+    parts = raw_language.split("_", 1)
+    if len(parts) == 2 and len(parts[0]) == 3 and parts[0].islower():
+        return LanguageSpec(raw_language, parts[0])
+
+    raise SubtatixError(
+        f"Unsupported language '{language}'. Use a Whisper language code like 'es' "
+        "or a full NLLB language code like 'spa_Latn'."
+    )
+
+
+def resolve_translation_output_path(
+    subtitle_path: Path,
+    target_language: LanguageSpec,
+) -> Path:
+    if subtitle_path.suffix != ".srt":
+        raise SubtatixError(f"Expected an .srt subtitle file, got: {subtitle_path}")
+    return subtitle_path.with_name(
+        f"{subtitle_path.stem}.{target_language.suffix}{subtitle_path.suffix}"
+    )
+
+
+def get_translation_backend(
+    model_name: str = DEFAULT_TRANSLATION_MODEL,
+) -> tuple[object, object, str]:
+    if model_name not in _TRANSLATION_MODELS:
+        release_memory()
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+        device = get_device()
+        _TRANSLATION_MODELS[model_name] = (tokenizer, model.to(device), device)
+    return _TRANSLATION_MODELS[model_name]
+
+
+def get_available_nllb_languages(
+    model_name: str = DEFAULT_TRANSLATION_MODEL,
+) -> tuple[str, ...]:
+    global _NLLB_LANGUAGE_CODES
+    if _NLLB_LANGUAGE_CODES is None:
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        _NLLB_LANGUAGE_CODES = tuple(sorted(tokenizer.additional_special_tokens))
+    return _NLLB_LANGUAGE_CODES
+
+
+def translate_batch(
+    texts: list[str],
+    src_lang: str,
+    tgt_lang: str,
+    model_name: str = DEFAULT_TRANSLATION_MODEL,
+    max_length: int = 400,
+) -> list[str]:
+    tokenizer, model, device = get_translation_backend(model_name)
+    tokenizer.src_lang = src_lang
+    inputs = tokenizer(
+        texts,
+        return_tensors="pt",
+        padding=True,
+        truncation=True,
+    ).to(device)
+    translated_tokens = model.generate(
+        **inputs,
+        forced_bos_token_id=tokenizer.convert_tokens_to_ids(tgt_lang),
+        max_length=max_length,
+    )
+    return tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)
+
+
+def translate_subtitles(
+    document: SubtitleDocument,
+    target_language: str,
+    model_name: str = DEFAULT_TRANSLATION_MODEL,
+    batch_size: int = DEFAULT_TRANSLATION_BATCH_SIZE,
+    max_length: int = 400,
+    progress_callback: Callable[[int, int], None] | None = None,
+) -> Path:
+    source = resolve_language(document.source_language)
+    target = resolve_language(target_language)
+    output_path = resolve_translation_output_path(document.subtitle_path, target)
+    cues = document.cues
+    total_batches = max(1, math.ceil(len(cues) / batch_size))
+
+    translated_texts: list[str] = []
+    for batch_index, start in enumerate(range(0, len(cues), batch_size), start=1):
+        batch = cues[start : start + batch_size]
+        translated_texts.extend(
+            translate_batch(
+                [cue.text for cue in batch],
+                src_lang=source.nllb_code,
+                tgt_lang=target.nllb_code,
+                model_name=model_name,
+                max_length=max_length,
+            )
+        )
+        if progress_callback is not None:
+            progress_callback(batch_index, total_batches)
+
+    translated_cues = [
+        SubtitleCue(start=cue.start, end=cue.end, text=text)
+        for cue, text in zip(cues, translated_texts, strict=True)
+    ]
+    write_srt(output_path, translated_cues)
+    return output_path