Delete .ipynb_checkpoints

.ipynb_checkpoints/caption_folder-checkpoint.py

@@ -1,206 +0,0 @@
-"""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

@@ -1,771 +0,0 @@
-"""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

@@ -1,519 +0,0 @@
-"""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

@@ -1,19 +0,0 @@
-# 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

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

.ipynb_checkpoints/setup_vllm-checkpoint.sh

@@ -1,22 +0,0 @@
-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

@@ -1,254 +0,0 @@
-"""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))