Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+model/tokenizer.json filter=lfs diff=lfs merge=lfs -text
+processor/tokenizer.json filter=lfs diff=lfs merge=lfs -text

.ipynb_checkpoints/caption_folder-checkpoint.py

@@ -0,0 +1,206 @@
+"""Caption every video in a folder with the packaged vLLM model.
+
+Edit the variables below, then run:
+
+    cd /workspace/polished_model_v3
+    source .venv/bin/activate
+    python caption_folder.py
+"""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+SCRIPT_DIR = Path(__file__).resolve().parent
+sys.path.insert(0, str(SCRIPT_DIR / "vllm"))
+
+from vllm_caption_runtime import build_vllm_request, load_llm, load_processor, sampling_params
+
+
+# =============================================================================
+# Paths
+# =============================================================================
+
+INPUT_FOLDER = "/workspace/videos_to_caption"
+OUTPUT_FOLDER = "/workspace/captions_out"
+MODEL_DIR = str(SCRIPT_DIR / "model")
+
+RECURSIVE = True
+OVERWRITE_EXISTING = False
+OUTPUT_EXTENSION = ".txt"
+ERROR_EXTENSION = ".error.txt"
+VIDEO_EXTENSIONS = (".mp4", ".mov", ".mkv", ".webm", ".avi", ".m4v")
+
+
+# =============================================================================
+# Prompt Settings
+# =============================================================================
+
+PROMPT_OVERRIDE = ""
+CAPTION_LENGTH = "very large"
+INCLUDE_WATERMARK_INFO = False
+HAS_THINKING = True
+
+VULGARITY = "low"
+UNCERTAINTY = "low"
+CHARACTER_NAMES = "none"
+FLUFF = "none"
+HAS_REPETITION = False
+SPECULATION = "low"
+TEMPORAL_DETAIL = "medium"
+VISUAL_SPECIFICITY = "moderate"
+CAMERA_DETAIL = "medium"
+CAPTION_STYLE = "plain"
+
+
+# =============================================================================
+# vLLM / Generation Hyperparameters
+# =============================================================================
+
+NUM_FRAMES = 12
+SAMPLING_RATE = 16_000
+MAX_MODEL_LEN = 4096
+MAX_NUM_SEQS = 20
+BATCH_SIZE = 20
+GPU_MEMORY_UTILIZATION = 0.88
+DTYPE = "bfloat16"
+ENFORCE_EAGER = False
+ENABLE_PREFIX_CACHING = False
+TRUST_REMOTE_CODE = False
+
+MAX_TOKENS = 1200
+TEMPERATURE = 0.0
+TOP_P = 0.9
+REPETITION_PENALTY = 1.1
+
+
+def prompt_settings() -> dict[str, object]:
+    return {
+        "caption_length": CAPTION_LENGTH,
+        "include_watermark_info": INCLUDE_WATERMARK_INFO,
+        "has_thinking": HAS_THINKING,
+        "vulgarity": VULGARITY,
+        "uncertainty": UNCERTAINTY,
+        "character_names": CHARACTER_NAMES,
+        "fluff": FLUFF,
+        "has_repetition": HAS_REPETITION,
+        "speculation": SPECULATION,
+        "temporal_detail": TEMPORAL_DETAIL,
+        "visual_specificity": VISUAL_SPECIFICITY,
+        "camera_detail": CAMERA_DETAIL,
+        "caption_style": CAPTION_STYLE,
+    }
+
+
+def find_videos(input_folder: Path) -> list[Path]:
+    iterator = input_folder.rglob("*") if RECURSIVE else input_folder.glob("*")
+    videos = [
+        path
+        for path in iterator
+        if path.is_file() and path.suffix.lower() in VIDEO_EXTENSIONS
+    ]
+    return sorted(videos)
+
+
+def output_path_for(video_path: Path, input_folder: Path, output_folder: Path) -> Path:
+    relative = video_path.relative_to(input_folder)
+    return (output_folder / relative).with_suffix(OUTPUT_EXTENSION)
+
+
+def write_text(path: Path, text: str) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(text.rstrip() + "\n", encoding="utf-8")
+
+
+def chunks(items: list[Path], size: int) -> list[list[Path]]:
+    return [items[index : index + size] for index in range(0, len(items), size)]
+
+
+def main() -> None:
+    input_folder = Path(INPUT_FOLDER)
+    output_folder = Path(OUTPUT_FOLDER)
+    if not input_folder.is_dir():
+        raise RuntimeError(f"INPUT_FOLDER does not exist or is not a directory: {input_folder}")
+
+    videos = find_videos(input_folder)
+    if not videos:
+        print(f"No videos found in {input_folder}")
+        return
+
+    processor = load_processor(MODEL_DIR)
+    llm = load_llm(
+        model_dir=MODEL_DIR,
+        max_model_len=MAX_MODEL_LEN,
+        max_num_seqs=MAX_NUM_SEQS,
+        gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
+        dtype=DTYPE,
+        enforce_eager=ENFORCE_EAGER,
+        enable_prefix_caching=ENABLE_PREFIX_CACHING,
+        trust_remote_code=TRUST_REMOTE_CODE,
+    )
+    params = sampling_params(
+        temperature=TEMPERATURE,
+        max_tokens=MAX_TOKENS,
+        top_p=TOP_P,
+        repetition_penalty=REPETITION_PENALTY,
+    )
+
+    done = 0
+    skipped = 0
+    failed = 0
+    settings = prompt_settings()
+
+    for batch in chunks(videos, BATCH_SIZE):
+        requests = []
+        request_videos = []
+        for video_path in batch:
+            output_path = output_path_for(video_path, input_folder, output_folder)
+            if output_path.exists() and not OVERWRITE_EXISTING:
+                skipped += 1
+                continue
+            try:
+                requests.append(
+                    build_vllm_request(
+                        processor=processor,
+                        model_dir=MODEL_DIR,
+                        video_path=video_path,
+                        num_frames=NUM_FRAMES,
+                        sampling_rate=SAMPLING_RATE,
+                        prompt_override=PROMPT_OVERRIDE,
+                        prompt_settings=settings,
+                    )
+                )
+                request_videos.append(video_path)
+            except Exception as exc:
+                failed += 1
+                error_path = output_path.with_suffix(ERROR_EXTENSION)
+                write_text(error_path, f"{type(exc).__name__}: {exc}")
+                print(f"FAILED preprocess {video_path}: {exc}")
+
+        if not requests:
+            continue
+
+        try:
+            outputs = llm.generate(requests, sampling_params=params)
+        except Exception as exc:
+            failed += len(request_videos)
+            for video_path in request_videos:
+                output_path = output_path_for(video_path, input_folder, output_folder)
+                error_path = output_path.with_suffix(ERROR_EXTENSION)
+                write_text(error_path, f"{type(exc).__name__}: {exc}")
+            print(f"FAILED batch starting {request_videos[0]}: {exc}")
+            continue
+
+        for video_path, output in zip(request_videos, outputs, strict=True):
+            output_path = output_path_for(video_path, input_folder, output_folder)
+            text = output.outputs[0].text if output.outputs else ""
+            write_text(output_path, text)
+            done += 1
+            print(f"WROTE {output_path}")
+
+    print(f"Done. wrote={done} skipped={skipped} failed={failed}")
+
+
+if __name__ == "__main__":
+    main()

.ipynb_checkpoints/caption_model_runtime-checkpoint.py

@@ -0,0 +1,771 @@
+"""Runtime helpers for the condensed Gemma 4 Parakeet caption model."""
+
+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}."
+    )

.ipynb_checkpoints/infer-checkpoint.py

@@ -0,0 +1,519 @@
+"""Run Gemma 4 31B Parakeet-hybrid caption inference on one video.
+
+Edit the variables below, then run on the Vast instance:
+
+    source /venv/main/bin/activate
+    cd /workspace/polished_model_v3
+    python infer.py
+
+This uses the condensed inference package:
+- video + text prompt are passed through the Gemma processor
+- audio is supplied as a sidecar WAV
+- processor-created audio features are replaced with Parakeet-native features
+- the language LoRA is already merged into model/
+- the trained audio projector is loaded from model/embed_audio.safetensors
+"""
+
+from __future__ import annotations
+
+import contextlib
+import io
+import json
+import logging
+import subprocess
+import tempfile
+import warnings
+from pathlib import Path
+from typing import Any
+
+import torch
+from transformers import AutoModelForMultimodalLM, AutoProcessor
+
+SCRIPT_DIR = Path(__file__).resolve().parent
+MODEL_ROOT = SCRIPT_DIR
+
+from caption_model_runtime import (
+    DEFAULT_PARAKEET_MODEL_ID,
+    CAPTION_LENGTH_LABELS,
+    CAPTION_SETTING_FIELD_CHOICES,
+    DEFAULT_CAPTION_SETTING_VALUES,
+    format_caption_settings_prompt,
+    gemma_core,
+    install_parakeet_audio_bridge,
+    load_state_file,
+    replace_batch_audio_features,
+)
+
+
+# Edit these.
+VIDEO_PATH = "/workspace/test7.mp4"  # Options: path to the video you want to caption.
+MODEL_PATH = str(MODEL_ROOT / "model")  # Options: merged model path.
+PROCESSOR_PATH = str(MODEL_ROOT / "processor")  # Options: processor path.
+AUDIO_PROJECTOR_PATH = str(MODEL_ROOT / "model" / "embed_audio.safetensors")  # Options: trained audio projector path.
+
+# Parakeet hybrid audio bridge settings. These should normally match the packaged model.
+PARAKEET_MODEL_ID = str(MODEL_ROOT / "parakeet")  # Options: "nvidia/parakeet-tdt-0.6b-v3" or compatible local/HF path.
+PARAKEET_BRIDGE_MODE = "tdt_token_embeddings_with_encoder_context"  # Options: "encoder", "tdt_tokens", "tdt_token_embeddings", "encoder_soft_tdt_token_embeddings", "tdt_token_embeddings_with_encoder_context".
+PARAKEET_NATIVE_FEATURES = True  # Options: True to replace Gemma audio features with Parakeet features, False for debugging only.
+PARAKEET_TDT_FILTER_BLANK_TOKENS = True  # Options: True or False.
+PARAKEET_TDT_FILTER_SPECIAL_TOKEN_IDS = True  # Options: True or False.
+PROJECTOR_INTERMEDIATE_SIZE = 4096  # Options: integer; the packaged model uses 4096.
+PROJECTOR_DROPOUT = 0.0  # Options: float; inference should normally be 0.0.
+HYBRID_ENCODER_GATE_INIT = 0.0  # Options: float; saved audio projector weights override the initial gate.
+
+# Prompt settings. Empty PROMPT_OVERRIDE builds the standard dynamic prompt.
+PROMPT_OVERRIDE = ""  # Options: "" or any full custom prompt string.
+CAPTION_SETTINGS_JSON_PATH = ""  # Options: "" or a JSON path under /workspace/dataset_jsons to override the settings below.
+CAPTION_LENGTH = "very large"  # Options: "very small", "small", "medium", "large", "very large".
+INCLUDE_WATERMARK_INFO = False  # Options: True or False.
+VULGARITY = "low"  # Options: "none", "low", "medium", "high".
+UNCERTAINTY = "low"  # Options: "none", "low", "medium", "high".
+CHARACTER_NAMES = "none"  # Options: "none", "ambiguous", "single", "multiple".
+FLUFF = "none"  # Options: "none", "low", "medium", "high".
+HAS_REPETITION = False  # Options: True or False.
+SPECULATION = "low"  # Options: "none", "low", "medium", "high".
+TEMPORAL_DETAIL = "medium"  # Options: "static", "low", "medium", "high".
+VISUAL_SPECIFICITY = "moderate"  # Options: "generic", "moderate", "detailed", "excessive".
+CAMERA_DETAIL = "medium"  # Options: "none", "low", "medium", "high".
+CAPTION_STYLE = "plain"  # Options: "plain", "verbose", "ornate", "robotic".
+HAS_THINKING = True  # Options: True to request thought JSON plus final caption, False to request only the final caption.
+
+# Media settings. Training used separate sidecar audio and random frame counts.
+NUM_FRAMES = 12  # Options: None for processor default, or an integer frame count.
+FPS = None  # Options: None for processor default, or a float such as 1.0.
+SAMPLING_RATE = 16_000  # Options: normally 16000.
+AUDIO_MAX_LENGTH_SAMPLES = 0  # Options: 0 keeps full audio; positive integer truncates Parakeet audio.
+MAX_AUDIO_SECONDS = 0.0  # Options: 0.0 keeps full audio; positive float caps extracted sidecar audio.
+
+# Generation settings.
+MAX_NEW_TOKENS = 1200  # Options: positive integer token cap.
+TEMPERATURE = 0.0  # Options: 0.0 for greedy decoding, >0.0 for sampling.
+TOP_P = 0.9  # Options: float in (0, 1], used only when TEMPERATURE > 0.
+REPETITION_PENALTY = 1.1  # Options: 1.0 disables the penalty, >1.0 penalizes repetition.
+PRINT_INPUT_STATS = False  # Options: True or False.
+QUIET_MODEL_LOAD = True  # Options: True hides noisy missing-key load reports; False prints full loader output.
+
+# Usually leave these alone.
+LOCAL_FILES_ONLY = True  # Options: True to use cached files only, False to allow downloads.
+DTYPE = torch.bfloat16  # Options: torch.bfloat16, torch.float16, torch.float32.
+DEVICE_MAP = "auto"  # Options: "auto", "cuda", or another Transformers device_map value.
+ATTN_IMPLEMENTATION = "sdpa"  # Options: "sdpa", "flash_attention_2", None.
+
+warnings.filterwarnings(
+    "ignore",
+    message=r"RNN module weights are not part of single contiguous chunk of memory.*",
+    category=UserWarning,
+)
+
+
+@contextlib.contextmanager
+def quiet_model_load() -> Any:
+    if not QUIET_MODEL_LOAD:
+        yield
+        return
+    load_report_logger = logging.getLogger("transformers.utils.loading_report")
+    old_level = load_report_logger.level
+    load_report_logger.setLevel(logging.ERROR)
+    patched_modules: list[tuple[Any, Any]] = []
+    try:
+        import transformers.modeling_utils as modeling_utils
+        import transformers.utils.loading_report as loading_report
+
+        original_report = modeling_utils.log_state_dict_report
+
+        def quiet_report(
+            model: Any,
+            pretrained_model_name_or_path: str,
+            ignore_mismatched_sizes: bool,
+            loading_info: Any,
+            logger: logging.Logger | None = None,
+        ) -> None:
+            has_fatal_issue = bool(getattr(loading_info, "error_msgs", None)) or bool(
+                getattr(loading_info, "conversion_errors", None)
+            )
+            if not ignore_mismatched_sizes and bool(getattr(loading_info, "mismatched_keys", None)):
+                has_fatal_issue = True
+            if has_fatal_issue:
+                original_report(
+                    model,
+                    pretrained_model_name_or_path,
+                    ignore_mismatched_sizes,
+                    loading_info,
+                    logger=logger,
+                )
+
+        for module in (loading_report, modeling_utils):
+            patched_modules.append((module, module.log_state_dict_report))
+            module.log_state_dict_report = quiet_report
+    except Exception:
+        patched_modules = []
+    try:
+        with contextlib.redirect_stdout(io.StringIO()), contextlib.redirect_stderr(io.StringIO()):
+            yield
+    finally:
+        for module, original in patched_modules:
+            module.log_state_dict_report = original
+        load_report_logger.setLevel(old_level)
+
+
+def load_parakeet_projector(model: torch.nn.Module) -> None:
+    state_path = Path(AUDIO_PROJECTOR_PATH)
+    state = load_state_file(state_path)
+    module = gemma_core(model).embed_audio
+    module.load_state_dict(state, strict=True)
+
+
+def video_has_audio_stream(video_path: Path) -> bool:
+    cmd = [
+        "ffprobe",
+        "-v",
+        "error",
+        "-select_streams",
+        "a:0",
+        "-show_entries",
+        "stream=index",
+        "-of",
+        "csv=p=0",
+        str(video_path),
+    ]
+    result = subprocess.run(cmd, check=True, capture_output=True, text=True)
+    return bool(result.stdout.strip())
+
+
+def probe_video_duration_seconds(video_path: Path) -> float:
+    cmd = [
+        "ffprobe",
+        "-v",
+        "error",
+        "-show_entries",
+        "format=duration",
+        "-of",
+        "default=noprint_wrappers=1:nokey=1",
+        str(video_path),
+    ]
+    result = subprocess.run(cmd, check=True, capture_output=True, text=True)
+    duration = float(result.stdout.strip())
+    if duration <= 0:
+        raise RuntimeError(f"Video duration must be positive: {video_path}")
+    return duration
+
+
+def extract_audio(video_path: Path, audio_path: Path) -> None:
+    cmd = [
+        "ffmpeg",
+        "-hide_banner",
+        "-loglevel",
+        "error",
+        "-y",
+        "-i",
+        str(video_path),
+        "-vn",
+        "-ac",
+        "1",
+        "-ar",
+        str(SAMPLING_RATE),
+    ]
+    if MAX_AUDIO_SECONDS > 0:
+        cmd.extend(["-t", f"{MAX_AUDIO_SECONDS:.6f}"])
+    cmd.extend(["-c:a", "pcm_s16le", str(audio_path)])
+    subprocess.run(cmd, check=True)
+
+
+def create_silent_audio(audio_path: Path, duration_seconds: float) -> None:
+    if MAX_AUDIO_SECONDS > 0:
+        duration_seconds = min(duration_seconds, MAX_AUDIO_SECONDS)
+    cmd = [
+        "ffmpeg",
+        "-hide_banner",
+        "-loglevel",
+        "error",
+        "-y",
+        "-f",
+        "lavfi",
+        "-i",
+        f"anullsrc=channel_layout=mono:sample_rate={SAMPLING_RATE}",
+        "-t",
+        f"{duration_seconds:.6f}",
+        "-ac",
+        "1",
+        "-ar",
+        str(SAMPLING_RATE),
+        "-c:a",
+        "pcm_s16le",
+        str(audio_path),
+    ]
+    subprocess.run(cmd, check=True)
+
+
+def prepare_sidecar_audio(video_path: Path, tmpdir: Path) -> tuple[Path, bool]:
+    audio_path = tmpdir / "sidecar_audio.wav"
+    if video_has_audio_stream(video_path):
+        extract_audio(video_path, audio_path)
+        return audio_path, True
+
+    duration_seconds = probe_video_duration_seconds(video_path)
+    print(f"input_video_has_audio=false; creating_silent_sidecar_audio duration_seconds={duration_seconds:.3f}", flush=True)
+    create_silent_audio(audio_path, duration_seconds)
+    return audio_path, False
+
+
+def bool_from_json(value: Any, field_name: str) -> bool:
+    if isinstance(value, bool):
+        return value
+    if isinstance(value, str):
+        lowered = value.strip().lower()
+        if lowered in {"1", "true", "yes", "y", "on"}:
+            return True
+        if lowered in {"0", "false", "no", "n", "off"}:
+            return False
+    raise ValueError(f"{field_name} must be boolean-like, got {value!r}")
+
+
+def load_prompt_settings_json() -> dict[str, Any]:
+    if not CAPTION_SETTINGS_JSON_PATH.strip():
+        return {}
+    path = Path(CAPTION_SETTINGS_JSON_PATH)
+    data = json.loads(path.read_text(encoding="utf-8"))
+    if not isinstance(data, dict):
+        raise ValueError(f"CAPTION_SETTINGS_JSON_PATH must point to a JSON object: {path}")
+    return data
+
+
+def build_prompt() -> str:
+    if PROMPT_OVERRIDE.strip():
+        return PROMPT_OVERRIDE.strip()
+
+    settings: dict[str, Any] = {
+        "caption_length": CAPTION_LENGTH,
+        "include_watermark_info": INCLUDE_WATERMARK_INFO,
+        **DEFAULT_CAPTION_SETTING_VALUES,
+        "vulgarity": VULGARITY,
+        "uncertainty": UNCERTAINTY,
+        "character_names": CHARACTER_NAMES,
+        "fluff": FLUFF,
+        "has_repetition": HAS_REPETITION,
+        "speculation": SPECULATION,
+        "temporal_detail": TEMPORAL_DETAIL,
+        "visual_specificity": VISUAL_SPECIFICITY,
+        "camera_detail": CAMERA_DETAIL,
+        "caption_style": CAPTION_STYLE,
+        "has_thinking": HAS_THINKING,
+    }
+    settings.update(load_prompt_settings_json())
+    settings["has_thinking"] = HAS_THINKING
+
+    settings["caption_length"] = str(settings["caption_length"]).strip().lower()
+    if settings["caption_length"] not in CAPTION_LENGTH_LABELS:
+        raise ValueError(f"caption_length must be one of {CAPTION_LENGTH_LABELS}, got {settings['caption_length']!r}")
+    settings["include_watermark_info"] = bool_from_json(settings["include_watermark_info"], "include_watermark_info")
+    settings["has_repetition"] = bool_from_json(settings["has_repetition"], "has_repetition")
+    settings["has_thinking"] = bool_from_json(settings["has_thinking"], "has_thinking")
+    for field_name, allowed in CAPTION_SETTING_FIELD_CHOICES.items():
+        value = str(settings[field_name]).strip().lower()
+        if value not in allowed:
+            raise ValueError(f"{field_name} must be one of {allowed}, got {value!r}")
+        settings[field_name] = value
+
+    return format_caption_settings_prompt(settings)
+
+
+def build_messages(video_path: Path, audio_path: Path, prompt: str) -> list[dict[str, Any]]:
+    return [
+        {
+            "role": "user",
+            "content": [
+                {"type": "video", "path": str(video_path)},
+                {"type": "text", "text": prompt},
+                {"type": "audio", "path": str(audio_path)},
+            ],
+        },
+        {"role": "assistant", "content": [{"type": "text", "text": ""}]},
+    ]
+
+
+def trim_empty_assistant_terminator(inputs: dict[str, torch.Tensor], processor: Any) -> dict[str, torch.Tensor]:
+    eos_tail = processor.tokenizer.encode("<turn|>\n", add_special_tokens=False)
+    if not eos_tail:
+        return inputs
+
+    tail_len = len(eos_tail)
+    input_ids = inputs["input_ids"][0]
+    if input_ids[-tail_len:].tolist() != eos_tail:
+        return inputs
+
+    trimmed = {}
+    for key, value in inputs.items():
+        if isinstance(value, torch.Tensor) and value.ndim >= 2 and value.shape[1] == input_ids.shape[0]:
+            trimmed[key] = value[:, :-tail_len]
+        else:
+            trimmed[key] = value
+    return trimmed
+
+
+def tensor_stats(tensor: torch.Tensor | None, mask: torch.Tensor | None = None) -> dict[str, Any]:
+    if tensor is None:
+        return {"present": False}
+    stats_tensor = tensor.detach().float().cpu()
+    result: dict[str, Any] = {
+        "present": True,
+        "shape": list(tensor.shape),
+        "mean": round(float(stats_tensor.mean().item()), 8),
+        "std": round(float(stats_tensor.std().item()), 8),
+        "abs_mean": round(float(stats_tensor.abs().mean().item()), 8),
+    }
+    if mask is not None:
+        result["mask_shape"] = list(mask.shape)
+        result["mask_sum"] = int(mask.detach().cpu().sum().item())
+    return result
+
+
+def print_input_stats(inputs: dict[str, torch.Tensor], label: str) -> None:
+    stats = {
+        "input_ids_shape": list(inputs["input_ids"].shape),
+        "input_features": tensor_stats(inputs.get("input_features"), inputs.get("input_features_mask")),
+        "keys": sorted(inputs.keys()),
+    }
+    print(f"{label}=" + json.dumps(stats, sort_keys=True), flush=True)
+
+
+def move_inputs_to_model_device(inputs: dict[str, Any], model: torch.nn.Module) -> dict[str, Any]:
+    device = getattr(model, "device", None)
+    if device is None:
+        try:
+            device = next(model.parameters()).device
+        except StopIteration:
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    moved = {}
+    for key, value in inputs.items():
+        moved[key] = value.to(device) if isinstance(value, torch.Tensor) else value
+    return moved
+
+
+def generation_kwargs() -> dict[str, Any]:
+    kwargs: dict[str, Any] = {
+        "max_new_tokens": MAX_NEW_TOKENS,
+        "do_sample": TEMPERATURE > 0,
+        "repetition_penalty": REPETITION_PENALTY,
+        "use_cache": True,
+    }
+    if TEMPERATURE > 0:
+        kwargs["temperature"] = TEMPERATURE
+        kwargs["top_p"] = TOP_P
+    return kwargs
+
+
+def prepare_inputs(
+    processor: Any,
+    parakeet_processor: Any,
+    video_path: Path,
+    audio_path: Path,
+    prompt: str,
+) -> dict[str, torch.Tensor]:
+    processor_kwargs: dict[str, Any] = {
+        "padding": True,
+        "truncation": False,
+        "sampling_rate": SAMPLING_RATE,
+    }
+    if NUM_FRAMES is not None:
+        processor_kwargs["num_frames"] = NUM_FRAMES
+    if FPS is not None:
+        processor_kwargs["fps"] = FPS
+
+    inputs = processor.apply_chat_template(
+        build_messages(video_path, audio_path, prompt),
+        tokenize=True,
+        return_dict=True,
+        return_tensors="pt",
+        load_audio_from_video=False,
+        processor_kwargs=processor_kwargs,
+    )
+    inputs = trim_empty_assistant_terminator(inputs, processor)
+    if PRINT_INPUT_STATS:
+        print_input_stats(inputs, "input_stats_before_parakeet_swap")
+    if PARAKEET_NATIVE_FEATURES:
+        replace_batch_audio_features(
+            inputs,
+            audio_paths=[str(audio_path)],
+            parakeet_processor=parakeet_processor,
+            sampling_rate=SAMPLING_RATE,
+            max_length_samples=AUDIO_MAX_LENGTH_SAMPLES,
+        )
+        if PRINT_INPUT_STATS:
+            print_input_stats(inputs, "input_stats_after_parakeet_swap")
+    return inputs
+
+
+def load_model(processor_path: Path) -> tuple[Any, torch.nn.Module, Any]:
+    processor = AutoProcessor.from_pretrained(str(processor_path), local_files_only=LOCAL_FILES_ONLY)
+    parakeet_processor = AutoProcessor.from_pretrained(PARAKEET_MODEL_ID, local_files_only=LOCAL_FILES_ONLY)
+
+    model_kwargs: dict[str, Any] = {
+        "local_files_only": LOCAL_FILES_ONLY,
+        "dtype": DTYPE,
+        "low_cpu_mem_usage": True,
+        "device_map": DEVICE_MAP,
+    }
+    if ATTN_IMPLEMENTATION:
+        model_kwargs["attn_implementation"] = ATTN_IMPLEMENTATION
+
+    with quiet_model_load():
+        model = AutoModelForMultimodalLM.from_pretrained(MODEL_PATH, **model_kwargs)
+    bridge_args = type(
+        "BridgeArgs",
+        (),
+        {
+            "parakeet_model_id": PARAKEET_MODEL_ID,
+            "parakeet_bridge_mode": PARAKEET_BRIDGE_MODE,
+            "local_files_only": LOCAL_FILES_ONLY,
+            "projector_intermediate_size": PROJECTOR_INTERMEDIATE_SIZE,
+            "projector_dropout": PROJECTOR_DROPOUT,
+            "hybrid_encoder_gate_init": HYBRID_ENCODER_GATE_INIT,
+            "parakeet_tdt_filter_blank_tokens": PARAKEET_TDT_FILTER_BLANK_TOKENS,
+            "parakeet_tdt_filter_special_token_ids": PARAKEET_TDT_FILTER_SPECIAL_TOKEN_IDS,
+        },
+    )()
+    with quiet_model_load():
+        install_parakeet_audio_bridge(model, bridge_args)
+    gemma_core(model)
+    load_parakeet_projector(model)
+
+    model.eval()
+    return processor, model, parakeet_processor
+
+
+def main() -> None:
+    video_path = Path(VIDEO_PATH)
+    model_path = Path(MODEL_PATH)
+    processor_path = Path(PROCESSOR_PATH)
+    audio_projector_path = Path(AUDIO_PROJECTOR_PATH)
+    parakeet_path = Path(PARAKEET_MODEL_ID)
+
+    for label, path in (
+        ("VIDEO_PATH", video_path),
+        ("MODEL_PATH", model_path),
+        ("processor", processor_path),
+        ("audio_projector", audio_projector_path),
+        ("parakeet", parakeet_path),
+    ):
+        if not path.exists():
+            raise FileNotFoundError(f"{label} does not exist: {path}")
+
+    prompt = build_prompt()
+    processor, model, parakeet_processor = load_model(processor_path)
+    with tempfile.TemporaryDirectory(prefix="gemma4_caption_inference_") as tmpdir_raw:
+        tmpdir = Path(tmpdir_raw)
+        audio_path, _had_audio = prepare_sidecar_audio(video_path, tmpdir)
+        inputs = prepare_inputs(processor, parakeet_processor, video_path, audio_path, prompt)
+        moved_inputs = move_inputs_to_model_device(inputs, model)
+        input_len = moved_inputs["input_ids"].shape[-1]
+        with torch.inference_mode():
+            output_ids = model.generate(**moved_inputs, **generation_kwargs())
+
+    new_tokens = output_ids[0, input_len:]
+    response = processor.decode(new_tokens, skip_special_tokens=True).strip()
+    print(response, flush=True)
+
+
+if __name__ == "__main__":
+    main()

.ipynb_checkpoints/requirements-checkpoint.txt

@@ -0,0 +1,19 @@
+# Inference dependencies for polished_model_v3.
+# ffmpeg/ffprobe are required as system binaries for video/audio extraction.
+
+--extra-index-url https://download.pytorch.org/whl/cu130
+
+accelerate==1.14.0
+av==17.1.0
+librosa==0.11.0
+numpy==2.4.4
+opencv-python-headless>=4.10.0
+safetensors==0.8.0
+scipy==1.18.0
+soundfile==0.14.0
+torch==2.12.1+cu130
+torchcodec==0.14.0
+torchvision==0.27.1+cu130
+transformers==5.12.1
+
+-e ./vllm

.ipynb_checkpoints/run_inference-checkpoint.sh

@@ -0,0 +1,5 @@
+set -euo pipefail
+
+cd "$(dirname "$0")"
+source /venv/main/bin/activate
+python infer.py

.ipynb_checkpoints/setup_vllm-checkpoint.sh

@@ -0,0 +1,22 @@
+set -euo pipefail
+
+cd "$(dirname "$0")"
+
+if [ ! -d vllm ]; then
+  if [ ! -f vllm.zip ]; then
+    echo "Missing vllm/ and vllm.zip" >&2
+    exit 1
+  fi
+  unzip -q vllm.zip
+fi
+
+python3 -m venv .venv
+source .venv/bin/activate
+
+pip install --upgrade pip setuptools wheel
+pip install -r requirements.txt
+
+python - <<'PY'
+import vllm
+print("vLLM import OK:", getattr(vllm, "__version__", "unknown"))
+PY

.ipynb_checkpoints/vllm_caption_runtime-checkpoint.py

@@ -0,0 +1,254 @@
+"""Shared helpers for the polished vLLM caption examples."""
+
+from __future__ import annotations
+
+import subprocess
+import sys
+from pathlib import Path
+from typing import Any
+
+import cv2
+import numpy as np
+from transformers import AutoProcessor
+from vllm import LLM, SamplingParams
+
+CAPTION_LENGTH_LABELS = ("very small", "small", "medium", "large", "very large")
+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"),
+}
+
+
+def bool_from_value(value: Any, field_name: str) -> bool:
+    if isinstance(value, bool):
+        return value
+    if isinstance(value, str):
+        lowered = value.strip().lower()
+        if lowered in {"1", "true", "yes", "y", "on"}:
+            return True
+        if lowered in {"0", "false", "no", "n", "off"}:
+            return False
+    raise ValueError(f"{field_name} must be boolean-like, got {value!r}")
+
+
+def format_caption_settings_prompt(settings: dict[str, Any]) -> str:
+    caption_length = str(settings["caption_length"]).strip().lower()
+    if caption_length not in CAPTION_LENGTH_LABELS:
+        raise ValueError(f"caption_length must be one of {CAPTION_LENGTH_LABELS}, got {caption_length!r}")
+
+    include_watermark_info = bool_from_value(
+        settings["include_watermark_info"], "include_watermark_info"
+    )
+    has_repetition = bool_from_value(settings["has_repetition"], "has_repetition")
+    has_thinking = bool_from_value(settings["has_thinking"], "has_thinking")
+
+    normalized: dict[str, str] = {}
+    for field_name, choices in CAPTION_SETTING_FIELD_CHOICES.items():
+        value = str(settings[field_name]).strip().lower()
+        if value not in choices:
+            raise ValueError(f"{field_name} must be one of {choices}, got {value!r}")
+        normalized[field_name] = value
+
+    watermark_instruction = (
+        "Include watermark info." if include_watermark_info else "Do not include watermark info."
+    )
+    thinking_instruction = (
+        "Output thought JSON before the final caption."
+        if has_thinking
+        else "Do not output thought JSON; output only the caption."
+    )
+    setting_text = (
+        f"vulgarity={normalized['vulgarity']}; "
+        f"uncertainty={normalized['uncertainty']}; "
+        f"character_names={normalized['character_names']}; "
+        f"fluff={normalized['fluff']}; "
+        f"has_repetition={str(has_repetition).lower()}; "
+        f"has_thinking={str(has_thinking).lower()}; "
+        f"speculation={normalized['speculation']}; "
+        f"temporal_detail={normalized['temporal_detail']}; "
+        f"visual_specificity={normalized['visual_specificity']}; "
+        f"camera_detail={normalized['camera_detail']}; "
+        f"caption_style={normalized['caption_style']}"
+    )
+    return (
+        f"Write a {caption_length} caption for this clip using both the visuals and the audio. "
+        f"{watermark_instruction} {thinking_instruction} Match these caption settings: {setting_text}."
+    )
+
+
+def load_video_frames(path: str | Path, num_frames: int) -> tuple[np.ndarray, dict[str, Any]]:
+    video_path = str(path)
+    cap = cv2.VideoCapture(video_path)
+    if not cap.isOpened():
+        raise RuntimeError(f"Could not open video: {video_path}")
+
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0)
+    fps = float(cap.get(cv2.CAP_PROP_FPS) or 24.0)
+    if total_frames <= 0:
+        cap.release()
+        raise RuntimeError(f"Could not determine frame count: {video_path}")
+
+    indices = np.linspace(0, max(0, total_frames - 1), num_frames).round().astype(int)
+    frames = []
+    for idx in indices:
+        cap.set(cv2.CAP_PROP_POS_FRAMES, int(idx))
+        ok, frame_bgr = cap.read()
+        if ok:
+            frames.append(cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB))
+    cap.release()
+
+    if not frames:
+        raise RuntimeError(f"Could not read frames: {video_path}")
+
+    metadata = {
+        "fps": fps,
+        "duration": total_frames / fps if fps > 0 else 0.0,
+        "total_num_frames": total_frames,
+        "frames_indices": [int(x) for x in indices[: len(frames)]],
+        "video_backend": "opencv",
+        "do_sample_frames": False,
+    }
+    return np.stack(frames, axis=0), metadata
+
+
+def load_audio(path: str | Path, sampling_rate: int) -> tuple[np.ndarray, int]:
+    video_path = str(path)
+    cmd = [
+        "ffmpeg",
+        "-hide_banner",
+        "-loglevel",
+        "error",
+        "-i",
+        video_path,
+        "-vn",
+        "-ac",
+        "1",
+        "-ar",
+        str(sampling_rate),
+        "-f",
+        "f32le",
+        "-",
+    ]
+    result = subprocess.run(cmd, check=False, capture_output=True)
+    if result.returncode == 0 and result.stdout:
+        return np.frombuffer(result.stdout, dtype=np.float32), sampling_rate
+
+    frames, metadata = load_video_frames(video_path, 2)
+    del frames
+    duration = max(1.0, float(metadata.get("duration") or 1.0))
+    return np.zeros(max(1, int(duration * sampling_rate)), dtype=np.float32), sampling_rate
+
+
+def build_prompt(
+    *,
+    processor: AutoProcessor,
+    model_dir: str | Path,
+    video_path: str | Path,
+    prompt_override: str,
+    prompt_settings: dict[str, Any],
+) -> str:
+    del model_dir
+    prompt_text = prompt_override.strip() or format_caption_settings_prompt(prompt_settings)
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "video", "path": str(video_path)},
+                {"type": "text", "text": prompt_text},
+                {"type": "audio", "path": "/tmp/polished_model_v3_audio.wav"},
+            ],
+        },
+        {"role": "assistant", "content": [{"type": "text", "text": ""}]},
+    ]
+    prompt = processor.apply_chat_template(messages, tokenize=False)
+    assistant_tail = "<turn|>\n"
+    if prompt.endswith(assistant_tail):
+        prompt = prompt[: -len(assistant_tail)]
+    return prompt
+
+
+def build_vllm_request(
+    *,
+    processor: AutoProcessor,
+    model_dir: str | Path,
+    video_path: str | Path,
+    num_frames: int,
+    sampling_rate: int,
+    prompt_override: str,
+    prompt_settings: dict[str, Any],
+) -> dict[str, Any]:
+    video, video_metadata = load_video_frames(video_path, num_frames)
+    audio, sr = load_audio(video_path, sampling_rate)
+    prompt = build_prompt(
+        processor=processor,
+        model_dir=model_dir,
+        video_path=video_path,
+        prompt_override=prompt_override,
+        prompt_settings=prompt_settings,
+    )
+    return {
+        "prompt": prompt,
+        "multi_modal_data": {
+            "video": [(video, video_metadata)],
+            "audio": (audio, sr),
+        },
+    }
+
+
+def load_processor(model_dir: str | Path) -> AutoProcessor:
+    return AutoProcessor.from_pretrained(str(model_dir), local_files_only=True)
+
+
+def load_llm(
+    *,
+    model_dir: str | Path,
+    max_model_len: int,
+    max_num_seqs: int,
+    gpu_memory_utilization: float,
+    dtype: str,
+    enforce_eager: bool,
+    enable_prefix_caching: bool,
+    trust_remote_code: bool,
+) -> LLM:
+    return LLM(
+        model=str(model_dir),
+        tokenizer=str(model_dir),
+        max_model_len=max_model_len,
+        max_num_seqs=max_num_seqs,
+        gpu_memory_utilization=gpu_memory_utilization,
+        limit_mm_per_prompt={"video": 1, "audio": 1, "image": 0},
+        enforce_eager=enforce_eager,
+        enable_prefix_caching=enable_prefix_caching,
+        trust_remote_code=trust_remote_code,
+        dtype=dtype,
+    )
+
+
+def sampling_params(
+    *,
+    temperature: float,
+    max_tokens: int,
+    top_p: float,
+    repetition_penalty: float,
+) -> SamplingParams:
+    kwargs: dict[str, Any] = {
+        "temperature": temperature,
+        "max_tokens": max_tokens,
+        "repetition_penalty": repetition_penalty,
+    }
+    if temperature > 0:
+        kwargs["top_p"] = top_p
+    return SamplingParams(**kwargs)
+
+
+def ensure_local_vllm_source(script_dir: Path) -> None:
+    local_vllm = script_dir / "vllm"
+    if local_vllm.is_dir():
+        sys.path.insert(0, str(local_vllm))

README.md

@@ -1,3 +1,165 @@
 ---
 license: apache-2.0
+language:
+  - en
+library_name: transformers
+pipeline_tag: image-to-text
+tags:
+  - video-captioning
+  - multimodal
+  - gemma
+  - parakeet
 ---
+
+# Gemma 4 31B Parakeet Video Caption Model
+
+This repository packages a Gemma 4 31B video captioning model with a Parakeet-based audio bridge and a patched vLLM runtime. It is intended for local or server-side inference on CUDA GPUs.
+
+The model produces raw caption responses. When `HAS_THINKING = True`, the response includes the model's thought JSON followed by the final caption. When `HAS_THINKING = False`, the prompt asks for the final caption only.
+
+> This package uses custom runtime code. Stock vLLM does not know how to load this Parakeet audio bridge unless you unzip and use the bundled patched `vllm/` folder.
+
+## Repository Contents
+
+- `model/`: Gemma 4 31B caption model weights. The caption LoRA is already merged.
+- `model/embed_audio.safetensors`: trained Parakeet-to-Gemma audio projector.
+- `model/processor_config.json`, `model/tokenizer.json`, `model/tokenizer_config.json`, `model/chat_template.jinja`: processor/tokenizer files copied into `model/` so vLLM can point directly at `model/`.
+- `processor/`: original processor/tokenizer files, kept for compatibility.
+- `parakeet/`: bundled NVIDIA Parakeet TDT 0.6B audio model.
+- `vllm.zip`: zipped patched vLLM source tree with Gemma4 Parakeet bridge support. Unzip this before installing.
+- `example_run.py`: edit variables at the top, point it at one video, print the raw model output.
+- `caption_folder.py`: edit variables at the top, point it at an input folder and output folder, write one raw output per video.
+- `vllm_caption_runtime.py`: shared helper code for the vLLM examples.
+- `infer.py` and `caption_model_runtime.py`: Transformers inference path from v2, kept for compatibility.
+- `requirements.txt`: inference dependencies plus editable install of the unzipped bundled `vllm/`.
+- `setup_vllm.sh`: convenience setup script.
+- `assets/decode_throughput.png`: vLLM decode-throughput stress test graph.
+
+## Installation
+
+You need `ffmpeg`, `ffprobe`, and `unzip` available on `PATH`.
+
+```bash
+cd /path/to/this/repo
+unzip -q vllm.zip
+python3 -m venv .venv
+source .venv/bin/activate
+pip install --upgrade pip setuptools wheel
+pip install -r requirements.txt
+```
+
+Or run the helper:
+
+```bash
+cd /path/to/this/repo
+./setup_vllm.sh
+source .venv/bin/activate
+```
+
+## Single-Video Inference
+
+Edit `VIDEO_PATH` in `example_run.py`:
+
+```python
+VIDEO_PATH = "/path/to/video.mp4"
+```
+
+Then run:
+
+```bash
+cd /path/to/this/repo
+source .venv/bin/activate
+python example_run.py
+```
+
+The script prints the raw model output.
+
+## Folder Captioning
+
+Edit these variables in `caption_folder.py`:
+
+```python
+INPUT_FOLDER = "/path/to/videos"
+OUTPUT_FOLDER = "/path/to/captions"
+```
+
+Then run:
+
+```bash
+cd /path/to/this/repo
+source .venv/bin/activate
+python caption_folder.py
+```
+
+The script mirrors the input folder structure into `OUTPUT_FOLDER` and writes one `.txt` file per video. It loads the model once and batches videos using `BATCH_SIZE`.
+
+## Prompt Settings
+
+Both `example_run.py` and `caption_folder.py` are variable-based. They do not use command-line args.
+
+Prompt settings are grouped under `Prompt Settings`, including:
+
+- `HAS_THINKING`
+- `CAPTION_LENGTH`
+- `INCLUDE_WATERMARK_INFO`
+- `VULGARITY`
+- `UNCERTAINTY`
+- `CHARACTER_NAMES`
+- `FLUFF`
+- `HAS_REPETITION`
+- `SPECULATION`
+- `TEMPORAL_DETAIL`
+- `VISUAL_SPECIFICITY`
+- `CAMERA_DETAIL`
+- `CAPTION_STYLE`
+
+Runtime and generation settings are grouped under `vLLM / Generation Hyperparameters`, including:
+
+- `NUM_FRAMES`
+- `MAX_MODEL_LEN`
+- `MAX_NUM_SEQS`
+- `BATCH_SIZE` in `caption_folder.py`
+- `GPU_MEMORY_UTILIZATION`
+- `MAX_TOKENS`
+- `TEMPERATURE`
+- `TOP_P`
+- `REPETITION_PENALTY`
+
+## Performance
+
+This graph shows measured vLLM decode throughput on an RTX PRO 6000 using 12 video frames, `max_tokens = 96`, `max_model_len = 4096`, prefix cache disabled, and unique requests:
+
+![Decode Throughput](assets/decode_throughput.png)
+
+For short caption requests on the RTX PRO 6000 test machine, `MAX_NUM_SEQS = 20` was the best throughput point. vLLM reported about 24k KV-cache tokens, roughly 5.9 full 4096-token requests, so long-context traffic should use lower concurrency.
+
+## Architecture
+
+The base language and vision model is Gemma 4 31B. The caption LoRA has already been merged into the Gemma weights in `model/`, so there is no separate LoRA adapter to load.
+
+Audio is handled by a Parakeet bridge:
+
+1. The prompt uses Gemma's multimodal chat format with video and audio placeholders.
+2. vLLM's Gemma4 multimodal processor builds the text/video/audio request structure.
+3. The patched runtime detects this package layout by finding `model/embed_audio.safetensors` and sibling folder `parakeet/`.
+4. Instead of using Gemma's original audio tower, the runtime extracts waveform audio and converts it with the bundled Parakeet processor.
+5. The frozen Parakeet TDT model produces audio representations. The packaged bridge mode is `tdt_token_embeddings_with_encoder_context`, combining Parakeet TDT token embeddings with encoder context.
+6. The trained projector in `model/embed_audio.safetensors` maps the Parakeet-side audio stream into Gemma's 5376-dimensional hidden space.
+7. Those projected audio embeddings are inserted into the Gemma prompt where the audio soft tokens live.
+8. Gemma generates the raw thought/caption output using video frames, text prompt settings, and projected Parakeet audio features.
+
+In short: Parakeet is frozen and used as the audio recognizer/encoder, the learned bridge maps Parakeet audio features into Gemma's embedding space, and the merged Gemma caption model writes the final output.
+
+## Limitations
+
+- This is a custom multimodal package, not a stock Transformers pipeline.
+- It requires the bundled patched vLLM runtime for the vLLM examples.
+- Very long videos or long-context requests should use lower concurrency than short caption jobs.
+- Audio captioning depends on the Parakeet bridge and may still miss unclear, quiet, overlapping, or non-speech sounds.
+- The scripts intentionally print or save the raw model output rather than parsing it.
+
+## Notes
+
+- Videos with no audio are handled by inserting a silent audio array so the model still receives the expected audio input shape.
+- Keep `model/`, `parakeet/`, `processor/`, `vllm.zip`, the Python scripts, and `assets/` together when moving or uploading this repository.
+- If uploading to Hugging Face, upload this README as `README.md` at the repository root.

caption_folder.py

@@ -0,0 +1,206 @@
+"""Caption every video in a folder with the packaged vLLM model.
+
+Edit the variables below, then run:
+
+    cd /workspace/polished_model_v3
+    source .venv/bin/activate
+    python caption_folder.py
+"""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+SCRIPT_DIR = Path(__file__).resolve().parent
+sys.path.insert(0, str(SCRIPT_DIR / "vllm"))
+
+from vllm_caption_runtime import build_vllm_request, load_llm, load_processor, sampling_params
+
+
+# =============================================================================
+# Paths
+# =============================================================================
+
+INPUT_FOLDER = "/workspace/videos_to_caption"
+OUTPUT_FOLDER = "/workspace/captions_out"
+MODEL_DIR = str(SCRIPT_DIR / "model")
+
+RECURSIVE = True
+OVERWRITE_EXISTING = False
+OUTPUT_EXTENSION = ".txt"
+ERROR_EXTENSION = ".error.txt"
+VIDEO_EXTENSIONS = (".mp4", ".mov", ".mkv", ".webm", ".avi", ".m4v")
+
+
+# =============================================================================
+# Prompt Settings
+# =============================================================================
+
+PROMPT_OVERRIDE = ""
+CAPTION_LENGTH = "very large"
+INCLUDE_WATERMARK_INFO = False
+HAS_THINKING = True
+
+VULGARITY = "low"
+UNCERTAINTY = "low"
+CHARACTER_NAMES = "none"
+FLUFF = "none"
+HAS_REPETITION = False
+SPECULATION = "low"
+TEMPORAL_DETAIL = "medium"
+VISUAL_SPECIFICITY = "moderate"
+CAMERA_DETAIL = "medium"
+CAPTION_STYLE = "plain"
+
+
+# =============================================================================
+# vLLM / Generation Hyperparameters
+# =============================================================================
+
+NUM_FRAMES = 12
+SAMPLING_RATE = 16_000
+MAX_MODEL_LEN = 4096
+MAX_NUM_SEQS = 20
+BATCH_SIZE = 20
+GPU_MEMORY_UTILIZATION = 0.88
+DTYPE = "bfloat16"
+ENFORCE_EAGER = False
+ENABLE_PREFIX_CACHING = False
+TRUST_REMOTE_CODE = False
+
+MAX_TOKENS = 1200
+TEMPERATURE = 0.0
+TOP_P = 0.9
+REPETITION_PENALTY = 1.1
+
+
+def prompt_settings() -> dict[str, object]:
+    return {
+        "caption_length": CAPTION_LENGTH,
+        "include_watermark_info": INCLUDE_WATERMARK_INFO,
+        "has_thinking": HAS_THINKING,
+        "vulgarity": VULGARITY,
+        "uncertainty": UNCERTAINTY,
+        "character_names": CHARACTER_NAMES,
+        "fluff": FLUFF,
+        "has_repetition": HAS_REPETITION,
+        "speculation": SPECULATION,
+        "temporal_detail": TEMPORAL_DETAIL,
+        "visual_specificity": VISUAL_SPECIFICITY,
+        "camera_detail": CAMERA_DETAIL,
+        "caption_style": CAPTION_STYLE,
+    }
+
+
+def find_videos(input_folder: Path) -> list[Path]:
+    iterator = input_folder.rglob("*") if RECURSIVE else input_folder.glob("*")
+    videos = [
+        path
+        for path in iterator
+        if path.is_file() and path.suffix.lower() in VIDEO_EXTENSIONS
+    ]
+    return sorted(videos)
+
+
+def output_path_for(video_path: Path, input_folder: Path, output_folder: Path) -> Path:
+    relative = video_path.relative_to(input_folder)
+    return (output_folder / relative).with_suffix(OUTPUT_EXTENSION)
+
+
+def write_text(path: Path, text: str) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    path.write_text(text.rstrip() + "\n", encoding="utf-8")
+
+
+def chunks(items: list[Path], size: int) -> list[list[Path]]:
+    return [items[index : index + size] for index in range(0, len(items), size)]
+
+
+def main() -> None:
+    input_folder = Path(INPUT_FOLDER)
+    output_folder = Path(OUTPUT_FOLDER)
+    if not input_folder.is_dir():
+        raise RuntimeError(f"INPUT_FOLDER does not exist or is not a directory: {input_folder}")
+
+    videos = find_videos(input_folder)
+    if not videos:
+        print(f"No videos found in {input_folder}")
+        return
+
+    processor = load_processor(MODEL_DIR)
+    llm = load_llm(
+        model_dir=MODEL_DIR,
+        max_model_len=MAX_MODEL_LEN,
+        max_num_seqs=MAX_NUM_SEQS,
+        gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
+        dtype=DTYPE,
+        enforce_eager=ENFORCE_EAGER,
+        enable_prefix_caching=ENABLE_PREFIX_CACHING,
+        trust_remote_code=TRUST_REMOTE_CODE,
+    )
+    params = sampling_params(
+        temperature=TEMPERATURE,
+        max_tokens=MAX_TOKENS,
+        top_p=TOP_P,
+        repetition_penalty=REPETITION_PENALTY,
+    )
+
+    done = 0
+    skipped = 0
+    failed = 0
+    settings = prompt_settings()
+
+    for batch in chunks(videos, BATCH_SIZE):
+        requests = []
+        request_videos = []
+        for video_path in batch:
+            output_path = output_path_for(video_path, input_folder, output_folder)
+            if output_path.exists() and not OVERWRITE_EXISTING:
+                skipped += 1
+                continue
+            try:
+                requests.append(
+                    build_vllm_request(
+                        processor=processor,
+                        model_dir=MODEL_DIR,
+                        video_path=video_path,
+                        num_frames=NUM_FRAMES,
+                        sampling_rate=SAMPLING_RATE,
+                        prompt_override=PROMPT_OVERRIDE,
+                        prompt_settings=settings,
+                    )
+                )
+                request_videos.append(video_path)
+            except Exception as exc:
+                failed += 1
+                error_path = output_path.with_suffix(ERROR_EXTENSION)
+                write_text(error_path, f"{type(exc).__name__}: {exc}")
+                print(f"FAILED preprocess {video_path}: {exc}")
+
+        if not requests:
+            continue
+
+        try:
+            outputs = llm.generate(requests, sampling_params=params)
+        except Exception as exc:
+            failed += len(request_videos)
+            for video_path in request_videos:
+                output_path = output_path_for(video_path, input_folder, output_folder)
+                error_path = output_path.with_suffix(ERROR_EXTENSION)
+                write_text(error_path, f"{type(exc).__name__}: {exc}")
+            print(f"FAILED batch starting {request_videos[0]}: {exc}")
+            continue
+
+        for video_path, output in zip(request_videos, outputs, strict=True):
+            output_path = output_path_for(video_path, input_folder, output_folder)
+            text = output.outputs[0].text if output.outputs else ""
+            write_text(output_path, text)
+            done += 1
+            print(f"WROTE {output_path}")
+
+    print(f"Done. wrote={done} skipped={skipped} failed={failed}")
+
+
+if __name__ == "__main__":
+    main()

caption_model_runtime.py

@@ -0,0 +1,771 @@
+"""Runtime helpers for the condensed Gemma 4 Parakeet caption model."""
+
+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}."
+    )

example_run.py

@@ -0,0 +1,123 @@
+"""Run the packaged vLLM caption model on one video.
+
+Edit the variables below, then run:
+
+    cd /workspace/polished_model_v3
+    source .venv/bin/activate
+    python example_run.py
+"""
+
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+SCRIPT_DIR = Path(__file__).resolve().parent
+sys.path.insert(0, str(SCRIPT_DIR / "vllm"))
+
+from vllm_caption_runtime import build_vllm_request, load_llm, load_processor, sampling_params
+
+
+# =============================================================================
+# Paths
+# =============================================================================
+
+VIDEO_PATH = "/workspace/test7.mp4"
+MODEL_DIR = str(SCRIPT_DIR / "model")
+
+
+# =============================================================================
+# Prompt Settings
+# =============================================================================
+
+PROMPT_OVERRIDE = ""
+CAPTION_LENGTH = "very large"
+INCLUDE_WATERMARK_INFO = False
+HAS_THINKING = True
+
+VULGARITY = "low"
+UNCERTAINTY = "low"
+CHARACTER_NAMES = "none"
+FLUFF = "none"
+HAS_REPETITION = False
+SPECULATION = "low"
+TEMPORAL_DETAIL = "medium"
+VISUAL_SPECIFICITY = "moderate"
+CAMERA_DETAIL = "medium"
+CAPTION_STYLE = "plain"
+
+
+# =============================================================================
+# vLLM / Generation Hyperparameters
+# =============================================================================
+
+NUM_FRAMES = 12
+SAMPLING_RATE = 16_000
+MAX_MODEL_LEN = 4096
+MAX_NUM_SEQS = 1
+GPU_MEMORY_UTILIZATION = 0.88
+DTYPE = "bfloat16"
+ENFORCE_EAGER = False
+ENABLE_PREFIX_CACHING = False
+TRUST_REMOTE_CODE = False
+
+MAX_TOKENS = 1200
+TEMPERATURE = 0.0
+TOP_P = 0.9
+REPETITION_PENALTY = 1.1
+
+
+def prompt_settings() -> dict[str, object]:
+    return {
+        "caption_length": CAPTION_LENGTH,
+        "include_watermark_info": INCLUDE_WATERMARK_INFO,
+        "has_thinking": HAS_THINKING,
+        "vulgarity": VULGARITY,
+        "uncertainty": UNCERTAINTY,
+        "character_names": CHARACTER_NAMES,
+        "fluff": FLUFF,
+        "has_repetition": HAS_REPETITION,
+        "speculation": SPECULATION,
+        "temporal_detail": TEMPORAL_DETAIL,
+        "visual_specificity": VISUAL_SPECIFICITY,
+        "camera_detail": CAMERA_DETAIL,
+        "caption_style": CAPTION_STYLE,
+    }
+
+
+def main() -> None:
+    processor = load_processor(MODEL_DIR)
+    request = build_vllm_request(
+        processor=processor,
+        model_dir=MODEL_DIR,
+        video_path=VIDEO_PATH,
+        num_frames=NUM_FRAMES,
+        sampling_rate=SAMPLING_RATE,
+        prompt_override=PROMPT_OVERRIDE,
+        prompt_settings=prompt_settings(),
+    )
+
+    llm = load_llm(
+        model_dir=MODEL_DIR,
+        max_model_len=MAX_MODEL_LEN,
+        max_num_seqs=MAX_NUM_SEQS,
+        gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
+        dtype=DTYPE,
+        enforce_eager=ENFORCE_EAGER,
+        enable_prefix_caching=ENABLE_PREFIX_CACHING,
+        trust_remote_code=TRUST_REMOTE_CODE,
+    )
+    outputs = llm.generate(
+        [request],
+        sampling_params=sampling_params(
+            temperature=TEMPERATURE,
+            max_tokens=MAX_TOKENS,
+            top_p=TOP_P,
+            repetition_penalty=REPETITION_PENALTY,
+        ),
+    )
+    print(outputs[0].outputs[0].text)
+
+
+if __name__ == "__main__":
+    main()

infer.py

@@ -0,0 +1,519 @@
+"""Run Gemma 4 31B Parakeet-hybrid caption inference on one video.
+
+Edit the variables below, then run on the Vast instance:
+
+    source /venv/main/bin/activate
+    cd /workspace/polished_model_v3
+    python infer.py
+
+This uses the condensed inference package:
+- video + text prompt are passed through the Gemma processor
+- audio is supplied as a sidecar WAV
+- processor-created audio features are replaced with Parakeet-native features
+- the language LoRA is already merged into model/
+- the trained audio projector is loaded from model/embed_audio.safetensors
+"""
+
+from __future__ import annotations
+
+import contextlib
+import io
+import json
+import logging
+import subprocess
+import tempfile
+import warnings
+from pathlib import Path
+from typing import Any
+
+import torch
+from transformers import AutoModelForMultimodalLM, AutoProcessor
+
+SCRIPT_DIR = Path(__file__).resolve().parent
+MODEL_ROOT = SCRIPT_DIR
+
+from caption_model_runtime import (
+    DEFAULT_PARAKEET_MODEL_ID,
+    CAPTION_LENGTH_LABELS,
+    CAPTION_SETTING_FIELD_CHOICES,
+    DEFAULT_CAPTION_SETTING_VALUES,
+    format_caption_settings_prompt,
+    gemma_core,
+    install_parakeet_audio_bridge,
+    load_state_file,
+    replace_batch_audio_features,
+)
+
+
+# Edit these.
+VIDEO_PATH = "/workspace/test7.mp4"  # Options: path to the video you want to caption.
+MODEL_PATH = str(MODEL_ROOT / "model")  # Options: merged model path.
+PROCESSOR_PATH = str(MODEL_ROOT / "processor")  # Options: processor path.
+AUDIO_PROJECTOR_PATH = str(MODEL_ROOT / "model" / "embed_audio.safetensors")  # Options: trained audio projector path.
+
+# Parakeet hybrid audio bridge settings. These should normally match the packaged model.
+PARAKEET_MODEL_ID = str(MODEL_ROOT / "parakeet")  # Options: "nvidia/parakeet-tdt-0.6b-v3" or compatible local/HF path.
+PARAKEET_BRIDGE_MODE = "tdt_token_embeddings_with_encoder_context"  # Options: "encoder", "tdt_tokens", "tdt_token_embeddings", "encoder_soft_tdt_token_embeddings", "tdt_token_embeddings_with_encoder_context".
+PARAKEET_NATIVE_FEATURES = True  # Options: True to replace Gemma audio features with Parakeet features, False for debugging only.
+PARAKEET_TDT_FILTER_BLANK_TOKENS = True  # Options: True or False.
+PARAKEET_TDT_FILTER_SPECIAL_TOKEN_IDS = True  # Options: True or False.
+PROJECTOR_INTERMEDIATE_SIZE = 4096  # Options: integer; the packaged model uses 4096.
+PROJECTOR_DROPOUT = 0.0  # Options: float; inference should normally be 0.0.
+HYBRID_ENCODER_GATE_INIT = 0.0  # Options: float; saved audio projector weights override the initial gate.
+
+# Prompt settings. Empty PROMPT_OVERRIDE builds the standard dynamic prompt.
+PROMPT_OVERRIDE = ""  # Options: "" or any full custom prompt string.
+CAPTION_SETTINGS_JSON_PATH = ""  # Options: "" or a JSON path under /workspace/dataset_jsons to override the settings below.
+CAPTION_LENGTH = "very large"  # Options: "very small", "small", "medium", "large", "very large".
+INCLUDE_WATERMARK_INFO = False  # Options: True or False.
+VULGARITY = "low"  # Options: "none", "low", "medium", "high".
+UNCERTAINTY = "low"  # Options: "none", "low", "medium", "high".
+CHARACTER_NAMES = "none"  # Options: "none", "ambiguous", "single", "multiple".
+FLUFF = "none"  # Options: "none", "low", "medium", "high".
+HAS_REPETITION = False  # Options: True or False.
+SPECULATION = "low"  # Options: "none", "low", "medium", "high".
+TEMPORAL_DETAIL = "medium"  # Options: "static", "low", "medium", "high".
+VISUAL_SPECIFICITY = "moderate"  # Options: "generic", "moderate", "detailed", "excessive".
+CAMERA_DETAIL = "medium"  # Options: "none", "low", "medium", "high".
+CAPTION_STYLE = "plain"  # Options: "plain", "verbose", "ornate", "robotic".
+HAS_THINKING = True  # Options: True to request thought JSON plus final caption, False to request only the final caption.
+
+# Media settings. Training used separate sidecar audio and random frame counts.
+NUM_FRAMES = 12  # Options: None for processor default, or an integer frame count.
+FPS = None  # Options: None for processor default, or a float such as 1.0.
+SAMPLING_RATE = 16_000  # Options: normally 16000.
+AUDIO_MAX_LENGTH_SAMPLES = 0  # Options: 0 keeps full audio; positive integer truncates Parakeet audio.
+MAX_AUDIO_SECONDS = 0.0  # Options: 0.0 keeps full audio; positive float caps extracted sidecar audio.
+
+# Generation settings.
+MAX_NEW_TOKENS = 1200  # Options: positive integer token cap.
+TEMPERATURE = 0.0  # Options: 0.0 for greedy decoding, >0.0 for sampling.
+TOP_P = 0.9  # Options: float in (0, 1], used only when TEMPERATURE > 0.
+REPETITION_PENALTY = 1.1  # Options: 1.0 disables the penalty, >1.0 penalizes repetition.
+PRINT_INPUT_STATS = False  # Options: True or False.
+QUIET_MODEL_LOAD = True  # Options: True hides noisy missing-key load reports; False prints full loader output.
+
+# Usually leave these alone.
+LOCAL_FILES_ONLY = True  # Options: True to use cached files only, False to allow downloads.
+DTYPE = torch.bfloat16  # Options: torch.bfloat16, torch.float16, torch.float32.
+DEVICE_MAP = "auto"  # Options: "auto", "cuda", or another Transformers device_map value.
+ATTN_IMPLEMENTATION = "sdpa"  # Options: "sdpa", "flash_attention_2", None.
+
+warnings.filterwarnings(
+    "ignore",
+    message=r"RNN module weights are not part of single contiguous chunk of memory.*",
+    category=UserWarning,
+)
+
+
+@contextlib.contextmanager
+def quiet_model_load() -> Any:
+    if not QUIET_MODEL_LOAD:
+        yield
+        return
+    load_report_logger = logging.getLogger("transformers.utils.loading_report")
+    old_level = load_report_logger.level
+    load_report_logger.setLevel(logging.ERROR)
+    patched_modules: list[tuple[Any, Any]] = []
+    try:
+        import transformers.modeling_utils as modeling_utils
+        import transformers.utils.loading_report as loading_report
+
+        original_report = modeling_utils.log_state_dict_report
+
+        def quiet_report(
+            model: Any,
+            pretrained_model_name_or_path: str,
+            ignore_mismatched_sizes: bool,
+            loading_info: Any,
+            logger: logging.Logger | None = None,
+        ) -> None:
+            has_fatal_issue = bool(getattr(loading_info, "error_msgs", None)) or bool(
+                getattr(loading_info, "conversion_errors", None)
+            )
+            if not ignore_mismatched_sizes and bool(getattr(loading_info, "mismatched_keys", None)):
+                has_fatal_issue = True
+            if has_fatal_issue:
+                original_report(
+                    model,
+                    pretrained_model_name_or_path,
+                    ignore_mismatched_sizes,
+                    loading_info,
+                    logger=logger,
+                )
+
+        for module in (loading_report, modeling_utils):
+            patched_modules.append((module, module.log_state_dict_report))
+            module.log_state_dict_report = quiet_report
+    except Exception:
+        patched_modules = []
+    try:
+        with contextlib.redirect_stdout(io.StringIO()), contextlib.redirect_stderr(io.StringIO()):
+            yield
+    finally:
+        for module, original in patched_modules:
+            module.log_state_dict_report = original
+        load_report_logger.setLevel(old_level)
+
+
+def load_parakeet_projector(model: torch.nn.Module) -> None:
+    state_path = Path(AUDIO_PROJECTOR_PATH)
+    state = load_state_file(state_path)
+    module = gemma_core(model).embed_audio
+    module.load_state_dict(state, strict=True)
+
+
+def video_has_audio_stream(video_path: Path) -> bool:
+    cmd = [
+        "ffprobe",
+        "-v",
+        "error",
+        "-select_streams",
+        "a:0",
+        "-show_entries",
+        "stream=index",
+        "-of",
+        "csv=p=0",
+        str(video_path),
+    ]
+    result = subprocess.run(cmd, check=True, capture_output=True, text=True)
+    return bool(result.stdout.strip())
+
+
+def probe_video_duration_seconds(video_path: Path) -> float:
+    cmd = [
+        "ffprobe",
+        "-v",
+        "error",
+        "-show_entries",
+        "format=duration",
+        "-of",
+        "default=noprint_wrappers=1:nokey=1",
+        str(video_path),
+    ]
+    result = subprocess.run(cmd, check=True, capture_output=True, text=True)
+    duration = float(result.stdout.strip())
+    if duration <= 0:
+        raise RuntimeError(f"Video duration must be positive: {video_path}")
+    return duration
+
+
+def extract_audio(video_path: Path, audio_path: Path) -> None:
+    cmd = [
+        "ffmpeg",
+        "-hide_banner",
+        "-loglevel",
+        "error",
+        "-y",
+        "-i",
+        str(video_path),
+        "-vn",
+        "-ac",
+        "1",
+        "-ar",
+        str(SAMPLING_RATE),
+    ]
+    if MAX_AUDIO_SECONDS > 0:
+        cmd.extend(["-t", f"{MAX_AUDIO_SECONDS:.6f}"])
+    cmd.extend(["-c:a", "pcm_s16le", str(audio_path)])
+    subprocess.run(cmd, check=True)
+
+
+def create_silent_audio(audio_path: Path, duration_seconds: float) -> None:
+    if MAX_AUDIO_SECONDS > 0:
+        duration_seconds = min(duration_seconds, MAX_AUDIO_SECONDS)
+    cmd = [
+        "ffmpeg",
+        "-hide_banner",
+        "-loglevel",
+        "error",
+        "-y",
+        "-f",
+        "lavfi",
+        "-i",
+        f"anullsrc=channel_layout=mono:sample_rate={SAMPLING_RATE}",
+        "-t",
+        f"{duration_seconds:.6f}",
+        "-ac",
+        "1",
+        "-ar",
+        str(SAMPLING_RATE),
+        "-c:a",
+        "pcm_s16le",
+        str(audio_path),
+    ]
+    subprocess.run(cmd, check=True)
+
+
+def prepare_sidecar_audio(video_path: Path, tmpdir: Path) -> tuple[Path, bool]:
+    audio_path = tmpdir / "sidecar_audio.wav"
+    if video_has_audio_stream(video_path):
+        extract_audio(video_path, audio_path)
+        return audio_path, True
+
+    duration_seconds = probe_video_duration_seconds(video_path)
+    print(f"input_video_has_audio=false; creating_silent_sidecar_audio duration_seconds={duration_seconds:.3f}", flush=True)
+    create_silent_audio(audio_path, duration_seconds)
+    return audio_path, False
+
+
+def bool_from_json(value: Any, field_name: str) -> bool:
+    if isinstance(value, bool):
+        return value
+    if isinstance(value, str):
+        lowered = value.strip().lower()
+        if lowered in {"1", "true", "yes", "y", "on"}:
+            return True
+        if lowered in {"0", "false", "no", "n", "off"}:
+            return False
+    raise ValueError(f"{field_name} must be boolean-like, got {value!r}")
+
+
+def load_prompt_settings_json() -> dict[str, Any]:
+    if not CAPTION_SETTINGS_JSON_PATH.strip():
+        return {}
+    path = Path(CAPTION_SETTINGS_JSON_PATH)
+    data = json.loads(path.read_text(encoding="utf-8"))
+    if not isinstance(data, dict):
+        raise ValueError(f"CAPTION_SETTINGS_JSON_PATH must point to a JSON object: {path}")
+    return data
+
+
+def build_prompt() -> str:
+    if PROMPT_OVERRIDE.strip():
+        return PROMPT_OVERRIDE.strip()
+
+    settings: dict[str, Any] = {
+        "caption_length": CAPTION_LENGTH,
+        "include_watermark_info": INCLUDE_WATERMARK_INFO,
+        **DEFAULT_CAPTION_SETTING_VALUES,
+        "vulgarity": VULGARITY,
+        "uncertainty": UNCERTAINTY,
+        "character_names": CHARACTER_NAMES,
+        "fluff": FLUFF,
+        "has_repetition": HAS_REPETITION,
+        "speculation": SPECULATION,
+        "temporal_detail": TEMPORAL_DETAIL,
+        "visual_specificity": VISUAL_SPECIFICITY,
+        "camera_detail": CAMERA_DETAIL,
+        "caption_style": CAPTION_STYLE,
+        "has_thinking": HAS_THINKING,
+    }
+    settings.update(load_prompt_settings_json())
+    settings["has_thinking"] = HAS_THINKING
+
+    settings["caption_length"] = str(settings["caption_length"]).strip().lower()
+    if settings["caption_length"] not in CAPTION_LENGTH_LABELS:
+        raise ValueError(f"caption_length must be one of {CAPTION_LENGTH_LABELS}, got {settings['caption_length']!r}")
+    settings["include_watermark_info"] = bool_from_json(settings["include_watermark_info"], "include_watermark_info")
+    settings["has_repetition"] = bool_from_json(settings["has_repetition"], "has_repetition")
+    settings["has_thinking"] = bool_from_json(settings["has_thinking"], "has_thinking")
+    for field_name, allowed in CAPTION_SETTING_FIELD_CHOICES.items():
+        value = str(settings[field_name]).strip().lower()
+        if value not in allowed:
+            raise ValueError(f"{field_name} must be one of {allowed}, got {value!r}")
+        settings[field_name] = value
+
+    return format_caption_settings_prompt(settings)
+
+
+def build_messages(video_path: Path, audio_path: Path, prompt: str) -> list[dict[str, Any]]:
+    return [
+        {
+            "role": "user",
+            "content": [
+                {"type": "video", "path": str(video_path)},
+                {"type": "text", "text": prompt},
+                {"type": "audio", "path": str(audio_path)},
+            ],
+        },
+        {"role": "assistant", "content": [{"type": "text", "text": ""}]},
+    ]
+
+
+def trim_empty_assistant_terminator(inputs: dict[str, torch.Tensor], processor: Any) -> dict[str, torch.Tensor]:
+    eos_tail = processor.tokenizer.encode("<turn|>\n", add_special_tokens=False)
+    if not eos_tail:
+        return inputs
+
+    tail_len = len(eos_tail)
+    input_ids = inputs["input_ids"][0]
+    if input_ids[-tail_len:].tolist() != eos_tail:
+        return inputs
+
+    trimmed = {}
+    for key, value in inputs.items():
+        if isinstance(value, torch.Tensor) and value.ndim >= 2 and value.shape[1] == input_ids.shape[0]:
+            trimmed[key] = value[:, :-tail_len]
+        else:
+            trimmed[key] = value
+    return trimmed
+
+
+def tensor_stats(tensor: torch.Tensor | None, mask: torch.Tensor | None = None) -> dict[str, Any]:
+    if tensor is None:
+        return {"present": False}
+    stats_tensor = tensor.detach().float().cpu()
+    result: dict[str, Any] = {
+        "present": True,
+        "shape": list(tensor.shape),
+        "mean": round(float(stats_tensor.mean().item()), 8),
+        "std": round(float(stats_tensor.std().item()), 8),
+        "abs_mean": round(float(stats_tensor.abs().mean().item()), 8),
+    }
+    if mask is not None:
+        result["mask_shape"] = list(mask.shape)
+        result["mask_sum"] = int(mask.detach().cpu().sum().item())
+    return result
+
+
+def print_input_stats(inputs: dict[str, torch.Tensor], label: str) -> None:
+    stats = {
+        "input_ids_shape": list(inputs["input_ids"].shape),
+        "input_features": tensor_stats(inputs.get("input_features"), inputs.get("input_features_mask")),
+        "keys": sorted(inputs.keys()),
+    }
+    print(f"{label}=" + json.dumps(stats, sort_keys=True), flush=True)
+
+
+def move_inputs_to_model_device(inputs: dict[str, Any], model: torch.nn.Module) -> dict[str, Any]:
+    device = getattr(model, "device", None)
+    if device is None:
+        try:
+            device = next(model.parameters()).device
+        except StopIteration:
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    moved = {}
+    for key, value in inputs.items():
+        moved[key] = value.to(device) if isinstance(value, torch.Tensor) else value
+    return moved
+
+
+def generation_kwargs() -> dict[str, Any]:
+    kwargs: dict[str, Any] = {
+        "max_new_tokens": MAX_NEW_TOKENS,
+        "do_sample": TEMPERATURE > 0,
+        "repetition_penalty": REPETITION_PENALTY,
+        "use_cache": True,
+    }
+    if TEMPERATURE > 0:
+        kwargs["temperature"] = TEMPERATURE
+        kwargs["top_p"] = TOP_P
+    return kwargs
+
+
+def prepare_inputs(
+    processor: Any,
+    parakeet_processor: Any,
+    video_path: Path,
+    audio_path: Path,
+    prompt: str,
+) -> dict[str, torch.Tensor]:
+    processor_kwargs: dict[str, Any] = {
+        "padding": True,
+        "truncation": False,
+        "sampling_rate": SAMPLING_RATE,
+    }
+    if NUM_FRAMES is not None:
+        processor_kwargs["num_frames"] = NUM_FRAMES
+    if FPS is not None:
+        processor_kwargs["fps"] = FPS
+
+    inputs = processor.apply_chat_template(
+        build_messages(video_path, audio_path, prompt),
+        tokenize=True,
+        return_dict=True,
+        return_tensors="pt",
+        load_audio_from_video=False,
+        processor_kwargs=processor_kwargs,
+    )
+    inputs = trim_empty_assistant_terminator(inputs, processor)
+    if PRINT_INPUT_STATS:
+        print_input_stats(inputs, "input_stats_before_parakeet_swap")
+    if PARAKEET_NATIVE_FEATURES:
+        replace_batch_audio_features(
+            inputs,
+            audio_paths=[str(audio_path)],
+            parakeet_processor=parakeet_processor,
+            sampling_rate=SAMPLING_RATE,
+            max_length_samples=AUDIO_MAX_LENGTH_SAMPLES,
+        )
+        if PRINT_INPUT_STATS:
+            print_input_stats(inputs, "input_stats_after_parakeet_swap")
+    return inputs
+
+
+def load_model(processor_path: Path) -> tuple[Any, torch.nn.Module, Any]:
+    processor = AutoProcessor.from_pretrained(str(processor_path), local_files_only=LOCAL_FILES_ONLY)
+    parakeet_processor = AutoProcessor.from_pretrained(PARAKEET_MODEL_ID, local_files_only=LOCAL_FILES_ONLY)
+
+    model_kwargs: dict[str, Any] = {
+        "local_files_only": LOCAL_FILES_ONLY,
+        "dtype": DTYPE,
+        "low_cpu_mem_usage": True,
+        "device_map": DEVICE_MAP,
+    }
+    if ATTN_IMPLEMENTATION:
+        model_kwargs["attn_implementation"] = ATTN_IMPLEMENTATION
+
+    with quiet_model_load():
+        model = AutoModelForMultimodalLM.from_pretrained(MODEL_PATH, **model_kwargs)
+    bridge_args = type(
+        "BridgeArgs",
+        (),
+        {
+            "parakeet_model_id": PARAKEET_MODEL_ID,
+            "parakeet_bridge_mode": PARAKEET_BRIDGE_MODE,
+            "local_files_only": LOCAL_FILES_ONLY,
+            "projector_intermediate_size": PROJECTOR_INTERMEDIATE_SIZE,
+            "projector_dropout": PROJECTOR_DROPOUT,
+            "hybrid_encoder_gate_init": HYBRID_ENCODER_GATE_INIT,
+            "parakeet_tdt_filter_blank_tokens": PARAKEET_TDT_FILTER_BLANK_TOKENS,
+            "parakeet_tdt_filter_special_token_ids": PARAKEET_TDT_FILTER_SPECIAL_TOKEN_IDS,
+        },
+    )()
+    with quiet_model_load():
+        install_parakeet_audio_bridge(model, bridge_args)
+    gemma_core(model)
+    load_parakeet_projector(model)
+
+    model.eval()
+    return processor, model, parakeet_processor
+
+
+def main() -> None:
+    video_path = Path(VIDEO_PATH)
+    model_path = Path(MODEL_PATH)
+    processor_path = Path(PROCESSOR_PATH)
+    audio_projector_path = Path(AUDIO_PROJECTOR_PATH)
+    parakeet_path = Path(PARAKEET_MODEL_ID)
+
+    for label, path in (
+        ("VIDEO_PATH", video_path),
+        ("MODEL_PATH", model_path),
+        ("processor", processor_path),
+        ("audio_projector", audio_projector_path),
+        ("parakeet", parakeet_path),
+    ):
+        if not path.exists():
+            raise FileNotFoundError(f"{label} does not exist: {path}")
+
+    prompt = build_prompt()
+    processor, model, parakeet_processor = load_model(processor_path)
+    with tempfile.TemporaryDirectory(prefix="gemma4_caption_inference_") as tmpdir_raw:
+        tmpdir = Path(tmpdir_raw)
+        audio_path, _had_audio = prepare_sidecar_audio(video_path, tmpdir)
+        inputs = prepare_inputs(processor, parakeet_processor, video_path, audio_path, prompt)
+        moved_inputs = move_inputs_to_model_device(inputs, model)
+        input_len = moved_inputs["input_ids"].shape[-1]
+        with torch.inference_mode():
+            output_ids = model.generate(**moved_inputs, **generation_kwargs())
+
+    new_tokens = output_ids[0, input_len:]
+    response = processor.decode(new_tokens, skip_special_tokens=True).strip()
+    print(response, flush=True)
+
+
+if __name__ == "__main__":
+    main()

