lattifai 1.0.5__py3-none-any.whl → 1.2.0__py3-none-any.whl

lattifai/audio2.py

@@ -36,7 +36,7 @@ class AudioData(namedtuple("AudioData", ["sampling_rate", "ndarray", "path", "st
     @property
     def streaming_mode(self) -> bool:
         """Indicates whether streaming mode is enabled based on streaming_chunk_secs."""
-        if self.streaming_chunk_secs is not None:
+        if self.streaming_chunk_secs:
             return self.duration > self.streaming_chunk_secs * 1.1
         return False
 

lattifai/caption/caption.py

@@ -4,17 +4,19 @@ import json
 import re
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, TypeVar
 
 from lhotse.supervision import AlignmentItem
 from lhotse.utils import Pathlike
 from tgt import TextGrid
 
-from ..config.caption import InputCaptionFormat, OutputCaptionFormat
+from ..config.caption import InputCaptionFormat, OutputCaptionFormat  # noqa: F401
 from .supervision import Supervision
 from .text_parser import normalize_text as normalize_text_fn
 from .text_parser import parse_speaker_text, parse_timestamp_text
 
+DiarizationOutput = TypeVar("DiarizationOutput")
+
 
 @dataclass
 class Caption:
@@ -40,7 +42,7 @@ class Caption:
     # Audio Event Detection results
     audio_events: Optional[TextGrid] = None
     # Speaker Diarization results
-    speaker_diarization: Optional[TextGrid] = None
+    speaker_diarization: Optional[DiarizationOutput] = None
     # Alignment results
     alignments: List[Supervision] = field(default_factory=list)
 
@@ -272,7 +274,7 @@ class Caption:
         cls,
         transcription: List[Supervision],
         audio_events: Optional[TextGrid] = None,
-        speaker_diarization: Optional[TextGrid] = None,
+        speaker_diarization: Optional[DiarizationOutput] = None,
         language: Optional[str] = None,
         source_path: Optional[Pathlike] = None,
         metadata: Optional[Dict[str, str]] = None,
@@ -283,7 +285,7 @@ class Caption:
         Args:
             transcription: List of transcription supervision segments
             audio_events: Optional TextGrid with audio event detection results
-            speaker_diarization: Optional TextGrid with speaker diarization results
+            speaker_diarization: Optional DiarizationOutput with speaker diarization results
             language: Language code
             source_path: Source file path
             metadata: Additional metadata
@@ -384,9 +386,9 @@ class Caption:
         """
         Read speaker diarization TextGrid from file.
         """
-        from tgt import read_textgrid
+        from lattifai_core.diarization import DiarizationOutput
 
-        self.speaker_diarization = read_textgrid(path)
+        self.speaker_diarization = DiarizationOutput.read(path)
         return self.speaker_diarization
 
     def write_speaker_diarization(
@@ -399,9 +401,7 @@ class Caption:
         if not self.speaker_diarization:
             raise ValueError("No speaker diarization data to write.")
 
-        from tgt import write_to_file
-
-        write_to_file(self.speaker_diarization, path, format="long")
+        self.speaker_diarization.write(path)
         return path
 
     @staticmethod
@@ -451,7 +451,10 @@ class Caption:
                     else:
                         if include_speaker_in_text and sup.speaker is not None:
                             # Use [SPEAKER]: format for consistency with parsing
-                            text = f"[{sup.speaker}]: {sup.text}"
+                            if not sup.has_custom("original_speaker") or sup.custom["original_speaker"]:
+                                text = f"[{sup.speaker}]: {sup.text}"
+                            else:
+                                text = f"{sup.text}"
                         else:
                             text = sup.text
                         f.write(f"[{sup.start:.2f}-{sup.end:.2f}] {text}\n")
@@ -471,7 +474,12 @@ class Caption:
             tg = TextGrid()
             supervisions, words, scores = [], [], {"utterances": [], "words": []}
             for supervision in sorted(alignments, key=lambda x: x.start):
-                if include_speaker_in_text and supervision.speaker is not None:
+                # Respect `original_speaker` custom flag: default to include speaker when missing
+                if (
+                    include_speaker_in_text
+                    and supervision.speaker is not None
+                    and (not supervision.has_custom("original_speaker") or supervision.custom["original_speaker"])
+                ):
                     text = f"{supervision.speaker} {supervision.text}"
                 else:
                     text = supervision.text
@@ -526,7 +534,10 @@ class Caption:
                         )
                 else:
                     if include_speaker_in_text and sup.speaker is not None:
-                        text = f"{sup.speaker} {sup.text}"
+                        if not sup.has_custom("original_speaker") or sup.custom["original_speaker"]:
+                            text = f"{sup.speaker} {sup.text}"
+                        else:
+                            text = f"{sup.text}"
                     else:
                         text = sup.text
                     subs.append(
@@ -830,7 +841,16 @@ class Caption:
             if cls._is_youtube_vtt_with_word_timestamps(content):
                 return cls._parse_youtube_vtt_with_word_timestamps(content, normalize_text)
 
-        if format == "gemini" or str(caption).endswith("Gemini.md") or str(caption).endswith("Gemini3.md"):
+        # Match Gemini format: explicit format, or files ending with Gemini.md/Gemini3.md,
+        # or files containing "gemini" in the name with .md extension
+        caption_str = str(caption).lower()
+        is_gemini_format = (
+            format == "gemini"
+            or str(caption).endswith("Gemini.md")
+            or str(caption).endswith("Gemini3.md")
+            or ("gemini" in caption_str and caption_str.endswith(".md"))
+        )
+        if is_gemini_format:
             from .gemini_reader import GeminiReader
 
             supervisions = GeminiReader.extract_for_alignment(caption)
@@ -1242,7 +1262,11 @@ class Caption:
             if include_speaker_in_text:
                 file.write("speaker\tstart\tend\ttext\n")
                 for supervision in alignments:
-                    speaker = supervision.speaker or ""
+                    # Respect `original_speaker` custom flag: default to True when missing
+                    include_speaker = supervision.speaker and (
+                        not supervision.has_custom("original_speaker") or supervision.custom["original_speaker"]
+                    )
+                    speaker = supervision.speaker if include_speaker else ""
                     start_ms = round(1000 * supervision.start)
                     end_ms = round(1000 * supervision.end)
                     text = supervision.text.strip().replace("\t", " ")
@@ -1280,7 +1304,10 @@ class Caption:
                 writer = csv.writer(file)
                 writer.writerow(["speaker", "start", "end", "text"])
                 for supervision in alignments:
-                    speaker = supervision.speaker or ""
+                    include_speaker = supervision.speaker and (
+                        not supervision.has_custom("original_speaker") or supervision.custom["original_speaker"]
+                    )
+                    speaker = supervision.speaker if include_speaker else ""
                     start_ms = round(1000 * supervision.start)
                     end_ms = round(1000 * supervision.end)
                     text = supervision.text.strip()
@@ -1318,7 +1345,12 @@ class Caption:
                 end = supervision.end
                 text = supervision.text.strip().replace("\t", " ")
 
-                if include_speaker_in_text and supervision.speaker:
+                # Respect `original_speaker` custom flag when adding speaker prefix
+                if (
+                    include_speaker_in_text
+                    and supervision.speaker
+                    and (not supervision.has_custom("original_speaker") or supervision.custom["original_speaker"])
+                ):
                     text = f"[[{supervision.speaker}]]{text}"
 
                 file.write(f"{start}\t{end}\t{text}\n")
@@ -1364,9 +1396,13 @@ class Caption:
                 # Write timestamp line
                 file.write(f"{start_time},{end_time}\n")
 
-                # Write text (with optional speaker)
+                # Write text (with optional speaker). Respect `original_speaker` custom flag.
                 text = supervision.text.strip()
-                if include_speaker_in_text and supervision.speaker:
+                if (
+                    include_speaker_in_text
+                    and supervision.speaker
+                    and (not supervision.has_custom("original_speaker") or supervision.custom["original_speaker"])
+                ):
                     text = f"{supervision.speaker}: {text}"
 
                 file.write(f"{text}\n")

lattifai/cli/__init__.py

@@ -5,12 +5,14 @@ import nemo_run as run  # noqa: F401
 # Import and re-export entrypoints at package level so NeMo Run can find them
 from lattifai.cli.alignment import align
 from lattifai.cli.caption import convert
+from lattifai.cli.diarization import diarize
 from lattifai.cli.transcribe import transcribe, transcribe_align
 from lattifai.cli.youtube import youtube
 
 __all__ = [
     "align",
     "convert",
+    "diarize",
     "transcribe",
     "transcribe_align",
     "youtube",

lattifai/cli/caption.py

@@ -14,7 +14,7 @@ from lattifai.utils import safe_print
 def convert(
     input_path: Pathlike,
     output_path: Pathlike,
-    include_speaker_in_text: bool = True,
+    include_speaker_in_text: bool = False,
     normalize_text: bool = False,
 ):
     """

lattifai/cli/diarization.py

@@ -0,0 +1,110 @@
+"""Speaker diarization CLI entry point with nemo_run."""
+
+from pathlib import Path
+from typing import Optional
+
+import colorful
+import nemo_run as run
+from typing_extensions import Annotated
+
+from lattifai.client import LattifAI
+from lattifai.config import AlignmentConfig, CaptionConfig, ClientConfig, DiarizationConfig, MediaConfig
+from lattifai.utils import safe_print
+
+__all__ = ["diarize"]
+
+
+@run.cli.entrypoint(name="run", namespace="diarization")
+def diarize(
+    input_media: Optional[str] = None,
+    input_caption: Optional[str] = None,
+    output_caption: Optional[str] = None,
+    media: Annotated[Optional[MediaConfig], run.Config[MediaConfig]] = None,
+    caption: Annotated[Optional[CaptionConfig], run.Config[CaptionConfig]] = None,
+    client: Annotated[Optional[ClientConfig], run.Config[ClientConfig]] = None,
+    alignment: Annotated[Optional[AlignmentConfig], run.Config[AlignmentConfig]] = None,
+    diarization: Annotated[Optional[DiarizationConfig], run.Config[DiarizationConfig]] = None,
+):
+    """Run speaker diarization on aligned captions and audio."""
+
+    media_config = media or MediaConfig()
+    caption_config = caption or CaptionConfig()
+    diarization_config = diarization or DiarizationConfig()
+
+    if input_media and media_config.input_path:
+        raise ValueError("Cannot specify both positional input_media and media.input_path.")
+    if input_media:
+        media_config.set_input_path(input_media)
+    if not media_config.input_path:
+        raise ValueError("Input media path must be provided via positional input_media or media.input_path.")
+
+    if input_caption and caption_config.input_path:
+        raise ValueError("Cannot specify both positional input_caption and caption.input_path.")
+    if input_caption:
+        caption_config.set_input_path(input_caption)
+    if not caption_config.input_path:
+        raise ValueError("Input caption path must be provided via positional input_caption or caption.input_path.")
+
+    if output_caption and caption_config.output_path:
+        raise ValueError("Cannot specify both positional output_caption and caption.output_path.")
+    if output_caption:
+        caption_config.set_output_path(output_caption)
+
+    diarization_config.enabled = True
+
+    client_instance = LattifAI(
+        client_config=client,
+        alignment_config=alignment,
+        caption_config=caption_config,
+        diarization_config=diarization_config,
+    )
+
+    safe_print(colorful.cyan("🎧 Loading media for diarization..."))
+    media_audio = client_instance.audio_loader(
+        media_config.input_path,
+        channel_selector=media_config.channel_selector,
+        streaming_chunk_secs=media_config.streaming_chunk_secs,
+    )
+
+    safe_print(colorful.cyan("📖 Loading caption segments..."))
+    caption_obj = client_instance._read_caption(
+        caption_config.input_path,
+        input_caption_format=None if caption_config.input_format == "auto" else caption_config.input_format,
+        verbose=False,
+    )
+
+    if not caption_obj.alignments:
+        caption_obj.alignments = caption_obj.supervisions
+
+    if not caption_obj.alignments:
+        raise ValueError("Caption does not contain segments for diarization.")
+
+    if caption_config.output_path:
+        output_path = caption_config.output_path
+    else:
+        from datetime import datetime
+
+        input_caption_path = Path(caption_config.input_path)
+        timestamp = datetime.now().strftime("%Y%m%d_%H")
+        default_output = (
+            input_caption_path.parent / f"{input_caption_path.stem}.diarized.{timestamp}.{caption_config.output_format}"
+        )
+        caption_config.set_output_path(default_output)
+        output_path = caption_config.output_path
+
+    safe_print(colorful.cyan("🗣️ Performing speaker diarization..."))
+    diarized_caption = client_instance.speaker_diarization(
+        input_media=media_audio,
+        caption=caption_obj,
+        output_caption_path=output_path,
+    )
+
+    return diarized_caption
+
+
+def main():
+    run.cli.main(diarize)
+
+
+if __name__ == "__main__":
+    main()

lattifai/cli/transcribe.py

@@ -117,7 +117,9 @@ def transcribe(
 
     # Create transcriber
     if not transcription_config.lattice_model_path:
-        transcription_config.lattice_model_path = _resolve_model_path("LattifAI/Lattice-1")
+        transcription_config.lattice_model_path = _resolve_model_path(
+            "LattifAI/Lattice-1", getattr(transcription_config, "model_hub", "huggingface")
+        )
     transcriber = create_transcriber(transcription_config=transcription_config)
 
     safe_print(colorful.cyan(f"🎤 Starting transcription with {transcriber.name}..."))

lattifai/cli/youtube.py

@@ -25,6 +25,7 @@ def youtube(
     caption: Annotated[Optional[CaptionConfig], run.Config[CaptionConfig]] = None,
     transcription: Annotated[Optional[TranscriptionConfig], run.Config[TranscriptionConfig]] = None,
     diarization: Annotated[Optional[DiarizationConfig], run.Config[DiarizationConfig]] = None,
+    use_transcription: bool = False,
 ):
     """
     Download media from YouTube (when needed) and align captions.
@@ -55,6 +56,11 @@ def youtube(
             Fields: gemini_api_key, model_name, language, device
         diarization: Speaker diarization configuration.
             Fields: enabled, num_speakers, min_speakers, max_speakers, device
+        use_transcription: If True, skip YouTube caption download and directly use
+            transcription.model_name to transcribe. If False (default), first try to
+            download YouTube captions; if download fails (no captions available or
+            errors like HTTP 429), automatically fallback to transcription if
+            transcription.model_name is configured.
 
     Examples:
         # Download from YouTube and align (positional argument)
@@ -108,7 +114,11 @@ def youtube(
         transcription_config=transcription,
         diarization_config=diarization,
     )
+
     # Call the client's youtube method
+    # If use_transcription=True, skip YouTube caption download and use transcription directly.
+    # If use_transcription=False (default), try YouTube captions first; on failure,
+    # automatically fallback to transcription if transcription.model_name is configured.
     return lattifai_client.youtube(
         url=media_config.input_path,
         output_dir=media_config.output_dir,
@@ -118,6 +128,7 @@ def youtube(
         split_sentence=caption_config.split_sentence,
         channel_selector=media_config.channel_selector,
         streaming_chunk_secs=media_config.streaming_chunk_secs,
+        use_transcription=use_transcription,
     )
 
 

lattifai/client.py

@@ -56,6 +56,7 @@ class LattifAI(LattifAIClientMixin, SyncAPIClient):
 
         # Initialize base API client
         super().__init__(config=client_config)
+        self.config = client_config
 
         # Initialize all configs with defaults
         alignment_config, transcription_config, diarization_config = self._init_configs(
@@ -106,7 +107,13 @@ class LattifAI(LattifAIClientMixin, SyncAPIClient):
                 )
 
             if not input_caption:
-                caption = self._transcribe(media_audio, source_lang=self.caption_config.source_lang, is_async=False)
+                output_dir = None
+                if output_caption_path:
+                    output_dir = Path(str(output_caption_path)).parent
+                    output_dir.mkdir(parents=True, exist_ok=True)
+                caption = self._transcribe(
+                    media_audio, source_lang=self.caption_config.source_lang, is_async=False, output_dir=output_dir
+                )
             else:
                 caption = self._read_caption(input_caption, input_caption_format)
 
@@ -260,18 +267,13 @@ class LattifAI(LattifAIClientMixin, SyncAPIClient):
             caption.supervisions = supervisions
             caption.alignments = alignments
 
-            # Step 5: Speaker diarization
-            if self.diarization_config.enabled and self.diarizer:
-                safe_print(colorful.cyan("🗣️  Performing speaker diarization..."))
-                caption = self.speaker_diarization(
-                    input_media=media_audio,
-                    caption=caption,
-                    output_caption_path=output_caption_path,
-                )
-            elif output_caption_path:
+            if output_caption_path:
                 self._write_caption(caption, output_caption_path)
 
-            return caption
+            # Profile if enabled
+            if self.config.profile:
+                self.aligner.profile()
+
         except (CaptionProcessingError, LatticeEncodingError, AlignmentError, LatticeDecodingError):
             # Re-raise our specific errors as-is
             raise
@@ -284,6 +286,17 @@ class LattifAI(LattifAIClientMixin, SyncAPIClient):
                 context={"original_error": str(e), "error_type": e.__class__.__name__},
             )
 
+        # Step 5: Speaker diarization
+        if self.diarization_config.enabled and self.diarizer:
+            safe_print(colorful.cyan("🗣️  Performing speaker diarization..."))
+            caption = self.speaker_diarization(
+                input_media=media_audio,
+                caption=caption,
+                output_caption_path=output_caption_path,
+            )
+
+        return caption
+
     def speaker_diarization(
         self,
         input_media: AudioData,
@@ -315,7 +328,14 @@ class LattifAI(LattifAIClientMixin, SyncAPIClient):
                 caption.read_speaker_diarization(diarization_file)
 
         diarization, alignments = self.diarizer.diarize_with_alignments(
-            input_media, caption.alignments, diarization=caption.speaker_diarization
+            input_media,
+            caption.alignments,
+            diarization=caption.speaker_diarization,
+            alignment_fn=self.aligner.alignment,
+            transcribe_fn=self.transcriber.transcribe_numpy if self.transcriber else None,
+            separate_fn=self.aligner.separate if self.aligner.worker.separator_ort else None,
+            debug=self.diarizer.config.debug,
+            output_path=output_caption_path,
         )
         caption.alignments = alignments
         caption.speaker_diarization = diarization
@@ -324,105 +344,6 @@ class LattifAI(LattifAIClientMixin, SyncAPIClient):
         if output_caption_path:
             self._write_caption(caption, output_caption_path)
 
-            if self.diarizer.config.debug:
-                # debug
-                from tgt import Interval, IntervalTier, TextGrid, write_to_file
-
-                debug_tg = TextGrid()
-                transcript_tier = IntervalTier(
-                    start_time=0,
-                    end_time=input_media.duration,
-                    name="transcript",
-                    objects=[Interval(sup.start, sup.end, sup.text) for sup in caption.alignments],
-                )
-                debug_tg.add_tier(transcript_tier)
-
-                speaker_tier = IntervalTier(
-                    start_time=0,
-                    end_time=input_media.duration,
-                    name="speaker",
-                    objects=[Interval(sup.start, sup.end, sup.speaker) for sup in caption.alignments],
-                )
-                debug_tg.add_tier(speaker_tier)
-
-                from collections import defaultdict
-
-                spk2intervals = defaultdict(lambda: [])
-                num_multispk = 0
-
-                segments, skipks = [], []
-                for k, supervision in enumerate(caption.alignments):  # TODO: alignments 本身存在 overlap, eg: [event]
-                    # supervision = caption.alignments[k]
-                    if supervision.custom.get("speaker", []):
-                        num_multispk += 1
-                    else:
-                        continue
-
-                    if k in skipks:
-                        continue
-
-                    for speaker in supervision.custom.get("speaker", []):
-                        for name, start_time, end_time in speaker:
-                            spk2intervals[name].append(Interval(start_time, end_time, name))
-
-                    _segments = []
-                    if k > 0:
-                        _segments.append(caption.alignments[k - 1])
-                    _segments.append(supervision)
-                    while k + 1 < len(caption.alignments):
-                        skipks.append(k + 1)
-                        next_sup = caption.alignments[k + 1]
-                        if not next_sup.custom.get("speaker", []):
-                            k += 1
-                            break
-                        _segments.append(next_sup)
-                        k += 1
-
-                    if segments:
-                        if _segments[0].start >= segments[-1][-1].end:
-                            segments.append(_segments)
-                        else:
-                            if _segments[1:]:
-                                segments.append(_segments[1:])
-                            else:
-                                pass
-                    else:
-                        segments.append(_segments)
-
-                print(
-                    f"Number of multi-speaker segments: {num_multispk}/{len(caption.alignments)} segments: {len(segments)}"
-                )
-
-                for speaker, intervals in sorted(spk2intervals.items(), key=lambda x: x[0]):
-                    speaker_tier = IntervalTier(
-                        start_time=0, end_time=input_media.duration, name=speaker, objects=intervals
-                    )
-                    debug_tg.add_tier(speaker_tier)
-
-                for tier in caption.speaker_diarization.tiers:
-                    tier.name = f"Diarization-{tier.name}"
-                    debug_tg.add_tier(tier)
-
-                tier = IntervalTier(
-                    start_time=0,
-                    end_time=input_media.duration,
-                    name="resegment",
-                    objects=[
-                        Interval(round(sup.start, 2), round(sup.end, 2), sup.text)
-                        for _segments in segments
-                        for sup in _segments
-                    ],
-                )
-                debug_tg.add_tier(tier)
-
-                # if caption.audio_events:
-                #     for tier in caption.audio_events.tiers:
-                #         # tier.name = f"{tier.name}"
-                #         debug_tg.add_tier(tier)
-
-                debug_tgt_file = Path(str(output_caption_path)).with_suffix(".DiarizationDebug.TextGrid")
-                write_to_file(debug_tg, debug_tgt_file, format="long")
-
         return caption
 
     def youtube(

lattifai/config/alignment.py

@@ -21,6 +21,9 @@ class AlignmentConfig:
     model_name: str = "LattifAI/Lattice-1"
     """Model identifier or path to local model directory (e.g., 'LattifAI/Lattice-1')."""
 
+    model_hub: Literal["huggingface", "modelscope"] = "huggingface"
+    """Which model hub to use when resolving remote model names: 'huggingface' or 'modelscope'."""
+
     device: Literal["cpu", "cuda", "mps", "auto"] = "auto"
     """Computation device: 'cpu' for CPU, 'cuda' for NVIDIA GPU, 'mps' for Apple Silicon."""
 
@@ -79,6 +82,17 @@ class AlignmentConfig:
     Default: 10000. Typical range: 1000-20000.
     """
 
+    # Alignment timing configuration
+    start_margin: float = 0.08
+    """Maximum start time margin (in seconds) to extend segment boundaries at the beginning.
+    Default: 0.08. Typical range: 0.0-0.5.
+    """
+
+    end_margin: float = 0.20
+    """Maximum end time margin (in seconds) to extend segment boundaries at the end.
+    Default: 0.20. Typical range: 0.0-0.5.
+    """
+
     client_wrapper: Optional["SyncAPIClient"] = field(default=None, repr=False)
     """Reference to the SyncAPIClient instance. Auto-set during client initialization."""
 

lattifai/config/client.py

@@ -26,6 +26,11 @@ class ClientConfig:
     default_headers: Optional[Dict[str, str]] = field(default=None)
     """Optional static headers to include in all requests."""
 
+    profile: bool = False
+    """Enable profiling of client operations tasks.
+    When True, prints detailed timing information for various stages of the process.
+    """
+
     def __post_init__(self):
         """Validate and auto-populate configuration after initialization."""
 

lattifai/config/transcription.py

@@ -12,6 +12,7 @@ if TYPE_CHECKING:
 SUPPORTED_TRANSCRIPTION_MODELS = Literal[
     "gemini-2.5-pro",
     "gemini-3-pro-preview",
+    "gemini-3-flash-preview",
     "nvidia/parakeet-tdt-0.6b-v3",
     "nvidia/canary-1b-v2",
     "iic/SenseVoiceSmall",
@@ -50,6 +51,9 @@ class TranscriptionConfig:
     lattice_model_path: Optional[str] = None
     """Path to local LattifAI model. Will be auto-set in LattifAI client."""
 
+    model_hub: Literal["huggingface", "modelscope"] = "huggingface"
+    """Which model hub to use when resolving lattice models for transcription."""
+
     client_wrapper: Optional["SyncAPIClient"] = field(default=None, repr=False)
     """Reference to the SyncAPIClient instance. Auto-set during client initialization."""