lattifai 1.2.1__py3-none-any.whl → 1.3.0__py3-none-any.whl

lattifai/config/event.py

@@ -0,0 +1,102 @@
+"""Audio Event Detection configuration for LattifAI."""
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Dict, List, Literal, Optional
+
+from ..utils import _select_device
+
+if TYPE_CHECKING:
+    from ..client import SyncAPIClient
+
+
+@dataclass
+class EventConfig:
+    """
+    Audio Event Detection configuration.
+
+    Settings for detecting audio events (Speech, Music, Male, Female...) in audio files using the AED model.
+
+    Event Matching:
+        When event_matching is enabled, the AED system will:
+        1. Parse [Event] markers from input captions (e.g., [Music], [Applause])
+        2. Match caption events to AED labels using semantic matching
+        3. Force detection of matched labels even if not in top_k
+        4. Update caption timestamps based on AED detection results
+
+        Event matching logic is implemented in lattifai_core.event.EventMatcher.
+    """
+
+    enabled: bool = False
+    """Enable audio event detection."""
+
+    device: Literal["cpu", "cuda", "mps", "auto"] = "auto"
+    """Computation device for Event Detection models."""
+
+    vad_chunk_size: float = 30.0
+    """VAD chunk size in seconds for speech segmentation."""
+
+    vad_max_gap: float = 2.0
+    """Maximum gap in seconds between VAD segments to merge."""
+
+    fast_mode: bool = True
+    """Enable fast mode (only detect top_k classes, skip others)."""
+
+    model_path: str = ""
+    """Path to pretrained model. If empty, uses default bundled model."""
+
+    event_matching: bool = True
+    """Whether update events in the alignment"""
+
+    extra_events: List[str] = field(default_factory=list)
+    """Additional event types to always detect, even if not in top_k.
+    Example: ["Applause", "Laughter", "Music"]
+    """
+
+    event_aliases: Dict[str, List[str]] = field(default_factory=dict)
+    """Custom aliases mapping [Event] markers to AED labels.
+
+    Core AED labels (14 types):
+        [Applause], [Baby cry], [Battle cry], [Bellow], [Children shouting],
+        [Laughter], [Music], [Shout], [Singing], [Sound effect],
+        [Speech], [Whoop], [Yell]
+
+    Custom aliases extend built-ins (not replace):
+        {"[Audience reaction]": ["[Applause]", "[Cheering]"]}
+    """
+
+    time_tolerance: float = 20.0
+    """Max time (seconds) non-Speech events can extend beyond supervision boundaries."""
+
+    update_timestamps: bool = True
+    """Whether to update caption event timestamps based on AED detections."""
+
+    duplicate_strategy: Literal["keep_all", "merge_first", "split"] = "merge_first"
+    """Strategy for handling multiple [Event] markers mapped to same AED interval.
+    - keep_all: Update all events to same time range (may cause overlapping)
+    - merge_first: Keep only first event per interval, skip duplicates
+    - split: Split interval at speech boundaries (not yet implemented)
+    """
+
+    client_wrapper: Optional["SyncAPIClient"] = field(default=None, repr=False)
+    """Reference to the SyncAPIClient instance. Auto-set during client initialization."""
+
+    def __post_init__(self):
+        """Validate and auto-populate configuration after initialization."""
+        # Validate device
+        if self.device not in ("cpu", "cuda", "mps", "auto") and not self.device.startswith("cuda:"):
+            raise ValueError(f"device must be one of ('cpu', 'cuda', 'mps', 'auto'), got '{self.device}'")
+
+        if self.device == "auto":
+            self.device = _select_device(self.device)
+
+        # Validate vad_chunk_size
+        if self.vad_chunk_size < 0:
+            raise ValueError("vad_chunk_size must be non-negative")
+
+        # Validate vad_max_gap
+        if self.vad_max_gap < 0:
+            raise ValueError("vad_max_gap must be non-negative")
+
+        # Validate time_tolerance
+        if self.time_tolerance < 0:
+            raise ValueError("time_tolerance must be non-negative")

