from __future__ import annotations import argparse import gc import re import sys import time from contextlib import ExitStack from pathlib import Path import ctranslate2 from faster_whisper import WhisperModel VIDEO_EXTENSIONS = {".avi", ".m4v", ".mkv", ".mov", ".mp4", ".webm"} def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="Transcribe videos to timestamped TXT and SRT files." ) parser.add_argument( "--input", type=Path, default=Path("input"), help="Folder containing video files (default: ./input).", ) parser.add_argument( "--output", type=Path, default=Path("output"), help="Folder for TXT and SRT files (default: ./output).", ) parser.add_argument("--model", default="large-v3") parser.add_argument( "--language", default="fa", help="Language code such as fa, en or de; use auto for detection.", ) parser.add_argument( "--device", choices=("auto", "cuda", "cpu"), default="auto", ) parser.add_argument( "--compute-type", default=None, help="Optional override such as float16, int8_float16 or int8.", ) parser.add_argument("--beam-size", type=int, default=5) parser.add_argument( "--model-cache", type=Path, default=None, help="Optional persistent folder for downloaded model files.", ) return parser.parse_args() def natural_key(path: Path) -> list[tuple[int, int | str]]: return [ (0, int(part)) if part.isdigit() else (1, part.casefold()) for part in re.split(r"(\d+)", path.name) ] def srt_time(seconds: float) -> str: total_ms = max(0, int(round(seconds * 1000))) hours, remainder = divmod(total_ms, 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 select_device(requested: str) -> tuple[str, str]: try: cuda_devices = ctranslate2.get_cuda_device_count() except Exception: cuda_devices = 0 if requested == "cuda": if cuda_devices < 1: raise RuntimeError("CUDA was requested, but no compatible GPU was detected.") return "cuda", "float16" if requested == "cpu": return "cpu", "int8" if cuda_devices > 0: return "cuda", "float16" return "cpu", "int8" def preserve_old_part(part_path: Path) -> None: if not part_path.exists(): return backup = part_path.with_name(f"{part_path.name}.previous-{time.time_ns()}") part_path.rename(backup) print(f"Preserved unfinished output: {backup.name}", flush=True) def transcribe_video( video_path: Path, output_dir: Path, model: WhisperModel, language: str | None, beam_size: int, ) -> None: txt_path = output_dir / f"{video_path.stem}.txt" srt_path = output_dir / f"{video_path.stem}.srt" txt_part = output_dir / f".{video_path.stem}.txt.part" srt_part = output_dir / f".{video_path.stem}.srt.part" need_txt = not txt_path.exists() need_srt = not srt_path.exists() if not need_txt and not need_srt: print(f"Skipped: {video_path.name} (TXT and SRT already exist)") return if not need_txt: print(f"Keeping existing file: {txt_path.name}") if not need_srt: print(f"Keeping existing file: {srt_path.name}") if need_txt: preserve_old_part(txt_part) if need_srt: preserve_old_part(srt_part) print(f"\nStarting: {video_path.name}", flush=True) segments, info = model.transcribe( str(video_path), language=language, task="transcribe", beam_size=beam_size, vad_filter=True, vad_parameters={"min_silence_duration_ms": 500}, ) duration = float(getattr(info, "duration", 0.0) or 0.0) detected_language = getattr(info, "language", None) probability = getattr(info, "language_probability", None) if detected_language: message = f"Language: {detected_language}" if probability is not None: message += f" ({probability:.1%})" print(message, flush=True) subtitle_number = 0 next_progress = 0 with ExitStack() as stack: txt_file = ( stack.enter_context(txt_part.open("w", encoding="utf-8", newline="\n")) if need_txt else None ) srt_file = ( stack.enter_context(srt_part.open("w", encoding="utf-8", newline="\n")) if need_srt else None ) for segment in segments: text = segment.text.strip() if not text: continue subtitle_number += 1 start = srt_time(segment.start) end = srt_time(segment.end) if txt_file: txt_file.write(f"[{start} --> {end}]\n{text}\n\n") if srt_file: srt_file.write( f"{subtitle_number}\n" f"{start} --> {end}\n" f"{text}\n\n" ) if duration > 0: progress = min(100, int(segment.end / duration * 100)) if progress >= next_progress: print( f"Progress: {progress:3d}% ({subtitle_number} segments)", flush=True, ) next_progress = (progress // 5 + 1) * 5 if txt_file: txt_file.flush() if srt_file: srt_file.flush() if need_txt: if txt_path.exists(): raise FileExistsError(f"Final output appeared during processing: {txt_path}") txt_part.rename(txt_path) print(f"Created: {txt_path.name}") if need_srt: if srt_path.exists(): raise FileExistsError(f"Final output appeared during processing: {srt_path}") srt_part.rename(srt_path) print(f"Created: {srt_path.name}") print(f"Finished: {video_path.name}", flush=True) def find_videos(input_dir: Path) -> list[Path]: videos = sorted( ( path for path in input_dir.iterdir() if path.is_file() and path.suffix.casefold() in VIDEO_EXTENSIONS ), key=natural_key, ) stems: dict[str, list[str]] = {} for video in videos: stems.setdefault(video.stem.casefold(), []).append(video.name) collisions = [names for names in stems.values() if len(names) > 1] if collisions: examples = "; ".join(", ".join(names) for names in collisions) raise ValueError( "Video filenames must have unique names before the extension. " f"Conflicts: {examples}" ) return videos def main() -> int: args = parse_args() input_dir = args.input.expanduser().resolve() output_dir = args.output.expanduser().resolve() if not input_dir.is_dir(): print(f"ERROR: Input folder does not exist: {input_dir}", file=sys.stderr) return 2 if args.beam_size < 1: print("ERROR: --beam-size must be at least 1.", file=sys.stderr) return 2 output_dir.mkdir(parents=True, exist_ok=True) try: videos = find_videos(input_dir) except ValueError as exc: print(f"ERROR: {exc}", file=sys.stderr) return 2 if not videos: formats = ", ".join(sorted(VIDEO_EXTENSIONS)) print(f"No supported video files found in: {input_dir}") print(f"Supported extensions: {formats}") return 0 try: device, default_compute_type = select_device(args.device) compute_type = args.compute_type or default_compute_type model_options: dict[str, str] = { "device": device, "compute_type": compute_type, } if args.model_cache: model_cache = args.model_cache.expanduser().resolve() model_cache.mkdir(parents=True, exist_ok=True) model_options["download_root"] = str(model_cache) print(f"Files found: {len(videos)}") print(f"Device: {device}") print(f"Compute type: {compute_type}") print(f"Loading model: {args.model}", flush=True) model = WhisperModel(args.model, **model_options) except Exception as exc: print(f"ERROR while loading the model: {exc}", file=sys.stderr) return 2 language = None if args.language.casefold() == "auto" else args.language failures: list[tuple[str, str]] = [] for video_path in videos: try: transcribe_video( video_path=video_path, output_dir=output_dir, model=model, language=language, beam_size=args.beam_size, ) except Exception as exc: failures.append((video_path.name, str(exc))) print( f"\nERROR while processing {video_path.name}: {exc}", file=sys.stderr, flush=True, ) finally: gc.collect() if failures: print("\nFailed files:", file=sys.stderr) for filename, error in failures: print(f"- {filename}: {error}", file=sys.stderr) return 1 print("\nAll files completed successfully.") return 0 if __name__ == "__main__": raise SystemExit(main())