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

lattifai/config/media.py

@@ -52,12 +52,23 @@ class MediaConfig:
     sample_rate: Optional[int] = None
     """Audio sample rate in Hz (e.g., 16000, 44100)."""
 
-    channels: Optional[int] = None
-    """Number of audio channels (1=mono, 2=stereo)."""
-
     channel_selector: Optional[str | int] = "average"
     """Audio channel selection strategy: 'average', 'left', 'right', or channel index."""
 
+    # Audio Streaming Configuration
+    streaming_chunk_secs: Optional[float] = 600.0
+    """Duration in seconds of each audio chunk for streaming mode.
+    When set to a value (e.g., 600.0), enables streaming mode for processing very long audio files (>1 hour).
+    Audio is processed in chunks to keep memory usage low (<4GB peak), suitable for 20+ hour files.
+    When None, disables streaming and loads entire audio into memory.
+    Valid range: 1-1800 seconds (minimum 1 second, maximum 30 minutes).
+    Default: 600 seconds (10 minutes).
+    Recommended: Use 60 seconds or larger for optimal performance.
+    - Smaller chunks: Lower memory usage, more frequent I/O
+    - Larger chunks: Better alignment context, higher memory usage
+    Note: Streaming may add slight processing overhead but enables handling arbitrarily long files.
+    """
+
     # Output / download configuration
     output_dir: Path = field(default_factory=lambda: Path.cwd())
     """Directory for output files (default: current working directory)."""
@@ -87,12 +98,21 @@ class MediaConfig:
         self._normalize_media_format()
         self._process_input_path()
         self._process_output_path()
+        self._validate_streaming_config()
 
     def _setup_output_directory(self) -> None:
         """Ensure output directory exists and is valid."""
         resolved_output_dir = self._ensure_dir(self.output_dir)
         self.output_dir = resolved_output_dir
 
