srtforge 0.1.0__py3-none-any.whl

srtforge/__init__.py

@@ -0,0 +1,3 @@
+"""srtforge: generate SRT subtitles from video using local Whisper (MLX)."""
+
+__version__ = "0.1.0"

srtforge/audio.py

@@ -0,0 +1,57 @@
+"""Extract audio from a video file using ffmpeg."""
+
+from __future__ import annotations
+
+import shutil
+import subprocess
+import tempfile
+from pathlib import Path
+
+
+class FFmpegError(RuntimeError):
+    """Raised when ffmpeg is missing or fails."""
+
+
+def ensure_ffmpeg() -> str:
+    """Return the ffmpeg executable path or raise if not found."""
+    ffmpeg = shutil.which("ffmpeg")
+    if not ffmpeg:
+        raise FFmpegError(
+            "ffmpeg not found on PATH. Install it (e.g. `nb install ffmpeg`)."
+        )
+    return ffmpeg
+
+
+def extract_audio(video_path: Path, sample_rate: int = 16000) -> Path:
+    """Extract mono PCM WAV at the given sample rate into a temp file.
+
+    Returns the path to the WAV file. Caller is responsible for deletion.
+    """
+    ffmpeg = ensure_ffmpeg()
+    if not video_path.exists():
+        raise FileNotFoundError(f"Video not found: {video_path}")
+
+    tmp = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
+    tmp.close()
+    out_path = Path(tmp.name)
+
+    cmd = [
+        ffmpeg,
+        "-y",
+        "-i", str(video_path),
+        "-vn",                       # drop video
+        "-ac", "1",                  # mono
+        "-ar", str(sample_rate),     # sample rate
+        "-c:a", "pcm_s16le",         # 16-bit PCM
+        str(out_path),
+    ]
+    proc = subprocess.run(cmd, capture_output=True, text=True)
+    if proc.returncode != 0:
+        out_path.unlink(missing_ok=True)
+        raise FFmpegError(
+            f"ffmpeg failed (exit {proc.returncode}):\n{proc.stderr.strip()[-2000:]}"
+        )
+    if out_path.stat().st_size == 0:
+        out_path.unlink(missing_ok=True)
+        raise FFmpegError("ffmpeg produced no audio. Does the file have an audio track?")
+    return out_path

srtforge/cli.py

@@ -0,0 +1,237 @@
+"""Command-line interface for srtforge."""
+
+from __future__ import annotations
+
+import argparse
+import re
+import sys
+from pathlib import Path
+
+from . import __version__
+from .audio import FFmpegError, extract_audio
+from .merge import burn_subtitles
+from .segment import Opts, fit_cues, resegment
+from .srt import segments_to_srt
+from .transcribe import DEFAULT_MODEL, transcribe
+from .translate import DEFAULT_TRANSLATE_MODEL, TranslationError, translate_segments
+
+
+def _log(msg: str) -> None:
+    print(msg, file=sys.stderr, flush=True)
+
+
+def _positive_int(value: str) -> int:
+    parsed = int(value)
+    if parsed <= 0:
+        raise argparse.ArgumentTypeError("must be greater than 0")
+    return parsed
+
+
+def _positive_float(value: str) -> float:
+    parsed = float(value)
+    if parsed <= 0:
+        raise argparse.ArgumentTypeError("must be greater than 0")
+    return parsed
+
+
+def _crf(value: str) -> int:
+    parsed = int(value)
+    if not 0 <= parsed <= 51:
+        raise argparse.ArgumentTypeError("must be between 0 and 51")
+    return parsed
+
+
+def _translated_path(output: Path, language: str) -> Path:
+    """Insert a language tag before the suffix: video.srt -> video.<lang>.srt."""
+    tag = re.sub(r"[^\w.-]+", "-", language.strip().lower()).strip(".-_")
+    tag = tag or "translated"
+    return output.with_suffix(f".{tag}{output.suffix}")
+
+
+def build_parser() -> argparse.ArgumentParser:
+    p = argparse.ArgumentParser(
+        prog="srtforge",
+        description="Generate SRT subtitles from a video using local Whisper (MLX).",
+        epilog="Also: `srtforge merge VIDEO SUBS.srt -o OUT.mp4` to burn subtitles "
+               "into a video. Run `srtforge merge -h` for details.",
+    )
+    p.add_argument("video", type=Path, help="Path to the input video (e.g. MP4).")
+    p.add_argument(
+        "-o", "--output", type=Path, default=None,
+        help="Output .srt path (default: alongside the video).",
+    )
+    p.add_argument(
+        "-m", "--model", default=DEFAULT_MODEL,
+        help=f"Whisper model / HF repo (default: {DEFAULT_MODEL}).",
+    )
+    p.add_argument(
+        "-l", "--language", default=None,
+        help="Language code (e.g. en, es). Default: auto-detect.",
+    )
+    p.add_argument(
+        "--word-timestamps", action="store_true",
+        help="Compute word-level timestamps (slower).",
+    )
+    p.add_argument(
+        "--no-resegment", dest="resegment", action="store_false",
+        help="Disable sentence-aware re-cueing; keep raw Whisper segments.",
+    )
+    p.add_argument(
+        "--max-line-length", type=_positive_int, default=Opts.max_cpl, metavar="N",
+        help=f"Max characters per subtitle line (default: {Opts.max_cpl}).",
+    )
+    p.add_argument(
+        "--max-lines", type=_positive_int, default=Opts.max_lines, metavar="N",
+        help=f"Max lines per subtitle cue (default: {Opts.max_lines}).",
+    )
+    p.add_argument(
+        "--reading-speed", type=_positive_float, default=Opts.reading_cps, metavar="CPS",
+        help=f"Reading speed in characters/second (default: {Opts.reading_cps}).",
+    )
+    p.add_argument(
+        "-t", "--translate", default=None, metavar="LANG",
+        help="Also translate subtitles into LANG (e.g. es, Spanish, Japanese) "
+             "using a local MLX LLM. Writes a second .<lang>.srt file.",
+    )
+    p.add_argument(
+        "--translate-model", default=DEFAULT_TRANSLATE_MODEL,
+        help=f"MLX LLM used for translation (default: {DEFAULT_TRANSLATE_MODEL}).",
+    )
+    p.add_argument("--version", action="version", version=f"srtforge {__version__}")
+    return p
+
+
+def _gen_main(argv: list[str] | None = None) -> int:
+    args = build_parser().parse_args(argv)
+
+    output = args.output or args.video.with_suffix(".srt")
+    translating = bool(args.translate)
+    total = 4 if translating else 3
+    opts = Opts(
+        max_cpl=args.max_line_length,
+        max_lines=args.max_lines,
+        reading_cps=args.reading_speed,
+    )
+    # Sentence-aware re-cueing needs word-level timestamps.
+    need_words = args.word_timestamps or args.resegment
+    wav_path = None
+    try:
+        _log(f"[1/{total}] Extracting audio from {args.video} ...")
+        wav_path = extract_audio(args.video)
+
+        _log(f"[2/{total}] Transcribing with {args.model} (downloads model on first run) ...")
+        segments = transcribe(
+            wav_path,
+            model=args.model,
+            language=args.language,
+            word_timestamps=need_words,
+        )
+
+        if args.resegment:
+            segments = resegment(segments, opts)
+
+        _log(f"[3/{total}] Writing {output} ...")
+        srt_text = segments_to_srt(segments)
+        output.write_text(srt_text, encoding="utf-8")
+        if not srt_text:
+            _log("Warning: no speech detected; wrote an empty SRT file.")
+        else:
+            _log(f"  {srt_text.count('-->')} subtitle entries written.")
+
+        if translating:
+            if not srt_text:
+                _log("Skipping translation: no speech to translate.")
+            else:
+                t_out = _translated_path(output, args.translate)
+                _log(
+                    f"[4/{total}] Translating to {args.translate} with "
+                    f"{args.translate_model} (downloads model on first run) ..."
+                )
+                t_segments = translate_segments(
+                    segments,
+                    target_language=args.translate,
+                    model=args.translate_model,
+                )
+                if args.resegment:
+                    # Translated text length differs; fit/split to line limits.
+                    t_segments = fit_cues(t_segments, opts)
+                t_out.write_text(segments_to_srt(t_segments), encoding="utf-8")
+                _log(f"  Translated subtitles written to {t_out}")
+
+        _log("Done.")
+        return 0
+
+    except (FFmpegError, FileNotFoundError, TranslationError) as e:
+        _log(f"Error: {e}")
+        return 1
+    except KeyboardInterrupt:
+        _log("Interrupted.")
+        return 130
+    finally:
+        if wav_path is not None:
+            wav_path.unlink(missing_ok=True)
+
+
+def build_merge_parser() -> argparse.ArgumentParser:
+    p = argparse.ArgumentParser(
+        prog="srtforge merge",
+        description="Burn a subtitle file permanently into a video (hard subs).",
+    )
+    p.add_argument("video", type=Path, help="Input video file.")
+    p.add_argument("subtitles", type=Path, help="Subtitle file (.srt) to burn in.")
+    p.add_argument(
+        "-o", "--output", type=Path, default=None,
+        help="Output video path (default: <video>.subbed<ext>).",
+    )
+    p.add_argument(
+        "--crf", type=_crf, default=18, metavar="N",
+        help="x264 quality, lower = better/larger (default: 18, ~visually lossless).",
+    )
+    p.add_argument(
+        "--preset", default="slow",
+        help="x264 speed/efficiency preset (default: slow).",
+    )
+    p.add_argument(
+        "--font-size", type=_positive_int, default=None, metavar="N",
+        help="Override subtitle font size.",
+    )
+    p.add_argument(
+        "--font-name", default=None, metavar="NAME",
+        help="Override subtitle font name.",
+    )
+    return p
+
+
+def _merge_main(argv: list[str]) -> int:
+    args = build_merge_parser().parse_args(argv)
+    output = args.output or args.video.with_suffix(f".subbed{args.video.suffix}")
+    try:
+        _log(f"Burning {args.subtitles} into {args.video} (crf={args.crf}, preset={args.preset}) ...")
+        burn_subtitles(
+            args.video,
+            args.subtitles,
+            output,
+            crf=args.crf,
+            preset=args.preset,
+            font_size=args.font_size,
+            font_name=args.font_name,
+        )
+        _log(f"Done. Wrote {output}")
+        return 0
+    except (FFmpegError, FileNotFoundError) as e:
+        _log(f"Error: {e}")
+        return 1
+    except KeyboardInterrupt:
+        _log("Interrupted.")
+        return 130
+
+
+def main(argv: list[str] | None = None) -> int:
+    argv = list(sys.argv[1:] if argv is None else argv)
+    if argv and argv[0] == "merge":
+        return _merge_main(argv[1:])
+    return _gen_main(argv)
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

srtforge/merge.py

@@ -0,0 +1,129 @@
+"""Burn subtitles permanently into a video using ffmpeg (libass)."""
+
+from __future__ import annotations
+
+import glob
+import re
+import shutil
+import subprocess
+from pathlib import Path
+
+from .audio import FFmpegError
+
+# Where fuller ffmpeg builds (with libass) tend to live.
+_FFMPEG_GLOBS = (
+    "/opt/nanobrew/prefix/Cellar/ffmpeg-full/*/bin/ffmpeg",
+    "/opt/nanobrew/prefix/Cellar/ffmpeg/*/bin/ffmpeg",
+    "/opt/homebrew/Cellar/ffmpeg*/*/bin/ffmpeg",
+    "/usr/local/Cellar/ffmpeg*/*/bin/ffmpeg",
+)
+
+
+def _has_subtitles_filter(binary: str) -> bool:
+    """True if the ffmpeg ``binary`` exposes the libass ``subtitles`` filter."""
+    try:
+        out = subprocess.run(
+            [binary, "-hide_banner", "-filters"],
+            capture_output=True, text=True, timeout=30,
+        ).stdout
+    except Exception:
+        return False
+    return re.search(r"(?m)^\s*\S+\s+subtitles\s", out) is not None
+
+
+def _resolve_ffmpeg_with_subtitles() -> str:
+    """Find an ffmpeg that supports subtitle burn-in (libass).
+
+    Prefers the one on PATH; otherwise searches common install locations for a
+    fuller build (e.g. ``ffmpeg-full``).
+    """
+    candidates: list[str] = []
+    on_path = shutil.which("ffmpeg")
+    if on_path:
+        candidates.append(on_path)
+    for pattern in _FFMPEG_GLOBS:
+        candidates.extend(sorted(glob.glob(pattern), reverse=True))
+
+    seen: set[str] = set()
+    for c in candidates:
+        if c in seen:
+            continue
+        seen.add(c)
+        if _has_subtitles_filter(c):
+            return c
+    raise FFmpegError(
+        "No ffmpeg with the 'subtitles' filter (libass) was found. Your default "
+        "ffmpeg lacks libass. Install a full build, e.g. `nb install ffmpeg-full`."
+    )
+
+
+def _escape_filter_value(value: str) -> str:
+    """Escape a value embedded in an ffmpeg filtergraph option.
+
+    The filtergraph parser treats these characters as syntax even when the
+    whole filtergraph is passed as a single subprocess argument.
+    """
+    special = "\\':,[];"
+    return "".join(f"\\{ch}" if ch in special else ch for ch in value)
+
+
+def _escape_sub_path(path: Path) -> str:
+    """Escape a subtitle path for the ffmpeg ``subtitles`` filter value."""
+    return _escape_filter_value(str(path))
+
+
+def _force_style(font_size: int | None, font_name: str | None) -> str:
+    parts: list[str] = []
+    if font_name:
+        parts.append(f"FontName={font_name}")
+    if font_size:
+        parts.append(f"FontSize={font_size}")
+    return ",".join(parts)
+
+
+def burn_subtitles(
+    video: Path,
+    srt: Path,
+    output: Path,
+    crf: int = 18,
+    preset: str = "slow",
+    font_size: int | None = None,
+    font_name: str | None = None,
+) -> Path:
+    """Render ``srt`` onto ``video`` and write ``output``.
+
+    Framerate and resolution are preserved; audio is stream-copied. Video is
+    re-encoded with libx264 at the given CRF/preset (CRF 18 ~ visually lossless).
+    Returns the output path.
+    """
+    ffmpeg = _resolve_ffmpeg_with_subtitles()
+    if not video.exists():
+        raise FileNotFoundError(f"Video not found: {video}")
+    if not srt.exists():
+        raise FileNotFoundError(f"Subtitle file not found: {srt}")
+
+    vf = f"subtitles=filename={_escape_sub_path(srt)}"
+    style = _force_style(font_size, font_name)
+    if style:
+        vf += f":force_style={_escape_filter_value(style)}"
+
+    cmd = [
+        ffmpeg,
+        "-y",
+        "-i", str(video),
+        "-vf", vf,
+        "-c:v", "libx264",
+        "-crf", str(crf),
+        "-preset", preset,
+        "-pix_fmt", "yuv420p",
+        "-c:a", "copy",          # keep original audio losslessly
+        "-movflags", "+faststart",
+        str(output),
+    ]
+    proc = subprocess.run(cmd, capture_output=True, text=True)
+    if proc.returncode != 0:
+        raise FFmpegError(
+            f"ffmpeg failed to burn subtitles (exit {proc.returncode}):\n"
+            f"{proc.stderr.strip()[-2000:]}"
+        )
+    return output