subtitle-engine 0.1.0__py3-none-any.whl

subtitle_engine/__init__.py

@@ -0,0 +1,3 @@
+"""Subtitle Engine — generate SRT files with WhisperX."""
+
+__version__ = "0.1.0"

subtitle_engine/captioner.py

@@ -0,0 +1,78 @@
+"""Generate captions from transcripts using a local Ollama instance."""
+
+import json
+from typing import Optional
+
+import requests
+
+
+def _default_prompt(transcript: str) -> str:
+    """Build the prompt sent to the LLM."""
+    return (
+        "Create a short, engaging caption (1-2 sentences) for a video based on the following transcript. "
+        "Write the caption in the same language as the transcript. "
+        "Answer directly with the caption only, without any thinking or explanation.\n\n"
+        f"Transcript:\n{transcript}"
+    )
+
+
+def generate_caption(
+    transcript: str,
+    *,
+    model: str,
+    host: str = "http://localhost:11434",
+    prompt: Optional[str] = None,
+) -> str:
+    """Generate a caption from a transcript via Ollama.
+
+    Parameters
+    ----------
+    transcript:
+        The transcript text to summarize.
+    model:
+        Name of the Ollama model to use.
+    host:
+        Base URL of the Ollama API.
+    prompt:
+        Custom prompt. A default prompt is used if omitted.
+
+    Returns
+    -------
+    The generated caption string.
+    """
+    if not transcript.strip():
+        raise ValueError("Cannot generate a caption from an empty transcript")
+
+    url = f"{host.rstrip('/')}/api/generate"
+    payload = {
+        "model": model,
+        "prompt": prompt or _default_prompt(transcript),
+        "stream": False,
+    }
+
+    try:
+        response = requests.post(url, json=payload, timeout=300)
+    except requests.ConnectionError as exc:
+        raise ConnectionError(
+            f"Could not connect to Ollama at {host}. Is Ollama running?"
+        ) from exc
+
+    response.raise_for_status()
+    data = response.json()
+    caption = data.get("response", "").strip()
+
+    if not caption:
+        raise ValueError(
+            "Ollama returned an empty caption. "
+            "This can happen with some models or languages — try a different --ollama-model."
+        )
+
+    return caption
+
+
+def list_models(host: str = "http://localhost:11434") -> list[str]:
+    """Return the names of models available in the local Ollama instance."""
+    url = f"{host.rstrip('/')}/api/tags"
+    response = requests.get(url, timeout=30)
+    response.raise_for_status()
+    return [model["name"] for model in response.json().get("models", [])]

subtitle_engine/cli.py