+    def _validate_streaming_config(self) -> None:
+        """Validate streaming configuration parameters."""
+        if self.streaming_chunk_secs is not None:
+            if not 1.0 <= self.streaming_chunk_secs <= 1800.0:
+                raise ValueError(
+                    f"streaming_chunk_secs must be between 1 and 1800 seconds (1 second to 30 minutes), got {self.streaming_chunk_secs}. Recommended: 60 seconds or larger."
+                )
+
     def _validate_default_formats(self) -> None:
         """Validate default audio and video formats."""
         self.default_audio_format = self._normalize_format(self.default_audio_format)

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."""
 

lattifai/diarization/lattifai.py

@@ -1,11 +1,12 @@
 """LattifAI speaker diarization implementation."""
 
 import logging
-from collections import defaultdict
-from typing import List, Optional, Tuple
+from dataclasses import dataclass
+from typing import Callable, List, Optional, Tuple
 
-import torch
-from tgt import Interval, IntervalTier, TextGrid
+import numpy as np
+from lattifai_core.diarization import DiarizationOutput
+from tgt import TextGrid
 
 from lattifai.audio2 import AudioData
 from lattifai.caption import Supervision
@@ -60,7 +61,7 @@ class LattifAIDiarizer:
         num_speakers: Optional[int] = None,
         min_speakers: Optional[int] = None,
         max_speakers: Optional[int] = None,
-    ) -> TextGrid:
+    ) -> DiarizationOutput:
         """Perform speaker diarization on the input audio."""
         return self.diarizer.diarize(
             input_media,
@@ -73,11 +74,16 @@ class LattifAIDiarizer:
         self,
         input_media: AudioData,
         alignments: List[Supervision],
-        diarization: Optional[TextGrid] = None,
+        diarization: Optional[DiarizationOutput] = None,
         num_speakers: Optional[int] = None,
         min_speakers: Optional[int] = None,
         max_speakers: Optional[int] = None,
-    ) -> Tuple[TextGrid, List[Supervision]]:
+        alignment_fn: Optional[Callable] = None,
+        transcribe_fn: Optional[Callable] = None,
+        separate_fn: Optional[Callable] = None,
+        debug: bool = False,
+        output_path: Optional[str] = None,
+    ) -> Tuple[DiarizationOutput, List[Supervision]]:
         """Diarize the given media input and return alignments with refined speaker labels."""
         return self.diarizer.diarize_with_alignments(
             input_media,
@@ -86,4 +92,9 @@ class LattifAIDiarizer:
             num_speakers=num_speakers,
             min_speakers=min_speakers,
             max_speakers=max_speakers,
+            alignment_fn=alignment_fn,
+            transcribe_fn=transcribe_fn,
+            separate_fn=separate_fn,
+            debug=debug,
+            output_path=output_path,
         )

lattifai/errors.py

@@ -11,11 +11,15 @@ LATTICE_DECODING_FAILURE_HELP = (
     "1) Media(Audio/Video) and text content mismatch:\n"
     "   - The transcript/caption does not accurately match the media content\n"
     "   - Text may be from a different version or section of the media\n"
-    "   ⚠️  Note: Gemini transcription may occasionally skip large segments of media, causing alignment failures.\n"
-    "       We will detect and fix this issue in the next version.\n\n"
-    "2) Unsupported media type:\n"
+    "2) Text formatting issues:\n"
+    "   - Special characters, HTML entities, or unusual punctuation may cause alignment failures\n"
+    "   - Text normalization is enabled by default (caption.normalize_text=True)\n"
+    "     If you disabled it, try re-enabling: caption.normalize_text=True\n"
+    "3) Unsupported media type:\n"
     "   - Singing is not yet supported, this will be optimized in future versions\n\n"
     "💡 Troubleshooting tips:\n"
+    "   • Text normalization is enabled by default to handle special characters\n"
+    "     (no action needed unless you explicitly set caption.normalize_text=False)\n"
     "   • Verify the transcript matches the media by listening to a few segments\n"
     "   • For YouTube videos, manually check if auto-generated transcript are accurate\n"
     "       • Consider using a different transcription source if Gemini results are incomplete"

lattifai/mixin.py

@@ -10,6 +10,7 @@ from lhotse.utils import Pathlike
 from lattifai.audio2 import AudioData
 from lattifai.caption import Caption
 from lattifai.errors import CaptionProcessingError
+from lattifai.utils import safe_print
 
 if TYPE_CHECKING:
     from .config import AlignmentConfig, CaptionConfig, ClientConfig, DiarizationConfig, TranscriptionConfig
@@ -183,7 +184,9 @@ class LattifAIClientMixin:
         from lattifai.utils import _resolve_model_path
 
         if transcription_config is not None:
-            transcription_config.lattice_model_path = _resolve_model_path(alignment_config.model_name)
+            transcription_config.lattice_model_path = _resolve_model_path(
+                alignment_config.model_name, getattr(alignment_config, "model_hub", "huggingface")
+            )
 
         # Set client_wrapper for all configs
         alignment_config.client_wrapper = self
@@ -278,7 +281,7 @@ class LattifAIClientMixin:
 
         try:
             if verbose:
-                print(colorful.cyan(f"📖 Step 1: Reading caption file from {input_caption}"))
+                safe_print(colorful.cyan(f"📖 Step 1: Reading caption file from {input_caption}"))
             caption = Caption.read(
                 input_caption,
                 format=input_caption_format,
@@ -287,18 +290,18 @@ class LattifAIClientMixin:
             diarization_file = Path(str(input_caption)).with_suffix(".SpkDiar")
             if diarization_file.exists():
                 if verbose:
-                    print(colorful.cyan(f"📖 Step 1b: Reading speaker diarization from {diarization_file}"))
+                    safe_print(colorful.cyan(f"📖 Step 1b: Reading speaker diarization from {diarization_file}"))
                 caption.read_speaker_diarization(diarization_file)
             events_file = Path(str(input_caption)).with_suffix(".AED")
             if events_file.exists():
                 if verbose:
-                    print(colorful.cyan(f"📖 Step 1c: Reading audio events from {events_file}"))
+                    safe_print(colorful.cyan(f"📖 Step 1c: Reading audio events from {events_file}"))
                 from tgt import read_textgrid
 
                 caption.audio_events = read_textgrid(events_file)
 
             if verbose:
-                print(colorful.green(f"         ✓ Parsed {len(caption)} caption segments"))
+                safe_print(colorful.green(f"         ✓ Parsed {len(caption)} caption segments"))
             return caption
         except Exception as e:
             raise CaptionProcessingError(
@@ -332,10 +335,10 @@ class LattifAIClientMixin:
             )
             diarization_file = Path(str(output_caption_path)).with_suffix(".SpkDiar")
             if not diarization_file.exists() and caption.speaker_diarization:
-                print(colorful.green(f"    Writing speaker diarization to: {diarization_file}"))
+                safe_print(colorful.green(f"    Writing speaker diarization to: {diarization_file}"))
                 caption.write_speaker_diarization(diarization_file)
 
-            print(colorful.green(f"🎉🎉🎉🎉🎉 Caption file written to: {output_caption_path}"))
+            safe_print(colorful.green(f"🎉🎉🎉🎉🎉 Caption file written to: {output_caption_path}"))
             return result
         except Exception as e:
             raise CaptionProcessingError(
@@ -352,14 +355,14 @@ class LattifAIClientMixin:
         force_overwrite: bool,
     ) -> str:
         """Download media from YouTube (async implementation)."""
-        print(colorful.cyan("📥 Downloading media from YouTube..."))
+        safe_print(colorful.cyan("📥 Downloading media from YouTube..."))
         media_file = await self.downloader.download_media(
             url=url,
             output_dir=str(output_dir),
             media_format=media_format,
             force_overwrite=force_overwrite,
         )
-        print(colorful.green(f"    ✓ Media downloaded: {media_file}"))
+        safe_print(colorful.green(f"    ✓ Media downloaded: {media_file}"))
         return media_file
 
     def _download_media_sync(
@@ -379,6 +382,7 @@ class LattifAIClientMixin:
         media_file: Union[str, Path, AudioData],
         source_lang: Optional[str],
         is_async: bool = False,
+        output_dir: Optional[Path] = None,
     ) -> Caption:
         """
         Get captions by downloading or transcribing.
