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