lattifai/config/media.py

@@ -91,6 +91,26 @@ class MediaConfig:
     force_overwrite: bool = False
     """Overwrite existing output files without prompting."""
 
+    audio_track_id: Optional[str] = "original"
+    """Audio track ID for multi-language YouTube videos.
+    - "original": Select the original audio track (default)
+    - Language code (e.g., "en", "ja", "fr"): Select by language
+    - Format ID (e.g., "251-drc", "140-0"): Select specific format
+    - None: No filtering, use yt-dlp default selection
+    """
+
+    quality: str = "best"
+    """Media quality for YouTube downloads.
+    For audio:
+    - "best": Highest bitrate (default)
+    - "medium": ~128 kbps
+    - "low": ~50 kbps
+    - Numeric string (e.g., "128"): Target bitrate in kbps
+    For video:
+    - "best": Highest resolution (default)
+    - "1080", "720", "480", "360": Target resolution
+    """
+
     def __post_init__(self) -> None:
         """Validate configuration and normalize paths/formats."""
         self._setup_output_directory()

lattifai/config/transcription.py

@@ -7,7 +7,7 @@ from typing import TYPE_CHECKING, Literal, Optional
 from ..utils import _select_device
 
 if TYPE_CHECKING:
-    from ..base_client import SyncAPIClient
+    from ..client import SyncAPIClient
 
 SUPPORTED_TRANSCRIPTION_MODELS = Literal[
     "gemini-2.5-pro",
@@ -48,6 +48,30 @@ class TranscriptionConfig:
     language: Optional[str] = None
     """Target language code for transcription (e.g., 'en', 'zh', 'ja')."""
 
+    prompt: Optional[str] = None
+    """Custom prompt text or path to prompt file for transcription.
+    If the value is an existing file path, the file contents will be used.
+    Otherwise, the value is used directly as the prompt text."""
+
+    description: Optional[str] = None
+    """Media description from platforms like YouTube, Xiaoyuzhou (小宇宙), etc.
+    Used to provide context for transcription."""
+
+    thinking: bool = True
+    """Enable Gemini's thinking mode (Gemini models only). Set to False to disable thinking."""
+
+    include_thoughts: bool = False
+    """Include Gemini's thinking process in the output (Gemini models only). Requires thinking=True."""
+
+    temperature: Optional[float] = None
+    """Sampling temperature for generation. Higher values increase randomness."""
+
+    top_k: Optional[float] = None
+    """Top-k sampling parameter. Limits token selection to top k candidates."""
+
+    top_p: Optional[float] = None
+    """Nucleus sampling parameter. Limits token selection by cumulative probability."""
+
     lattice_model_path: Optional[str] = None
     """Path to local LattifAI model. Will be auto-set in LattifAI client."""
 

lattifai/data/__init__.py

@@ -0,0 +1,8 @@
+"""Data types for LattifAI.
+
+Provides extended Caption class with transcription/alignment/diarization support.
+"""
+
+from .caption import Caption
+
+__all__ = ["Caption"]

lattifai/data/caption.py

@@ -0,0 +1,228 @@
+"""Extended Caption class with transcription, alignment, and diarization support."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, TypeVar
+
+from lattifai.caption import Caption as BaseCaption
+from lattifai.caption import Pathlike, Supervision
+
+if TYPE_CHECKING:
+    from lattifai_core.event import LEDOutput
+
+DiarizationOutput = TypeVar("DiarizationOutput")
+
+
+@dataclass
+class Caption(BaseCaption):
+    """
+    Extended Caption with transcription, alignment, and diarization support.
+
+    Inherits from BaseCaption and adds fields for:
+    - alignments: Post-alignment results
+    - transcription: ASR results
+    - event: LattifAI Event Detection results (LEDOutput)
+    - diarization: Speaker diarization results
+
+
+    These fields are used in the LattifAI pipeline for:
+    - Forced alignment results
+    - Storing intermediate transcription results
+    - LattifAI Event Detection (music, applause, speech, etc.)
+    - Speaker identification and separation
+
+    """
+
+    # Alignment results
+    alignments: List[Supervision] = field(default_factory=list)
+
+    # Transcription results
+    transcription: List[Supervision] = field(default_factory=list)
+
+    # LattifAI Event Detection results
+    event: Optional["LEDOutput"] = None
+
+    # Speaker Diarization results
+    diarization: Optional[DiarizationOutput] = None
+
+    def __len__(self) -> int:
+        """Return the number of supervision segments."""
+        return len(self.supervisions or self.transcription)
+
+    def __repr__(self) -> str:
+        """String representation of Caption."""
+        lang = f"lang={self.language}" if self.language else "lang=unknown"
+        kind_str = f"kind={self.kind}" if self.kind else ""
+        parts = [f"Caption({len(self.supervisions or self.transcription)} segments", lang]
+        if kind_str:
+            parts.append(kind_str)
+        if self.duration:
+            parts.append(f"duration={self.duration:.2f}s")
+        return ", ".join(parts) + ")"
+
+    def with_margins(
+        self,
+        start_margin: float = 0.08,
+        end_margin: float = 0.20,
+        min_gap: float = 0.08,
+        collision_mode: str = "trim",
+    ) -> "Caption":
+        """
+        Create a new Caption with segment boundaries adjusted based on word-level alignment.
+
+        Uses supervision.alignment['word'] to recalculate segment start/end times
+        with the specified margins applied around the actual speech boundaries.
+
+        Prefers alignments > supervisions > transcription as source.
+
+        Args:
+            start_margin: Seconds to extend before the first word (default: 0.08)
+            end_margin: Seconds to extend after the last word (default: 0.20)
+            min_gap: Minimum gap between segments for collision handling (default: 0.08)
+            collision_mode: How to handle segment overlap - 'trim' or 'gap' (default: 'trim')
+
+        Returns:
+            New Caption instance with adjusted timestamps
+
+        Note:
+            Segments without alignment data will keep their original timestamps.
+        """
+        from lattifai.caption.standardize import apply_margins_to_captions
+
+        # Determine which supervisions to use (priority: alignments > supervisions > transcription)
+        if self.alignments:
+            source_sups = self.alignments
+        elif self.supervisions:
+            source_sups = self.supervisions
+        else:
+            source_sups = self.transcription
+
+        adjusted_sups = apply_margins_to_captions(
+            source_sups,
+            start_margin=start_margin,
+            end_margin=end_margin,
+            min_gap=min_gap,
+            collision_mode=collision_mode,
+        )
+
+        return Caption(
+            supervisions=adjusted_sups,
+            transcription=self.transcription,
+            event=self.event,
+            diarization=self.diarization,
+            alignments=[],  # Clear alignments since we've applied them
+            language=self.language,
+            kind=self.kind,
+            source_format=self.source_format,
+            source_path=self.source_path,
+            metadata=self.metadata.copy() if self.metadata else {},
+        )
+
+    def write(
+        self,
+        path=None,
+        output_format: Optional[str] = None,
+        include_speaker_in_text: bool = True,
+        word_level: bool = False,
+        karaoke_config=None,
+        metadata: Optional[Dict[str, Any]] = None,
+    ):
+        """
+        Write caption to file or return as bytes.
+
+        Prefers alignments > supervisions > transcription as source.
+
+        Args:
+            path: Path to output caption file, BytesIO object, or None to return bytes
+            output_format: Output format (e.g., 'srt', 'vtt', 'ass')
+            include_speaker_in_text: Whether to include speaker labels in text
+            word_level: Use word-level output format if supported
+            karaoke_config: Karaoke configuration
+            metadata: Optional metadata dict to pass to writer
+
+        Returns:
+            Path to the written file if path is a file path, or bytes if path is BytesIO/None
+        """
+        # Temporarily swap supervisions to use the priority order
+        original_supervisions = self.supervisions
+
+        if self.alignments:
+            self.supervisions = self.alignments
+        elif not self.supervisions and self.transcription:
+            self.supervisions = self.transcription
+
+        try:
+            result = super().write(
+                path=path,
+                output_format=output_format,
+                include_speaker_in_text=include_speaker_in_text,
+                word_level=word_level,
+                karaoke_config=karaoke_config,
+                metadata=metadata,
+            )
+        finally:
+            # Restore original supervisions
+            self.supervisions = original_supervisions
+
+        return result
+
+    @classmethod
+    def from_transcription_results(
+        cls,
+        transcription: List[Supervision],
+        event: Optional["LEDOutput"] = None,
+        diarization: Optional[DiarizationOutput] = None,
+        language: Optional[str] = None,
+        source_path: Optional[Pathlike] = None,
+        metadata: Optional[Dict[str, str]] = None,
+    ) -> "Caption":
+        """
+        Create Caption from transcription results including audio events and diarization.
+
+        Args:
+            transcription: List of transcription supervision segments
+            event: Optional LEDOutput with event detection results
+            diarization: Optional DiarizationOutput with speaker diarization results
+            language: Language code
+            source_path: Source file path
+            metadata: Additional metadata
+
+        Returns:
+            New Caption instance with transcription data
+        """
+        return cls(
+            transcription=transcription,
+            event=event,
+            diarization=diarization,
+            language=language,
+            kind="transcription",
+            source_format="asr",
+            source_path=source_path,
+            metadata=metadata or {},
+        )
+
+    def read_diarization(
+        self,
+        path: Pathlike,
+    ) -> "DiarizationOutput":
+        """
+        Read speaker diarization TextGrid from file.
+        """
+        from lattifai_core.diarization import DiarizationOutput
+
+        self.diarization = DiarizationOutput.read(path)
+        return self.diarization
+
+    def write_diarization(
+        self,
+        path: Pathlike,
+    ) -> Pathlike:
+        """
+        Write speaker diarization TextGrid to file.
+        """
+        if not self.diarization:
+            raise ValueError("No speaker diarization data to write.")
+
+        self.diarization.write(path)
+        return path

lattifai/diarization/__init__.py

@@ -1,4 +1,44 @@
-"""Speaker diarization module for LattifAI."""
+"""Speaker diarization module for LattifAI.
+
+This module provides multi-speaker identification and labeling capabilities
+using pyannote.audio-based diarization models. It can identify who spoke
+when in an audio file and optionally match detected speakers with existing
+speaker labels from input captions.
+
+Key Components:
+    LattifAIDiarizer: Main diarization class that wraps pyannote.audio
+        pipelines for speaker segmentation and clustering.
+
+Features:
+    - Automatic speaker detection with configurable min/max speaker counts
+    - Speaker label preservation from input captions (e.g., "Alice:", ">> Bob:")
+    - Integration with alignment results to assign speakers to words/segments
+    - Support for pre-computed diarization results (avoid reprocessing)
+
+Configuration:
+    Use DiarizationConfig to control:
+    - enabled: Whether to run diarization
+    - min_speakers/max_speakers: Constrain speaker count detection
+    - device: GPU/CPU device selection
+    - debug: Enable verbose output
+
+Example:
+    >>> from lattifai import LattifAI
+    >>> from lattifai.config import DiarizationConfig
+    >>> client = LattifAI(diarization_config=DiarizationConfig(enabled=True))
+    >>> caption = client.alignment(audio="speech.wav", input_caption="transcript.srt")
+    >>> for seg in caption.supervisions:
+    ...     print(f"{seg.speaker}: {seg.text}")
+
+Performance Notes:
+    - Diarization adds ~10-30% processing time to alignment
+    - GPU acceleration recommended for longer audio files
+    - Results are cached when output_path is provided
+
+See Also:
+    - lattifai.config.DiarizationConfig: Configuration options
+    - lattifai.client.LattifAI.speaker_diarization: Direct diarization method
+"""
 
 from .lattifai import LattifAIDiarizer