@@ -400,14 +404,23 @@ class LattifAIClientMixin:
             # Transcription mode: use Transcriber to transcribe
             self._validate_transcription_setup()
 
-            print(colorful.cyan(f"🎤 Transcribing({self.transcriber.name}) media: {str(media_file)} ..."))
+            safe_print(colorful.cyan(f"🎤 Transcribing({self.transcriber.name}) media: {str(media_file)} ..."))
             transcription = await self.transcriber.transcribe_file(media_file, language=source_lang)
-            print(colorful.green("         ✓ Transcription completed."))
+            safe_print(colorful.green("         ✓ Transcription completed."))
 
             if "gemini" in self.transcriber.name.lower():
+                safe_print(colorful.yellow("🔍 Gemini raw output:"))
+                safe_print(colorful.yellow(f"{transcription[:1000]}..."))  # Print first 1000 chars
+
                 # write to temp file and use Caption read
-                with tempfile.NamedTemporaryFile(suffix=self.transcriber.file_suffix, delete=True) as tmp_file:
-                    tmp_path = Path(tmp_file.name)
+                # On Windows, we need to close the file before writing to it
+                tmp_file = tempfile.NamedTemporaryFile(
+                    suffix=self.transcriber.file_suffix, delete=False, mode="w", encoding="utf-8"
+                )
+                tmp_path = Path(tmp_file.name)
+                tmp_file.close()  # Close file before writing
+
+                try:
                     await asyncio.to_thread(
                         self.transcriber.write,
                         transcription,
@@ -417,6 +430,22 @@ class LattifAIClientMixin:
                     transcription = self._read_caption(
                         tmp_path, input_caption_format="gemini", normalize_text=False, verbose=False
                     )
+                finally:
+                    # Clean up temp file
+                    if tmp_path.exists():
+                        tmp_path.unlink()
+            else:
+                safe_print(colorful.yellow(f"🔍 {self.transcriber.name} raw output:"))
+                if isinstance(transcription, Caption):
+                    safe_print(colorful.yellow(f"Caption with {len(transcription.transcription)} segments"))
+                    if transcription.transcription:
+                        safe_print(colorful.yellow(f"First segment: {transcription.transcription[0].text}"))
+
+            if output_dir:
+                # Generate transcript file path
+                transcript_file = output_dir / f"{Path(str(media_file)).stem}_{self.transcriber.file_name}"
+                await asyncio.to_thread(self.transcriber.write, transcription, transcript_file, encoding="utf-8")
+                safe_print(colorful.green(f"         ✓ Transcription saved to: {transcript_file}"))
 
             return transcription
 
@@ -459,7 +488,7 @@ class LattifAIClientMixin:
             if self.caption_config.input_path:
                 caption_path = Path(self.caption_config.input_path)
                 if caption_path.exists():
-                    print(colorful.green(f"📄 Using provided caption file: {caption_path}"))
+                    safe_print(colorful.green(f"📄 Using provided caption file: {caption_path}"))
                     return str(caption_path)
                 else:
                     raise FileNotFoundError(f"Provided caption path does not exist: {caption_path}")
@@ -496,7 +525,7 @@ class LattifAIClientMixin:
 
                     # elif choice == "overwrite": continue to transcribe below
 
-                print(colorful.cyan(f"🎤 Transcribing media with {transcriber_name}..."))
+                safe_print(colorful.cyan(f"🎤 Transcribing media with {transcriber_name}..."))
                 if self.transcriber.supports_url:
                     transcription = await self.transcriber.transcribe(url, language=source_lang)
                 else:
@@ -508,7 +537,7 @@ class LattifAIClientMixin:
                     caption_file = transcription
                 else:
                     caption_file = str(transcript_file)
-                print(colorful.green(f"         ✓ Transcription completed: {caption_file}"))
+                safe_print(colorful.green(f"         ✓ Transcription completed: {caption_file}"))
             else:
                 # Download YouTube captions
                 caption_file = await self.downloader.download_captions(

lattifai/server/app.py

@@ -232,7 +232,7 @@ async def align_files(
     normalize_text: bool = Form(False),
     output_format: str = Form("srt"),
     transcription_model: str = Form("nvidia/parakeet-tdt-0.6b-v3"),
-    alignment_model: str = Form("Lattifai/Lattice-1"),
+    alignment_model: str = Form("LattifAI/Lattice-1"),
 ):
     # Check if LATTIFAI_API_KEY is set
     if not os.environ.get("LATTIFAI_API_KEY"):
@@ -423,4 +423,5 @@ def process_alignment(
             input_caption=str(caption_path) if caption_path else None,
             output_caption_path=str(output_caption_path) if output_caption_path else None,
             split_sentence=split_sentence,
+            streaming_chunk_secs=None,  # Server API default: no streaming
         )

lattifai/transcription/__init__.py

@@ -70,7 +70,7 @@ def create_transcriber(
         raise ValueError(
             f"Cannot determine transcriber for model_name='{transcription_config.model_name}'. "
             f"Supported patterns: \n"
-            f"  - Gemini API models: 'gemini-2.5-pro', 'gemini-3-pro-preview'\n"
+            f"  - Gemini API models: 'gemini-2.5-pro', 'gemini-3-pro-preview', 'gemini-3-flash-preview'\n"
             f"  - Local HF models: 'nvidia/parakeet-*', 'iic/SenseVoiceSmall', etc.\n"
             f"Please specify a valid model_name."
         )

lattifai/transcription/base.py

@@ -2,10 +2,12 @@
 
 from abc import ABC, abstractmethod
 from pathlib import Path
-from typing import Optional, Union
+from typing import List, Optional, Union
+
+import numpy as np
 
 from lattifai.audio2 import AudioData
-from lattifai.caption import Caption
+from lattifai.caption import Caption, Supervision
 from lattifai.config import TranscriptionConfig
 from lattifai.logging import get_logger
 
@@ -96,6 +98,23 @@ class BaseTranscriber(ABC):
             language: Optional language code for transcription.
         """
 
+    @abstractmethod
+    def transcribe_numpy(
+        self,
+        audio: Union[np.ndarray, List[np.ndarray]],
+        language: Optional[str] = None,
+    ) -> Union[Supervision, List[Supervision]]:
+        """
+        Transcribe audio from a numpy array and return Supervision.
+
+        Args:
+            audio_array: Audio data as numpy array (shape: [samples]).
+            language: Optional language code for transcription.
+
+        Returns:
+            Supervision object with transcription info.
+        """
+
     @abstractmethod
     def write(self, transcript: Union[str, Caption], output_file: Path, encoding: str = "utf-8") -> Path:
         """

lattifai/transcription/gemini.py

@@ -2,12 +2,14 @@
 
 import asyncio
 from pathlib import Path
-from typing import Optional, Union
+from typing import List, Optional, Union
 
+import numpy as np
 from google import genai
 from google.genai.types import GenerateContentConfig, Part, ThinkingConfig
 
 from lattifai.audio2 import AudioData
+from lattifai.caption import Supervision
 from lattifai.config import TranscriptionConfig
 from lattifai.transcription.base import BaseTranscriber
 from lattifai.transcription.prompts import get_prompt_loader
@@ -118,6 +120,130 @@ class GeminiTranscriber(BaseTranscriber):
             self.logger.error(f"Gemini transcription failed: {str(e)}")
             raise RuntimeError(f"Gemini transcription failed: {str(e)}")
 
+    def transcribe_numpy(
+        self,
+        audio: Union[np.ndarray, List[np.ndarray]],
+        language: Optional[str] = None,
+    ) -> Union[Supervision, List[Supervision]]:
+        """
+        Transcribe audio from a numpy array (or list of arrays) and return Supervision.
+
+        Note: Gemini API does not support word-level alignment. The returned
+        Supervision will contain only the full transcription text without alignment.
+
+        Args:
+            audio: Audio data as numpy array (shape: [samples]),
+                   or a list of such arrays for batch processing.
+            language: Optional language code for transcription.
+
+        Returns:
+            Supervision object (or list of Supervision objects) with transcription text (no alignment).
+
+        Raises:
+            ValueError: If API key not provided
+            RuntimeError: If transcription fails
+        """
+        # Handle batch processing
+        if isinstance(audio, list):
+            return [self.transcribe_numpy(arr, language=language) for arr in audio]
+
+        audio_array = audio
+        # Use default sample rate of 16000 Hz
+        sample_rate = 16000
+
+        if self.config.verbose:
+            self.logger.info(f"🎤 Starting Gemini transcription for numpy array (sample_rate={sample_rate})")
+
+        # Ensure audio is in the correct shape
+        if audio_array.ndim == 1:
+            audio_array = audio_array.reshape(1, -1)
+        elif audio_array.ndim > 2:
+            raise ValueError(f"Audio array must be 1D or 2D, got shape {audio_array.shape}")
+
+        # Save numpy array to temporary file
+        import tempfile
+
+        import soundfile as sf
+
+        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp_file:
+            # Transpose to (samples, channels) for soundfile
+            sf.write(tmp_file.name, audio_array.T, sample_rate)
+            tmp_path = Path(tmp_file.name)
+
+        try:
+            # Transcribe using simple ASR prompt
+            import asyncio
+
+            transcript = asyncio.run(self._transcribe_with_simple_prompt(tmp_path, language=language))
+
+            # Create Supervision object from transcript
+            duration = audio_array.shape[-1] / sample_rate
+            supervision = Supervision(
+                id="gemini-transcription",
+                recording_id="numpy-array",
+                start=0.0,
+                duration=duration,
+                text=transcript,
+                speaker=None,
+                alignment=None,  # Gemini does not provide word-level alignment
+            )
+
+            return supervision
+
+        finally:
+            # Clean up temporary file
+            if tmp_path.exists():
+                tmp_path.unlink()
+
+    async def _transcribe_with_simple_prompt(self, media_file: Path, language: Optional[str] = None) -> str:
+        """
+        Transcribe audio using a simple ASR prompt instead of complex instructions.
+
+        Args:
+            media_file: Path to audio file
+            language: Optional language code
+
+        Returns:
+            Transcribed text
+        """
+        client = self._get_client()
+
+        # Upload audio file
+        if self.config.verbose:
+            self.logger.info("📤 Uploading audio file to Gemini...")
+        uploaded_file = client.files.upload(file=str(media_file))
+
+        # Simple ASR prompt
+        system_prompt = "Transcribe the audio."
+        if language:
+            system_prompt = f"Transcribe the audio in {language}."
+
+        # Create simple generation config
+        simple_config = GenerateContentConfig(
+            system_instruction=system_prompt,
+            response_modalities=["TEXT"],
+        )
+
+        contents = Part.from_uri(file_uri=uploaded_file.uri, mime_type=uploaded_file.mime_type)
+        response = await asyncio.get_event_loop().run_in_executor(
+            None,
+            lambda: client.models.generate_content(
+                model=self.config.model_name,
+                contents=contents,
+                config=simple_config,
+            ),
+        )
+
+        if not response.text:
+            raise RuntimeError("Empty response from Gemini API")
+
+        transcript = response.text.strip()
+
+        if self.config.verbose:
+            self.logger.info(f"✅ Transcription completed: {len(transcript)} characters")
+
+        return transcript
+
     def _get_transcription_prompt(self) -> str:
         """Get (and cache) transcription system prompt from prompts module."""
         if self._system_prompt is not None:

lattifai/transcription/lattifai.py

@@ -1,10 +1,12 @@
 """Transcription module with config-driven architecture."""
 
 from pathlib import Path
-from typing import Optional, Union
+from typing import List, Optional, Union
+
+import numpy as np
 
 from lattifai.audio2 import AudioData
-from lattifai.caption import Caption
+from lattifai.caption import Caption, Supervision
 from lattifai.config import TranscriptionConfig
 from lattifai.transcription.base import BaseTranscriber
 from lattifai.transcription.prompts import get_prompt_loader  # noqa: F401
@@ -74,6 +76,32 @@ class LattifAITranscriber(BaseTranscriber):
 
         return caption
 
+    def transcribe_numpy(
+        self,
+        audio: Union[np.ndarray, List[np.ndarray]],
+        language: Optional[str] = None,
+    ) -> Union[Supervision, List[Supervision]]:
+        """
+        Transcribe audio from a numpy array (or list of arrays) and return Supervision.
+
+        Args:
+            audio: Audio data as numpy array (shape: [samples]),
+                   or a list of such arrays for batch processing.
+            language: Optional language code for transcription.
+
+        Returns:
+            Supervision object (or list of Supervision objects) with transcription and alignment info.
+        """
+        if self._transcriber is None:
+            from lattifai_core.transcription import LattifAITranscriber as CoreLattifAITranscriber
+
+            self._transcriber = CoreLattifAITranscriber.from_pretrained(model_config=self.config)
+
+        # Delegate to core transcriber which handles both single arrays and lists
+        return self._transcriber.transcribe(
+            audio, language=language, return_hypotheses=True, progress_bar=False, timestamps=True
+        )[0]
+
     def write(
         self, transcript: Caption, output_file: Path, encoding: str = "utf-8", cache_audio_events: bool = True
     ) -> Path: