aisrt 0.1.0__py3-none-any.whl

aisrt/__init__.py

@@ -0,0 +1,3 @@
+"""Ultimate SRT Generator."""
+
+__version__ = "0.1.0"

aisrt/assembly.py

@@ -0,0 +1,192 @@
+"""Broadcast-quality SubRip (SRT) formatting and Atomic File I/O."""
+
+import os
+from pathlib import Path
+from typing import Any
+
+from loguru import logger
+
+
+def _format_timestamp(seconds: float) -> str:
+    """Format a timestamp (in float seconds) to SRT standard: HH:MM:SS,mmm."""
+    hours = int(seconds // 3600)
+    minutes = int((seconds % 3600) // 60)
+    secs = int(seconds % 60)
+    millis = int(round((seconds - int(seconds)) * 1000))
+
+    if millis == 1000:
+        secs += 1
+        millis = 0
+        if secs == 60:
+            secs = 0
+            minutes += 1
+            if minutes == 60:
+                minutes = 0
+                hours += 1
+
+    return f"{hours:02d}:{minutes:02d}:{secs:02d},{millis:03d}"
+
+
+class SRTFormatter:
+    """Chunks Whisper words into broadcast-standard SRT format."""
+
+    def __init__(self, max_chars_per_line: int = 42, max_lines: int = 2) -> None:
+        """Initialize the SRT chunker.
+
+        Args:
+            max_chars_per_line: Maximum characters before wrapping a line.
+            max_lines: Maximum lines per subtitle block.
+        """
+        self.max_chars_per_line = max_chars_per_line
+        self.max_lines = max_lines
+        self.terminal_punctuation = {".", "?", "!", "。", "？", "！"}
+
+    def format_segments(self, segments: Any) -> str:
+        """Iterate over faster-whisper Segment/Word objects and yield SRT blocks.
+
+        Requires word_timestamps=True in the Whisper model transcribe() call.
+
+        Args:
+            segments: A generator of faster-whisper Segment objects.
+
+        Returns:
+            The complete SRT file content as a string.
+        """
+        self._srt_blocks: list[str] = []
+        self._block_idx = 1
+
+        for segment in segments:
+            if not getattr(segment, "words", None):
+                self._format_raw_segment(segment)
+            else:
+                self._format_word_segment(segment)
+
+        return "\n".join(self._srt_blocks)
+
+    def _format_raw_segment(self, segment: Any) -> None:
+        """Fallback formatter for segments without word timestamps."""
+        text = segment.text.strip()
+        if text:
+            start = _format_timestamp(segment.start)
+            end = _format_timestamp(segment.end)
+            self._srt_blocks.append(f"{self._block_idx}\n{start} --> {end}\n{text}\n")
+            self._block_idx += 1
+
+    def _format_word_segment(self, segment: Any) -> None:
+        """Advanced formatter that chunks based on character count and punctuation."""
+        current_words: list[str] = []
+        current_start: float | None = None
+        current_end: float = 0.0
+        char_count = 0
+        line_count = 1
+
+        for word_obj in segment.words:
+            word = word_obj.word.strip()
+            if not word:
+                continue
+
+            # Temporal gap check: flush if silence > 1.5s
+            if current_end > 0.0 and (word_obj.start - current_end) > 1.5:
+                if current_words and current_start is not None:
+                    self._flush_words(current_words, current_start, current_end)
+                    current_words = []
+                    current_start = None
+                    char_count = 0
+                    line_count = 1
+
+            if current_start is None:
+                current_start = word_obj.start
+
+            current_words.append(word_obj.word)
+            current_end = word_obj.end
+            char_count += len(word)
+
+            is_terminal = any(word.endswith(p) for p in self.terminal_punctuation)
+
+            # If appending this word exceeds the line length, wrap BEFORE adding it
+            if char_count > self.max_chars_per_line and line_count < self.max_lines:
+                # Insert newline before the current word
+                current_words.pop()  # Remove the word we just added
+                current_words.append("\n")
+                current_words.append(word_obj.word.lstrip())
+                char_count = len(word)
+                line_count += 1
+
+            is_too_long = char_count >= self.max_chars_per_line
+
+            if is_terminal or (is_too_long and line_count >= self.max_lines):
+                self._flush_words(current_words, current_start, current_end)
+                current_words = []
+                current_start = None
+                char_count = 0
+                line_count = 1
+
+        if current_words and current_start is not None:
+            self._flush_words(current_words, current_start, current_end)
+
+    def _flush_words(self, words: list[str], start_time: float, end_time: float) -> None:
+        """Write the aggregated words to the block list."""
+        text = "".join(words).strip()
+        if text:
+            start_str = _format_timestamp(start_time)
+            end_str = _format_timestamp(end_time)
+            self._srt_blocks.append(f"{self._block_idx}\n{start_str} --> {end_str}\n{text}\n")
+            self._block_idx += 1
+
+
+class AtomicWriter:
+    """Handles cross-device POSIX atomic file writing and metadata inheritance."""
+
+    @staticmethod
+    def write_srt(source_video: Path, srt_content: str, language_code: str = "en") -> Path:
+        """Write the SRT securely, inheriting the permissions of the source video.
+
+        Args:
+            source_video: The original MKV/MP4 file.
+            srt_content: The fully formatted SRT text block.
+            language_code: The locale suffix for the subtitle (e.g., 'en', 'eng').
+
+        Returns:
+            The Path to the finalized, atomically committed SRT file.
+        """
+        final_srt_path = source_video.with_suffix(f".{language_code}.srt")
+        temp_srt_path = source_video.with_name(f".{source_video.stem}.srt.tmp")
+
+        logger.debug(f"Assembling atomic SRT chunks in {temp_srt_path}")
+
+        try:
+            # 1. Write to hidden temp file in the same directory
+            temp_srt_path.write_text(srt_content, encoding="utf-8")
+
+            # 2. Inherit metadata from the source video
+            stat = source_video.stat()
+
+            try:
+                os.chown(temp_srt_path, stat.st_uid, stat.st_gid)
+            except PermissionError:
+                # Running as non-root over SMB/NFS might restrict chown
+                logger.debug(
+                    f"Insufficient permissions to chown {temp_srt_path} to "
+                    f"UID:{stat.st_uid}/GID:{stat.st_gid}. Proceeding anyway."
+                )
+
+            try:
+                os.chmod(temp_srt_path, stat.st_mode)
+            except PermissionError:
+                logger.debug(f"Insufficient permissions to chmod {temp_srt_path}")
+
+            # 3. Cross-device safe Atomic Rename
+            # os.replace is atomic on POSIX if both files are on the same filesystem.
+            # We write the temp file in the same folder to guarantee this and prevent EXDEV errors.
+            os.replace(temp_srt_path, final_srt_path)
+            logger.info(f"Successfully generated and committed {final_srt_path.name}")
+            return final_srt_path
+
+        except Exception as e:
+            # Clean up the temp file if the atomic commit fails
+            if temp_srt_path.exists():
+                try:
+                    temp_srt_path.unlink()
+                except OSError:
+                    pass
+            raise RuntimeError(f"Atomic subtitle write failed for {source_video.name}: {e}") from e

aisrt/cli.py

@@ -0,0 +1,177 @@
+"""CLI commands for the SRT Generator."""
+
+import asyncio
+import sys
+from pathlib import Path
+from typing import Annotated
+
+import typer
+from loguru import logger
+from rich.console import Console
+from rich.table import Table
+
+from aisrt.config import AppConfig, FilterConfig, HardwareConfig
+from aisrt.discovery import DiscoveryEngine
+from aisrt.hardware import HardwareProfiler, ModelRouter, setup_thread_safety
+from aisrt.state import StateTracker
+
+app = typer.Typer(help="Ultimate SRT Generator", add_completion=False)
+console = Console()
+
+
+def configure_logging(verbose: bool) -> None:
+    """Configure Loguru to output cleanly via Rich."""
+    logger.remove()
+    log_level = "DEBUG" if verbose else "INFO"
+    logger.add(sys.stderr, level=log_level, colorize=True)
+
+
+@app.command()
+def scan(
+    media_dir: Annotated[Path, typer.Argument(help="Root directory containing media files")],
+    min_age_mins: Annotated[int, typer.Option(help="Minimum file age in minutes")] = 15,
+    force_device: Annotated[str | None, typer.Option(help="Force specific device")] = None,
+    force_model: Annotated[str | None, typer.Option(help="Force specific model")] = None,
+    verbose: Annotated[bool, typer.Option("--verbose", "-v", help="Enable debug logging")] = False,
+) -> None:
+    """Perform a dry-run scan of the media directory and profile hardware."""
+    configure_logging(verbose)
+
+    # Compile the configuration
+    hw_config = HardwareConfig(force_device=force_device, force_model=force_model)
+    flt_config = FilterConfig(min_age_mins=min_age_mins)
+    config = AppConfig(
+        media_dir=media_dir,
+        dry_run=True,
+        hardware=hw_config,
+        filters=flt_config,
+    )
+
+    # 1. Profile Hardware
+    setup_thread_safety()
+    console.print("\n[bold cyan]1. Profiling Hardware...[/bold cyan]")
+    profile = HardwareProfiler.profile()
+    _ = ModelRouter.get_config(profile, config.hardware)
+
+    # 2. Run the Async Discovery Engine
+    console.print(f"\n[bold cyan]2. Scanning Directory: {config.media_dir}[/bold cyan]")
+    asyncio.run(_run_scan(config))
+
+
+async def _run_scan(config: AppConfig) -> None:
+    """Execute the asynchronous scanning process."""
+    table = Table(title="Media File Discovery Report", show_lines=True)
+    table.add_column("File Path", style="dim", no_wrap=False)
+    table.add_column("Size (MB)", justify="right", style="green")
+    table.add_column("Action", style="magenta")
+    table.add_column("Reason", style="yellow")
+
+    async with StateTracker(config.db_path) as tracker:
+        engine = DiscoveryEngine(config.media_dir, config.filters, tracker)
+
+        process_count = 0
+        skip_count = 0
+
+        async for media_file, action_str in engine.scan():
+            size_mb = media_file.size / (1024 * 1024)
+            path_str = str(media_file.path.relative_to(config.media_dir))
+
+            if action_str == "PROCESS":
+                table.add_row(
+                    path_str,
+                    f"{size_mb:.1f}",
+                    "[bold green]PROCESS[/bold green]",
+                    "Needs Subtitle",
+                )
+                process_count += 1
+            else:
+                reason = action_str.replace("SKIP: ", "")
+                table.add_row(path_str, f"{size_mb:.1f}", "[dim]SKIP[/dim]", reason)
+                skip_count += 1
+
+    console.print(table)
+    console.print(
+        f"\n[bold]Summary:[/bold] {process_count} files to process, {skip_count} files skipped."
+    )
+
+
+@app.command()
+def run(
+    media_dir: Annotated[Path, typer.Argument(help="Root directory containing media files")],
+    min_age_mins: Annotated[int, typer.Option(help="Minimum file age in minutes")] = 15,
+    translate: Annotated[
+        bool, typer.Option("--translate", help="Enable AI translation to English")
+    ] = False,
+    watch: Annotated[bool, typer.Option("--watch", help="Run continuously in daemon mode")] = False,
+    watch_interval: Annotated[
+        int, typer.Option("--watch-interval", help="Minutes between scans in watch mode")
+    ] = 60,
+    force_device: Annotated[str | None, typer.Option(help="Force specific device")] = None,
+    force_model: Annotated[str | None, typer.Option(help="Force specific model")] = None,
+    verbose: Annotated[bool, typer.Option("--verbose", "-v", help="Enable debug logging")] = False,
+) -> None:
+    """Run the live SRT generation pipeline."""
+    configure_logging(verbose)
+
+    # Compile the configuration
+    hw_config = HardwareConfig(force_device=force_device, force_model=force_model)
+    flt_config = FilterConfig(min_age_mins=min_age_mins)
+    config = AppConfig(
+        media_dir=media_dir,
+        dry_run=False,
+        translate=translate,
+        watch=watch,
+        watch_interval_mins=watch_interval,
+        hardware=hw_config,
+        filters=flt_config,
+    )
+
+    setup_thread_safety()
+
+    console.print("\n[bold cyan]1. Profiling Hardware & Initializing Models...[/bold cyan]")
+    profile = HardwareProfiler.profile()
+    model_cfg = ModelRouter.get_config(profile, config.hardware)
+
+    # Initialize the STT singleton before starting async loop
+    from aisrt.stt import STTWorker
+
+    stt_worker = STTWorker()
+    stt_worker.initialize(model_cfg)
+
+    console.print(f"\n[bold cyan]2. Starting Async Pipeline on {config.media_dir}[/bold cyan]")
+
+    # We define a wrapper to inject the db context manager and the pipeline
+    async def _execute_pipeline() -> None:
+        from aisrt.pipeline import Pipeline
+
+        async with StateTracker(config.db_path) as tracker:
+            while True:
+                engine = DiscoveryEngine(config.media_dir, config.filters, tracker)
+                pipeline = Pipeline(
+                    engine, cpu_cores=profile.physical_cores, translate=config.translate
+                )
+                await pipeline.run()
+
+                if not config.watch:
+                    break
+
+                console.print(
+                    f"\n[bold yellow]Sleeping for {config.watch_interval_mins} "
+                    f"minutes...[/bold yellow]"
+                )
+                await asyncio.sleep(config.watch_interval_mins * 60)
+                console.print(
+                    f"\n[bold cyan]Waking up and scanning Directory: {config.media_dir}[/bold cyan]"
+                )
+
+    try:
+        asyncio.run(_execute_pipeline())
+        console.print("\n[bold green]Pipeline finished successfully.[/bold green]")
+    except KeyboardInterrupt:
+        console.print("\n[bold red]Pipeline interrupted by user.[/bold red]")
+    finally:
+        stt_worker.close()
+
+
+if __name__ == "__main__":
+    app()

aisrt/config.py

@@ -0,0 +1,80 @@
+"""Configuration schemas for the SRT Generator."""
+
+from pathlib import Path
+
+from pydantic import BaseModel, Field
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+
+class HardwareConfig(BaseModel):
+    """Configuration for hardware acceleration and inference limits."""
+
+    force_device: str | None = Field(
+        default=None,
+        description="Force a specific compute device (e.g., 'cuda'). Auto-detect if None.",
+    )
+    force_compute_type: str | None = Field(
+        default=None,
+        description="Force compute type (e.g., 'float16', 'int8'). Auto-detect if None.",
+    )
+    force_model: str | None = Field(
+        default=None,
+        description="Force a Whisper model (e.g., 'large-v3-turbo'). Auto-detect if None.",
+    )
+
+
+class FilterConfig(BaseModel):
+    """Configuration for filtering media files during discovery."""
+
+    min_age_mins: int = Field(
+        default=15,
+        description="Minimum file age in minutes to avoid processing active downloads.",
+    )
+    extensions: list[str] = Field(
+        default_factory=lambda: [".mkv", ".mp4", ".avi", ".webm"],
+        description="List of valid media file extensions to process.",
+    )
+    exclude_patterns: list[str] = Field(
+        default_factory=lambda: ["*sample*", "*extras*", "*featurettes*"],
+        description="Glob patterns for directories or files to ignore.",
+    )
+    target_languages: list[str] = Field(
+        default_factory=lambda: ["eng", "en"],
+        description="Target subtitle languages to generate/check for.",
+    )
+
+
+class AppConfig(BaseSettings):
+    """Main application configuration."""
+
+    media_dir: Path = Field(
+        description="The root directory containing media to scan.",
+    )
+    db_path: Path = Field(
+        default_factory=lambda: Path.home() / ".config" / "aisrt" / "state.db",
+        description="Path to the local SQLite state database.",
+    )
+    dry_run: bool = Field(
+        default=False,
+        description="If True, only scan and report what would be done (no execution).",
+    )
+    translate: bool = Field(
+        default=False,
+        description="If True, translates foreign audio to English using Whisper's translate task.",
+    )
+    watch: bool = Field(
+        default=False,
+        description="If True, runs the pipeline continuously in daemon mode.",
+    )
+    watch_interval_mins: int = Field(
+        default=60,
+        description="Interval in minutes between scans when running in watch mode.",
+    )
+    hardware: HardwareConfig = Field(default_factory=HardwareConfig)
+    filters: FilterConfig = Field(default_factory=FilterConfig)
+
+    model_config = SettingsConfigDict(
+        env_prefix="AISRT_",
+        env_nested_delimiter="__",
+        extra="ignore",
+    )

aisrt/discovery.py

@@ -0,0 +1,195 @@
+"""NAS-Safe File Discovery Engine."""
+
+import asyncio
+import json
+import os
+import time
+from collections.abc import AsyncGenerator
+from dataclasses import dataclass
+from pathlib import Path
+
+from loguru import logger
+
+from aisrt.config import FilterConfig
+from aisrt.state import StateTracker
+
+
+@dataclass
+class MediaFile:
+    """Represents a discovered media file pending processing."""
+
+    path: Path
+    size: int
+    mtime: float
+    inode: int
+
+
+class DiscoveryEngine:
+    """Safely crawls a media directory and filters files based on state and config."""
+
+    def __init__(self, media_dir: Path, config: FilterConfig, state_tracker: StateTracker) -> None:
+        """Initialize the discovery engine.
+
+        Args:
+            media_dir: The root directory to scan.
+            config: The filtering configuration rules.
+            state_tracker: The active SQLite state tracker.
+        """
+        self.media_dir = media_dir
+        self.config = config
+        self.state_tracker = state_tracker
+
+    async def scan(self) -> AsyncGenerator[tuple[MediaFile, str], None]:
+        """Scan the media directory and yield files with their action status.
+
+        Yields:
+            A tuple of (MediaFile, action_string).
+            action_string is 'PROCESS' if the file needs STT, or a 'SKIP: <reason>' string.
+        """
+        loop = asyncio.get_running_loop()
+
+        def _walk(directory: Path) -> list[Path]:
+            paths = []
+            try:
+                for entry in os.scandir(directory):
+                    path = Path(entry.path)
+
+                    if any(path.match(p) for p in self.config.exclude_patterns):
+                        continue
+
+                    if entry.is_dir(follow_symlinks=False):
+                        paths.extend(_walk(path))
+                    elif entry.is_file(follow_symlinks=False):
+                        if path.suffix.lower() in self.config.extensions:
+                            paths.append(path)
+            except PermissionError:
+                logger.warning(f"Permission denied: {directory}")
+            return paths
+
+        logger.info(f"Starting directory scan at {self.media_dir}...")
+        all_files = await loop.run_in_executor(None, _walk, self.media_dir)
+        logger.info(f"Found {len(all_files)} potential media files. Analyzing...")
+
+        current_time = time.time()
+
+        for file_path in all_files:
+            try:
+                stat = file_path.stat()
+                media_file = MediaFile(
+                    path=file_path,
+                    size=stat.st_size,
+                    mtime=stat.st_mtime,
+                    inode=stat.st_ino,
+                )
+            except OSError as e:
+                logger.warning(f"Could not stat {file_path}: {e}")
+                continue
+
+            action_str = await self._analyze_file(media_file, current_time)
+            yield media_file, action_str
+
+    async def _analyze_file(self, media_file: MediaFile, current_time: float) -> str:
+        """Determine if a single file should be processed or skipped."""
+        min_age_seconds = self.config.min_age_mins * 60
+
+        if (current_time - media_file.mtime) < min_age_seconds:
+            return f"SKIP: Modified recently (< {self.config.min_age_mins}m)"
+
+        if self._has_sibling_subtitle(media_file.path):
+            return "SKIP: External sibling subtitle exists"
+
+        db_state = await self.state_tracker.get_state(str(media_file.path))
+        if db_state and db_state.status == "COMPLETED" and db_state.size == media_file.size:
+            return "SKIP: Already processed (Database)"
+
+        is_hardlink = await self.state_tracker.check_hardlink_processed(
+            media_file.inode, media_file.size
+        )
+        if is_hardlink:
+            return "SKIP: Hardlink to already processed file"
+
+        if db_state and db_state.status == "EMBEDDED_EXISTS":
+            return "SKIP: Embedded English subtitle exists (Database)"
+
+        has_embedded = await self._check_embedded_subtitles(media_file.path)
+        if has_embedded:
+            await self.state_tracker.update_state(
+                file_path=str(media_file.path),
+                inode=media_file.inode,
+                mtime=media_file.mtime,
+                size=media_file.size,
+                status="EMBEDDED_EXISTS",
+            )
+            return "SKIP: Embedded English subtitle detected"
+
+        return "PROCESS"
+
+    def _has_sibling_subtitle(self, video_path: Path) -> bool:
+        """Check if an external SRT file exists next to the video."""
+        base_name = video_path.stem
+        dir_name = video_path.parent
+
+        check_suffixes = [".srt"]
+        for lang in self.config.target_languages:
+            check_suffixes.append(f".{lang}.srt")
+
+        for suffix in check_suffixes:
+            if (dir_name / f"{base_name}{suffix}").exists():
+                return True
+
+        return False
+
+    async def _check_embedded_subtitles(self, video_path: Path) -> bool:
+        """Run ffprobe to check if an embedded target-language subtitle exists."""
+        cmd = [
+            "ffprobe",
+            "-v",
+            "error",
+            "-select_streams",
+            "s",
+            "-show_entries",
+            "stream=index,codec_name:stream_tags=language",
+            "-of",
+            "json",
+            str(video_path),
+        ]
+
+        try:
+            process = await asyncio.create_subprocess_exec(
+                *cmd,
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.PIPE,
+            )
+            stdout, _ = await process.communicate()
+
+            if process.returncode != 0:
+                logger.warning(f"ffprobe failed on {video_path}")
+                return False
+
+            data = json.loads(stdout.decode("utf-8"))
+            streams = data.get("streams", [])
+
+            for stream in streams:
+                codec = stream.get("codec_name", "").lower()
+                tags = stream.get("tags", {})
+                lang = tags.get("language", "").lower()
+
+                # Only skip if we find a text-based subtitle track in the target language.
+                # Image-based subs (hdmv_pgs_subtitle) force transcodes on many players.
+                if lang in self.config.target_languages:
+                    if codec in ["subrip", "ass", "mov_text", "webvtt"]:
+                        return True
+                    else:
+                        logger.debug(
+                            f"Ignoring embedded {codec} subtitle in {video_path} (forces transcode)"
+                        )
+
+        except FileNotFoundError:
+            logger.error("ffprobe not found. Please ensure FFmpeg is installed and in PATH.")
+            raise
+        except json.JSONDecodeError:
+            logger.warning(f"Failed to parse ffprobe JSON for {video_path}")
+        except Exception as e:
+            logger.warning(f"Error checking embedded streams for {video_path}: {e}")
+
+        return False