lattifai 1.2.2__py3-none-any.whl → 1.3.1__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/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,8 +6,8 @@ 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.config.caption import KaraokeConfig
 from lattifai.utils import safe_print
 
 
@@ -68,7 +68,7 @@ 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
@@ -121,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()
@@ -176,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()
@@ -235,9 +235,9 @@ def diff(
     """
     from pathlib import Path
 
-    from lattifai.alignment.sentence_splitter import SentenceSplitter
     from lattifai.alignment.text_align import align_supervisions_and_transcription
-    from lattifai.caption import Caption
+    from lattifai.caption import SentenceSplitter
+    from lattifai.data import Caption
 
     ref_path = Path(ref_path).expanduser()
     hyp_path = Path(hyp_path).expanduser()

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.")
@@ -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}"))
 

lattifai/cli/youtube.py

@@ -11,6 +11,7 @@ from lattifai.config import (
     CaptionConfig,
     ClientConfig,
     DiarizationConfig,
+    EventConfig,
     MediaConfig,
     TranscriptionConfig,
 )
@@ -25,6 +26,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,
+    event: Annotated[Optional[EventConfig], run.Config[EventConfig]] = None,
     use_transcription: bool = False,
 ):
     """
@@ -114,6 +116,7 @@ def youtube(
         caption_config=caption_config,
         transcription_config=transcription,
         diarization_config=diarization,
+        event_config=event,
     )
 
     # Call the client's youtube method

lattifai/client.py

@@ -7,10 +7,18 @@ import colorful
 from lattifai_core.client import SyncAPIClient
 from lhotse.utils import Pathlike
 
-from lattifai.alignment import Lattice1Aligner, Segmenter, align_supervisions_and_transcription
+from lattifai.alignment import Lattice1Aligner, Segmenter
 from lattifai.audio2 import AudioData, AudioLoader
-from lattifai.caption import Caption, InputCaptionFormat
-from lattifai.config import AlignmentConfig, CaptionConfig, ClientConfig, DiarizationConfig, TranscriptionConfig
+from lattifai.caption import InputCaptionFormat
+from lattifai.config import (
+    AlignmentConfig,
+    CaptionConfig,
+    ClientConfig,
+    DiarizationConfig,
+    EventConfig,
+    TranscriptionConfig,
+)
+from lattifai.data import Caption
 from lattifai.errors import (
     AlignmentError,
     CaptionProcessingError,
@@ -22,6 +30,7 @@ from lattifai.utils import safe_print
 
 if TYPE_CHECKING:
     from lattifai.diarization import LattifAIDiarizer  # noqa: F401
+    from lattifai.event import LattifAIEventDetector  # noqa: F401
 
 
 class LattifAI(LattifAIClientMixin, SyncAPIClient):
@@ -41,6 +50,7 @@ class LattifAI(LattifAIClientMixin, SyncAPIClient):
         caption_config: Optional[CaptionConfig] = None,
         transcription_config: Optional[TranscriptionConfig] = None,
         diarization_config: Optional[DiarizationConfig] = None,
+        event_config: Optional[EventConfig] = None,
     ) -> None:
         __doc__ = LattifAIClientMixin._INIT_DOC.format(
             client_class="LattifAI",
@@ -59,8 +69,8 @@ class LattifAI(LattifAIClientMixin, SyncAPIClient):
         self.config = client_config
 
         # Initialize all configs with defaults
-        alignment_config, transcription_config, diarization_config = self._init_configs(
-            alignment_config, transcription_config, diarization_config
+        alignment_config, transcription_config, diarization_config, event_config = self._init_configs(
+            alignment_config, transcription_config, diarization_config, event_config
         )
 
         # Store configs
@@ -82,6 +92,14 @@ class LattifAI(LattifAIClientMixin, SyncAPIClient):
 
             self.diarizer = LattifAIDiarizer(config=self.diarization_config)
 
+        # Initialize event detector if enabled
+        self.event_config = event_config
+        self.event_detector: Optional["LattifAIEventDetector"] = None
+        if self.event_config.enabled:
+            from lattifai.event import LattifAIEventDetector  # noqa: F811
+
+            self.event_detector = LattifAIEventDetector(config=self.event_config)
+
         # Initialize shared components (transcriber, downloader)
         self._init_shared_components(transcription_config)
 
@@ -126,6 +144,8 @@ class LattifAI(LattifAIClientMixin, SyncAPIClient):
                 safe_print(colorful.cyan(f"🔄   Using segmented alignment strategy: {alignment_strategy}"))
 
                 if caption.supervisions and alignment_strategy == "transcription":
+                    from lattifai.alignment.text_align import align_supervisions_and_transcription
+
                     if "gemini" in self.transcriber.name.lower():
                         raise ValueError(
                             f"Transcription-based alignment is not supported for {self.transcriber.name} "
@@ -139,7 +159,7 @@ class LattifAI(LattifAIClientMixin, SyncAPIClient):
                             output_dir=Path(str(output_caption_path)).parent if output_caption_path else None,
                         )
                         caption.transcription = transcript.supervisions or transcript.transcription
-                        caption.audio_events = transcript.audio_events
+                        caption.event = transcript.event
                     if not caption.transcription:
                         raise ValueError("Transcription is empty after transcription step.")
 
@@ -234,13 +254,15 @@ class LattifAI(LattifAIClientMixin, SyncAPIClient):
             if self.config.profile:
                 self.aligner.profile()
 
-        except (CaptionProcessingError, LatticeEncodingError, AlignmentError, LatticeDecodingError):
+        except (CaptionProcessingError, LatticeEncodingError) as e:
             # Re-raise our specific errors as-is
-            raise
+            raise e
+        except LatticeDecodingError as e:
+            raise e
         except Exception as e:
             # Catch any unexpected errors and wrap them
             raise AlignmentError(
-                "Unexpected error during alignment process",
+                message="Unexpected error during alignment process",
                 media_path=str(input_media),
                 caption_path=str(input_caption),
                 context={"original_error": str(e), "error_type": e.__class__.__name__},
@@ -255,6 +277,13 @@ class LattifAI(LattifAIClientMixin, SyncAPIClient):
                 output_caption_path=output_caption_path,
             )
 
+        # Step 6: Event detection
+        if self.event_config.enabled and self.event_detector:
+            safe_print(colorful.cyan("🔊 Performing audio event detection..."))
+            caption = self.event_detector.detect_and_update_caption(caption, media_audio)
+            if output_caption_path:
+                self._write_caption(caption, output_caption_path)
+
         return caption
 
     def speaker_diarization(
@@ -285,12 +314,12 @@ class LattifAI(LattifAIClientMixin, SyncAPIClient):
             diarization_file = Path(str(output_caption_path)).with_suffix(".SpkDiar")
             if diarization_file.exists():
                 safe_print(colorful.cyan(f"Reading existing speaker diarization from {diarization_file}"))
-                caption.read_speaker_diarization(diarization_file)
+                caption.read_diarization(diarization_file)
 
         diarization, alignments = self.diarizer.diarize_with_alignments(
             input_media,
             caption.alignments,
-            diarization=caption.speaker_diarization,
+            diarization=caption.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,
@@ -298,7 +327,7 @@ class LattifAI(LattifAIClientMixin, SyncAPIClient):
             output_path=output_caption_path,
         )
         caption.alignments = alignments
-        caption.speaker_diarization = diarization
+        caption.diarization = diarization
 
         # Write output if requested
         if output_caption_path: