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

lattifai/alignment/lattice1_aligner.py

@@ -35,7 +35,8 @@ class Lattice1Aligner(object):
             raise ValueError("AlignmentConfig.client_wrapper is not set. It must be initialized by the client.")
 
         client_wrapper = config.client_wrapper
-        model_path = _resolve_model_path(config.model_name)
+        # Resolve model path using configured model hub
+        model_path = _resolve_model_path(config.model_name, getattr(config, "model_hub", "huggingface"))
 
         self.tokenizer = _load_tokenizer(client_wrapper, model_path, config.model_name, config.device)
         self.worker = _load_worker(model_path, config.device, config)
@@ -53,6 +54,29 @@ class Lattice1Aligner(object):
         """
         return self.worker.emission(ndarray)
 
+    def separate(self, audio: np.ndarray) -> np.ndarray:
+        """Separate audio using separator model.
+
+        Args:
+            audio: np.ndarray object containing the audio to separate, shape (1, T)
+
+        Returns:
+            Separated audio as numpy array
+
+        Raises:
+            RuntimeError: If separator model is not available
+        """
+        if self.worker.separator_ort is None:
+            raise RuntimeError("Separator model not available. separator.onnx not found in model path.")
+
+        # Run separator model
+        separator_output = self.worker.separator_ort.run(
+            None,
+            {"audio": audio},
+        )
+
+        return separator_output[0]
+
     def alignment(
         self,
         audio: AudioData,
@@ -120,7 +144,12 @@ class Lattice1Aligner(object):
                 safe_print(colorful.cyan("🎯 Step 4: Decoding lattice results to aligned segments"))
             try:
                 alignments = self.tokenizer.detokenize(
-                    lattice_id, lattice_results, supervisions=supervisions, return_details=return_details
+                    lattice_id,
+                    lattice_results,
+                    supervisions=supervisions,
+                    return_details=return_details,
+                    start_margin=self.config.start_margin,
+                    end_margin=self.config.end_margin,
                 )
                 if verbose:
                     safe_print(colorful.green(f"         ✓ Successfully aligned {len(alignments)} segments"))

lattifai/alignment/lattice1_worker.py

@@ -1,6 +1,7 @@
 import json
 import time
 from collections import defaultdict
+from pathlib import Path
 from typing import Any, Dict, Optional, Tuple
 
 import numpy as np
@@ -73,6 +74,19 @@ class Lattice1Worker:
         else:
             self.extractor = None  # ONNX model includes feature extractor
 
+        # Initialize separator if available
+        separator_model_path = Path(model_path) / "separator.onnx"
+        if separator_model_path.exists():
+            try:
+                self.separator_ort = ort.InferenceSession(
+                    str(separator_model_path),
+                    providers=providers + ["CPUExecutionProvider"],
+                )
+            except Exception as e:
+                raise ModelLoadError(f"separator model from {model_path}", original_error=e)
+        else:
+            self.separator_ort = None
+
         self.device = torch.device(device)
         self.timings = defaultdict(lambda: 0.0)
 

lattifai/alignment/tokenizer.py

@@ -214,7 +214,7 @@ class LatticeTokenizer:
             else:
                 with open(words_model_path, "rb") as f:
                     data = pickle.load(f)
-        except pickle.UnpicklingError as e:
+        except Exception as e:
             del e
             import msgpack
 
@@ -434,6 +434,8 @@ class LatticeTokenizer:
         lattice_results: Tuple[torch.Tensor, Any, Any, float, float],
         supervisions: List[Supervision],
         return_details: bool = False,
+        start_margin: float = 0.08,
+        end_margin: float = 0.20,
     ) -> List[Supervision]:
         emission, results, labels, frame_shift, offset, channel = lattice_results  # noqa: F841
         response = self.client_wrapper.post(
@@ -448,6 +450,8 @@ class LatticeTokenizer:
                 "channel": channel,
                 "return_details": False if return_details is None else return_details,
                 "destroy_lattice": True,
+                "start_margin": start_margin,
+                "end_margin": end_margin,
             },
         )
         if response.status_code == 400:
@@ -538,12 +542,9 @@ def _load_tokenizer(
     tokenizer_cls: Type[LatticeTokenizer] = LatticeTokenizer,
 ) -> LatticeTokenizer:
     """Instantiate tokenizer with consistent error handling."""
-    try:
-        return tokenizer_cls.from_pretrained(
-            client_wrapper=client_wrapper,
-            model_path=model_path,
-            model_name=model_name,
-            device=device,
-        )
-    except Exception as e:
-        raise ModelLoadError(f"tokenizer from {model_path}", original_error=e)
+    return tokenizer_cls.from_pretrained(
+        client_wrapper=client_wrapper,
+        model_path=model_path,
+        model_name=model_name,
+        device=device,
+    )

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,108 @@
+"""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 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,
+    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,
+        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/client.py

@@ -106,7 +106,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 +266,9 @@ 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
         except (CaptionProcessingError, LatticeEncodingError, AlignmentError, LatticeDecodingError):
             # Re-raise our specific errors as-is
             raise
@@ -284,6 +281,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 +323,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 +339,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/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."""