caption_model_runtime.py

from __future__ import annotations

import math
import types
from collections import OrderedDict
from pathlib import Path
from typing import Any

import numpy as np
import soundfile as sf
import torch
from torch import nn
from transformers import AutoModelForTDT

try:
    from scipy.signal import resample_poly
except Exception:  # pragma: no cover - only needed for non-16k audio.
    resample_poly = None


def unwrap_parallel(model: torch.nn.Module) -> torch.nn.Module:
    while hasattr(model, "module"):
        model = model.module
    return model


def gemma_core(model: torch.nn.Module) -> torch.nn.Module:
    base = unwrap_parallel(model)
    core = getattr(base, "model", None)
    if core is not None and hasattr(core, "audio_tower") and hasattr(core, "embed_audio"):
        return core
    if hasattr(base, "audio_tower") and hasattr(base, "embed_audio"):
        return base
    raise AttributeError("Could not locate Gemma4Model core with audio_tower/embed_audio")


def load_state_file(path: Path) -> dict[str, torch.Tensor]:
    if path.suffix == ".safetensors":
        from safetensors.torch import load_file

        return load_file(str(path))
    return torch.load(path, map_location="cpu")


DEFAULT_PARAKEET_MODEL_ID = "nvidia/parakeet-tdt-0.6b-v3"

class FrozenParakeetAudioTower(nn.Module):
    """Gemma audio tower replacement backed by a frozen Parakeet encoder.

    Gemma's processor creates one audio soft token per roughly four 10ms feature
    frames. Parakeet's encoder subsamples by eight, so the tower upsamples the
    Parakeet sequence back to the Gemma audio-token count before Gemma scatters
    the projected features into the prompt.
    """

    def __init__(
        self,
        model_id: str,
        *,
        local_files_only: bool,
        dtype: torch.dtype,
        expected_subsample_factor: int = 4,
    ) -> None:
        super().__init__()
        parakeet = AutoModelForTDT.from_pretrained(
            model_id,
            local_files_only=local_files_only,
            dtype=dtype,
            low_cpu_mem_usage=True,
        )
        self.encoder = parakeet.encoder
        self.hidden_size = int(parakeet.config.encoder_config.hidden_size)
        self.token_hidden_size = int(getattr(parakeet.config, "decoder_hidden_size", 640))
        self.model_id = model_id
        self.expected_subsample_factor = int(expected_subsample_factor)
        self.register_buffer("_parakeet_bridge_marker", torch.ones(1), persistent=True)
        for parameter in self.encoder.parameters():
            parameter.requires_grad = False
        self._disable_decode_expert_switching()
        self.encoder.eval()
        del parakeet

    def _disable_decode_expert_switching(self) -> None:
        def get_correct_experts_implementation(encoder: nn.Module, implementation: Any = None) -> Any:
            del encoder
            return implementation

        def set_experts_implementation(encoder: nn.Module, implementation: Any = None) -> None:
            del encoder, implementation
            return None

        self.encoder.get_correct_experts_implementation = types.MethodType(
            get_correct_experts_implementation,
            self.encoder,
        )
        self.encoder.set_experts_implementation = types.MethodType(
            set_experts_implementation,
            self.encoder,
        )

    def train(self, mode: bool = True) -> "FrozenParakeetAudioTower":
        super().train(mode)
        self.encoder.eval()
        return self

    def state_dict(self, *args: Any, **kwargs: Any) -> OrderedDict[str, torch.Tensor]:
        prefix = kwargs.get("prefix", "")
        destination = kwargs.get("destination")
        if destination is None:
            destination = OrderedDict()
        destination[prefix + "_parakeet_bridge_marker"] = self._parakeet_bridge_marker.detach().cpu()
        return destination

    def load_state_dict(self, state_dict: dict[str, torch.Tensor], strict: bool = True, assign: bool = False):
        del assign
        marker = state_dict.get("_parakeet_bridge_marker")
        if marker is not None:
            self._parakeet_bridge_marker.copy_(marker.to(self._parakeet_bridge_marker.device))
        missing = [] if marker is not None or not strict else ["_parakeet_bridge_marker"]
        unexpected = [key for key in state_dict if key != "_parakeet_bridge_marker"]
        if strict and (missing or unexpected):
            raise RuntimeError(f"Parakeet audio tower state mismatch: missing={missing} unexpected={unexpected}")
        return missing, unexpected

    @staticmethod
    def _gemma_audio_mask(input_features_mask: torch.Tensor, target_length: int) -> torch.Tensor:
        mask = input_features_mask
        while mask.shape[1] > target_length:
            mask = mask[:, ::2]
        if mask.shape[1] > target_length:
            mask = mask[:, :target_length]
        if mask.shape[1] < target_length:
            pad = torch.zeros(
                (mask.shape[0], target_length - mask.shape[1]),
                dtype=mask.dtype,
                device=mask.device,
            )
            mask = torch.cat([mask, pad], dim=1)
        return mask.bool()

    def forward(
        self,
        input_features: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        **kwargs: Any,
    ) -> Any:
        del kwargs
        if attention_mask is None:
            attention_mask = torch.ones(
                input_features.shape[:2],
                dtype=torch.long,
                device=input_features.device,
            )
        target_length = (input_features.shape[1] + self.expected_subsample_factor - 1) // self.expected_subsample_factor
        encoder_dtype = next(self.encoder.parameters()).dtype
        with torch.no_grad():
            encoded = self.encoder(
                input_features=input_features.to(dtype=encoder_dtype),
                attention_mask=attention_mask.long(),
                output_attention_mask=True,
            )
            hidden = encoded.last_hidden_state

        if hidden.shape[1] != target_length:
            hidden = torch.nn.functional.interpolate(
                hidden.transpose(1, 2).float(),
                size=target_length,
                mode="linear",
                align_corners=False,
            ).transpose(1, 2).to(dtype=encoder_dtype)

        output_mask = self._gemma_audio_mask(attention_mask, target_length)
        return type(
            "ParakeetAudioTowerOutput",
            (),
            {
                "last_hidden_state": hidden,
                "attention_mask": output_mask,
                "pooler_output": None,
            },
        )()


class FrozenParakeetTDTTokenAudioTower(nn.Module):
    """Gemma audio tower that exposes Parakeet's audio-derived token stream.

    This still does not insert transcript text into the Gemma prompt. Parakeet
    runs from audio features to its own TDT token/duration sequence internally,
    then the decoder hidden states for that sequence become Gemma audio soft
    tokens after the trainable projector.
    """

    def __init__(
        self,
        model_id: str,
        *,
        local_files_only: bool,
        dtype: torch.dtype,
        expected_subsample_factor: int = 4,
        min_token_repeats: int = 1,
        token_feature_source: str = "decoder_states",
        filter_blank_tokens: bool = True,
        filter_special_token_ids: bool = True,
    ) -> None:
        super().__init__()
        self.tdt = AutoModelForTDT.from_pretrained(
            model_id,
            local_files_only=local_files_only,
            dtype=dtype,
            low_cpu_mem_usage=True,
        )
        self.hidden_size = int(self.tdt.config.decoder_hidden_size)
        self.blank_token_id = int(self.tdt.config.blank_token_id)
        self.model_id = model_id
        self.expected_subsample_factor = int(expected_subsample_factor)
        self.min_token_repeats = max(1, int(min_token_repeats))
        self.filter_blank_tokens = bool(filter_blank_tokens)
        self.filter_special_token_ids = bool(filter_special_token_ids)
        self.special_token_ids = {0, 2, 3}
        if token_feature_source not in {"decoder_states", "token_embeddings"}:
            raise ValueError(f"Unsupported token_feature_source={token_feature_source}")
        self.token_feature_source = token_feature_source
        self.register_buffer("_parakeet_tdt_token_bridge_marker", torch.ones(1), persistent=True)
        for parameter in self.tdt.parameters():
            parameter.requires_grad = False
        self._disable_decode_expert_switching()
        self.tdt.eval()

    def _disable_decode_expert_switching(self) -> None:
        def get_correct_experts_implementation(encoder: nn.Module, implementation: Any = None) -> Any:
            del encoder
            return implementation

        def set_experts_implementation(encoder: nn.Module, implementation: Any = None) -> None:
            del encoder, implementation
            return None

        self.tdt.encoder.get_correct_experts_implementation = types.MethodType(
            get_correct_experts_implementation,
            self.tdt.encoder,
        )
        self.tdt.encoder.set_experts_implementation = types.MethodType(
            set_experts_implementation,
            self.tdt.encoder,
        )
        self.tdt.get_correct_experts_implementation = types.MethodType(
            get_correct_experts_implementation,
            self.tdt,
        )
        self.tdt.set_experts_implementation = types.MethodType(
            set_experts_implementation,
            self.tdt,
        )

    def train(self, mode: bool = True) -> "FrozenParakeetTDTTokenAudioTower":
        super().train(mode)
        self.tdt.eval()
        return self

    def state_dict(self, *args: Any, **kwargs: Any) -> OrderedDict[str, torch.Tensor]:
        prefix = kwargs.get("prefix", "")
        destination = kwargs.get("destination")
        if destination is None:
            destination = OrderedDict()
        destination[prefix + "_parakeet_tdt_token_bridge_marker"] = (
            self._parakeet_tdt_token_bridge_marker.detach().cpu()
        )
        return destination

    def load_state_dict(self, state_dict: dict[str, torch.Tensor], strict: bool = True, assign: bool = False):
        del assign
        marker = state_dict.get("_parakeet_tdt_token_bridge_marker")
        if marker is not None:
            self._parakeet_tdt_token_bridge_marker.copy_(marker.to(self._parakeet_tdt_token_bridge_marker.device))
        missing = [] if marker is not None or not strict else ["_parakeet_tdt_token_bridge_marker"]
        unexpected = [key for key in state_dict if key != "_parakeet_tdt_token_bridge_marker"]
        if strict and (missing or unexpected):
            raise RuntimeError(f"Parakeet TDT token audio tower state mismatch: missing={missing} unexpected={unexpected}")
        return missing, unexpected

    def _target_length(self, input_features: torch.Tensor) -> int:
        return (input_features.shape[1] + self.expected_subsample_factor - 1) // self.expected_subsample_factor

    def _sequence_to_target_length(
        self,
        token_ids: torch.Tensor,
        decoder_states: torch.Tensor,
        durations: torch.Tensor,
        target_length: int,
    ) -> torch.Tensor:
        keep_mask = torch.ones_like(token_ids, dtype=torch.bool)
        if self.filter_blank_tokens:
            keep_mask &= token_ids.ne(self.blank_token_id)
        if self.filter_special_token_ids:
            for special_id in self.special_token_ids:
                keep_mask &= token_ids.ne(special_id)
        if keep_mask.any():
            decoder_states = decoder_states[keep_mask]
            durations = durations[keep_mask]
        repeats = durations.long().clamp_min(self.min_token_repeats)
        expanded = decoder_states.repeat_interleave(repeats, dim=0)
        if expanded.numel() == 0:
            expanded = decoder_states[:1]
        if expanded.shape[0] != target_length:
            expanded = torch.nn.functional.interpolate(
                expanded.transpose(0, 1).unsqueeze(0).float(),
                size=target_length,
                mode="linear",
                align_corners=False,
            ).squeeze(0).transpose(0, 1).to(dtype=decoder_states.dtype)
        return expanded

    def forward(
        self,
        input_features: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        **kwargs: Any,
    ) -> Any:
        del kwargs
        if attention_mask is None:
            attention_mask = torch.ones(
                input_features.shape[:2],
                dtype=torch.long,
                device=input_features.device,
            )
        target_length = self._target_length(input_features)
        model_dtype = next(self.tdt.parameters()).dtype
        with torch.no_grad():
            generated = self.tdt.generate(
                input_features=input_features.to(dtype=model_dtype),
                attention_mask=attention_mask.long(),
                max_new_tokens=target_length,
            )
            token_ids = generated.sequences.to(device=input_features.device)
            durations = generated.durations.to(device=input_features.device)
            if self.token_feature_source == "token_embeddings":
                decoder_states = self.tdt.decoder.embedding(token_ids)
            else:
                decoder_states = self.tdt.decoder(token_ids)

        projected_inputs: list[torch.Tensor] = []
        output_masks: list[torch.Tensor] = []
        for batch_index in range(decoder_states.shape[0]):
            expanded = self._sequence_to_target_length(
                token_ids[batch_index],
                decoder_states[batch_index],
                durations[batch_index],
                target_length,
            )
            projected_inputs.append(expanded)
            output_masks.append(torch.ones(target_length, dtype=torch.bool, device=input_features.device))

        hidden = torch.stack(projected_inputs, dim=0)
        output_mask = torch.stack(output_masks, dim=0)
        return type(
            "ParakeetTDTTokenAudioTowerOutput",
            (),
            {
                "last_hidden_state": hidden,
                "attention_mask": output_mask,
                "pooler_output": None,
            },
        )()