model/audio_config.json

@@ -0,0 +1,25 @@
+{
+  "attention_chunk_size": 12,
+  "attention_context_left": 13,
+  "attention_context_right": 0,
+  "attention_invalid_logits_value": -1000000000.0,
+  "attention_logit_cap": 50.0,
+  "conv_kernel_size": 5,
+  "dtype": "bfloat16",
+  "gradient_clipping": 10000000000.0,
+  "hidden_act": "silu",
+  "hidden_size": 1024,
+  "initializer_range": 0.02,
+  "model_type": "gemma4_audio",
+  "num_attention_heads": 8,
+  "num_hidden_layers": 12,
+  "output_proj_dims": 1536,
+  "residual_weight": 0.5,
+  "rms_norm_eps": 1e-06,
+  "subsampling_conv_channels": [
+    128,
+    32
+  ],
+  "transformers_version": "5.12.1",
+  "use_clipped_linears": true
+}

model/chat_template.jinja

@@ -0,0 +1,363 @@
+{%- macro format_parameters(properties, required, filter_keys=false) -%}
+    {%- set standard_keys = ['description', 'type', 'properties', 'required', 'nullable'] -%}
+    {%- set ns = namespace(found_first=false) -%}
+    {%- for key, value in properties | dictsort -%}
+        {%- set add_comma = false -%}
+        {%- if not filter_keys or key not in standard_keys -%}
+            {%- if ns.found_first %},{% endif -%}
+            {%- set ns.found_first = true -%}
+            {{ key }}:{
+            {%- if value['description'] -%}
+                description:<|"|>{{ value['description'] }}<|"|>
+                {%- set add_comma = true -%}
+            {%- endif -%}
+            {%- if value['type'] | upper == 'STRING' -%}
+                {%- if value['enum'] -%}
+                    {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                    enum:{{ format_argument(value['enum']) }}
+                {%- endif -%}
+            {%- elif value['type'] | upper == 'ARRAY' -%}
+                {%- if value['items'] is mapping and value['items'] -%}
+                    {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                    items:{
+                    {%- set ns_items = namespace(found_first=false) -%}
+                    {%- for item_key, item_value in value['items'] | dictsort -%}
+                        {%- if item_value is not none -%}
+                            {%- if ns_items.found_first %},{% endif -%}
+                            {%- set ns_items.found_first = true -%}
+                            {%- if item_key == 'properties' -%}
+                                properties:{
+                                {%- if item_value is mapping -%}
+                                    {{- format_parameters(item_value, value['items']['required'] | default([])) -}}
+                                {%- endif -%}
+                                }
+                            {%- elif item_key == 'required' -%}
+                                required:[
+                                {%- for req_item in item_value -%}
+                                    <|"|>{{- req_item -}}<|"|>
+                                    {%- if not loop.last %},{% endif -%}
+                                {%- endfor -%}
+                                ]
+                            {%- elif item_key == 'type' -%}
+                                {%- if item_value is string -%}
+                                    type:{{ format_argument(item_value | upper) }}
+                                {%- else -%}
+                                    type:{{ format_argument(item_value | map('upper') | list) }}
+                                {%- endif -%}
+                            {%- else -%}
+                                {{ item_key }}:{{ format_argument(item_value) }}
+                            {%- endif -%}
+                        {%- endif -%}
+                    {%- endfor -%}
+                    }
+                {%- endif -%}
+            {%- endif -%}
+            {%- if value['nullable'] %}
+                {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                nullable:true
+            {%- endif -%}
+            {%- if value['type'] | upper == 'OBJECT' -%}
+                {%- if value['properties'] is defined and value['properties'] is mapping -%}
+                    {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                    properties:{
+                    {{- format_parameters(value['properties'], value['required'] | default([])) -}}
+                    }
+                {%- elif value is mapping -%}
+                    {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                    properties:{
+                    {{- format_parameters(value, value['required'] | default([]), filter_keys=true) -}}
+                    }
+                {%- endif -%}
+                {%- if value['required'] -%}
+                    {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                    required:[
+                    {%- for item in value['required'] | default([]) -%}
+                        <|"|>{{- item -}}<|"|>
+                        {%- if not loop.last %},{% endif -%}
+                    {%- endfor -%}
+                    ]
+                {%- endif -%}
+            {%- endif -%}
+            {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+            type:<|"|>{{ value['type'] | upper }}<|"|>}
+        {%- endif -%}
+    {%- endfor -%}
+{%- endmacro -%}
+{%- macro format_function_declaration(tool_data) -%}
+    declaration:{{- tool_data['function']['name'] -}}{description:<|"|>{{- tool_data['function']['description'] -}}<|"|>
+    {%- set params = tool_data['function']['parameters'] -%}
+    {%- if params -%}
+        ,parameters:{
+        {%- if params['properties'] -%}
+            properties:{ {{- format_parameters(params['properties'], params['required']) -}} },
+        {%- endif -%}
+        {%- if params['required'] -%}
+            required:[
+            {%- for item in params['required'] -%}
+                <|"|>{{- item -}}<|"|>
+                {{- ',' if not loop.last -}}
+            {%- endfor -%}
+            ],
+        {%- endif -%}
+        {%- if params['type'] -%}
+            type:<|"|>{{- params['type'] | upper -}}<|"|>}
+        {%- endif -%}
+    {%- endif -%}
+    {%- if 'response' in tool_data['function'] -%}
+        {%- set response_declaration = tool_data['function']['response'] -%}
+        ,response:{
+        {%- if response_declaration['description'] -%}
+            description:<|"|>{{- response_declaration['description'] -}}<|"|>,
+        {%- endif -%}
+        {%- if response_declaration['type'] | upper == 'OBJECT' -%}
+            type:<|"|>{{- response_declaration['type'] | upper -}}<|"|>}
+        {%- endif -%}
+    {%- endif -%}
+    }
+{%- endmacro -%}
+{%- macro format_argument(argument, escape_keys=True) -%}
+    {%- if argument is string -%}
+        {{- '<|"|>' + argument + '<|"|>' -}}
+    {%- elif argument is boolean -%}
+        {{- 'true' if argument else 'false' -}}
+    {%- elif argument is mapping -%}
+        {{- '{' -}}
+        {%- set ns = namespace(found_first=false) -%}
+        {%- for key, value in argument | dictsort -%}
+            {%- if ns.found_first %},{% endif -%}
+            {%- set ns.found_first = true -%}
+            {%- if escape_keys -%}
+                {{- '<|"|>' + key + '<|"|>' -}}
+            {%- else -%}
+                {{- key -}}
+            {%- endif -%}
+            :{{- format_argument(value, escape_keys=escape_keys) -}}
+        {%- endfor -%}
+        {{- '}' -}}
+    {%- elif argument is sequence -%}
+        {{- '[' -}}
+        {%- for item in argument -%}
+            {{- format_argument(item, escape_keys=escape_keys) -}}
+            {%- if not loop.last %},{% endif -%}
+        {%- endfor -%}
+        {{- ']' -}}
+    {%- else -%}
+        {{- argument -}}
+    {%- endif -%}
+{%- endmacro -%}
+{%- macro strip_thinking(text) -%}
+    {%- set ns = namespace(result='') -%}
+    {%- for part in text.split('<channel|>') -%}
+        {%- if '<|channel>' in part -%}
+            {%- set ns.result = ns.result + part.split('<|channel>')[0] -%}
+        {%- else -%}
+            {%- set ns.result = ns.result + part -%}
+        {%- endif -%}
+    {%- endfor -%}
+    {{- ns.result | trim -}}
+{%- endmacro -%}
+
+{%- macro format_tool_response_block(tool_name, response) -%}
+    {{- '<|tool_response>' -}}
+    {%- if response is mapping -%}
+        {{- 'response:' + tool_name + '{' -}}
+        {%- for key, value in response | dictsort -%}
+            {{- key -}}:{{- format_argument(value, escape_keys=False) -}}
+            {%- if not loop.last %},{% endif -%}
+        {%- endfor -%}
+        {{- '}' -}}
+    {%- else -%}
+        {{- 'response:' + tool_name + '{value:' + format_argument(response, escape_keys=False) + '}' -}}
+    {%- endif -%}
+    {{- '<tool_response|>' -}}
+{%- endmacro -%}
+
+{%- set ns = namespace(prev_message_type=None) -%}
+{%- set loop_messages = messages -%}
+{{- bos_token -}}
+{#- Handle System/Tool Definitions Block -#}
+{%- if (enable_thinking is defined and enable_thinking) or tools or messages[0]['role'] in ['system', 'developer'] -%}
+    {{- '<|turn>system\n' -}}
+    {#- Inject Thinking token at the very top of the FIRST system turn -#}
+    {%- if enable_thinking is defined and enable_thinking -%}
+        {{- '<|think|>\n' -}}
+        {%- set ns.prev_message_type = 'think' -%}
+    {%- endif -%}
+    {%- if messages[0]['role'] in ['system', 'developer'] -%}
+        {%- if messages[0]['content'] is string -%}
+            {{- messages[0]['content'] | trim -}}
+        {%- elif messages[0]['content'] is sequence -%}
+            {%- for item in messages[0]['content'] -%}
+                {{- item['text'] | trim + ' '-}}
+            {%- endfor -%}
+        {%- endif -%}
+        {%- set loop_messages = messages[1:] -%}
+    {%- endif -%}
+    {%- if tools -%}
+        {%- for tool in tools %}
+            {{- '<|tool>' -}}
+            {{- format_function_declaration(tool) | trim -}}
+            {{- '<tool|>' -}}
+        {%- endfor %}
+        {%- set ns.prev_message_type = 'tool' -%}
+    {%- endif -%}
+    {{- '<turn|>\n' -}}
+{%- endif %}
+
+{#- Pre-scan: find last user message index for reasoning guard -#}
+{%- set ns_turn = namespace(last_user_idx=-1) -%}
+{%- for i in range(loop_messages | length) -%}
+    {%- if loop_messages[i]['role'] == 'user' -%}
+        {%- set ns_turn.last_user_idx = i -%}
+    {%- endif -%}
+{%- endfor -%}
+
+{#- Loop through messages -#}
+{%- for message in loop_messages -%}
+    {%- if message['role'] != 'tool' -%}
+    {%- set ns.prev_message_type = None -%}
+    {%- set role = 'model' if message['role'] == 'assistant' else message['role'] -%}
+    {#- Detect continuation: suppress duplicate <|turn>model when previous non-tool message was also assistant -#}
+    {%- set prev_nt = namespace(role=None, found=false) -%}
+    {%- if loop.index0 > 0 -%}
+        {%- for j in range(loop.index0 - 1, -1, -1) -%}
+            {%- if not prev_nt.found -%}
+                {%- if loop_messages[j]['role'] != 'tool' -%}
+                    {%- set prev_nt.role = loop_messages[j]['role'] -%}
+                    {%- set prev_nt.found = true -%}
+                {%- endif -%}
+            {%- endif -%}
+        {%- endfor -%}
+    {%- endif -%}
+    {%- set continue_same_model_turn = (role == 'model' and prev_nt.role == 'assistant') -%}
+    {%- if not continue_same_model_turn -%}
+        {{- '<|turn>' + role + '\n' }}
+    {%- endif -%}
+
+    {#- Render reasoning/reasoning_content as thinking channel -#}
+    {%- set thinking_text = message.get('reasoning') or message.get('reasoning_content') -%}
+    {%- if thinking_text and loop.index0 > ns_turn.last_user_idx and message.get('tool_calls') -%}
+        {{- '<|channel>thought\n' + thinking_text + '\n<channel|>' -}}
+    {%- endif -%}
+
+            {%- if message['tool_calls'] -%}
+                {%- for tool_call in message['tool_calls'] -%}
+                    {%- set function = tool_call['function'] -%}
+                    {{- '<|tool_call>call:' + function['name'] + '{' -}}
+                    {%- if function['arguments'] is mapping -%}
+                        {%- set ns_args = namespace(found_first=false) -%}
+                        {%- for key, value in function['arguments'] | dictsort -%}
+                            {%- if ns_args.found_first %},{% endif -%}
+                            {%- set ns_args.found_first = true -%}
+                            {{- key -}}:{{- format_argument(value, escape_keys=False) -}}
+                        {%- endfor -%}
+                    {%- elif function['arguments'] is string -%}
+                        {{- function['arguments'] -}}
+                    {%- endif -%}
+                    {{- '}<tool_call|>' -}}
+                {%- endfor -%}
+                {%- set ns.prev_message_type = 'tool_call' -%}
+            {%- endif -%}
+
+            {%- set ns_tr_out = namespace(flag=false) -%}
+            {%- if message.get('tool_responses') -%}
+                {#- Legacy: tool_responses embedded on the assistant message (Google/Gemma native) -#}
+                {%- for tool_response in message['tool_responses'] -%}
+                    {{- format_tool_response_block(tool_response['name'] | default('unknown'), tool_response['response']) -}}
+                    {%- set ns_tr_out.flag = true -%}
+                    {%- set ns.prev_message_type = 'tool_response' -%}
+                {%- endfor -%}
+            {%- elif message.get('tool_calls') -%}
+                {#- OpenAI Chat Completions: forward-scan consecutive role:tool messages -#}
+                {%- set ns_tool_scan = namespace(stopped=false) -%}
+                {%- for k in range(loop.index0 + 1, loop_messages | length) -%}
+                    {%- if ns_tool_scan.stopped -%}
+                    {%- elif loop_messages[k]['role'] != 'tool' -%}
+                        {%- set ns_tool_scan.stopped = true -%}
+                    {%- else -%}
+                        {%- set follow = loop_messages[k] -%}
+                        {#- Resolve tool_call_id to function name -#}
+                        {%- set ns_tname = namespace(name=follow.get('name') | default('unknown')) -%}
+                        {%- for tc in message['tool_calls'] -%}
+                            {%- if tc.get('id') == follow.get('tool_call_id') -%}
+                                {%- set ns_tname.name = tc['function']['name'] -%}
+                            {%- endif -%}
+                        {%- endfor -%}
+                        {#- Handle content as string or content-parts array -#}
+                        {%- set tool_body = follow.get('content') -%}
+                        {%- if tool_body is string -%}
+                            {{- format_tool_response_block(ns_tname.name, tool_body) -}}
+                        {%- elif tool_body is sequence and tool_body is not string -%}
+                            {%- set ns_txt = namespace(s='') -%}
+                            {%- for part in tool_body -%}
+                                {%- if part.get('type') == 'text' -%}
+                                    {%- set ns_txt.s = ns_txt.s + (part.get('text') | default('')) -%}
+                                {%- endif -%}
+                            {%- endfor -%}
+                            {{- format_tool_response_block(ns_tname.name, ns_txt.s) -}}
+                            {%- for part in tool_body -%}
+                                {%- if part.get('type') == 'image' -%}
+                                    {{- '<|image|>' -}}
+                                {%- elif part.get('type') == 'audio' -%}
+                                    {{- '<|audio|>' -}}
+                                {%- elif part.get('type') == 'video' -%}
+                                    {{- '<|video|>' -}}
+                                {%- endif -%}
+                            {%- endfor -%}
+                        {%- else -%}
+                            {{- format_tool_response_block(ns_tname.name, tool_body) -}}
+                        {%- endif -%}
+                        {%- set ns_tr_out.flag = true -%}
+                        {%- set ns.prev_message_type = 'tool_response' -%}
+                    {%- endif -%}
+                {%- endfor -%}
+            {%- endif -%}
+
+            {%- set captured_content -%}
+            {%- if message['content'] is string -%}
+                {%- if role == 'model' -%}
+                    {{- strip_thinking(message['content']) -}}
+                {%- else -%}
+                    {{- message['content'] | trim -}}
+                {%- endif -%}
+            {%- elif message['content'] is sequence -%}
+                {%- for item in message['content'] -%}
+                    {%- if item['type'] == 'text' -%}
+                        {%- if role == 'model' -%}
+                            {{- strip_thinking(item['text']) -}}
+                        {%- else -%}
+                            {{- item['text'] | trim -}}
+                        {%- endif -%}
+                    {%- elif item['type'] == 'image' -%}
+                        {{- '<|image|>' -}}
+                        {%- set ns.prev_message_type = 'image' -%}
+                    {%- elif item['type'] == 'audio' -%}
+                        {{- '<|audio|>' -}}
+                        {%- set ns.prev_message_type = 'audio' -%}
+                    {%- elif item['type'] == 'video' -%}
+                        {{- '<|video|>' -}}
+                        {%- set ns.prev_message_type = 'video' -%}
+                    {%- endif -%}
+                {%- endfor -%}
+            {%- endif -%}
+            {%- endset -%}
+
+            {{- captured_content -}}
+            {%- set has_content = captured_content | trim | length > 0 -%}
+
+        {%- if ns.prev_message_type == 'tool_call' and not ns_tr_out.flag -%}
+            {{- '<|tool_response>' -}}
+        {%- elif not (ns_tr_out.flag and not has_content) -%}
+            {{- '<turn|>\n' -}}
+        {%- endif -%}
+    {%- endif -%}
+{%- endfor -%}
+
+{%- if add_generation_prompt -%}
+    {%- if ns.prev_message_type != 'tool_response' and ns.prev_message_type != 'tool_call' -%}
+        {{- '<|turn>model\n' -}}
+        {%- if not enable_thinking | default(false) -%}
+            {{- '<|channel>thought\n<channel|>' -}}
+        {%- endif -%}
+    {%- endif -%}
+{%- endif -%}

model/config.json

@@ -0,0 +1,216 @@
+{
+  "architectures": [
+    "Gemma4ForConditionalGeneration"
+  ],
+  "audio_config": {
+    "_name_or_path": "",
+    "architectures": null,
+    "attention_chunk_size": 12,
+    "attention_context_left": 13,
+    "attention_context_right": 0,
+    "attention_invalid_logits_value": -1000000000.0,
+    "attention_logit_cap": 50.0,
+    "chunk_size_feed_forward": 0,
+    "conv_kernel_size": 5,
+    "dtype": "bfloat16",
+    "gradient_clipping": 10000000000.0,
+    "hidden_act": "silu",
+    "hidden_size": 1024,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "initializer_range": 0.02,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "model_type": "gemma4_audio",
+    "num_attention_heads": 8,
+    "num_hidden_layers": 12,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_proj_dims": 1536,
+    "problem_type": null,
+    "residual_weight": 0.5,
+    "return_dict": true,
+    "rms_norm_eps": 1e-06,
+    "subsampling_conv_channels": [
+      128,
+      32
+    ],
+    "use_clipped_linears": true
+  },
+  "audio_token_id": 258881,
+  "boa_token_id": 256000,
+  "boi_token_id": 255999,
+  "dtype": "bfloat16",
+  "eoa_token_id": 258883,
+  "eoa_token_index": 258883,
+  "eoi_token_id": 258882,
+  "eos_token_id": [
+    1,
+    106
+  ],
+  "image_token_id": 258880,
+  "initializer_range": 0.02,
+  "model_type": "gemma4",
+  "text_config": {
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "attention_k_eq_v": true,
+    "bos_token_id": 2,
+    "dtype": "bfloat16",
+    "enable_moe_block": false,
+    "eos_token_id": 1,
+    "expert_intermediate_size": null,
+    "final_logit_softcapping": 30.0,
+    "global_head_dim": 512,
+    "head_dim": 256,
+    "hidden_activation": "gelu_pytorch_tanh",
+    "hidden_size": 5376,
+    "hidden_size_per_layer_input": 0,
+    "initializer_range": 0.02,
+    "intermediate_size": 21504,
+    "layer_types": [
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention"
+    ],
+    "max_position_embeddings": 262144,
+    "model_type": "gemma4_text",
+    "moe_intermediate_size": null,
+    "num_attention_heads": 32,
+    "num_experts": null,
+    "num_global_key_value_heads": 4,
+    "num_hidden_layers": 60,
+    "num_key_value_heads": 16,
+    "num_kv_shared_layers": 0,
+    "pad_token_id": 0,
+    "rms_norm_eps": 1e-06,
+    "rope_parameters": {
+      "full_attention": {
+        "partial_rotary_factor": 0.25,
+        "rope_theta": 1000000.0,
+        "rope_type": "proportional"
+      },
+      "sliding_attention": {
+        "rope_theta": 10000.0,
+        "rope_type": "default"
+      }
+    },
+    "sliding_window": 1024,
+    "tie_word_embeddings": true,
+    "top_k_experts": null,
+    "use_bidirectional_attention": "vision",
+    "use_cache": true,
+    "use_double_wide_mlp": false,
+    "vocab_size": 262144,
+    "vocab_size_per_layer_input": 262144
+  },
+  "tie_word_embeddings": true,
+  "transformers_version": "5.12.1",
+  "video_token_id": 258884,
+  "vision_config": {
+    "_name_or_path": "",
+    "architectures": null,
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "chunk_size_feed_forward": 0,
+    "default_output_length": 280,
+    "dtype": "bfloat16",
+    "global_head_dim": 72,
+    "head_dim": 72,
+    "hidden_activation": "gelu_pytorch_tanh",
+    "hidden_size": 1152,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "initializer_range": 0.02,
+    "intermediate_size": 4304,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "max_position_embeddings": 131072,
+    "model_type": "gemma4_vision",
+    "num_attention_heads": 16,
+    "num_hidden_layers": 27,
+    "num_key_value_heads": 16,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "patch_size": 16,
+    "pooling_kernel_size": 3,
+    "position_embedding_size": 10240,
+    "problem_type": null,
+    "return_dict": true,
+    "rms_norm_eps": 1e-06,
+    "rope_parameters": {
+      "rope_theta": 100.0,
+      "rope_type": "default"
+    },
+    "standardize": true,
+    "use_clipped_linears": false
+  },
+  "vision_soft_tokens_per_image": 280
+}

model/embed_audio.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1f961b58fe76eb58c936f7b48cfb45d87ae431a0d820b2a0f995bd22a936b19f
+size 101801106

model/gemma31b_with_audio_config.json

@@ -0,0 +1,216 @@
+{
+  "architectures": [
+    "Gemma4ForConditionalGeneration"
+  ],
+  "audio_config": {
+    "_name_or_path": "",
+    "architectures": null,
+    "attention_chunk_size": 12,
+    "attention_context_left": 13,
+    "attention_context_right": 0,
+    "attention_invalid_logits_value": -1000000000.0,
+    "attention_logit_cap": 50.0,
+    "chunk_size_feed_forward": 0,
+    "conv_kernel_size": 5,
+    "dtype": "bfloat16",
+    "gradient_clipping": 10000000000.0,
+    "hidden_act": "silu",
+    "hidden_size": 1024,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "initializer_range": 0.02,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "model_type": "gemma4_audio",
+    "num_attention_heads": 8,
+    "num_hidden_layers": 12,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_proj_dims": 1536,
+    "problem_type": null,
+    "residual_weight": 0.5,
+    "return_dict": true,
+    "rms_norm_eps": 1e-06,
+    "subsampling_conv_channels": [
+      128,
+      32
+    ],
+    "use_clipped_linears": true
+  },
+  "audio_token_id": 258881,
+  "boa_token_id": 256000,
+  "boi_token_id": 255999,
+  "bos_token_id": 2,
+  "dtype": "bfloat16",
+  "eoa_token_id": 258883,
+  "eoa_token_index": 258883,
+  "eoi_token_id": 258882,
+  "eos_token_id": 1,
+  "image_token_id": 258880,
+  "initializer_range": 0.02,
+  "model_type": "gemma4",
+  "pad_token_id": 0,
+  "text_config": {
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "attention_k_eq_v": true,
+    "bos_token_id": 2,
+    "dtype": "bfloat16",
+    "enable_moe_block": false,
+    "eos_token_id": 1,
+    "expert_intermediate_size": null,
+    "final_logit_softcapping": 30.0,
+    "global_head_dim": 512,
+    "head_dim": 256,
+    "hidden_activation": "gelu_pytorch_tanh",
+    "hidden_size": 5376,
+    "hidden_size_per_layer_input": 0,
+    "initializer_range": 0.02,
+    "intermediate_size": 21504,
+    "layer_types": [
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "sliding_attention",
+      "full_attention"
+    ],
+    "max_position_embeddings": 262144,
+    "model_type": "gemma4_text",
+    "moe_intermediate_size": null,
+    "num_attention_heads": 32,
+    "num_experts": null,
+    "num_global_key_value_heads": 4,
+    "num_hidden_layers": 60,
+    "num_key_value_heads": 16,
+    "num_kv_shared_layers": 0,
+    "pad_token_id": 0,
+    "rms_norm_eps": 1e-06,
+    "rope_parameters": {
+      "full_attention": {
+        "partial_rotary_factor": 0.25,
+        "rope_theta": 1000000.0,
+        "rope_type": "proportional"
+      },
+      "sliding_attention": {
+        "rope_theta": 10000.0,
+        "rope_type": "default"
+      }
+    },
+    "sliding_window": 1024,
+    "tie_word_embeddings": true,
+    "top_k_experts": null,
+    "use_bidirectional_attention": "vision",
+    "use_cache": true,
+    "use_double_wide_mlp": false,
+    "vocab_size": 262144,
+    "vocab_size_per_layer_input": 262144
+  },
+  "tie_word_embeddings": true,
+  "transformers_version": "5.12.1",
+  "use_cache": false,
+  "video_token_id": 258884,
+  "vision_config": {
+    "_name_or_path": "",
+    "architectures": null,
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "chunk_size_feed_forward": 0,
+    "default_output_length": 280,
+    "dtype": "bfloat16",
+    "global_head_dim": 72,
+    "head_dim": 72,
+    "hidden_activation": "gelu_pytorch_tanh",
+    "hidden_size": 1152,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "initializer_range": 0.02,
+    "intermediate_size": 4304,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "max_position_embeddings": 131072,
+    "model_type": "gemma4_vision",
+    "num_attention_heads": 16,
+    "num_hidden_layers": 27,
+    "num_key_value_heads": 16,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "patch_size": 16,
+    "pooling_kernel_size": 3,
+    "position_embedding_size": 10240,
+    "problem_type": null,
+    "return_dict": true,
+    "rms_norm_eps": 1e-06,
+    "rope_parameters": {
+      "rope_theta": 100.0,
+      "rope_type": "default"
+    },
+    "standardize": true,
+    "use_clipped_linears": false
+  },
+  "vision_soft_tokens_per_image": 280
+}

model/generation_config.json

@@ -0,0 +1,14 @@
+{
+  "bos_token_id": 2,
+  "do_sample": true,
+  "eos_token_id": [
+    1,
+    106,
+    50
+  ],
+  "pad_token_id": 0,
+  "temperature": 1.0,
+  "top_k": 64,
+  "top_p": 0.95,
+  "transformers_version": "5.12.1"
+}

model/model-00001-of-00007.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d6b0656ed09013386c59ad75c974d4ece51e3e8cd6fce0b05a49405b723be47a
+size 9777274982

model/model-00002-of-00007.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ad6ee3a08ec001a62cec0af7fe5ca8b62a9edaa2079211792f359cf3c3caa536
+size 9964586468

model/model-00003-of-00007.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d10cff174725b889e3a8197119c01076bad608030a853ddf359c46ad8189a0f3
+size 9975564812

model/model-00004-of-00007.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:732ce9afb9c62c56af850674a2139681f58a1cf314599810bea06897b329079a
+size 9843454454

model/model-00005-of-00007.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:25a6ce60aaf1b2ea49562d1c98d07c9d5b9ea7dce6d7cf987b82207046534d6d
+size 9920512908

model/model-00006-of-00007.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b25ed49d8e9c8c5208f9ecd1250172fc274b1a9c9ef4b5fef301c4bb81c465e2
+size 9799403148

model/model-00007-of-00007.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9fa309158e5fbc7626c7080c0e86b62bb875b2405d2615c63fd9df247f212b29
+size 3891800932

model/processor_config.json

@@ -0,0 +1,75 @@
+{
+  "audio_ms_per_token": 40,
+  "audio_seq_length": 750,
+  "feature_extractor": {
+    "dither": 0.0,
+    "feature_extractor_type": "Gemma4AudioFeatureExtractor",
+    "feature_size": 128,
+    "fft_length": 512,
+    "fft_overdrive": false,
+    "frame_length": 320,
+    "hop_length": 160,
+    "input_scale_factor": 1.0,
+    "max_frequency": 8000.0,
+    "mel_floor": 0.001,
+    "min_frequency": 0.0,
+    "padding_side": "right",
+    "padding_value": 0.0,
+    "per_bin_mean": null,
+    "per_bin_stddev": null,
+    "preemphasis": 0.0,
+    "preemphasis_htk_flavor": true,
+    "return_attention_mask": true,
+    "sampling_rate": 16000
+  },
+  "image_processor": {
+    "do_convert_rgb": true,
+    "do_normalize": false,
+    "do_rescale": true,
+    "do_resize": true,
+    "image_mean": [
+      0.0,
+      0.0,
+      0.0
+    ],
+    "image_processor_type": "Gemma4ImageProcessor",
+    "image_seq_length": 280,
+    "image_std": [
+      1.0,
+      1.0,
+      1.0
+    ],
+    "max_soft_tokens": 280,
+    "patch_size": 16,
+    "pooling_kernel_size": 3,
+    "resample": 3,
+    "rescale_factor": 0.00392156862745098
+  },
+  "image_seq_length": 280,
+  "processor_class": "Gemma4Processor",
+  "video_processor": {
+    "do_convert_rgb": true,
+    "do_normalize": true,
+    "do_rescale": true,
+    "do_resize": true,
+    "do_sample_frames": true,
+    "image_mean": [
+      0.0,
+      0.0,
+      0.0
+    ],
+    "image_std": [
+      1.0,
+      1.0,
+      1.0
+    ],
+    "max_soft_tokens": 70,
+    "num_frames": 32,
+    "patch_size": 16,
+    "pooling_kernel_size": 3,
+    "resample": 3,
+    "rescale_factor": 0.00392156862745098,
+    "return_metadata": false,
+    "video_processor_type": "Gemma4VideoProcessor"
+  }
+}

model/tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a2619fe11b50dbed06ac443c51d757b354d0b62d64baa514404d4e84e6713519
+size 32169780

model/tokenizer_config.json

@@ -0,0 +1,99 @@
+{
+  "audio_token": "<|audio|>",
+  "backend": "tokenizers",
+  "boa_token": "<|audio>",
+  "boi_token": "<|image>",
+  "bos_token": "<bos>",
+  "eoa_token": "<audio|>",
+  "eoc_token": "<channel|>",
+  "eoi_token": "<image|>",
+  "eos_token": "<eos>",
+  "eot_token": "<turn|>",
+  "escape_token": "<|\"|>",
+  "etc_token": "<tool_call|>",
+  "etd_token": "<tool|>",
+  "etr_token": "<tool_response|>",
+  "extra_special_tokens": [
+    "<|video|>"
+  ],
+  "image_token": "<|image|>",
+  "is_local": true,
+  "local_files_only": true,
+  "mask_token": "<mask>",
+  "max_length": null,
+  "model_max_length": 1000000000000000019884624838656,
+  "model_specific_special_tokens": {
+    "audio_token": "<|audio|>",
+    "boa_token": "<|audio>",
+    "boi_token": "<|image>",
+    "eoa_token": "<audio|>",
+    "eoc_token": "<channel|>",
+    "eoi_token": "<image|>",
+    "eot_token": "<turn|>",
+    "escape_token": "<|\"|>",
+    "etc_token": "<tool_call|>",
+    "etd_token": "<tool|>",
+    "etr_token": "<tool_response|>",
+    "image_token": "<|image|>",
+    "soc_token": "<|channel>",
+    "sot_token": "<|turn>",
+    "stc_token": "<|tool_call>",
+    "std_token": "<|tool>",
+    "str_token": "<|tool_response>",
+    "think_token": "<|think|>"
+  },
+  "pad_to_multiple_of": null,
+  "pad_token": "<pad>",
+  "pad_token_type_id": 0,
+  "padding_side": "left",
+  "processor_class": "Gemma4Processor",
+  "response_schema": {
+    "properties": {
+      "content": {
+        "type": "string"
+      },
+      "role": {
+        "const": "assistant"
+      },
+      "thinking": {
+        "type": "string"
+      },
+      "tool_calls": {
+        "items": {
+          "properties": {
+            "function": {
+              "properties": {
+                "arguments": {
+                  "additionalProperties": {},
+                  "type": "object",
+                  "x-parser": "gemma4-tool-call"
+                },
+                "name": {
+                  "type": "string"
+                }
+              },
+              "type": "object",
+              "x-regex": "call\\:(?P<name>\\w+)(?P<arguments>\\{.*\\})"
+            },
+            "type": {
+              "const": "function"
+            }
+          },
+          "type": "object"
+        },
+        "type": "array",
+        "x-regex-iterator": "<\\|tool_call>(.*?)<tool_call\\|>"
+      }
+    },
+    "type": "object",
+    "x-regex": "(\\<\\|channel\\>thought\\n(?P<thinking>.*?)\\<channel\\|\\>)?(?P<tool_calls>\\<\\|tool_call\\>.*\\<tool_call\\|\\>)?(?P<content>(?:(?!\\<turn\\|\\>)(?!\\<\\|tool_response\\>).)+)?(?:\\<turn\\|\\>|\\<\\|tool_response\\>)?"
+  },
+  "soc_token": "<|channel>",
+  "sot_token": "<|turn>",
+  "stc_token": "<|tool_call>",
+  "std_token": "<|tool>",
+  "str_token": "<|tool_response>",
+  "think_token": "<|think|>",
+  "tokenizer_class": "GemmaTokenizer",
+  "unk_token": "<unk>"
+}

parakeet/config.json

@@ -0,0 +1,49 @@
+{
+  "architectures": [
+    "ParakeetForTDT"
+  ],
+  "blank_token_id": 8192,
+  "decoder_hidden_size": 640,
+  "dtype": "float32",
+  "durations": [
+    0,
+    1,
+    2,
+    3,
+    4
+  ],
+  "encoder_config": {
+    "activation_dropout": 0.1,
+    "attention_bias": false,
+    "attention_dropout": 0.1,
+    "conv_kernel_size": 9,
+    "convolution_bias": false,
+    "dropout": 0.1,
+    "dropout_positions": 0.0,
+    "hidden_act": "silu",
+    "hidden_size": 1024,
+    "initializer_range": 0.02,
+    "intermediate_size": 4096,
+    "layerdrop": 0.1,
+    "max_position_embeddings": 5000,
+    "model_type": "parakeet_encoder",
+    "num_attention_heads": 8,
+    "num_hidden_layers": 24,
+    "num_key_value_heads": 8,
+    "num_mel_bins": 128,
+    "scale_input": false,
+    "subsampling_conv_channels": 256,
+    "subsampling_conv_kernel_size": 3,
+    "subsampling_conv_stride": 2,
+    "subsampling_factor": 8
+  },
+  "hidden_act": "relu",
+  "initializer_range": 0.02,
+  "is_encoder_decoder": true,
+  "max_symbols_per_step": 10,
+  "model_type": "parakeet_tdt",
+  "num_decoder_layers": 2,
+  "pad_token_id": 2,
+  "transformers_version": "5.6.0.dev0",
+  "vocab_size": 8193
+}

parakeet/generation_config.json

@@ -0,0 +1,16 @@
+{
+  "_from_model_config": true,
+  "decoder_start_token_id": 8192,
+  "eos_token_id": 3,
+  "output_attentions": false,
+  "output_hidden_states": false,
+  "pad_token_id": 2,
+  "suppress_tokens": [
+    8193,
+    8194,
+    8195,
+    8196,
+    8197
+  ],
+  "transformers_version": "5.6.0.dev0"
+}

parakeet/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3a2026366188c8c68598edbbff92f8d11590a08e0ae2e6775544e7b07d6a5e11
+size 2508311120

parakeet/processor_config.json

@@ -0,0 +1,16 @@
+{
+  "blank_token": "<blank>",
+  "feature_extractor": {
+    "feature_extractor_type": "ParakeetFeatureExtractor",
+    "feature_size": 128,
+    "hop_length": 160,
+    "n_fft": 512,
+    "padding_side": "right",
+    "padding_value": 0.0,
+    "preemphasis": 0.97,
+    "return_attention_mask": true,
+    "sampling_rate": 16000,
+    "win_length": 400
+  },
+  "processor_class": "ParakeetProcessor"
+}

parakeet/tokenizer_config.json

@@ -0,0 +1,10 @@
+{
+  "backend": "tokenizers",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|endoftext|>",
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "<pad>",
+  "processor_class": "ParakeetProcessor",
+  "tokenizer_class": "ParakeetTokenizer",
+  "unk_token": "<unk>"
+}

processor/chat_template.jinja

@@ -0,0 +1,363 @@
+{%- macro format_parameters(properties, required, filter_keys=false) -%}
+    {%- set standard_keys = ['description', 'type', 'properties', 'required', 'nullable'] -%}
+    {%- set ns = namespace(found_first=false) -%}
+    {%- for key, value in properties | dictsort -%}
+        {%- set add_comma = false -%}
+        {%- if not filter_keys or key not in standard_keys -%}
+            {%- if ns.found_first %},{% endif -%}
+            {%- set ns.found_first = true -%}
+            {{ key }}:{
+            {%- if value['description'] -%}
+                description:<|"|>{{ value['description'] }}<|"|>
+                {%- set add_comma = true -%}
+            {%- endif -%}
+            {%- if value['type'] | upper == 'STRING' -%}
+                {%- if value['enum'] -%}
+                    {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                    enum:{{ format_argument(value['enum']) }}
+                {%- endif -%}
+            {%- elif value['type'] | upper == 'ARRAY' -%}
+                {%- if value['items'] is mapping and value['items'] -%}
+                    {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                    items:{
+                    {%- set ns_items = namespace(found_first=false) -%}
+                    {%- for item_key, item_value in value['items'] | dictsort -%}
+                        {%- if item_value is not none -%}
+                            {%- if ns_items.found_first %},{% endif -%}
+                            {%- set ns_items.found_first = true -%}
+                            {%- if item_key == 'properties' -%}
+                                properties:{
+                                {%- if item_value is mapping -%}
+                                    {{- format_parameters(item_value, value['items']['required'] | default([])) -}}
+                                {%- endif -%}
+                                }
+                            {%- elif item_key == 'required' -%}
+                                required:[
+                                {%- for req_item in item_value -%}
+                                    <|"|>{{- req_item -}}<|"|>
+                                    {%- if not loop.last %},{% endif -%}
+                                {%- endfor -%}
+                                ]
+                            {%- elif item_key == 'type' -%}
+                                {%- if item_value is string -%}
+                                    type:{{ format_argument(item_value | upper) }}
+                                {%- else -%}
+                                    type:{{ format_argument(item_value | map('upper') | list) }}
+                                {%- endif -%}
+                            {%- else -%}
+                                {{ item_key }}:{{ format_argument(item_value) }}
+                            {%- endif -%}
+                        {%- endif -%}
+                    {%- endfor -%}
+                    }
+                {%- endif -%}
+            {%- endif -%}
+            {%- if value['nullable'] %}
+                {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                nullable:true
+            {%- endif -%}
+            {%- if value['type'] | upper == 'OBJECT' -%}
+                {%- if value['properties'] is defined and value['properties'] is mapping -%}
+                    {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                    properties:{
+                    {{- format_parameters(value['properties'], value['required'] | default([])) -}}
+                    }
+                {%- elif value is mapping -%}
+                    {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                    properties:{
+                    {{- format_parameters(value, value['required'] | default([]), filter_keys=true) -}}
+                    }
+                {%- endif -%}
+                {%- if value['required'] -%}
+                    {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+                    required:[
+                    {%- for item in value['required'] | default([]) -%}
+                        <|"|>{{- item -}}<|"|>
+                        {%- if not loop.last %},{% endif -%}
+                    {%- endfor -%}
+                    ]
+                {%- endif -%}
+            {%- endif -%}
+            {%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
+            type:<|"|>{{ value['type'] | upper }}<|"|>}
+        {%- endif -%}
+    {%- endfor -%}
+{%- endmacro -%}
+{%- macro format_function_declaration(tool_data) -%}
+    declaration:{{- tool_data['function']['name'] -}}{description:<|"|>{{- tool_data['function']['description'] -}}<|"|>
+    {%- set params = tool_data['function']['parameters'] -%}
+    {%- if params -%}
+        ,parameters:{
+        {%- if params['properties'] -%}
+            properties:{ {{- format_parameters(params['properties'], params['required']) -}} },
+        {%- endif -%}
+        {%- if params['required'] -%}
+            required:[
+            {%- for item in params['required'] -%}
+                <|"|>{{- item -}}<|"|>
+                {{- ',' if not loop.last -}}
+            {%- endfor -%}
+            ],
+        {%- endif -%}
+        {%- if params['type'] -%}
+            type:<|"|>{{- params['type'] | upper -}}<|"|>}
+        {%- endif -%}
+    {%- endif -%}
+    {%- if 'response' in tool_data['function'] -%}
+        {%- set response_declaration = tool_data['function']['response'] -%}
+        ,response:{
+        {%- if response_declaration['description'] -%}
+            description:<|"|>{{- response_declaration['description'] -}}<|"|>,
+        {%- endif -%}
+        {%- if response_declaration['type'] | upper == 'OBJECT' -%}
+            type:<|"|>{{- response_declaration['type'] | upper -}}<|"|>}
+        {%- endif -%}
+    {%- endif -%}
+    }
+{%- endmacro -%}
+{%- macro format_argument(argument, escape_keys=True) -%}
+    {%- if argument is string -%}
+        {{- '<|"|>' + argument + '<|"|>' -}}
+    {%- elif argument is boolean -%}
+        {{- 'true' if argument else 'false' -}}
+    {%- elif argument is mapping -%}
+        {{- '{' -}}
+        {%- set ns = namespace(found_first=false) -%}
+        {%- for key, value in argument | dictsort -%}
+            {%- if ns.found_first %},{% endif -%}
+            {%- set ns.found_first = true -%}
+            {%- if escape_keys -%}
+                {{- '<|"|>' + key + '<|"|>' -}}
+            {%- else -%}
+                {{- key -}}
+            {%- endif -%}
+            :{{- format_argument(value, escape_keys=escape_keys) -}}
+        {%- endfor -%}
+        {{- '}' -}}
+    {%- elif argument is sequence -%}
+        {{- '[' -}}
+        {%- for item in argument -%}
+            {{- format_argument(item, escape_keys=escape_keys) -}}
+            {%- if not loop.last %},{% endif -%}
+        {%- endfor -%}
+        {{- ']' -}}
+    {%- else -%}
+        {{- argument -}}
+    {%- endif -%}
+{%- endmacro -%}
+{%- macro strip_thinking(text) -%}
+    {%- set ns = namespace(result='') -%}
+    {%- for part in text.split('<channel|>') -%}
+        {%- if '<|channel>' in part -%}
+            {%- set ns.result = ns.result + part.split('<|channel>')[0] -%}
+        {%- else -%}
+            {%- set ns.result = ns.result + part -%}
+        {%- endif -%}
+    {%- endfor -%}
+    {{- ns.result | trim -}}
+{%- endmacro -%}
+
+{%- macro format_tool_response_block(tool_name, response) -%}
+    {{- '<|tool_response>' -}}
+    {%- if response is mapping -%}
+        {{- 'response:' + tool_name + '{' -}}
+        {%- for key, value in response | dictsort -%}
+            {{- key -}}:{{- format_argument(value, escape_keys=False) -}}
+            {%- if not loop.last %},{% endif -%}
+        {%- endfor -%}
+        {{- '}' -}}
+    {%- else -%}
+        {{- 'response:' + tool_name + '{value:' + format_argument(response, escape_keys=False) + '}' -}}
+    {%- endif -%}
+    {{- '<tool_response|>' -}}
+{%- endmacro -%}
+
+{%- set ns = namespace(prev_message_type=None) -%}
+{%- set loop_messages = messages -%}
+{{- bos_token -}}
+{#- Handle System/Tool Definitions Block -#}
+{%- if (enable_thinking is defined and enable_thinking) or tools or messages[0]['role'] in ['system', 'developer'] -%}
+    {{- '<|turn>system\n' -}}
+    {#- Inject Thinking token at the very top of the FIRST system turn -#}
+    {%- if enable_thinking is defined and enable_thinking -%}
+        {{- '<|think|>\n' -}}
+        {%- set ns.prev_message_type = 'think' -%}
+    {%- endif -%}
+    {%- if messages[0]['role'] in ['system', 'developer'] -%}
+        {%- if messages[0]['content'] is string -%}
+            {{- messages[0]['content'] | trim -}}
+        {%- elif messages[0]['content'] is sequence -%}
+            {%- for item in messages[0]['content'] -%}
+                {{- item['text'] | trim + ' '-}}
+            {%- endfor -%}
+        {%- endif -%}
+        {%- set loop_messages = messages[1:] -%}
+    {%- endif -%}
+    {%- if tools -%}
+        {%- for tool in tools %}
+            {{- '<|tool>' -}}
+            {{- format_function_declaration(tool) | trim -}}
+            {{- '<tool|>' -}}
+        {%- endfor %}
+        {%- set ns.prev_message_type = 'tool' -%}
+    {%- endif -%}
+    {{- '<turn|>\n' -}}
+{%- endif %}
+
+{#- Pre-scan: find last user message index for reasoning guard -#}
+{%- set ns_turn = namespace(last_user_idx=-1) -%}
+{%- for i in range(loop_messages | length) -%}
+    {%- if loop_messages[i]['role'] == 'user' -%}
+        {%- set ns_turn.last_user_idx = i -%}
+    {%- endif -%}
+{%- endfor -%}
+
+{#- Loop through messages -#}
+{%- for message in loop_messages -%}
+    {%- if message['role'] != 'tool' -%}
+    {%- set ns.prev_message_type = None -%}
+    {%- set role = 'model' if message['role'] == 'assistant' else message['role'] -%}
+    {#- Detect continuation: suppress duplicate <|turn>model when previous non-tool message was also assistant -#}
+    {%- set prev_nt = namespace(role=None, found=false) -%}
+    {%- if loop.index0 > 0 -%}
+        {%- for j in range(loop.index0 - 1, -1, -1) -%}
+            {%- if not prev_nt.found -%}
+                {%- if loop_messages[j]['role'] != 'tool' -%}
+                    {%- set prev_nt.role = loop_messages[j]['role'] -%}
+                    {%- set prev_nt.found = true -%}
+                {%- endif -%}
+            {%- endif -%}
+        {%- endfor -%}
+    {%- endif -%}
+    {%- set continue_same_model_turn = (role == 'model' and prev_nt.role == 'assistant') -%}
+    {%- if not continue_same_model_turn -%}
+        {{- '<|turn>' + role + '\n' }}
+    {%- endif -%}
+
+    {#- Render reasoning/reasoning_content as thinking channel -#}
+    {%- set thinking_text = message.get('reasoning') or message.get('reasoning_content') -%}
+    {%- if thinking_text and loop.index0 > ns_turn.last_user_idx and message.get('tool_calls') -%}
+        {{- '<|channel>thought\n' + thinking_text + '\n<channel|>' -}}
+    {%- endif -%}
+
+            {%- if message['tool_calls'] -%}
+                {%- for tool_call in message['tool_calls'] -%}
+                    {%- set function = tool_call['function'] -%}
+                    {{- '<|tool_call>call:' + function['name'] + '{' -}}
+                    {%- if function['arguments'] is mapping -%}
+                        {%- set ns_args = namespace(found_first=false) -%}
+                        {%- for key, value in function['arguments'] | dictsort -%}
+                            {%- if ns_args.found_first %},{% endif -%}
+                            {%- set ns_args.found_first = true -%}
+                            {{- key -}}:{{- format_argument(value, escape_keys=False) -}}
+                        {%- endfor -%}
+                    {%- elif function['arguments'] is string -%}
+                        {{- function['arguments'] -}}
+                    {%- endif -%}
+                    {{- '}<tool_call|>' -}}
+                {%- endfor -%}
+                {%- set ns.prev_message_type = 'tool_call' -%}
+            {%- endif -%}
+
+            {%- set ns_tr_out = namespace(flag=false) -%}
+            {%- if message.get('tool_responses') -%}
+                {#- Legacy: tool_responses embedded on the assistant message (Google/Gemma native) -#}
+                {%- for tool_response in message['tool_responses'] -%}
+                    {{- format_tool_response_block(tool_response['name'] | default('unknown'), tool_response['response']) -}}
+                    {%- set ns_tr_out.flag = true -%}
+                    {%- set ns.prev_message_type = 'tool_response' -%}
+                {%- endfor -%}
+            {%- elif message.get('tool_calls') -%}
+                {#- OpenAI Chat Completions: forward-scan consecutive role:tool messages -#}
+                {%- set ns_tool_scan = namespace(stopped=false) -%}
+                {%- for k in range(loop.index0 + 1, loop_messages | length) -%}
+                    {%- if ns_tool_scan.stopped -%}
+                    {%- elif loop_messages[k]['role'] != 'tool' -%}
+                        {%- set ns_tool_scan.stopped = true -%}
+                    {%- else -%}
+                        {%- set follow = loop_messages[k] -%}
+                        {#- Resolve tool_call_id to function name -#}
+                        {%- set ns_tname = namespace(name=follow.get('name') | default('unknown')) -%}
+                        {%- for tc in message['tool_calls'] -%}
+                            {%- if tc.get('id') == follow.get('tool_call_id') -%}
+                                {%- set ns_tname.name = tc['function']['name'] -%}
+                            {%- endif -%}
+                        {%- endfor -%}
+                        {#- Handle content as string or content-parts array -#}
+                        {%- set tool_body = follow.get('content') -%}
+                        {%- if tool_body is string -%}
+                            {{- format_tool_response_block(ns_tname.name, tool_body) -}}
+                        {%- elif tool_body is sequence and tool_body is not string -%}
+                            {%- set ns_txt = namespace(s='') -%}
+                            {%- for part in tool_body -%}
+                                {%- if part.get('type') == 'text' -%}
+                                    {%- set ns_txt.s = ns_txt.s + (part.get('text') | default('')) -%}
+                                {%- endif -%}
+                            {%- endfor -%}
+                            {{- format_tool_response_block(ns_tname.name, ns_txt.s) -}}
+                            {%- for part in tool_body -%}
+                                {%- if part.get('type') == 'image' -%}
+                                    {{- '<|image|>' -}}
+                                {%- elif part.get('type') == 'audio' -%}
+                                    {{- '<|audio|>' -}}
+                                {%- elif part.get('type') == 'video' -%}
+                                    {{- '<|video|>' -}}
+                                {%- endif -%}
+                            {%- endfor -%}
+                        {%- else -%}
+                            {{- format_tool_response_block(ns_tname.name, tool_body) -}}
+                        {%- endif -%}
+                        {%- set ns_tr_out.flag = true -%}
+                        {%- set ns.prev_message_type = 'tool_response' -%}
+                    {%- endif -%}
+                {%- endfor -%}
+            {%- endif -%}
+
+            {%- set captured_content -%}
+            {%- if message['content'] is string -%}
+                {%- if role == 'model' -%}
+                    {{- strip_thinking(message['content']) -}}
+                {%- else -%}
+                    {{- message['content'] | trim -}}
+                {%- endif -%}
+            {%- elif message['content'] is sequence -%}
+                {%- for item in message['content'] -%}
+                    {%- if item['type'] == 'text' -%}
+                        {%- if role == 'model' -%}
+                            {{- strip_thinking(item['text']) -}}
+                        {%- else -%}
+                            {{- item['text'] | trim -}}
+                        {%- endif -%}
+                    {%- elif item['type'] == 'image' -%}
+                        {{- '<|image|>' -}}
+                        {%- set ns.prev_message_type = 'image' -%}
+                    {%- elif item['type'] == 'audio' -%}
+                        {{- '<|audio|>' -}}
+                        {%- set ns.prev_message_type = 'audio' -%}
+                    {%- elif item['type'] == 'video' -%}
+                        {{- '<|video|>' -}}
+                        {%- set ns.prev_message_type = 'video' -%}
+                    {%- endif -%}
+                {%- endfor -%}
+            {%- endif -%}
+            {%- endset -%}
+
+            {{- captured_content -}}
+            {%- set has_content = captured_content | trim | length > 0 -%}
+
+        {%- if ns.prev_message_type == 'tool_call' and not ns_tr_out.flag -%}
+            {{- '<|tool_response>' -}}
+        {%- elif not (ns_tr_out.flag and not has_content) -%}
+            {{- '<turn|>\n' -}}
+        {%- endif -%}
+    {%- endif -%}
+{%- endfor -%}
+
+{%- if add_generation_prompt -%}
+    {%- if ns.prev_message_type != 'tool_response' and ns.prev_message_type != 'tool_call' -%}
+        {{- '<|turn>model\n' -}}
+        {%- if not enable_thinking | default(false) -%}
+            {{- '<|channel>thought\n<channel|>' -}}
+        {%- endif -%}
+    {%- endif -%}
+{%- endif -%}

processor/processor_config.json

@@ -0,0 +1,75 @@
+{
+  "audio_ms_per_token": 40,
+  "audio_seq_length": 750,
+  "feature_extractor": {
+    "dither": 0.0,
+    "feature_extractor_type": "Gemma4AudioFeatureExtractor",
+    "feature_size": 128,
+    "fft_length": 512,
+    "fft_overdrive": false,
+    "frame_length": 320,
+    "hop_length": 160,
+    "input_scale_factor": 1.0,
+    "max_frequency": 8000.0,
+    "mel_floor": 0.001,
+    "min_frequency": 0.0,
+    "padding_side": "right",
+    "padding_value": 0.0,
+    "per_bin_mean": null,
+    "per_bin_stddev": null,
+    "preemphasis": 0.0,
+    "preemphasis_htk_flavor": true,
+    "return_attention_mask": true,
+    "sampling_rate": 16000
+  },
+  "image_processor": {
+    "do_convert_rgb": true,
+    "do_normalize": false,
+    "do_rescale": true,
+    "do_resize": true,
+    "image_mean": [
+      0.0,
+      0.0,
+      0.0
+    ],
+    "image_processor_type": "Gemma4ImageProcessor",
+    "image_seq_length": 280,
+    "image_std": [
+      1.0,
+      1.0,
+      1.0
+    ],
+    "max_soft_tokens": 280,
+    "patch_size": 16,
+    "pooling_kernel_size": 3,
+    "resample": 3,
+    "rescale_factor": 0.00392156862745098
+  },
+  "image_seq_length": 280,
+  "processor_class": "Gemma4Processor",
+  "video_processor": {
+    "do_convert_rgb": true,
+    "do_normalize": true,
+    "do_rescale": true,
+    "do_resize": true,
+    "do_sample_frames": true,
+    "image_mean": [
+      0.0,
+      0.0,
+      0.0
+    ],
+    "image_std": [
+      1.0,
+      1.0,
+      1.0
+    ],
+    "max_soft_tokens": 70,
+    "num_frames": 32,
+    "patch_size": 16,
+    "pooling_kernel_size": 3,
+    "resample": 3,
+    "rescale_factor": 0.00392156862745098,
+    "return_metadata": false,
+    "video_processor_type": "Gemma4VideoProcessor"
+  }
+}

processor/tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a2619fe11b50dbed06ac443c51d757b354d0b62d64baa514404d4e84e6713519
+size 32169780

processor/tokenizer_config.json

@@ -0,0 +1,99 @@
+{
+  "audio_token": "<|audio|>",
+  "backend": "tokenizers",
+  "boa_token": "<|audio>",
+  "boi_token": "<|image>",
+  "bos_token": "<bos>",
+  "eoa_token": "<audio|>",
+  "eoc_token": "<channel|>",
+  "eoi_token": "<image|>",
+  "eos_token": "<eos>",
+  "eot_token": "<turn|>",
+  "escape_token": "<|\"|>",
+  "etc_token": "<tool_call|>",
+  "etd_token": "<tool|>",
+  "etr_token": "<tool_response|>",
+  "extra_special_tokens": [
+    "<|video|>"
+  ],
+  "image_token": "<|image|>",
+  "is_local": true,
+  "local_files_only": true,
+  "mask_token": "<mask>",
+  "max_length": null,
+  "model_max_length": 1000000000000000019884624838656,
+  "model_specific_special_tokens": {
+    "audio_token": "<|audio|>",
+    "boa_token": "<|audio>",
+    "boi_token": "<|image>",
+    "eoa_token": "<audio|>",
+    "eoc_token": "<channel|>",
+    "eoi_token": "<image|>",
+    "eot_token": "<turn|>",
+    "escape_token": "<|\"|>",
+    "etc_token": "<tool_call|>",
+    "etd_token": "<tool|>",
+    "etr_token": "<tool_response|>",
+    "image_token": "<|image|>",
+    "soc_token": "<|channel>",
+    "sot_token": "<|turn>",
+    "stc_token": "<|tool_call>",
+    "std_token": "<|tool>",
+    "str_token": "<|tool_response>",
+    "think_token": "<|think|>"
+  },
+  "pad_to_multiple_of": null,
+  "pad_token": "<pad>",
+  "pad_token_type_id": 0,
+  "padding_side": "left",
+  "processor_class": "Gemma4Processor",
+  "response_schema": {
+    "properties": {
+      "content": {
+        "type": "string"
+      },
+      "role": {
+        "const": "assistant"
+      },
+      "thinking": {
+        "type": "string"
+      },
+      "tool_calls": {
+        "items": {
+          "properties": {
+            "function": {
+              "properties": {
+                "arguments": {
+                  "additionalProperties": {},
+                  "type": "object",
+                  "x-parser": "gemma4-tool-call"
+                },
+                "name": {
+                  "type": "string"
+                }
+              },
+              "type": "object",
+              "x-regex": "call\\:(?P<name>\\w+)(?P<arguments>\\{.*\\})"
+            },
+            "type": {
+              "const": "function"
+            }
+          },
+          "type": "object"
+        },
+        "type": "array",
+        "x-regex-iterator": "<\\|tool_call>(.*?)<tool_call\\|>"
+      }
+    },
+    "type": "object",
+    "x-regex": "(\\<\\|channel\\>thought\\n(?P<thinking>.*?)\\<channel\\|\\>)?(?P<tool_calls>\\<\\|tool_call\\>.*\\<tool_call\\|\\>)?(?P<content>(?:(?!\\<turn\\|\\>)(?!\\<\\|tool_response\\>).)+)?(?:\\<turn\\|\\>|\\<\\|tool_response\\>)?"
+  },
+  "soc_token": "<|channel>",
+  "sot_token": "<|turn>",
+  "stc_token": "<|tool_call>",
+  "std_token": "<|tool>",
+  "str_token": "<|tool_response>",
+  "think_token": "<|think|>",
+  "tokenizer_class": "GemmaTokenizer",
+  "unk_token": "<unk>"
+}

requirements.txt

@@ -0,0 +1,19 @@
+# Inference dependencies for polished_model_v3.
+# ffmpeg/ffprobe are required as system binaries for video/audio extraction.
+
+--extra-index-url https://download.pytorch.org/whl/cu130
+
+accelerate==1.14.0
+av==17.1.0
+librosa==0.11.0
+numpy==2.4.4
+opencv-python-headless>=4.10.0
+safetensors==0.8.0
+scipy==1.18.0
+soundfile==0.14.0
+torch==2.12.1+cu130
+torchcodec==0.14.0
+torchvision==0.27.1+cu130
+transformers==5.12.1
+
+-e ./vllm

run_inference.sh

@@ -0,0 +1,5 @@
+set -euo pipefail
+
+cd "$(dirname "$0")"
+source /venv/main/bin/activate
+python infer.py

setup_vllm.sh

@@ -0,0 +1,22 @@
+set -euo pipefail
+
+cd "$(dirname "$0")"
+
+if [ ! -d vllm ]; then
+  if [ ! -f vllm.zip ]; then
+    echo "Missing vllm/ and vllm.zip" >&2
+    exit 1
+  fi
+  unzip -q vllm.zip
+fi
+
+python3 -m venv .venv
+source .venv/bin/activate
+
+pip install --upgrade pip setuptools wheel
+pip install -r requirements.txt
+
+python - <<'PY'
+import vllm
+print("vLLM import OK:", getattr(vllm, "__version__", "unknown"))
+PY

vllm.zip

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:641359f12335daec713a57084ecdc9b222916660440cf07c8b795bfab5707673
+size 328568395

vllm_caption_runtime.py

@@ -0,0 +1,254 @@
+"""Shared helpers for the polished vLLM caption examples."""
+
+from __future__ import annotations
+
+import subprocess
+import sys
+from pathlib import Path
+from typing import Any
+
+import cv2
+import numpy as np
+from transformers import AutoProcessor
+from vllm import LLM, SamplingParams
+
+CAPTION_LENGTH_LABELS = ("very small", "small", "medium", "large", "very large")
+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"),
+}
+
+
+def bool_from_value(value: Any, field_name: str) -> bool:
+    if isinstance(value, bool):
+        return value
+    if isinstance(value, str):
+        lowered = value.strip().lower()
+        if lowered in {"1", "true", "yes", "y", "on"}:
+            return True
+        if lowered in {"0", "false", "no", "n", "off"}:
+            return False
+    raise ValueError(f"{field_name} must be boolean-like, got {value!r}")
+
+
+def format_caption_settings_prompt(settings: dict[str, Any]) -> str:
+    caption_length = str(settings["caption_length"]).strip().lower()
+    if caption_length not in CAPTION_LENGTH_LABELS:
+        raise ValueError(f"caption_length must be one of {CAPTION_LENGTH_LABELS}, got {caption_length!r}")
+
+    include_watermark_info = bool_from_value(
+        settings["include_watermark_info"], "include_watermark_info"
+    )
+    has_repetition = bool_from_value(settings["has_repetition"], "has_repetition")
+    has_thinking = bool_from_value(settings["has_thinking"], "has_thinking")
+
+    normalized: dict[str, str] = {}
+    for field_name, choices in CAPTION_SETTING_FIELD_CHOICES.items():
+        value = str(settings[field_name]).strip().lower()
+        if value not in choices:
+            raise ValueError(f"{field_name} must be one of {choices}, got {value!r}")
+        normalized[field_name] = value
+
+    watermark_instruction = (
+        "Include watermark info." if include_watermark_info else "Do not include watermark info."
+    )
+    thinking_instruction = (
+        "Output thought JSON before the final caption."
+        if has_thinking
+        else "Do not output thought JSON; output only the caption."
+    )
+    setting_text = (
+        f"vulgarity={normalized['vulgarity']}; "
+        f"uncertainty={normalized['uncertainty']}; "
+        f"character_names={normalized['character_names']}; "
+        f"fluff={normalized['fluff']}; "
+        f"has_repetition={str(has_repetition).lower()}; "
+        f"has_thinking={str(has_thinking).lower()}; "
+        f"speculation={normalized['speculation']}; "
+        f"temporal_detail={normalized['temporal_detail']}; "
+        f"visual_specificity={normalized['visual_specificity']}; "
+        f"camera_detail={normalized['camera_detail']}; "
+        f"caption_style={normalized['caption_style']}"
+    )
+    return (
+        f"Write a {caption_length} caption for this clip using both the visuals and the audio. "
+        f"{watermark_instruction} {thinking_instruction} Match these caption settings: {setting_text}."
+    )
+
+
+def load_video_frames(path: str | Path, num_frames: int) -> tuple[np.ndarray, dict[str, Any]]:
+    video_path = str(path)
+    cap = cv2.VideoCapture(video_path)
+    if not cap.isOpened():
+        raise RuntimeError(f"Could not open video: {video_path}")
+
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT) or 0)
+    fps = float(cap.get(cv2.CAP_PROP_FPS) or 24.0)
+    if total_frames <= 0:
+        cap.release()
+        raise RuntimeError(f"Could not determine frame count: {video_path}")
+
+    indices = np.linspace(0, max(0, total_frames - 1), num_frames).round().astype(int)
+    frames = []
+    for idx in indices:
+        cap.set(cv2.CAP_PROP_POS_FRAMES, int(idx))
+        ok, frame_bgr = cap.read()
+        if ok:
+            frames.append(cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB))
+    cap.release()
+
+    if not frames:
+        raise RuntimeError(f"Could not read frames: {video_path}")
+
+    metadata = {
+        "fps": fps,
+        "duration": total_frames / fps if fps > 0 else 0.0,
+        "total_num_frames": total_frames,
+        "frames_indices": [int(x) for x in indices[: len(frames)]],
+        "video_backend": "opencv",
+        "do_sample_frames": False,
+    }
+    return np.stack(frames, axis=0), metadata
+
+
+def load_audio(path: str | Path, sampling_rate: int) -> tuple[np.ndarray, int]:
+    video_path = str(path)
+    cmd = [
+        "ffmpeg",
+        "-hide_banner",
+        "-loglevel",
+        "error",
+        "-i",
+        video_path,
+        "-vn",
+        "-ac",
+        "1",
+        "-ar",
+        str(sampling_rate),
+        "-f",
+        "f32le",
+        "-",
+    ]
+    result = subprocess.run(cmd, check=False, capture_output=True)
+    if result.returncode == 0 and result.stdout:
+        return np.frombuffer(result.stdout, dtype=np.float32), sampling_rate
+
+    frames, metadata = load_video_frames(video_path, 2)
+    del frames
+    duration = max(1.0, float(metadata.get("duration") or 1.0))
+    return np.zeros(max(1, int(duration * sampling_rate)), dtype=np.float32), sampling_rate
+
+
+def build_prompt(
+    *,
+    processor: AutoProcessor,
+    model_dir: str | Path,
+    video_path: str | Path,
+    prompt_override: str,
+    prompt_settings: dict[str, Any],
+) -> str:
+    del model_dir
+    prompt_text = prompt_override.strip() or format_caption_settings_prompt(prompt_settings)
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "video", "path": str(video_path)},
+                {"type": "text", "text": prompt_text},
+                {"type": "audio", "path": "/tmp/polished_model_v3_audio.wav"},
+            ],
+        },
+        {"role": "assistant", "content": [{"type": "text", "text": ""}]},
+    ]
+    prompt = processor.apply_chat_template(messages, tokenize=False)
+    assistant_tail = "<turn|>\n"
+    if prompt.endswith(assistant_tail):
+        prompt = prompt[: -len(assistant_tail)]
+    return prompt
+
+
+def build_vllm_request(
+    *,
+    processor: AutoProcessor,
+    model_dir: str | Path,
+    video_path: str | Path,
+    num_frames: int,
+    sampling_rate: int,
+    prompt_override: str,
+    prompt_settings: dict[str, Any],
+) -> dict[str, Any]:
+    video, video_metadata = load_video_frames(video_path, num_frames)
+    audio, sr = load_audio(video_path, sampling_rate)
+    prompt = build_prompt(
+        processor=processor,
+        model_dir=model_dir,
+        video_path=video_path,
+        prompt_override=prompt_override,
+        prompt_settings=prompt_settings,
+    )
+    return {
+        "prompt": prompt,
+        "multi_modal_data": {
+            "video": [(video, video_metadata)],
+            "audio": (audio, sr),
+        },
+    }
+
+
+def load_processor(model_dir: str | Path) -> AutoProcessor:
+    return AutoProcessor.from_pretrained(str(model_dir), local_files_only=True)
+
+
+def load_llm(
+    *,
+    model_dir: str | Path,
+    max_model_len: int,
+    max_num_seqs: int,
+    gpu_memory_utilization: float,
+    dtype: str,
+    enforce_eager: bool,
+    enable_prefix_caching: bool,
+    trust_remote_code: bool,
+) -> LLM:
+    return LLM(
+        model=str(model_dir),
+        tokenizer=str(model_dir),
+        max_model_len=max_model_len,
+        max_num_seqs=max_num_seqs,
+        gpu_memory_utilization=gpu_memory_utilization,
+        limit_mm_per_prompt={"video": 1, "audio": 1, "image": 0},
+        enforce_eager=enforce_eager,
+        enable_prefix_caching=enable_prefix_caching,
+        trust_remote_code=trust_remote_code,
+        dtype=dtype,
+    )
+
+
+def sampling_params(
+    *,
+    temperature: float,
+    max_tokens: int,
+    top_p: float,
+    repetition_penalty: float,
+) -> SamplingParams:
+    kwargs: dict[str, Any] = {
+        "temperature": temperature,
+        "max_tokens": max_tokens,
+        "repetition_penalty": repetition_penalty,
+    }
+    if temperature > 0:
+        kwargs["top_p"] = top_p
+    return SamplingParams(**kwargs)
+
+
+def ensure_local_vllm_source(script_dir: Path) -> None:
+    local_vllm = script_dir / "vllm"
+    if local_vllm.is_dir():
+        sys.path.insert(0, str(local_vllm))