@@ -0,0 +1,170 @@
+"""Command-line interface for subtitle-engine."""
+
+from pathlib import Path
+from typing import Annotated, Optional
+
+import typer
+from rich.console import Console
+
+from subtitle_engine.captioner import generate_caption
+from subtitle_engine.srt_writer import write_srt
+from subtitle_engine.transcriber import transcribe
+from subtitle_engine.utils import resolve_output_path, validate_media_file
+
+app = typer.Typer(
+    help="Generate SRT subtitles from audio/video files using WhisperX",
+    no_args_is_help=True,
+)
+console = Console()
+
+
+@app.command()
+def main(
+    input_file: Annotated[
+        Path,
+        typer.Argument(
+            help="Audio or video file to transcribe",
+            exists=True,
+            file_okay=True,
+            dir_okay=False,
+            readable=True,
+        ),
+    ],
+    output: Annotated[
+        Optional[Path],
+        typer.Option(
+            "--output",
+            "-o",
+            help="Output SRT file (default: <input>.srt)",
+            file_okay=True,
+            dir_okay=False,
+        ),
+    ] = None,
+    model: Annotated[
+        str,
+        typer.Option(
+            "--model",
+            "-m",
+            help="WhisperX model: tiny, base, small, medium, large-v2, large-v3",
+        ),
+    ] = "small",
+    language: Annotated[
+        Optional[str],
+        typer.Option(
+            "--language",
+            "-l",
+            help="Language code, e.g. en, fi. Auto-detected if omitted.",
+        ),
+    ] = None,
+    device: Annotated[
+        Optional[str],
+        typer.Option(
+            "--device",
+            "-d",
+            help="Device: cpu, cuda or mps. Auto-detected if omitted.",
+        ),
+    ] = None,
+    batch_size: Annotated[
+        int,
+        typer.Option(
+            "--batch-size",
+            "-b",
+            help="WhisperX inference batch size",
+            min=1,
+        ),
+    ] = 16,
+    compute_type: Annotated[
+        Optional[str],
+        typer.Option(
+            "--compute-type",
+            "-c",
+            help="Compute type: int8 or float16. Auto-selected if omitted.",
+        ),
+    ] = None,
+    diarize: Annotated[
+        bool,
+        typer.Option(
+            "--diarize",
+            help="Run speaker diarization (requires --hf-token)",
+        ),
+    ] = False,
+    hf_token: Annotated[
+        Optional[str],
+        typer.Option(
+            "--hf-token",
+            help="Hugging Face token for diarization",
+            envvar="HF_TOKEN",
+        ),
+    ] = None,
+    caption: Annotated[
+        bool,
+        typer.Option(
+            "--caption",
+            help="Generate a caption from the transcript using Ollama",
+        ),
+    ] = False,
+    ollama_model: Annotated[
+        Optional[str],
+        typer.Option(
+            "--ollama-model",
+            help="Ollama model for caption generation. Required if --caption is set.",
+        ),
+    ] = None,
+    ollama_host: Annotated[
+        str,
+        typer.Option(
+            "--ollama-host",
+            help="Ollama API host",
+            envvar="OLLAMA_HOST",
+        ),
+    ] = "http://localhost:11434",
+) -> None:
+    """Generate SRT subtitles from a media file."""
+    try:
+        validate_media_file(input_file)
+        output_path = resolve_output_path(input_file, output)
+
+        if caption and not ollama_model:
+            raise ValueError("--ollama-model is required when using --caption")
+
+        console.print(f"[bold]Transcribing:[/bold] {input_file}")
+        console.print(f"[bold]Model:[/bold] {model}")
+        if language:
+            console.print(f"[bold]Language:[/bold] {language}")
+        if device:
+            console.print(f"[bold]Device:[/bold] {device}")
+
+        segments = transcribe(
+            input_file,
+            model_name=model,
+            language=language,
+            device=device,
+            batch_size=batch_size,
+            compute_type=compute_type,
+            diarize=diarize,
+            hf_token=hf_token,
+        )
+
+        write_srt(segments, output_path)
+        console.print(f"[green]Wrote subtitles to:[/green] {output_path}")
+
+        if caption:
+            transcript = " ".join(str(segment.get("text", "")).strip() for segment in segments)
+            caption_text = generate_caption(
+                transcript,
+                model=ollama_model,
+                host=ollama_host,
+            )
+            caption_path = output_path.with_suffix(".caption.txt")
+            caption_path.write_text(caption_text, encoding="utf-8")
+            console.print(f"[green]Wrote caption to:[/green] {caption_path}")
+    except (ValueError, FileNotFoundError, ConnectionError) as exc:
+        console.print(f"[red]Error:[/red] {exc}")
+        raise typer.Exit(code=1) from exc
+    except Exception as exc:  # noqa: BLE001
+        console.print(f"[red]Transcription failed:[/red] {exc}")
+        raise typer.Exit(code=1) from exc
+
+
+if __name__ == "__main__":
+    app()

subtitle_engine/srt_writer.py

