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

lattifai/audio2.py

@@ -162,168 +162,132 @@ class AudioLoader:
         sampling_rate: int,
         channel_selector: Optional[ChannelSelectorType],
     ) -> np.ndarray:
-        """Load audio from file or binary stream and resample to target rate.
+        """Load audio from file or binary stream and resample to target rate."""
+        audio_source: Union[str, BinaryIO] = audio
+        audio_path: Optional[Path] = None
 
-        Args:
-            audio: Path to audio file or binary stream.
-            sampling_rate: Target sampling rate.
-            channel_selector: How to select channels.
-
-        Returns:
-            Resampled audio as a NumPy array of shape (channels, samples).
-
-        Raises:
-            ImportError: If PyAV is needed but not installed.
-            ValueError: If no audio stream found.
-            RuntimeError: If audio loading fails.
-        """
         if isinstance(audio, Pathlike):
-            audio = str(Path(str(audio)).expanduser())
+            audio_path = Path(str(audio)).expanduser()
+            audio_source = str(audio_path)
+
+        if audio_path and audio_path.suffix.lower() in [".mp4", ".m4a", ".aac", ".mov", ".webm", ".avi", ".mkv"]:
+            return self._load_audio_with_av(audio_source, sampling_rate, channel_selector)
 
-        # load audio in chunks to reduce memory footprint for long files
         try:
-            # First check file duration to decide loading strategy
-            info = sf.info(audio)
-            duration = info.duration
+            return self._load_audio_with_soundfile(audio_source, sampling_rate, channel_selector)
+        except Exception as primary_error:
+            print(f"Primary error with soundfile: {primary_error}")
+            return self._load_audio_with_av(audio_source, sampling_rate, channel_selector, primary_error)
 
-            # For very long audio (>60 minutes), use chunk-based loading
-            if duration > 3600:  # 60 minutes
-                with sf.SoundFile(audio, "r") as f:
-                    sample_rate = f.samplerate
-                    total_frames = f.frames
+    def _load_audio_with_soundfile(
+        self,
+        audio: Union[str, BinaryIO],
+        sampling_rate: int,
+        channel_selector: Optional[ChannelSelectorType],
+    ) -> np.ndarray:
+        """Load audio via soundfile with chunking support for long inputs."""
+        info = sf.info(audio)
+        duration = info.duration
 
-                    # Pre-calculate output size to avoid list accumulation
-                    num_channels = 1 if channel_selector else f.channels
-                    expected_output_samples = int(total_frames * sampling_rate / sample_rate)
+        if duration > 3600:
+            with sf.SoundFile(audio, "r") as f:
+                sample_rate = f.samplerate
+                total_frames = f.frames
 
-                    # Pre-allocate output array
-                    waveform = np.zeros((num_channels, expected_output_samples), dtype=np.float32)
+                num_channels = 1 if channel_selector else f.channels
+                expected_output_samples = int(total_frames * sampling_rate / sample_rate)
+                waveform = np.zeros((num_channels, expected_output_samples), dtype=np.float32)
 
-                    # Use source sample rate for reading, not target
-                    chunk_frames = int(sample_rate * 1800)  # 30-minute chunks at source rate
-                    output_offset = 0
+                chunk_frames = int(sample_rate * 1800)
+                output_offset = 0
 
-                    while True:
-                        chunk = f.read(frames=chunk_frames, dtype="float32", always_2d=True)
-                        if chunk.size == 0:
-                            break
+                while True:
+                    chunk = f.read(frames=chunk_frames, dtype="float32", always_2d=True)
+                    if chunk.size == 0:
+                        break
 
-                        # Resample chunk -> (channels, samples)
-                        resampled_chunk = self._resample_audio(
-                            (chunk, sample_rate),
-                            sampling_rate,
-                            device=self.device,
-                            channel_selector=channel_selector,
-                        )
+                    resampled_chunk = self._resample_audio(
+                        (chunk, sample_rate),
+                        sampling_rate,
+                        device=self.device,
+                        channel_selector=channel_selector,
+                    )
 
-                        # Write directly to pre-allocated array
-                        chunk_length = resampled_chunk.shape[-1]
-                        waveform[..., output_offset : output_offset + chunk_length] = resampled_chunk
-                        output_offset += chunk_length
+                    chunk_length = resampled_chunk.shape[-1]
+                    waveform[..., output_offset : output_offset + chunk_length] = resampled_chunk
+                    output_offset += chunk_length
 
-                        # Clean up immediately
-                        del chunk, resampled_chunk
+                    del chunk, resampled_chunk
+
+                if output_offset < expected_output_samples:
+                    waveform = waveform[..., :output_offset]
 
-                    # Trim to actual size if needed (due to rounding in resampling)
-                    if output_offset < expected_output_samples:
-                        waveform = waveform[..., :output_offset]
+            return waveform
 
-                return waveform
+        waveform, sample_rate = sf.read(audio, always_2d=True, dtype="float32")
+        result = self._resample_audio(
+            (waveform, sample_rate),
+            sampling_rate,
+            device=self.device,
+            channel_selector=channel_selector,
+        )
+        del waveform
+        return result
+
+    def _load_audio_with_av(
+        self,
+        audio: Union[str, BinaryIO],
+        sampling_rate: int,
+        channel_selector: Optional[ChannelSelectorType],
+        primary_error: Optional[Exception] = None,
+    ) -> np.ndarray:
+        """Load audio via PyAV when soundfile is unavailable or unsuitable."""
+        try:
+            import av
+        except ImportError as exc:  # pragma: no cover
+            message = "PyAV (av) is required for loading certain audio formats. Install it with: pip install av"
+            if primary_error:
+                message = f"{message}\nPrimary error was: {primary_error}"
+            raise AudioLoadError(message) from exc
+
+        try:
+            container = av.open(audio)
+            audio_stream = next((s for s in container.streams if s.type == "audio"), None)
+
+            if audio_stream is None:
+                raise ValueError(f"No audio stream found in file: {audio}")
+
+            audio_stream.codec_context.format = av.AudioFormat("flt")
+            sample_rate = audio_stream.codec_context.sample_rate
+
+            duration_estimate = None
+            if audio_stream.duration and audio_stream.time_base:
+                duration_estimate = float(audio_stream.duration * audio_stream.time_base)
             else:
