diff --git a/README.md b/README.md index 4755a21..8e1e68b 100644 --- a/README.md +++ b/README.md @@ -92,3 +92,29 @@ Generate and burn subtitles into video without validating transcription: ```bash decipher transcribe -i video.mp4 --model small --subtitle_action burn ``` + +#### TwelveLabs Pegasus backend (optional) + +By default decipher transcribes locally with Whisper. You can optionally use +[TwelveLabs](https://twelvelabs.io) Pegasus instead — a video understanding +model that reads the video directly (not just the audio track), so it can use +on-screen context when transcribing. This is fully opt-in; the default +behaviour is unchanged. + +Install the extra and set your API key (a free key is available at +[twelvelabs.io](https://twelvelabs.io) with a generous free tier): + +```bash +pip install "decipher[pegasus]" +export TWELVELABS_API_KEY=tlk_... +``` + +Then pass `--backend pegasus`: + +```bash +decipher transcribe -i video.mp4 --backend pegasus +``` + +Transcription runs in the cloud, so no local GPU is required. The video is +uploaded to your TwelveLabs account; manage uploaded assets via the TwelveLabs +dashboard. diff --git a/decipher/__main__.py b/decipher/__main__.py index 8e4db7d..b853a1e 100644 --- a/decipher/__main__.py +++ b/decipher/__main__.py @@ -55,6 +55,14 @@ def cli(): choices=["add", "burn"], help="whether to perform subtitle add or burn action", ) + t.add_argument( + "--backend", + type=str, + default="whisper", + choices=["whisper", "pegasus"], + help="transcription backend: local 'whisper' (default) or TwelveLabs " + "'pegasus' (cloud, requires TWELVELABS_API_KEY)", + ) s = subparsers.add_parser("subtitle", help="subtitle a video") s.add_argument( @@ -106,6 +114,7 @@ def main(): args.task, args.batch_size, args.subtitle_action, + args.backend, ) elif args.action == "subtitle": output = subtitle( diff --git a/decipher/action.py b/decipher/action.py index 86747be..44fb90a 100644 --- a/decipher/action.py +++ b/decipher/action.py @@ -46,6 +46,7 @@ def transcribe( task="transcribe", batch_size=24, subtitle_action=None, + backend="whisper", ) -> PathStore: video_in = Path(video_in).absolute() assert video_in.exists(), f"File {video_in} does not exist" @@ -53,16 +54,27 @@ def transcribe( output_dir = Path(output_dir).absolute() output_dir.mkdir(parents=True, exist_ok=True) - with TemporaryDirectory() as _tempdir: - tempdir = Path(_tempdir) - audio_file = str(tempdir / "audio.aac") + if backend == "pegasus": + # Pegasus reads the video directly, so no audio extraction is needed. + from .pegasus import video_to_srt - ffprog( - ["ffmpeg", "-y", "-i", str(video_in), "-vn", "-c:a", "aac", audio_file], - desc=f"Extracting audio from video", - ) + pegasus_model = model if model.startswith("pegasus") else "pegasus1.5" + srt_text = video_to_srt(str(video_in), model=pegasus_model) + elif backend == "whisper": + with TemporaryDirectory() as _tempdir: + tempdir = Path(_tempdir) + audio_file = str(tempdir / "audio.aac") + + ffprog( + ["ffmpeg", "-y", "-i", str(video_in), "-vn", "-c:a", "aac", audio_file], + desc=f"Extracting audio from video", + ) - srt_text = audio_to_srt(audio_file, model, task, language, batch_size) + srt_text = audio_to_srt(audio_file, model, task, language, batch_size) + else: + raise ValueError( + f"Unknown backend {backend!r}, expected 'whisper' or 'pegasus'" + ) srt_file = output_dir / f"{video_in.stem}.srt" with open(srt_file, "w", encoding="utf-8") as f: diff --git a/decipher/pegasus.py b/decipher/pegasus.py new file mode 100644 index 0000000..d224c67 --- /dev/null +++ b/decipher/pegasus.py @@ -0,0 +1,95 @@ +"""TwelveLabs Pegasus transcription backend (optional). + +Pegasus is a video understanding model that reads the *video* directly, so it +can transcribe speech with on-screen context (speaker names, overlaid text, +non-speech cues) that an audio-only model never sees. This module is an opt-in +alternative to the default Whisper backend; nothing here runs unless you pass +``--backend pegasus``. + +Requires the ``twelvelabs`` package (``pip install decipher[pegasus]``) and a +``TWELVELABS_API_KEY`` environment variable. Grab a free key at +https://twelvelabs.io. +""" + +import os +import re +import time + +# Ask Pegasus to return ready-to-use SRT so we don't have to re-time anything. +_SRT_PROMPT = ( + "Transcribe all spoken words in this video into SubRip (SRT) subtitles. " + "Output ONLY the SRT: numbered cues, each with a " + "'HH:MM:SS,mmm --> HH:MM:SS,mmm' timestamp line followed by the spoken " + "text. Use real timestamps from the video, keep cues short (one or two " + "lines), do not add commentary, headings or code fences." +) + +_FENCE_RE = re.compile(r"^\s*```[a-zA-Z]*\s*|\s*```\s*$") +_CUE_RE = re.compile(r"\d{2}:\d{2}:\d{2},\d{3}\s*-->\s*\d{2}:\d{2}:\d{2},\d{3}") + + +def _strip_fences(text): + """Remove ``` code fences Pegasus sometimes wraps the SRT in.""" + return _FENCE_RE.sub("", text).strip() + + +def video_to_srt( + video_file, + model="pegasus1.5", + api_key=None, + poll_interval=5, + timeout=900, +): + """Transcribe a video to SRT text using TwelveLabs Pegasus. + + The video is uploaded as a TwelveLabs asset, then analysed with a prompt + that asks Pegasus for SRT-formatted output. Returns the SRT string. + """ + try: + from twelvelabs import TwelveLabs + except ImportError as e: # pragma: no cover - import-guard + raise ImportError( + "The pegasus backend needs the 'twelvelabs' package. " + "Install it with: pip install decipher[pegasus]" + ) from e + + api_key = api_key or os.environ.get("TWELVELABS_API_KEY") + if not api_key: + raise ValueError( + "TWELVELABS_API_KEY is not set. Get a free key at https://twelvelabs.io" + ) + + client = TwelveLabs(api_key=api_key) + + print("Uploading video to TwelveLabs ...") + with open(video_file, "rb") as f: + asset = client.assets.create(method="direct", file=f) + + # Wait for the asset to finish processing before analysing it. + deadline = time.monotonic() + timeout + status = asset.status + while status not in ("ready", "failed"): + if time.monotonic() > deadline: + raise TimeoutError( + f"Asset {asset.id} not ready after {timeout}s (status={status})" + ) + time.sleep(poll_interval) + status = client.assets.retrieve(asset.id).status + if status == "failed": + raise RuntimeError(f"TwelveLabs failed to process asset {asset.id}") + + print(f"Transcribing with TwelveLabs {model} ...") + result = client.analyze( + model_name=model, + video={"type": "asset_id", "asset_id": asset.id}, + prompt=_SRT_PROMPT, + max_tokens=2048, + ) + + srt_text = _strip_fences(result.data or "") + if not _CUE_RE.search(srt_text): + raise RuntimeError( + "Pegasus did not return SRT-formatted subtitles. Raw output:\n" + + srt_text[:500] + ) + return srt_text diff --git a/setup.py b/setup.py index cd1c545..22070ff 100644 --- a/setup.py +++ b/setup.py @@ -13,6 +13,9 @@ 'ffutils', 'stable-ts==2.17.3', ], + extras_require={ + 'pegasus': ['twelvelabs>=1.2.8'], + }, packages=find_packages(), entry_points={ 'console_scripts': [ diff --git a/tests/test_pegasus.py b/tests/test_pegasus.py new file mode 100644 index 0000000..58fd128 --- /dev/null +++ b/tests/test_pegasus.py @@ -0,0 +1,46 @@ +import os + +import pytest + +from decipher.pegasus import _CUE_RE, _strip_fences, video_to_srt + +SAMPLE_SRT = """1 +00:00:00,000 --> 00:00:02,500 +Hello world. + +2 +00:00:02,500 --> 00:00:05,000 +This is Decipher. +""" + + +def test_strip_fences_removes_code_block(): + fenced = "```srt\n" + SAMPLE_SRT + "\n```" + assert _strip_fences(fenced) == SAMPLE_SRT.strip() + + +def test_strip_fences_noop_when_plain(): + assert _strip_fences(SAMPLE_SRT) == SAMPLE_SRT.strip() + + +def test_cue_regex_matches_srt_timestamps(): + assert _CUE_RE.search(SAMPLE_SRT) + assert not _CUE_RE.search("no timestamps here") + + +def test_missing_api_key_raises(monkeypatch): + monkeypatch.delenv("TWELVELABS_API_KEY", raising=False) + with pytest.raises(ValueError, match="TWELVELABS_API_KEY"): + video_to_srt("does-not-matter.mp4") + + +@pytest.mark.skipif( + not os.environ.get("TWELVELABS_API_KEY"), + reason="requires TWELVELABS_API_KEY for a live TwelveLabs call", +) +def test_client_constructs_with_key(): + # Confirms the SDK is installed and the key is accepted (no network call). + from twelvelabs import TwelveLabs + + client = TwelveLabs(api_key=os.environ["TWELVELABS_API_KEY"]) + assert client is not None