"""AudioEmb audio-text embedding model.""" import math from torch.utils.checkpoint import checkpoint import wave from io import BytesIO from pathlib import Path from tempfile import NamedTemporaryFile from urllib.parse import urlparse from urllib.request import urlopen from .configuration_audio_emb import AudioEmbConfig import collections import collections.abc from dataclasses import dataclass import torch import torch.nn as nn import torchaudio.functional as F from torch import Tensor from torch.nn.functional import scaled_dot_product_attention from typing import Any, Dict, Callable, Iterable, List, Optional, Sequence, Tuple, Union, cast from transformers import PreTrainedModel, PreTrainedConfig, GenerationMixin from transformers import AutoTokenizer from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import ( Qwen2_5OmniTextConfig, ) from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import ( Qwen2_5OmniThinkerTextModel, ) from transformers.cache_utils import Cache from transformers.modeling_outputs import BaseModelOutputWithPast, ModelOutput from transformers.utils import can_return_tuple import copy try: import torchaudio except ImportError: torchaudio = None _Tuple2 = Union[int, Tuple[int, int], Sequence[int]] TARGET_SR = 16000 QUERY_INSTRUCTION = "Based on the question asked in the text query and context in the audio query, retrieve the relevant text document associated with that question." DOC_INSTRUCTION = "Represent the user's input." def _resolve_tuple2(x: _Tuple2) -> Tuple[int, int]: if isinstance(x, collections.abc.Sequence): assert len(x) == 2, ( f"Expected a sequence of length 2, got {x} with length {len(x)}" ) return cast(Tuple[int, int], tuple(x)) return (x, x) DASHENG_ARCH_CONFIG = { "audio_encoder_config": { "attn_drop_rate": 0.0, "center": True, "depth": 32, "drop_rate": 0.0, "embed_dim": 1280, "f_max": 8000.0, "f_min": 0.0, "hop_length": 160, "init_values": None, "input_channels": 1, "mlp_ratio": 4.0, "model_type": "midashenglm_dasheng_encoder", "n_fft": 512, "n_mels": 64, "num_heads": 16, "outputdim": 527, "patch_size": [ 64, 4 ], "patch_stride": [ 64, 4 ], "qkv_bias": True, "sample_rate": 16000, "target_length": 1008, "win_length": 512 }, "audio_projector_config": { "in_dim": 1280, "downsample_rate": 5, "out_dim": 3584, }, "text_config": { "attention_dropout": 0.0, "hidden_act": "silu", "hidden_size": 3584, "init_std": 0.02, "initializer_range": 0.02, "intermediate_size": 18944, "max_position_embeddings": 32768, "max_window_layers": 28, "model_type": "qwen2_5_omni_text", "num_attention_heads": 28, "num_hidden_layers": 28, "num_key_value_heads": 4, "rms_norm_eps": 1e-06, "rope_scaling": { "mrope_section": [ 16, 24, 24 ], "rope_type": "default", "type": "default" }, "rope_theta": 1000000.0, "sliding_window": 32768, "use_cache": True, "use_sliding_window": False, "vocab_size": 152064 }, "lite_random_decoder_config": { "attention_dropout": 0.0, "hidden_act": "silu", "hidden_size": 576, "init_std": 0.02, "initializer_range": 0.02, "intermediate_size": 1536, "max_position_embeddings": 2048, "max_window_layers": 12, "model_type": "qwen2_5_omni_text", "num_attention_heads": 8, "num_hidden_layers": 12, "num_key_value_heads": 4, "rms_norm_eps": 1e-06, "rope_scaling": { "mrope_section": [ 12, 12, 12 ], "rope_type": "default", "type": "default" }, "rope_theta": 1000000.0, "sliding_window": 2048, "use_cache": True, "use_sliding_window": False, "vocab_size": 152064 } } class DashengConfig(PreTrainedConfig): model_type = "midashenglm_dasheng_encoder" def __init__( self, embed_dim: int = 768, outputdim: int = 527, patch_size: Union[int, Tuple[int, int]] = 16, patch_stride: Union[int, Tuple[int, int]] = 16, input_channels: int = 1, target_length: int = 1012, depth: int = 12, num_heads: int = 12, mlp_ratio: float = 4.0, qkv_bias: bool = True, init_values: Optional[float] = None, drop_rate: float = 0.0, attn_drop_rate: float = 0.0, f_min: float = 0.0, f_max: float = 8000.0, center: bool = True, win_length: int = 512, hop_length: int = 160, sample_rate: int = 16000, n_fft: int = 512, n_mels: int = 64, **kwargs, ): self.embed_dim = embed_dim self.outputdim = outputdim self.patch_size = patch_size self.patch_stride = patch_stride self.input_channels = input_channels self.target_length = target_length self.depth = depth self.num_heads = num_heads self.mlp_ratio = mlp_ratio self.qkv_bias = qkv_bias self.init_values = init_values self.drop_rate = drop_rate self.attn_drop_rate = attn_drop_rate self.f_min = f_min self.f_max = f_max self.center = center self.win_length = win_length self.hop_length = hop_length self.sample_rate = sample_rate self.n_fft = n_fft self.n_mels = n_mels super().__init__(**kwargs) class AudioPatchEmbed(nn.Module): def __init__( self, input_size: _Tuple2 = 64, patch_size: _Tuple2 = 16, patch_stride: _Tuple2 = 16, in_chans: int = 1, embed_dim: int = 768, norm_layer: Optional[Callable] = None, flatten: bool = False, ): super().__init__() self.input_size = _resolve_tuple2(input_size) self.patch_size = _resolve_tuple2(patch_size) self.patch_stride = _resolve_tuple2(patch_stride) self.grid_size = ( self.input_size[0] // self.patch_stride[0], self.input_size[1] // self.patch_stride[1], ) self.num_patches = self.grid_size[0] * self.grid_size[1] self.flatten = flatten self.proj = nn.Conv2d( in_chans, embed_dim, kernel_size=self.patch_size, stride=self.patch_stride, ) self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity() def forward(self, x: torch.Tensor) -> torch.Tensor: x = self.proj(x) if self.flatten: x = torch.permute( torch.flatten(x, 2, 3), (0, 2, 1) ) # rearrange(x, "b c f t -> b (f t) c") x = self.norm(x) return x class LayerScale(nn.Module): def __init__(self, dim, init_values=1e-5, inplace=False): super().__init__() self.inplace = inplace self.gamma = nn.Parameter(init_values * torch.ones(dim)) def forward(self, x: torch.Tensor) -> torch.Tensor: return x.mul_(self.gamma) if self.inplace else x * self.gamma class DashengMlp(nn.Module): def __init__( self, in_features: int, hidden_features: Optional[int] = None, out_features: Optional[int] = None, drop: float = 0.0, ): super().__init__() out_features = out_features or in_features hidden_features = hidden_features or in_features self.fc1 = nn.Linear(in_features, hidden_features) self.act = nn.GELU() self.fc2 = nn.Linear(hidden_features, out_features) self.drop = nn.Dropout(drop) def forward(self, x: torch.Tensor) -> torch.Tensor: x = self.fc1(x) x = self.act(x) x = self.drop(x) x = self.fc2(x) x = self.drop(x) return x class DashengAttention(nn.Module): def __init__( self, dim: int, num_heads: int = 8, qkv_bias: bool = False, attn_drop: float = 0.0, proj_drop: float = 0.0, ): super().__init__() assert dim % num_heads == 0, "dim should be divisible by num_heads" self.num_heads = num_heads head_dim = dim // num_heads self.scale = head_dim**-0.5 self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) self.attn_drop = nn.Dropout(attn_drop) self.proj = nn.Linear(dim, dim) self.proj_drop = nn.Dropout(proj_drop) def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None): B, N, C = x.shape q, k, v = ( self.qkv(x) .reshape(B, N, 3, self.num_heads, C // self.num_heads) .permute(2, 0, 3, 1, 4) .unbind(0) ) x = scaled_dot_product_attention( q, k, v, attn_mask=mask[:, None, None, :] if mask is not None else None, ) x = x.transpose(1, 2).reshape(B, N, C) x = self.proj(x) x = self.proj_drop(x) return x class DashengBlock(nn.Module): def __init__( self, dim: int, num_heads: int, mlp_ratio: float = 4.0, qkv_bias: bool = False, drop: float = 0.0, attn_drop: float = 0.0, init_values: Optional[float] = None, ): super().__init__() self.norm1 = nn.LayerNorm(dim, eps=1e-6) self.attn = DashengAttention( dim, num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop, ) self.ls1 = ( LayerScale(dim, init_values=init_values) if init_values else nn.Identity() ) self.norm2 = nn.LayerNorm(dim, eps=1e-6) self.mlp = DashengMlp( in_features=dim, hidden_features=int(dim * mlp_ratio), drop=drop, ) self.ls2 = ( LayerScale(dim, init_values=init_values) if init_values else nn.Identity() ) # Kwargs usually has a mask parameter that is passed to Attention def forward( self, x: torch.Tensor, mask: Optional[torch.Tensor] = None, ) -> torch.Tensor: x = x + self.ls1(self.attn(self.norm1(x), mask)) x = x + self.ls2(self.mlp(self.norm2(x))) return x class DashengFrontend(nn.Module): def __init__(self, config: DashengConfig): super().__init__() self.config = config spectrogram_window, melscale_fbanks = self._build_frontend_buffers() self.register_buffer( "spectrogram_window", spectrogram_window, persistent=False, ) self.spectrogram_window: torch.Tensor self.register_buffer("melscale_fbanks", melscale_fbanks, persistent=False) self.melscale_fbanks: torch.Tensor def _build_frontend_buffers(self) -> tuple[torch.Tensor, torch.Tensor]: # Build on CPU explicitly: from_pretrained may construct modules on meta device. with torch.device("cpu"): spectrogram_window = torch.hann_window( self.config.win_length, dtype=torch.float32, ) melscale_fbanks = F.melscale_fbanks( n_freqs=self.config.n_fft // 2 + 1, f_min=self.config.f_min, f_max=self.config.f_max, n_mels=self.config.n_mels, sample_rate=self.config.sample_rate, ).to(torch.float32) return spectrogram_window, melscale_fbanks def ensure_frontend_buffers(self, device: torch.device) -> None: """Self-heal non-persistent audio frontend buffers if corrupted/uninitialized.""" expected_win_shape = (self.config.win_length,) expected_fb_shape = (self.config.n_fft // 2 + 1, self.config.n_mels) def _is_bad(name: str, tensor: torch.Tensor, expected_shape: tuple[int, ...]) -> bool: if tensor is None: return True if getattr(tensor, "is_meta", False): return True if tuple(tensor.shape) != expected_shape: return True t = tensor.detach().float() if not torch.isfinite(t).all().item(): return True if t.numel() > 0 and t.abs().max().item() > 1e6: return True return False win_bad = _is_bad("spectrogram_window", self.spectrogram_window, expected_win_shape) fb_bad = _is_bad("melscale_fbanks", self.melscale_fbanks, expected_fb_shape) if win_bad or fb_bad: new_win, new_fb = self._build_frontend_buffers() self.spectrogram_window = new_win.to(device=device) self.melscale_fbanks = new_fb.to(device=device) print( f"[WARN] Rebuilt frontend buffers (win_bad={win_bad}, fb_bad={fb_bad})", flush=True, ) else: if self.spectrogram_window.device != device: self.spectrogram_window = self.spectrogram_window.to(device=device) if self.melscale_fbanks.device != device: self.melscale_fbanks = self.melscale_fbanks.to(device=device) def forward(self, waveform: torch.Tensor) -> torch.Tensor: self.ensure_frontend_buffers(waveform.device) spectrogram = F.spectrogram( waveform=waveform.to(torch.float32), pad=0, window=self.spectrogram_window, n_fft=self.config.n_fft, hop_length=self.config.hop_length, win_length=self.config.win_length, power=2, normalized=False, center=self.config.center, ) mel_spectrogram = (spectrogram.mT @ self.melscale_fbanks.to(torch.float32)).mT # x has shape [batch, freq, time]. # F.amplitude_to_DB accepts inputs shaped as: # - [freq, time] # - [channel, freq, time] # - [..., channel, freq, time] # Here we insert a channel dimension of size 1 before calling it, # then remove that extra dimension afterward. log_mel_spectrogram = F.amplitude_to_DB( mel_spectrogram.unsqueeze(1), multiplier=10, amin=1e-10, db_multiplier=0, top_db=120, ).squeeze(1) return log_mel_spectrogram.to(waveform.dtype) class FixedAffine2d(nn.Module): """ Per-channel fixed affine transform: y = x * scale + bias where scale/bias are broadcast on (B, C, H, W). """ def __init__(self, scale: torch.Tensor, bias: torch.Tensor): super().__init__() self.register_buffer("scale", scale.reshape(1, -1, 1, 1)) self.register_buffer("bias", bias.reshape(1, -1, 1, 1)) @classmethod def from_batchnorm2d(cls, bn: nn.BatchNorm2d) -> "FixedAffine2d": if bn.running_mean is None or bn.running_var is None: raise ValueError("BatchNorm2d must have running stats to be converted.") if bn.affine: gamma = bn.weight.detach() beta = bn.bias.detach() else: gamma = torch.ones_like(bn.running_mean) beta = torch.zeros_like(bn.running_mean) running_mean = bn.running_mean.detach() running_var = bn.running_var.detach() scale = gamma / torch.sqrt(running_var + bn.eps) bias = beta - running_mean * scale return cls(scale=scale, bias=bias) def forward(self, x: torch.Tensor) -> torch.Tensor: return x * self.scale + self.bias class DashengAudioTransformer(PreTrainedModel): config_class = DashengConfig supports_gradient_checkpointing = True def __init__(self, config: DashengConfig): super().__init__(config) self.target_length = config.target_length self.embed_dim = config.embed_dim self.hop_length = config.hop_length self.gradient_checkpointing = False self.front_end = DashengFrontend(config) self.init_bn = nn.BatchNorm2d(config.n_mels, momentum=0.01) self.patch_embed = AudioPatchEmbed( input_size=(config.n_mels, config.target_length), embed_dim=config.embed_dim, in_chans=config.input_channels, patch_size=config.patch_size, flatten=False, patch_stride=config.patch_stride, ) self.time_pos_embed = nn.Parameter( torch.randn(1, config.embed_dim, 1, self.patch_embed.grid_size[1]) * 0.02 ) self.freq_pos_embed = nn.Parameter( torch.randn(1, config.embed_dim, self.patch_embed.grid_size[0], 1) * 0.02 ) self.pos_drop = nn.Dropout(p=config.drop_rate) self.blocks = nn.ModuleList( DashengBlock( dim=config.embed_dim, num_heads=config.num_heads, mlp_ratio=config.mlp_ratio, qkv_bias=config.qkv_bias, init_values=config.init_values, drop=config.drop_rate, attn_drop=config.attn_drop_rate, ) for _ in range(config.depth) ) self.norm = nn.LayerNorm(config.embed_dim, eps=1e-6) self.post_init() def replace_init_bn_with_fixed_affine(self): """ Call this after checkpoint is loaded and before inference/export. """ if isinstance(self.init_bn, nn.BatchNorm2d): self.init_bn.eval() self.init_bn = FixedAffine2d.from_batchnorm2d(self.init_bn) def forward_features( self, x: torch.Tensor, mask: Optional[torch.Tensor] = None, ) -> torch.Tensor: t = x.shape[-1] x = x + self.time_pos_embed[:, :, :, :t] x = ( x + self.freq_pos_embed[:, :, :, :] ) # Just to support __getitem__ in posembed x = torch.permute( torch.flatten(x, 2, 3), (0, 2, 1) ) # rearrange(x, "b c f t -> b (f t) c") x = self.pos_drop(x) for block in self.blocks: if self.gradient_checkpointing and self.training: x = self._gradient_checkpointing_func(block, x, mask) else: x = block(x, mask) x = self.norm(x) return x def _to_mask(self, lengths: torch.Tensor, max_length: int) -> torch.Tensor: batch_size = len(lengths) idx = torch.arange(max_length, device=lengths.device) idx = idx.repeat(batch_size).view(batch_size, max_length) mask = (idx < lengths.unsqueeze(-1)).bool() return mask def forward( self, x: torch.Tensor, x_length: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: x = self.front_end(x) target_length_in_patches = self.target_length // 4 x = x.unsqueeze(1) x = torch.permute(x, (0, 2, 1, 3)) x = self.init_bn(x) x = torch.permute(x, (0, 2, 1, 3)) x = self.patch_embed(x) t = x.shape[-1] input_splits = x.split(target_length_in_patches, dim=-1) if x_length is not None: assert len(x_length) == len(x), ( "batchsizes of input x and x_length need to be same" ) assert x_length.ndim == 1, "Lengths are of size (B,)" scaled_lengths = (x_length / (self.hop_length * 4)).long() mask = self._to_mask(max_length=t, lengths=scaled_lengths) split_masks = mask.split(target_length_in_patches, dim=-1) else: mask = None split_masks = [None] * len(input_splits) outputs = [] for split_x, split_mask in zip(input_splits, split_masks): split_x = self.forward_features(split_x, mask=split_mask) outputs.append(split_x) x = torch.cat(outputs, dim=1) return x, mask class AudioProjectorSubsample(nn.Module): def __init__( self, in_dim: int, out_dim: int, downsample_rate=5, dtype: Optional[torch.dtype] = None, ): super().__init__() self.k = downsample_rate self.out_dim = out_dim self.net = nn.Sequential( nn.Linear(in_dim * self.k, out_dim, dtype=dtype), nn.GELU(), nn.Linear(out_dim, out_dim, dtype=dtype), ) def forward(self, x, mask=None): batch_size, seq_len, dim = x.shape num_frames_to_discard = seq_len % self.k if num_frames_to_discard > 0: x = x[:, :-num_frames_to_discard, :] if mask is not None: mask = mask[:, :-num_frames_to_discard] if mask is None: mask = torch.ones(x.shape[:-1], dtype=torch.long, device=x.device) x = x.reshape( batch_size, -1, self.k * dim ) # rearrange(x, "b (s k) d -> b s (k d)", k=self.k) x = self.net(x) mask = mask.reshape( batch_size, -1, self.k ) # rearrange(mask, "b (s k) -> b s k", k=self.k) mask = mask.any(dim=-1).long() return x, mask @dataclass class Qwen25OmniTextModelOutput(ModelOutput): loss: Optional[torch.FloatTensor] = None logits: Optional[torch.FloatTensor] = None past_key_values: Optional[Cache] = None hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None attentions: Optional[Tuple[torch.FloatTensor, ...]] = None class Qwen25OmniThinkerTextOnlyDecoder(PreTrainedModel, GenerationMixin): config_class = Qwen2_5OmniTextConfig _supports_flash_attn_2 = True _supports_sdpa = True _supports_cache_class = True _supports_static_cache = True def __init__(self, config: Qwen2_5OmniTextConfig): super().__init__(config) self.model = Qwen2_5OmniThinkerTextModel._from_config(config) self.lm_head = nn.Linear( config.hidden_size, config.vocab_size, bias=False, ) self.post_init() @can_return_tuple def forward( self, input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, labels: Optional[torch.Tensor] = None, **kwargs, ) -> Union[Tuple, Qwen25OmniTextModelOutput]: if attention_mask is not None and position_ids is None: position_ids = ( attention_mask.long() .cumsum(dim=-1) .masked_fill_(attention_mask == 0, 1) - 1 ) outputs: BaseModelOutputWithPast = self.model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, cache_position=cache_position, return_dict=True, ) hidden_states = outputs.last_hidden_state logits = self.lm_head(hidden_states) loss = ( self.loss_function( logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs, ) if labels is not None else None ) return Qwen25OmniTextModelOutput( loss=loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, ) # Hardcoded architecture / training choices (not exposed in config.json). USE_LOGIT_SCALE = True HIDDEN_SIZE = 3584 class AudioEmbModel(PreTrainedModel): config_class = AudioEmbConfig def __init__(self, config: AudioEmbConfig): super().__init__(config) text_config = Qwen2_5OmniTextConfig(**DASHENG_ARCH_CONFIG["text_config"]) decoder = Qwen25OmniThinkerTextOnlyDecoder(text_config) self.model = decoder.model self.dasheng = DashengAudioTransformer( DashengConfig(**DASHENG_ARCH_CONFIG["audio_encoder_config"]) ) self.dasheng_down = AudioProjectorSubsample(**DASHENG_ARCH_CONFIG["audio_projector_config"]) self.dasheng_proj = nn.Identity() # Placeholder shapes match exported checkpoints (overwritten by load_state_dict). self.register_buffer("audio_start_token", torch.zeros(1, dtype=torch.long)) self.register_buffer("audio_end_token", torch.zeros(1, dtype=torch.long)) self.register_buffer("eos_token", torch.zeros(7, dtype=torch.long)) self.hidden_size = self.model.config.hidden_size self.use_checkpointing = False self.checkpoint_reentrant = False if USE_LOGIT_SCALE: init_value = math.log(1 / 0.07) self.register_buffer( "logit_scale", torch.tensor([init_value], dtype=torch.float32) ) else: self.logit_scale = None self.siglip_head = nn.Linear(self.hidden_size, self.hidden_size) self.dasheng.replace_init_bn_with_fixed_affine() self._tokenizer = None self.post_init() def set_tokenizer(self, tokenizer) -> None: """Attach a tokenizer for high-level encode APIs.""" self._tokenizer = tokenizer def _resolve_tokenizer(self): if self._tokenizer is not None: return self._tokenizer tokenizer = AutoTokenizer.from_pretrained(self.name_or_path, trust_remote_code=True) self._tokenizer = tokenizer return tokenizer @staticmethod def _to_list(x: Any) -> list[Any]: if x is None: return [] if isinstance(x, (list, tuple)): return list(x) return [x] @staticmethod def _build_prompt(instruction: str, text: Optional[str]) -> str: system_part = f"<|im_start|>system\n{instruction}<|im_end|>\n" if text is None: user_part = "<|im_start|>user\n" else: user_part = f"<|im_start|>user\n{text}" return system_part + user_part def _prepare_text_batch( self, tokenizer, texts: list[Optional[str]], *, task: str, instruction: Optional[str], device: torch.device, ) -> tuple[torch.Tensor, torch.Tensor]: if task not in ("query", "document"): raise ValueError(f"Unsupported task={task}. Use 'query' or 'document'.") default_instruction = QUERY_INSTRUCTION if task == "query" else DOC_INSTRUCTION instruction = instruction or default_instruction prompts = [self._build_prompt(instruction, text) for text in texts] encoded = tokenizer(prompts, padding=True, add_special_tokens=False, return_tensors="pt") text_ids = encoded["input_ids"].to(device) text_lens = encoded["attention_mask"].to(device).sum(dim=1) return text_ids, text_lens def _load_audio_path(self, path: Union[str, Path], target_sr: int) -> torch.Tensor: raw_path = str(path) parsed = urlparse(raw_path) is_remote_url = parsed.scheme in ("http", "https") local_path = None if is_remote_url else Path(path) suffix_source = Path(parsed.path) if is_remote_url else local_path suffix = suffix_source.suffix.lower() if is_remote_url: with urlopen(raw_path) as resp: audio_bytes = resp.read() source_for_torchaudio = NamedTemporaryFile( suffix=suffix or ".audio", delete=False, ) source_for_torchaudio.write(audio_bytes) source_for_torchaudio.flush() source_for_torchaudio.close() else: source_for_torchaudio = None path_for_wave = BytesIO(audio_bytes) if is_remote_url else str(local_path) try: if suffix in (".wav", ".wave"): with wave.open(path_for_wave, "rb") as wf: sr = wf.getframerate() n_channels = wf.getnchannels() sample_width = wf.getsampwidth() raw = wf.readframes(wf.getnframes()) if sample_width == 1: audio = torch.frombuffer(bytearray(raw), dtype=torch.uint8).float() audio = (audio - 128.0) / 128.0 elif sample_width == 2: audio = torch.frombuffer(bytearray(raw), dtype=torch.int16).float() / 32768.0 elif sample_width == 4: audio = torch.frombuffer(bytearray(raw), dtype=torch.int32).float() / 2147483648.0 else: raise ValueError(f"Unsupported WAV sample width: {sample_width}") if n_channels > 1: audio = audio.reshape(-1, n_channels).mean(dim=1) else: if torchaudio is None: raise ImportError("torchaudio is required for non-WAV audio paths.") load_target = ( source_for_torchaudio.name if is_remote_url else str(local_path) ) waveform, sr = torchaudio.load(load_target) if waveform.shape[0] > 1: waveform = waveform.mean(dim=0, keepdim=True) audio = waveform.squeeze(0) finally: if source_for_torchaudio is not None: try: Path(source_for_torchaudio.name).unlink(missing_ok=True) except OSError: pass if sr != target_sr: if torchaudio is None: raise ImportError("torchaudio is required for resampling.") audio = torchaudio.functional.resample(audio.unsqueeze(0), sr, target_sr).squeeze(0) return audio.float() def _prepare_audio_batch( self, audio_items: list[Optional[Union[str, Path, torch.Tensor]]], *, target_sr: int, device: torch.device, ) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor]]: tensor_list: list[torch.Tensor] = [] lens_list: list[int] = [] has_audio = False for item in audio_items: if item is None: tensor_list.append(torch.zeros(1, dtype=torch.float32)) lens_list.append(0) continue has_audio = True if isinstance(item, (str, Path)): wav = self._load_audio_path(item, target_sr=target_sr) elif isinstance(item, torch.Tensor): wav = item.detach().float().cpu() if wav.dim() == 2: wav = wav.mean(dim=0) elif wav.dim() != 1: raise ValueError("Audio tensor must be 1D waveform or 2D [channels, length].") else: raise TypeError(f"Unsupported audio item type: {type(item)}") if wav.numel() == 0: wav = torch.zeros(1, dtype=torch.float32) length = 0 else: length = int(wav.numel()) tensor_list.append(wav) lens_list.append(length) if not has_audio: return None, None max_len = max(t.numel() for t in tensor_list) padded = torch.zeros(len(tensor_list), max_len, dtype=torch.float32) for i, wav in enumerate(tensor_list): L = wav.numel() if L > 0: padded[i, :L] = wav return padded.to(device), torch.tensor(lens_list, dtype=torch.long, device=device) def encode( self, *, text: Optional[Union[str, list[str]]] = None, audio: Optional[Union[str, Path, torch.Tensor, list[Optional[Union[str, Path, torch.Tensor]]]]] = None, task: str = "document", instruction: Optional[str] = None, normalize: bool = True, device: Optional[Union[str, torch.device]] = None, ) -> torch.Tensor: """High-level embedding API. Accepts raw text/audio and returns embeddings.""" self.eval() tokenizer = self._resolve_tokenizer() device = torch.device(device) if device is not None else next(self.parameters()).device text_items = self._to_list(text) audio_items = self._to_list(audio) batch_size = max(len(text_items), len(audio_items)) if batch_size == 0: raise ValueError("At least one of text/audio must be provided.") if len(text_items) == 0: text_items = [None] * batch_size elif len(text_items) == 1 and batch_size > 1: text_items = text_items * batch_size elif len(text_items) != batch_size: raise ValueError("text and audio batch sizes must match (or be broadcastable length 1).") if len(audio_items) == 0: audio_items = [None] * batch_size elif len(audio_items) == 1 and batch_size > 1: audio_items = audio_items * batch_size elif len(audio_items) != batch_size: raise ValueError("text and audio batch sizes must match (or be broadcastable length 1).") text_ids, text_lens = self._prepare_text_batch( tokenizer, texts=text_items, task=task, instruction=instruction, device=device, ) audio_tensor, audio_lens = self._prepare_audio_batch( audio_items, target_sr=TARGET_SR, device=device, ) with torch.inference_mode(): emb = self( text_ids=text_ids, text_lens=text_lens, audio=audio_tensor, audio_lens=audio_lens, ) if normalize: emb = torch.nn.functional.normalize(emb.float(), dim=-1) return emb @staticmethod def _check_list_arg(name: str, value: Any, item_types: Tuple[type, ...]) -> None: if value is None: return if not isinstance(value, list): raise TypeError( f"`{name}` must be a list, got {type(value).__name__}." ) if len(value) == 0: raise ValueError(f"`{name}` must be a non-empty list when provided.") for i, item in enumerate(value): if not isinstance(item, item_types): allowed = ", ".join(t.__name__ for t in item_types) raise TypeError( f"`{name}[{i}]` must be one of ({allowed}), " f"got {type(item).__name__}." ) def encode_query( self, text: Optional[List[str]] = None, audio: Optional[List[Union[str, Path, torch.Tensor]]] = None, **kwargs, ) -> torch.Tensor: """Encode queries. Accepts text-only, audio-only, or both.""" self._check_list_arg("text", text, (str,)) self._check_list_arg("audio", audio, (str, Path, torch.Tensor)) if text is None and audio is None: raise ValueError( "encode_query requires at least one of `text` or `audio`." ) if text is not None and audio is not None and len(text) != len(audio): raise ValueError( f"encode_query: `text` (len={len(text)}) and `audio` (len={len(audio)}) " "must have the same length when both are provided." ) return self.encode(text=text, audio=audio, task="query", **kwargs) def encode_document( self, text: Optional[List[str]] = None, audio: Optional[List[Union[str, Path, torch.Tensor]]] = None, **kwargs, ) -> torch.Tensor: """Encode documents. Accepts exactly one of `text` or `audio`.""" self._check_list_arg("text", text, (str,)) self._check_list_arg("audio", audio, (str, Path, torch.Tensor)) if (text is None) == (audio is None): raise ValueError( "encode_document requires exactly one of `text` or `audio` " "(not both, not neither)." ) return self.encode(text=text, audio=audio, task="document", **kwargs) def forward(self, text_ids, text_lens, audio=None, audio_lens=None): device = text_ids.device B = text_ids.size(0) embed_dtype = self.model.embed_tokens.weight.dtype if audio is not None: audio_emb, audio_mask = self.dasheng( audio, audio_lens, ) audio_emb, audio_mask = self.dasheng_down(audio_emb, audio_mask) audio_lens = audio_mask.sum(dim=1) audio_emb = self.dasheng_proj(audio_emb) else: audio_emb = None text_emb = self.model.embed_tokens(text_ids) audio_start_emb = self.model.embed_tokens(self.audio_start_token.clone()) audio_end_emb = self.model.embed_tokens(self.audio_end_token.clone()) eos_emb = self.model.embed_tokens(self.eos_token.clone()) input_embeds = [] attention_masks = [] last_indices = [] for i in range(B): seq = [text_emb[i, : text_lens[i]]] if audio_emb is not None and audio_lens[i] > 0: seq.append(audio_start_emb) seq.append(audio_emb[i, : audio_lens[i]]) seq.append(audio_end_emb) seq.append(eos_emb) seq = torch.cat(seq, dim=0) input_embeds.append(seq) attention_masks.append(torch.ones(seq.size(0), device=device)) last_indices.append(seq.size(0) - 1) max_len = max(x.size(0) for x in input_embeds) padded_embeds = torch.zeros( B, max_len, self.hidden_size, device=device, dtype=embed_dtype ) padded_mask = torch.zeros(B, max_len, device=device) for i in range(B): L = input_embeds[i].size(0) padded_embeds[i, :L] = input_embeds[i] padded_mask[i, :L] = attention_masks[i] outputs = self.model( inputs_embeds=padded_embeds, attention_mask=padded_mask, ).last_hidden_state final_hidden = torch.stack([outputs[i, last_indices[i]] for i in range(B)]) out = self.siglip_head(final_hidden).squeeze(1) return out