class FrozenParakeetTDTTokenEncoderHybridAudioTower(FrozenParakeetTDTTokenAudioTower):
    """Expose both Parakeet TDT token embeddings and continuous encoder states.

    The token embedding stream carries the audio-derived ASR signal that already
    works. The continuous encoder stream preserves native acoustic information
    that does not survive the hard TDT token path.
    """

    def __init__(self, *args: Any, **kwargs: Any) -> None:
        super().__init__(*args, **kwargs)
        self.token_hidden_size = int(self.hidden_size)
        self.encoder_hidden_size = int(self.tdt.config.encoder_config.hidden_size)
        self.hidden_size = self.token_hidden_size + self.encoder_hidden_size

    @staticmethod
    def _match_length(hidden: torch.Tensor, target_length: int) -> torch.Tensor:
        if hidden.shape[1] == target_length:
            return hidden
        return torch.nn.functional.interpolate(
            hidden.transpose(1, 2).float(),
            size=target_length,
            mode="linear",
            align_corners=False,
        ).transpose(1, 2).to(dtype=hidden.dtype)

    def forward(
        self,
        input_features: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        **kwargs: Any,
    ) -> Any:
        del kwargs
        if attention_mask is None:
            attention_mask = torch.ones(
                input_features.shape[:2],
                dtype=torch.long,
                device=input_features.device,
            )
        target_length = self._target_length(input_features)
        model_dtype = next(self.tdt.parameters()).dtype
        model_features = input_features.to(dtype=model_dtype)
        with torch.no_grad():
            encoded = self.tdt.encoder(
                input_features=model_features,
                attention_mask=attention_mask.long(),
                output_attention_mask=True,
            )
            encoder_hidden = self._match_length(encoded.last_hidden_state, target_length)
            generated = self.tdt.generate(
                input_features=model_features,
                attention_mask=attention_mask.long(),
                max_new_tokens=target_length,
            )
            token_ids = generated.sequences.to(device=input_features.device)
            durations = generated.durations.to(device=input_features.device)
            if self.token_feature_source == "token_embeddings":
                decoder_states = self.tdt.decoder.embedding(token_ids)
            else:
                decoder_states = self.tdt.decoder(token_ids)

        token_inputs: list[torch.Tensor] = []
        for batch_index in range(decoder_states.shape[0]):
            token_inputs.append(
                self._sequence_to_target_length(
                    token_ids[batch_index],
                    decoder_states[batch_index],
                    durations[batch_index],
                    target_length,
                )
            )

        token_hidden = torch.stack(token_inputs, dim=0).to(dtype=model_dtype)
        encoder_hidden = encoder_hidden.to(dtype=model_dtype)
        hidden = torch.cat([token_hidden, encoder_hidden], dim=-1).to(dtype=model_dtype)
        output_mask = FrozenParakeetAudioTower._gemma_audio_mask(attention_mask, target_length)
        return type(
            "ParakeetTDTTokenEncoderHybridAudioTowerOutput",
            (),
            {
                "last_hidden_state": hidden,
                "attention_mask": output_mask,
                "pooler_output": None,
            },
        )()


class ParakeetToGemmaAudioProjector(nn.Module):
    def __init__(
        self,
        input_hidden_size: int,
        output_hidden_size: int,
        intermediate_size: int = 4096,
        dropout: float = 0.0,
    ) -> None:
        super().__init__()
        self.input_norm = nn.LayerNorm(input_hidden_size)
        self.up = nn.Linear(input_hidden_size, intermediate_size)
        self.act = nn.GELU()
        self.dropout = nn.Dropout(dropout)
        self.down = nn.Linear(intermediate_size, output_hidden_size)
        self.output_norm = nn.LayerNorm(output_hidden_size)

    def forward(self, inputs_embeds: torch.Tensor) -> torch.Tensor:
        output_dtype = inputs_embeds.dtype
        hidden = self.input_norm(inputs_embeds)
        hidden = self.up(hidden)
        hidden = self.act(hidden)
        hidden = self.dropout(hidden)
        hidden = self.down(hidden)
        return self.output_norm(hidden).to(dtype=output_dtype)


class ParakeetEncoderToTokenEmbeddingProjector(nn.Module):
    """Map continuous Parakeet encoder states through a speech-token-like space.

    The working TDT-token bridge proved that Gemma can use Parakeet's 640-dim
    token embedding space once it is projected into Gemma hidden size. This
    projector keeps the continuous encoder path, but gives it a trainable
    1024->640 bottleneck before the known-good 640->Gemma projector.
    """

    def __init__(
        self,
        input_hidden_size: int,
        token_hidden_size: int,
        output_hidden_size: int,
        intermediate_size: int = 4096,
        dropout: float = 0.0,
    ) -> None:
        super().__init__()
        self.encoder_to_token = nn.Sequential(
            nn.LayerNorm(input_hidden_size),
            nn.Linear(input_hidden_size, intermediate_size),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(intermediate_size, token_hidden_size),
            nn.LayerNorm(token_hidden_size),
        )
        self.token_projector = ParakeetToGemmaAudioProjector(
            input_hidden_size=token_hidden_size,
            output_hidden_size=output_hidden_size,
            intermediate_size=intermediate_size,
            dropout=dropout,
        )

    def forward(self, inputs_embeds: torch.Tensor) -> torch.Tensor:
        output_dtype = inputs_embeds.dtype
        token_like = self.encoder_to_token(inputs_embeds).to(dtype=output_dtype)
        return self.token_projector(token_like).to(dtype=output_dtype)