-                # For shorter audio, use standard loading
-                waveform, sample_rate = sf.read(audio, always_2d=True, dtype="float32")
-                # Resample and return directly to avoid double processing
-                result = self._resample_audio(
-                    (waveform, sample_rate),
-                    sampling_rate,
-                    device=self.device,
-                    channel_selector=channel_selector,
-                )
-                del waveform
-                return result
-        except Exception as primary_error:
-            print(f"Primary error with soundfile: {primary_error}")
-            # Fallback to PyAV for formats not supported by soundfile
-            try:
-                import av
-            except ImportError:
-                raise AudioLoadError(
-                    "PyAV (av) is required for loading certain audio formats. "
-                    f"Install it with: pip install av\n"
-                    f"Primary error was: {primary_error}"
-                )
-
-            try:
-                container = av.open(audio)
-                audio_stream = next((s for s in container.streams if s.type == "audio"), None)
-
-                if audio_stream is None:
-                    raise ValueError(f"No audio stream found in file: {audio}")
-
-                audio_stream.codec_context.format = av.AudioFormat("flt")  # 32-bit float
-                sample_rate = audio_stream.codec_context.sample_rate
-
-                # Estimate duration to decide processing strategy
-                duration_estimate = None
-                if audio_stream.duration and audio_stream.time_base:
-                    duration_estimate = float(audio_stream.duration * audio_stream.time_base)
-                else:
-                    print(f"WARNING: Failed to estimate duration for audio: {audio}")
-
-                # For very long audio (>30 minutes), process and resample in chunks
-                if duration_estimate and duration_estimate > 1800:
-                    # Estimate output size and pre-allocate with buffer
-                    num_channels = 1 if channel_selector else audio_stream.codec_context.channels
-                    estimated_samples = int(duration_estimate * sampling_rate * 1.1)  # 10% buffer
-                    waveform = np.zeros((num_channels, estimated_samples), dtype=np.float32)
-
-                    frames = []
-                    accumulated_samples = 0
-                    output_offset = 0
-                    chunk_sample_target = int(sample_rate * 600)  # 10 minutes at original rate
-
-                    for frame in container.decode(audio_stream):
-                        array = frame.to_ndarray()
-
-                        # Ensure shape is (samples, channels)
-                        if array.ndim == 1:
-                            array = array.reshape(-1, 1)
-                        elif array.ndim == 2 and array.shape[0] < array.shape[1]:
-                            array = array.T
-
-                        frames.append(array)
-                        accumulated_samples += array.shape[0]
-
-                        # Process chunk when accumulated enough samples
-                        if accumulated_samples >= chunk_sample_target:
-                            chunk = np.concatenate(frames, axis=0).astype(np.float32)
-                            del frames  # Free frames list before resampling
-                            # Resample chunk -> (channels, samples)
-                            resampled_chunk = self._resample_audio(
-                                (chunk, sample_rate),
-                                sampling_rate,
-                                device=self.device,
-                                channel_selector=channel_selector,
-                            )
-
-                            chunk_length = resampled_chunk.shape[-1]
-                            if output_offset + chunk_length > waveform.shape[-1]:
-                                print(
-                                    f"WARNING: Trimming resampled chunk from {chunk_length} to {waveform.shape[-1] - output_offset} samples to fit waveform buffer for audio: {audio}"  # noqa: E501
-                                )
-                                resampled_chunk = resampled_chunk[:, : waveform.shape[-1] - output_offset]
-
-                            # Write directly to array
-                            waveform[..., output_offset : output_offset + chunk_length] = resampled_chunk
-                            output_offset += chunk_length
-
-                            # Clean up immediately
-                            del chunk, resampled_chunk
-                            frames = []  # Create new list
-                            accumulated_samples = 0
-
-                    # Process remaining frames
-                    if frames:
+                print(f"WARNING: Failed to estimate duration for audio: {audio}")
+
+            if duration_estimate and duration_estimate > 1800:
+                num_channels = 1 if channel_selector else audio_stream.codec_context.channels
+                estimated_samples = int(duration_estimate * sampling_rate * 1.1)
+                waveform = np.zeros((num_channels, estimated_samples), dtype=np.float32)
+
+                frames = []
+                accumulated_samples = 0
+                output_offset = 0
+                chunk_sample_target = int(sample_rate * 600)
+
+                for frame in container.decode(audio_stream):
+                    array = frame.to_ndarray()
+
+                    if array.ndim == 1:
+                        array = array.reshape(-1, 1)
+                    elif array.ndim == 2 and array.shape[0] < array.shape[1]:
+                        array = array.T
+
+                    frames.append(array)
+                    accumulated_samples += array.shape[0]
+
+                    if accumulated_samples >= chunk_sample_target:
                         chunk = np.concatenate(frames, axis=0).astype(np.float32)
                         del frames
                         resampled_chunk = self._resample_audio(
@@ -335,53 +299,68 @@ class AudioLoader:
 
                         chunk_length = resampled_chunk.shape[-1]
                         if output_offset + chunk_length > waveform.shape[-1]:
-                            print(
-                                f"WARNING: Trimming resampled chunk from {chunk_length} to {waveform.shape[-1] - output_offset} samples to fit waveform buffer for audio: {audio}"  # noqa: E501
-                            )
+                            print("WARNING: Trimming resampled chunk to fit waveform buffer for audio: " f"{audio}")
                             resampled_chunk = resampled_chunk[:, : waveform.shape[-1] - output_offset]
 
                         waveform[..., output_offset : output_offset + chunk_length] = resampled_chunk
                         output_offset += chunk_length
-                        del chunk, resampled_chunk
-
-                    container.close()
 
-                    if output_offset == 0:
-                        raise ValueError(f"No audio data found in file: {audio}")
+                        del chunk, resampled_chunk
+                        frames = []
+                        accumulated_samples = 0
 
-                    # Trim to actual size
-                    waveform = waveform[..., :output_offset]
-                    return waveform
-                else:
-                    # For shorter audio, process in batches to reduce memory
-                    frames = []
-                    for frame in container.decode(audio_stream):
-                        array = frame.to_ndarray()
-                        # Ensure shape is (channels, samples)
-                        if array.ndim == 1:
-                            array = array.reshape(-1, 1)
-                        elif array.ndim == 2 and array.shape[0] < array.shape[1]:
-                            array = array.T
-                        frames.append(array)
-                    container.close()
-
-                    if not frames:
-                        raise ValueError(f"No audio data found in file: {audio}")
-
-                    # Concatenate remaining frames
-                    waveform = np.concatenate(frames, axis=0).astype(np.float32)
+                if frames:
+                    chunk = np.concatenate(frames, axis=0).astype(np.float32)
                     del frames
-                    # Resample and return directly
-                    result = self._resample_audio(
-                        (waveform, sample_rate),
+                    resampled_chunk = self._resample_audio(
+                        (chunk, sample_rate),
                         sampling_rate,
                         device=self.device,
                         channel_selector=channel_selector,
                     )
-                    del waveform
-                    return result
-            except Exception as e:
-                raise RuntimeError(f"Failed to load audio file {audio}: {e}")
+
+                    chunk_length = resampled_chunk.shape[-1]
+                    if output_offset + chunk_length > waveform.shape[-1]:
+                        print("WARNING: Trimming resampled chunk to fit waveform buffer for audio: " f"{audio}")
+                        resampled_chunk = resampled_chunk[:, : waveform.shape[-1] - output_offset]
+
+                    waveform[..., output_offset : output_offset + chunk_length] = resampled_chunk
+                    output_offset += chunk_length
+                    del chunk, resampled_chunk
+
+                container.close()
+
+                if output_offset == 0:
+                    raise ValueError(f"No audio data found in file: {audio}")
+
+                waveform = waveform[..., :output_offset]
+                return waveform
+
+            frames = []
+            for frame in container.decode(audio_stream):
+                array = frame.to_ndarray()
+                if array.ndim == 1:
+                    array = array.reshape(-1, 1)
+                elif array.ndim == 2 and array.shape[0] < array.shape[1]:
+                    array = array.T
+                frames.append(array)
+            container.close()
+
+            if not frames:
+                raise ValueError(f"No audio data found in file: {audio}")
+
+            waveform = np.concatenate(frames, axis=0).astype(np.float32)
+            del frames
+            result = self._resample_audio(
+                (waveform, sample_rate),
+                sampling_rate,
+                device=self.device,
+                channel_selector=channel_selector,
+            )
+            del waveform
+            return result
+        except Exception as exc:
+            raise RuntimeError(f"Failed to load audio file {audio}: {exc}")
 
     def __call__(
         self,

lattifai/cli/__init__.py

@@ -4,7 +4,7 @@ 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.caption import convert, diff
 from lattifai.cli.diarization import diarize
 from lattifai.cli.transcribe import transcribe, transcribe_align
 from lattifai.cli.youtube import youtube
@@ -12,6 +12,7 @@ from lattifai.cli.youtube import youtube
 __all__ = [
     "align",
     "convert",
+    "diff",
     "diarize",
     "transcribe",
     "transcribe_align",

lattifai/cli/alignment.py

@@ -1,5 +1,6 @@
 """Alignment CLI entry point with nemo_run."""
 
+import sys
 from typing import Optional
 
 import nemo_run as run
@@ -12,9 +13,11 @@ from lattifai.config import (
     CaptionConfig,
     ClientConfig,
     DiarizationConfig,
+    EventConfig,
     MediaConfig,
     TranscriptionConfig,
 )
+from lattifai.errors import LattifAIError
 
 __all__ = ["align"]
 
@@ -30,6 +33,7 @@ def align(
     alignment: Annotated[Optional[AlignmentConfig], run.Config[AlignmentConfig]] = None,
     transcription: Annotated[Optional[TranscriptionConfig], run.Config[TranscriptionConfig]] = None,
     diarization: Annotated[Optional[DiarizationConfig], run.Config[DiarizationConfig]] = None,
+    event: Annotated[Optional[EventConfig], run.Config[EventConfig]] = None,
 ):
     """
     Align audio/video with caption file.
@@ -121,6 +125,7 @@ def align(
         caption_config=caption_config,
         transcription_config=transcription,
         diarization_config=diarization,
+        event_config=event,
     )
 
     is_url = media_config.input_path.startswith(("http://", "https://"))

lattifai/cli/caption.py

@@ -6,6 +6,7 @@ import nemo_run as run
 from lhotse.utils import Pathlike
 from typing_extensions import Annotated
 
+from lattifai.caption.config import KaraokeConfig
 from lattifai.config import CaptionConfig
 from lattifai.utils import safe_print
 
@@ -16,6 +17,8 @@ def convert(
     output_path: Pathlike,
     include_speaker_in_text: bool = False,
     normalize_text: bool = False,
+    word_level: bool = False,
+    karaoke: bool = False,
 ):
     """
     Convert caption file to another format.
@@ -33,6 +36,11 @@ def convert(
         normalize_text: Whether to normalize caption text during conversion.
             This applies text cleaning such as removing HTML tags, decoding entities,
             collapsing whitespace, and standardizing punctuation.
+        word_level: Use word-level output format if supported.
+            When True without karaoke: outputs word-per-segment (each word as separate segment).
+            JSON format will include a 'words' field with word-level timestamps.
+        karaoke: Enable karaoke styling (requires word_level=True).
+            When True: outputs karaoke format (ASS \\kf tags, enhanced LRC, etc.).
 
     Examples:
         # Basic format conversion (positional arguments)
@@ -41,6 +49,15 @@ def convert(
         # Convert with text normalization
         lai caption convert input.srt output.json normalize_text=true
 
+        # Convert to word-per-segment output (if input has alignment)
+        lai caption convert input.json output.srt word_level=true
+
+        # Convert to karaoke format (ASS with \\kf tags)
+        lai caption convert input.json output.ass word_level=true karaoke=true
+
+        # Export JSON with word-level timestamps
+        lai caption convert input.srt output.json word_level=true
+
         # Mixing positional and keyword arguments
         lai caption convert input.srt output.vtt \\
             include_speaker_in_text=false \\
@@ -51,10 +68,18 @@ def convert(
             input_path=input.srt \\
             output_path=output.TextGrid
     """
-    from lattifai.caption import Caption
+    from lattifai.data import Caption
+
+    # Create karaoke_config if karaoke flag is set
+    karaoke_config = KaraokeConfig(enabled=True) if karaoke else None
 
     caption = Caption.read(input_path, normalize_text=normalize_text)
-    caption.write(output_path, include_speaker_in_text=include_speaker_in_text)
+    caption.write(
+        output_path,
+        include_speaker_in_text=include_speaker_in_text,
+        word_level=word_level,
+        karaoke_config=karaoke_config,
+    )
 
     safe_print(f"✅ Converted {input_path} -> {output_path}")
     return output_path
@@ -96,7 +121,7 @@ def normalize(
     """
     from pathlib import Path
 
-    from lattifai.caption import Caption
+    from lattifai.data import Caption
 
     input_path = Path(input_path).expanduser()
     output_path = Path(output_path).expanduser()
@@ -151,7 +176,7 @@ def shift(
     """
     from pathlib import Path
 
-    from lattifai.caption import Caption
+    from lattifai.data import Caption
 
     input_path = Path(input_path).expanduser()
     output_path = Path(output_path).expanduser()
@@ -178,6 +203,88 @@ def shift(
     return output_path
 
 
+@run.cli.entrypoint(name="diff", namespace="caption")
+def diff(
+    ref_path: Pathlike,
+    hyp_path: Pathlike,
+    split_sentence: bool = True,
+    verbose: bool = True,
+):
+    """
+    Compare and align caption supervisions with transcription segments.
+
+    This command reads a reference caption file and a hypothesis file, then performs
+    text alignment to show how they match up. It's useful for comparing
+    original subtitles against ASR (Automatic Speech Recognition) results.
+
+    Args:
+        ref_path: Path to reference caption file (ground truth)
+        hyp_path: Path to hypothesis file (e.g., ASR results)
+        split_sentence: Enable sentence splitting before alignment (default: True)
+        verbose: Enable verbose output to show detailed alignment info (default: True)
+
+    Examples:
+        # Compare reference with hypothesis (positional arguments)
+        lai caption diff subtitles.srt transcription.json
+
+        # Disable sentence splitting
+        lai caption diff subtitles.srt transcription.json split_sentence=false
+
+        # Disable verbose output
+        lai caption diff subtitles.srt transcription.json verbose=false
+    """
+    from pathlib import Path
+
+    from lattifai.alignment.text_align import align_supervisions_and_transcription
+    from lattifai.caption import SentenceSplitter
+    from lattifai.data import Caption
+
+    ref_path = Path(ref_path).expanduser()
+    hyp_path = Path(hyp_path).expanduser()
+
+    # Read reference caption (supervisions)
+    caption_obj = Caption.read(ref_path)
+
+    # Read hypothesis
+    hyp_obj = Caption.read(hyp_path)
+
+    # Apply sentence splitting if enabled
+    if split_sentence:
+        splitter = SentenceSplitter(device="cpu", lazy_init=True)
+        caption_obj.supervisions = splitter.split_sentences(caption_obj.supervisions)
+        hyp_obj.supervisions = splitter.split_sentences(hyp_obj.supervisions)
+
+    # Set transcription on caption object
+    caption_obj.transcription = hyp_obj.supervisions
+
+    safe_print(f"📖  Reference: {len(caption_obj.supervisions)} segments from {ref_path}")
+    safe_print(f"🎤 Hypothesis: {len(caption_obj.transcription)} segments from {hyp_path}")
+    if split_sentence:
+        safe_print("✂️  Sentence splitting: enabled")
+    safe_print("")
+
+    # Perform alignment
+    results = align_supervisions_and_transcription(
+        caption=caption_obj,
+        verbose=verbose,
+    )
+
+    # # Print summary
+    # safe_print("")
+    # safe_print("=" * 72)
+    # safe_print(f"📊 Alignment Summary: {len(results)} groups")
+    # for idx, (sub_align, asr_align, quality, timestamp, typing) in enumerate(results):
+    #     sub_count = len(sub_align) if sub_align else 0
+    #     asr_count = len(asr_align) if asr_align else 0
+    #     safe_print(f"  Group {idx + 1}: ref={sub_count}, hyp={asr_count}, {quality.info}, typing={typing}")
+
+    return results
+
+
+def main_diff():
+    run.cli.main(diff)
+
+
 def main_convert():
     run.cli.main(convert)
 

lattifai/cli/transcribe.py

@@ -92,10 +92,6 @@ def transcribe(
     client_wrapper = SyncAPIClient(config=client_config)
     transcription_config.client_wrapper = client_wrapper
 
-    # Initialize client wrapper to properly set client_wrapper
-    client_wrapper = SyncAPIClient(config=client_config)
-    transcription_config.client_wrapper = client_wrapper
-
     # Validate input is required
     if not input and not media_config.input_path:
         raise ValueError("Input is required. Provide input as positional argument or media.input_path.")
@@ -129,7 +125,7 @@ def transcribe(
         if is_url:
             # Download media first, then transcribe
             safe_print(colorful.cyan("    Downloading media from URL..."))
-            from lattifai.workflow.youtube import YouTubeDownloader
+            from lattifai.youtube import YouTubeDownloader
 
             downloader = YouTubeDownloader()
             input_path = asyncio.run(
@@ -170,7 +166,7 @@ def transcribe(
     safe_print(colorful.cyan(f"   Output: {final_output}"))
 
     # Write output
-    transcriber.write(transcript, final_output, encoding="utf-8", cache_audio_events=False)
+    transcriber.write(transcript, final_output, encoding="utf-8", cache_event=False)
 
     safe_print(colorful.green(f"🎉 Transcription completed: {final_output}"))