@@ -0,0 +1,45 @@
+"""Convert transcription segments to SRT format."""
+
+from pathlib import Path
+from typing import Iterable
+
+
+def _format_time(seconds: float) -> str:
+    """Convert seconds to SRT time format HH:MM:SS,mmm."""
+    total_millis = int(round(seconds * 1000))
+    hours = total_millis // 3_600_000
+    minutes = (total_millis % 3_600_000) // 60_000
+    secs = (total_millis % 60_000) // 1_000
+    millis = total_millis % 1_000
+
+    return f"{hours:02d}:{minutes:02d}:{secs:02d},{millis:03d}"
+
+
+def _format_segment(index: int, start: float, end: float, text: str) -> str:
+    """Format a single segment as an SRT block."""
+    cleaned_text = text.strip()
+    if not cleaned_text:
+        cleaned_text = "..."
+    return f"{index}\n{_format_time(start)} --> {_format_time(end)}\n{cleaned_text}\n"
+
+
+def segments_to_srt(segments: Iterable[dict]) -> str:
+    """Build an SRT string from WhisperX-style segments.
+
+    Each segment is expected to be a dict with keys:
+    ``start`` (float), ``end`` (float), and ``text`` (str).
+    """
+    blocks = []
+    for index, segment in enumerate(segments, start=1):
+        start = float(segment["start"])
+        end = float(segment["end"])
+        text = str(segment["text"])
+        blocks.append(_format_segment(index, start, end, text))
+    return "\n".join(blocks)
+
+
+def write_srt(segments: Iterable[dict], output_path: Path) -> None:
+    """Write segments to an SRT file."""
+    output_path = Path(output_path)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    output_path.write_text(segments_to_srt(segments), encoding="utf-8")

subtitle_engine/transcriber.py

@@ -0,0 +1,129 @@
+"""WhisperX transcription wrapper."""
+
+from pathlib import Path
+from typing import Optional
+
+import torch
+import whisperx
+
+
+VALID_MODELS = {"tiny", "base", "small", "medium", "large-v2", "large-v3"}
+VALID_DEVICES = {"cpu", "cuda", "mps"}
+
+
+def _detect_device(device: Optional[str]) -> str:
+    """Pick a device if none was specified."""
+    if device:
+        return device
+    if torch.cuda.is_available():
+        return "cuda"
+    if torch.backends.mps.is_available():
+        return "mps"
+    return "cpu"
+
+
+def _default_compute_type(device: str) -> str:
+    """Pick a safe compute type for the device."""
+    if device == "cpu":
+        return "int8"
+    return "float16"
+
+
+def _validate_model(model_name: str) -> None:
+    """Raise a ValueError if the model name is unknown."""
+    if model_name not in VALID_MODELS:
+        joined = ", ".join(sorted(VALID_MODELS))
+        raise ValueError(f"Unknown model '{model_name}'. Choose from: {joined}")
+
+
+def _validate_device(device: str) -> None:
+    """Raise a ValueError if the device name is unknown."""
+    if device not in VALID_DEVICES:
+        joined = ", ".join(sorted(VALID_DEVICES))
+        raise ValueError(f"Unknown device '{device}'. Choose from: {joined}")
+
+
+def transcribe(
+    audio_path: Path,
+    *,
+    model_name: str = "small",
+    language: Optional[str] = None,
+    device: Optional[str] = None,
+    batch_size: int = 16,
+    compute_type: Optional[str] = None,
+    diarize: bool = False,
+    hf_token: Optional[str] = None,
+) -> list[dict]:
+    """Transcribe an audio/video file and return SRT-ready segments.
+
+    Parameters
+    ----------
+    audio_path:
+        Path to the media file to transcribe.
+    model_name:
+        WhisperX model size. One of: tiny, base, small, medium, large-v2, large-v3.
+    language:
+        ISO language code, e.g. ``en`` or ``fi``. If ``None``, WhisperX auto-detects.
+    device:
+        ``cpu``, ``cuda`` or ``mps``. Auto-detected if omitted.
+    batch_size:
+        WhisperX batch size for transcription.
+    compute_type:
+        ``int8`` or ``float16``. Auto-selected per device if omitted.
+    diarize:
+        Whether to run speaker diarization.
+    hf_token:
+        Hugging Face token required for diarization.
+
+    Returns
+    -------
+    A list of segment dicts with ``start``, ``end`` and ``text`` keys.
+    """
+    _validate_model(model_name)
+
+    audio_path = Path(audio_path)
+    if not audio_path.exists():
+        raise FileNotFoundError(f"Audio file not found: {audio_path}")
+
+    device = _detect_device(device)
+    _validate_device(device)
+
+    if compute_type is None:
+        compute_type = _default_compute_type(device)
+
+    if diarize and not hf_token:
+        raise ValueError("--hf-token is required when using --diarize")
+
+    audio = whisperx.load_audio(str(audio_path))
+
+    model = whisperx.load_model(model_name, device, compute_type=compute_type)
+    result = model.transcribe(audio, batch_size=batch_size, language=language)
+
+    # Free transcription model memory before alignment
+    del model
+
+    detected_language = result.get("language")
+    if detected_language:
+        align_model, align_metadata = whisperx.load_align_model(
+            language_code=detected_language, device=device
+        )
+        result = whisperx.align(
+            result["segments"],
+            align_model,
+            align_metadata,
+            audio,
+            device,
+            return_char_alignments=False,
+        )
+        del align_model
+
+    if diarize:
+        diarize_model = whisperx.DiarizationPipeline(
+            model_name="pyannote/speaker-diarization-3.1",
+            use_auth_token=hf_token,
+            device=device,
+        )
+        diarize_segments = diarize_model(audio)
+        result = whisperx.assign_word_speakers(diarize_segments, result)
+
+    return result["segments"]

subtitle_engine/utils.py

@@ -0,0 +1,39 @@
+"""CLI helpers and path utilities."""
+
+from pathlib import Path
+from typing import Optional
+
+
+SUPPORTED_EXTENSIONS = {
+    ".mp3",
+    ".wav",
+    ".flac",
+    ".aac",
+    ".ogg",
+    ".m4a",
+    ".mp4",
+    ".mov",
+    ".mkv",
+    ".avi",
+    ".webm",
+}
+
+
+def resolve_output_path(input_path: Path, output: Optional[Path] = None) -> Path:
+    """Resolve the SRT output path.
+
+    If ``output`` is provided, use it. Otherwise create ``<input>.srt``
+    next to the input file.
+    """
+    if output:
+        return output
+    return input_path.with_suffix(".srt")
+
+
+def validate_media_file(path: Path) -> None:
+    """Raise a ValueError if the path does not look like a media file."""
+    if path.suffix.lower() not in SUPPORTED_EXTENSIONS:
+        joined = ", ".join(sorted(SUPPORTED_EXTENSIONS))
+        raise ValueError(
+            f"Unsupported file type '{path.suffix}'. Supported: {joined}"
+        )

subtitle_engine-0.1.0.dist-info/METADATA

@@ -0,0 +1,97 @@
+Metadata-Version: 2.4
+Name: subtitle-engine
+Version: 0.1.0
+Summary: Generate SRT subtitles from audio/video files using WhisperX
+Author: Leevi Puntanen
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/leevipuntanen/subtitle-engine
+Project-URL: Issues, https://github.com/leevipuntanen/subtitle-engine/issues
+Keywords: subtitles,srt,whisperx,transcription,asr
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: End Users/Desktop
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Multimedia :: Sound/Audio :: Speech
+Classifier: Topic :: Utilities
+Requires-Python: >=3.12
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: typer>=0.12.0
+Requires-Dist: whisperx>=3.8.0
+Requires-Dist: requests>=2.32.0
+Provides-Extra: dev
+Requires-Dist: pytest>=8.0.0; extra == "dev"
+Dynamic: license-file
+
+# subtitle-engine
+
+Generate `.srt` subtitle files from audio or video files using [WhisperX](https://github.com/m-bain/whisperX). Optionally generate a caption from the transcript with a local [Ollama](https://ollama.com/) LLM.
+
+## Installation
+
+Requires Python 3.12 or newer.
+
+```bash
+pip install subtitle-engine
+```
+
+Or install from source:
+
+```bash
+git clone https://github.com/leevipuntanen/subtitle-engine.git
+cd subtitle-engine
+python -m venv .venv
+source .venv/bin/activate
+pip install -e ".[dev]"
+```
+
+## Usage
+
+```bash
+# Basic usage — writes <input>.srt next to the source file
+subeng video.mp4
+
+# Specify output file
+subeng video.mp4 --output subtitles.srt
+
+# Use a different model or language
+subeng video.mp4 --model medium --language fi
+
+# Force CPU / CUDA / MPS
+subeng video.mp4 --device cpu
+
+# Speaker diarization (requires a Hugging Face token)
+subeng video.mp4 --diarize --hf-token $HF_TOKEN
+
+# Generate a caption from the transcript using Ollama
+subeng video.mp4 --caption --ollama-model qwen3.5:0.8b
+```
+
+## Options
+
+| Option | Description |
+|--------|-------------|
+| `--output`, `-o` | Output SRT file path |
+| `--model`, `-m` | WhisperX model: `tiny`, `base`, `small` (default), `medium`, `large-v2`, `large-v3` |
+| `--language`, `-l` | ISO language code, e.g. `en`, `fi`. Auto-detected if omitted. |
+| `--device`, `-d` | `cpu`, `cuda` or `mps`. Auto-detected if omitted. |
+| `--batch-size`, `-b` | Inference batch size (default: 16) |
+| `--compute-type`, `-c` | `int8` or `float16`. Auto-selected if omitted. |
+| `--diarize` | Enable speaker diarization |
+| `--hf-token` | Hugging Face token for diarization (or set `HF_TOKEN` env var) |
+| `--caption` | Generate a caption from the transcript via Ollama |
+| `--ollama-model` | Ollama model name (required with `--caption`) |
+| `--ollama-host` | Ollama API host (default: `http://localhost:11434`) |
+
+## Development
+
+Run the test suite:
+
+```bash
+pytest
+```
+
+## License
+
+MIT

subtitle_engine-0.1.0.dist-info/RECORD

@@ -0,0 +1,12 @@
+subtitle_engine/__init__.py,sha256=hQMR44TH4Q2t3giLJX1w4C48X7I1_BbLYvWv6UE193U,83
+subtitle_engine/captioner.py,sha256=YhddyIGe2Q05OHe9_VWjRgmnfu_z0fqHTeZGGuOVVYY,2273
+subtitle_engine/cli.py,sha256=I6mSqJ3CPMqmDYg0SXLnd8wVScajMKijUS5KmlnJOac,4923
+subtitle_engine/srt_writer.py,sha256=mCfYz5f_vezyKntJZknCDBfbr_wqN6tezOg3tdymTh4,1578
+subtitle_engine/transcriber.py,sha256=A9-iv0C69ZhwV2JeCAqOnBFVyymsVyXituMLCa0cbrw,3834
+subtitle_engine/utils.py,sha256=zhpyusQowoD2V0SGEt6CqgXiPzYfFSgrbdYc2tC5rcg,907
+subtitle_engine-0.1.0.dist-info/licenses/LICENSE,sha256=o1T82nY2oKB-YHPfmfFxSaovPjMXAO6zhl9gTwB7pxw,1071
+subtitle_engine-0.1.0.dist-info/METADATA,sha256=KGXISZd_1o-W6nbQGGtflIfOpALPrA0sXoEJ4Y9jCoc,2922
+subtitle_engine-0.1.0.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
+subtitle_engine-0.1.0.dist-info/entry_points.txt,sha256=DPgVIpA0GoIWbk4QkXbP_Ez4H7QYpFwqtC8QwPKV9v0,51
+subtitle_engine-0.1.0.dist-info/top_level.txt,sha256=VDgcRWJUpgMgxj4aQ_Xcby1C9FKQGEdVaC5GgRR4lYQ,16
+subtitle_engine-0.1.0.dist-info/RECORD,,

subtitle_engine-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (82.0.1)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

subtitle_engine-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+subeng = subtitle_engine.cli:app

subtitle_engine-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Leevi Puntanen
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

subtitle_engine-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+subtitle_engine