class ParakeetTDTTokenEncoderHybridProjector(nn.Module):
    """Project token embeddings plus continuous encoder states into Gemma space."""

    def __init__(
        self,
        token_hidden_size: int,
        encoder_hidden_size: int,
        output_hidden_size: int,
        intermediate_size: int = 4096,
        dropout: float = 0.0,
        encoder_gate_init: float = 0.0,
    ) -> None:
        super().__init__()
        self.token_hidden_size = int(token_hidden_size)
        self.encoder_hidden_size = int(encoder_hidden_size)
        self.token_projector = ParakeetToGemmaAudioProjector(
            input_hidden_size=token_hidden_size,
            output_hidden_size=output_hidden_size,
            intermediate_size=intermediate_size,
            dropout=dropout,
        )
        self.encoder_projector = ParakeetToGemmaAudioProjector(
            input_hidden_size=encoder_hidden_size,
            output_hidden_size=output_hidden_size,
            intermediate_size=intermediate_size,
            dropout=dropout,
        )
        self.encoder_gate = nn.Parameter(torch.tensor(float(encoder_gate_init), dtype=torch.float32))

    def forward(self, inputs_embeds: torch.Tensor) -> torch.Tensor:
        output_dtype = inputs_embeds.dtype
        token_hidden, encoder_hidden = torch.split(
            inputs_embeds,
            [self.token_hidden_size, self.encoder_hidden_size],
            dim=-1,
        )
        token_output = self.token_projector(token_hidden)
        encoder_output = self.encoder_projector(encoder_hidden)
        gate = torch.tanh(self.encoder_gate).to(dtype=output_dtype)
        return (token_output + gate * encoder_output).to(dtype=output_dtype)


def load_audio_array(path: str | Path, sampling_rate: int, max_length_samples: int) -> np.ndarray:
    audio, source_rate = sf.read(str(path), dtype="float32", always_2d=False)
    if audio.ndim > 1:
        audio = audio.mean(axis=1)
    if source_rate != sampling_rate:
        if resample_poly is None:
            raise RuntimeError(
                f"Audio {path} has sample rate {source_rate}, but scipy is unavailable for resampling"
            )
        divisor = math.gcd(int(source_rate), int(sampling_rate))
        audio = resample_poly(audio, sampling_rate // divisor, source_rate // divisor).astype("float32")
    if max_length_samples > 0 and audio.shape[0] > max_length_samples:
        audio = audio[:max_length_samples]
    return np.asarray(audio, dtype=np.float32)


def pad_or_trim_parakeet_features(
    input_features: torch.Tensor,
    attention_mask: torch.Tensor,
    target_length: int,
) -> tuple[torch.Tensor, torch.Tensor]:
    if input_features.shape[1] > target_length:
        input_features = input_features[:, :target_length]
        attention_mask = attention_mask[:, :target_length]
    if input_features.shape[1] < target_length:
        pad_length = target_length - input_features.shape[1]
        input_features = torch.nn.functional.pad(input_features, (0, 0, 0, pad_length), value=0.0)
        attention_mask = torch.nn.functional.pad(attention_mask, (0, pad_length), value=0)
    return input_features, attention_mask


def parakeet_feature_tensors(
    *,
    audio_paths: list[str],
    parakeet_processor: Any,
    sampling_rate: int,
    max_length_samples: int,
    target_length: int,
) -> tuple[torch.Tensor, torch.Tensor]:
    waveforms = [load_audio_array(path, sampling_rate, max_length_samples) for path in audio_paths]
    parakeet_batch = parakeet_processor(
        waveforms,
        sampling_rate=sampling_rate,
        return_tensors="pt",
        padding=True,
    )
    features = parakeet_batch["input_features"].float()
    mask = parakeet_batch.get("attention_mask")
    if mask is None:
        mask = torch.ones(features.shape[:2], dtype=torch.bool)
    features, mask = pad_or_trim_parakeet_features(features, mask.bool(), target_length)
    return features, mask


def replace_batch_audio_features(
    batch: dict[str, torch.Tensor],
    *,
    audio_paths: list[str],
    parakeet_processor: Any,
    sampling_rate: int,
    max_length_samples: int,
    prefix: str = "",
) -> None:
    feature_key = f"{prefix}input_features"
    mask_key = f"{prefix}input_features_mask"
    if feature_key not in batch:
        return
    original_features = batch[feature_key]
    original_mask = batch[mask_key]
    features, mask = parakeet_feature_tensors(
        audio_paths=audio_paths,
        parakeet_processor=parakeet_processor,
        sampling_rate=sampling_rate,
        max_length_samples=max_length_samples,
        target_length=int(original_features.shape[1]),
    )
    batch[feature_key] = features.to(dtype=original_features.dtype)
    batch[mask_key] = mask.to(dtype=original_mask.dtype)


def install_parakeet_audio_bridge(model: torch.nn.Module, args: argparse.Namespace) -> None:
    core = gemma_core(model)
    text_hidden_size = int(model.config.get_text_config().hidden_size)
    if args.parakeet_bridge_mode == "tdt_tokens":
        audio_tower = FrozenParakeetTDTTokenAudioTower(
            args.parakeet_model_id,
            local_files_only=args.local_files_only,
            dtype=torch.bfloat16,
            token_feature_source="decoder_states",
            filter_blank_tokens=getattr(args, "parakeet_tdt_filter_blank_tokens", True),
            filter_special_token_ids=getattr(args, "parakeet_tdt_filter_special_token_ids", True),
        )
    elif args.parakeet_bridge_mode in {"tdt_token_embeddings", "tdt_token_embeddings_with_encoder_context"}:
        tower_class = (
            FrozenParakeetTDTTokenEncoderHybridAudioTower
            if args.parakeet_bridge_mode == "tdt_token_embeddings_with_encoder_context"
            else FrozenParakeetTDTTokenAudioTower
        )
        audio_tower = tower_class(
            args.parakeet_model_id,
            local_files_only=args.local_files_only,
            dtype=torch.bfloat16,
            token_feature_source="token_embeddings",
            filter_blank_tokens=getattr(args, "parakeet_tdt_filter_blank_tokens", True),
            filter_special_token_ids=getattr(args, "parakeet_tdt_filter_special_token_ids", True),
        )
    else:
        audio_tower = FrozenParakeetAudioTower(
            args.parakeet_model_id,
            local_files_only=args.local_files_only,
            dtype=torch.bfloat16,
        )
    if args.parakeet_bridge_mode == "encoder_soft_tdt_token_embeddings":
        projector = ParakeetEncoderToTokenEmbeddingProjector(
            input_hidden_size=audio_tower.hidden_size,
            token_hidden_size=getattr(audio_tower, "token_hidden_size", 640),
            output_hidden_size=text_hidden_size,
            intermediate_size=args.projector_intermediate_size,
            dropout=args.projector_dropout,
        ).to(dtype=torch.bfloat16)
    elif args.parakeet_bridge_mode == "tdt_token_embeddings_with_encoder_context":
        projector = ParakeetTDTTokenEncoderHybridProjector(
            token_hidden_size=getattr(audio_tower, "token_hidden_size", 640),
            encoder_hidden_size=getattr(audio_tower, "encoder_hidden_size", 1024),
            output_hidden_size=text_hidden_size,
            intermediate_size=args.projector_intermediate_size,
            dropout=args.projector_dropout,
            encoder_gate_init=getattr(args, "hybrid_encoder_gate_init", 0.0),
        ).to(dtype=torch.bfloat16)
    else:
        projector = ParakeetToGemmaAudioProjector(
            input_hidden_size=audio_tower.hidden_size,
            output_hidden_size=text_hidden_size,
            intermediate_size=args.projector_intermediate_size,
            dropout=args.projector_dropout,
        ).to(dtype=torch.bfloat16)
    core.audio_tower = audio_tower
    core.embed_audio = projector
    target_device = getattr(model, "device", None)
    if isinstance(target_device, torch.device) and target_device.type != "cpu":
        core.audio_tower.to(device=target_device)
        core.embed_audio.to(device=target_device)
    print(
        "parakeet_audio_bridge_installed=true "
        f"bridge_mode={args.parakeet_bridge_mode} "
        f"parakeet_model_id={args.parakeet_model_id} "
        f"parakeet_hidden_size={audio_tower.hidden_size} "
        f"token_hidden_size={getattr(audio_tower, 'token_hidden_size', 'n/a')} "
        f"gemma_hidden_size={text_hidden_size} "
        f"projector_intermediate_size={args.projector_intermediate_size}",
        flush=True,
    )


CAPTION_LENGTH_LABELS = ("very small", "small", "medium", "large", "very large")
TAG_KEYS = ("tags", "tag_list", "tag_string", "danbooru_tags", "booru_tags")
CAPTION_SETTING_FIELD_CHOICES = {
    "vulgarity": ("none", "low", "medium", "high"),
    "uncertainty": ("none", "low", "medium", "high"),
    "character_names": ("none", "ambiguous", "single", "multiple"),
    "fluff": ("none", "low", "medium", "high"),
    "speculation": ("none", "low", "medium", "high"),
    "temporal_detail": ("static", "low", "medium", "high"),
    "visual_specificity": ("generic", "moderate", "detailed", "excessive"),
    "camera_detail": ("none", "low", "medium", "high"),
    "caption_style": ("plain", "verbose", "ornate", "robotic"),
}
CAPTION_SETTING_FIELDS = tuple(CAPTION_SETTING_FIELD_CHOICES)
DEFAULT_CAPTION_SETTING_VALUES = {
    "vulgarity": "none",
    "uncertainty": "none",
    "character_names": "none",
    "fluff": "none",
    "has_repetition": False,
    "has_thinking": True,
    "speculation": "none",
    "temporal_detail": "medium",
    "visual_specificity": "moderate",
    "camera_detail": "low",
    "caption_style": "plain",
}


def format_caption_settings_prompt(settings: dict[str, Any]) -> str:
    watermark_instruction = (
        "Include watermark info." if settings["include_watermark_info"] else "Do not include watermark info."
    )
    repetition_value = str(bool(settings["has_repetition"])).lower()
    thinking_value = str(bool(settings.get("has_thinking", True))).lower()
    thinking_instruction = (
        "Output thought JSON before the final caption."
        if settings.get("has_thinking", True)
        else "Do not output thought JSON; output only the caption."
    )
    setting_text = (
        f"vulgarity={settings['vulgarity']}; "
        f"uncertainty={settings['uncertainty']}; "
        f"character_names={settings['character_names']}; "
        f"fluff={settings['fluff']}; "
        f"has_repetition={repetition_value}; "
        f"has_thinking={thinking_value}; "
        f"speculation={settings['speculation']}; "
        f"temporal_detail={settings['temporal_detail']}; "
        f"visual_specificity={settings['visual_specificity']}; "
        f"camera_detail={settings['camera_detail']}; "
        f"caption_style={settings['caption_style']}"
    )
    return (
        f"Write a {settings['caption_length']} caption for this clip using both the visuals and the audio. "
        f"{watermark_instruction} {thinking_instruction} Match these caption settings: {setting_text}."
    )