Sentio model card, config and modeling code

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+resources/sentio_architecture.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,210 @@
+---
+license: apache-2.0
+pipeline_tag: feature-extraction
+base_model:
+- tencent/WeMM-Embedding-9B
+- Qwen/Qwen2.5-Omni-3B
+tags:
+- audio
+- multimodal-embedding
+- audio-retrieval
+- audio-classification
+- video-moment-retrieval
+---
+
+# Sentio
+
+![Sentio architecture](resources/sentio_architecture.png)
+
+Sentio is WeMM-Embedding-9B with ears, and with sharper eyes for video moments.
+
+WeMM-Embedding-9B is the strongest open embedding model for text, images, video and visual
+documents, but it cannot hear. [Wianmu-5B](https://huggingface.co/eduardo-terres2/Wianmu-5B) gave
+the 4B member of the family an audio tower by freezing both the embedding model and a pretrained
+audio encoder and training only a small connector between them, the recipe of Tianmu-Emb-Uni.
+Sentio does the same to the 9B, with two additions: before the audio stage, the 9B tower gets a
+small LoRA fine-tune for video moment retrieval, so the vision side also moves; and during the audio
+stage the audio encoder gets a LoRA of its own, trained together with the connector. The name is
+Latin for "I perceive": one model that hears, sees and reads into one space. Total 10.20B
+parameters, of which 152M (1.5%) were trained.
+
+## How it works
+
+Two experts, one that understands audio and one that understands everything else, and a trained
+piece in the middle that moves audio into the other model's embedding space.
+
+The audio side is the audio tower of Qwen2.5-Omni-3B (the encoder only, 638M parameters) with a
+LoRA on its attention and MLP projections. Everything else goes through WeMM-Embedding-9B. Audio
+lands in WeMM's 4096-d space, so an audio clip can be compared against a caption, a label name, an
+image or a video directly, and nothing that was indexed with the tower has to be re-indexed.
+
+```
+waveform
+  -> Whisper-style mel features (128 bins, up to 30 s)
+  -> Qwen2.5-Omni audio tower + LoRA (r=16), masked-mean pooled over frames  -> 2048-d
+  -> connector: 3-layer MLP, 2048 -> 4096 -> 8192 -> 4096, LayerNorm + GELU
+  -> modality adapter fusion
+  -> RMSNorm projection
+  -> L2 normalize  -> 4096-d embedding
+```
+
+**Two stages, three trained pieces.**
+
+1. *Stage 1, the tower.* WeMM-Embedding-9B gets a LoRA (r=16, α=32, 43.3M parameters) trained with
+   InfoNCE on video moment-retrieval pairs mixed with a replay of visual-document, video and image
+   pairs so that nothing already strong is lost. This is the only change to the non-audio path: a
+   text, image, video or document embedding from Sentio is the embedding of WeMM-9B with this LoRA
+   applied. It lifts MMEB-V2 from 80.58 to 80.68 and V-MRET from 58.46 to 61.53.
+2. *Stage 2, the ears.* With the tower frozen (LoRA included), the audio path is trained: a LoRA on
+   the audio tower (r=16, α=32, dropout 0.05, on q/k/v/out projections and both MLP layers of every
+   encoder block, 11.8M parameters) together with the audio head — connector, adapter fusion,
+   projection norm and prototypes, 96.8M parameters, almost all of it the connector (75.5M).
+
+**Adapter fusion.** One audio clip can be matched against very different things. For a caption you
+care about meaning, for another audio clip about acoustic detail, for a video about timing, and for
+a label about the category. So there are four small bottleneck adapters, one each for text, audio,
+video and label. A gate looks at the audio vector and mixes the four outputs with soft weights, with
+a small entropy bonus during training so the gate doesn't collapse onto one adapter.
+
+| Adapter | Input | Down projection | Activation | Bottleneck | Up projection | Output | Residual | Parameters |
+|---|---|---|---|---|---|---|---|---|
+| text  | 4096 | Linear 4096 → 512, no bias | GELU | 512 | Linear 512 → 4096, no bias | 4096 | yes | 4,194,304 |
+| audio | 4096 | Linear 4096 → 512, no bias | GELU | 512 | Linear 512 → 4096, no bias | 4096 | yes | 4,194,304 |
+| video | 4096 | Linear 4096 → 512, no bias | GELU | 512 | Linear 512 → 4096, no bias | 4096 | yes | 4,194,304 |
+| label | 4096 | Linear 4096 → 512, no bias | GELU | 512 | Linear 512 → 4096, no bias | 4096 | yes | 4,194,304 |
+| gate  | 4096 | Linear 4096 → 1024, no bias + LayerNorm(1024) | GELU | 1024 | Linear 1024 → 4, with bias | 4 (softmax weights) | no | 4,200,452 |
+
+**Shared prototypes.** 64 learnable vectors in the shared space. Each audio embedding and each target
+embedding gets a soft assignment over them, and training pushes the audio and its matching target
+to pick the same prototypes. They are only used for the loss; at inference you just use the embedding.
+
+**Parameter count** (measured from the weight files):
+
+| Component | Parameters | Trained |
+|---|---|---|
+| WeMM-Embedding-9B tower | 9,408M | frozen |
+| vdr LoRA on the tower (stage 1) | 43.3M | yes |
+| Qwen2.5-Omni-3B audio tower | 638M | frozen |
+| LoRA on the audio tower (stage 2) | 11.8M | yes |
+| Connector (3-layer MLP) | 75.5M | yes (stage 2) |
+| Adapter fusion (4 adapters + gate) | 21.0M | yes (stage 2) |
+| RMSNorm projection + 64 prototypes | 0.27M | yes (stage 2) |
+| **Total** | **10.20B** | **152M (1.5%)** |
+
+## Training
+
+*Stage 2 (audio).* The loss is symmetric in-batch InfoNCE between audio and target embeddings
+(temperature 0.07), which does most of the work. On top of it come the Tianmu prototype loss — a
+contrastive term over prototype assignments (0.03), a symmetric KL between the audio and target
+assignments (0.01) and a balance term so all prototypes get used (0.005) — and a gate-entropy bonus
+(0.01). For label datasets, batches are built so that each label appears at most once; otherwise two
+clips of "dog bark" in one batch count as negatives of each other.
+
+The data is 449k audio–target pairs over 37k distinct targets from seven public datasets, sampled
+in natural proportion: NSynth (instrument labels, the bulk of every epoch), SpeechCommands
+(spoken-word labels), FSD50K and RAVDESS (sound-event and emotion labels), Clotho (captions),
+Flickr8k Audio (spoken captions paired with the image they describe) and AVE (clips paired with
+their source video). Batch size 16 (smaller where a dataset has fewer distinct labels); learning
+rate 1e-5 for the head and 1e-4 for the audio-tower LoRA, constant after 500 warm-up steps;
+498,000 steps on one H200. Validation is the same in-batch InfoNCE on a held-out split every 2,500
+steps; the released checkpoint is the one with the lowest validation loss (step ~408,000).
+
+*Stage 1 (tower).* LoRA r=16, α=32 on the WeMM-9B tower, 600 steps of InfoNCE at temperature 0.045;
+8/11 of each batch are DiDeMo moment–query pairs and 1/11 each MSR-VTT video–text, MSCOCO
+text-to-image and ColPali page–query pairs, replayed so that video retrieval, image retrieval and
+visual documents keep their scores while V-MRET moves.
+
+## Evaluation
+
+All 190 MMEB-V3 tasks with the [Tencent WeMM evaluation harness](https://github.com/Tencent/WeMM-Embedding/tree/main/mmeb_v3_eval),
+with the Wianmu audio modifications for the 11 audio tasks; scored with the leaderboard's own code.
+For the two TUTSound tasks, whose queries carry both an audio clip and text, the model has no joint
+encoder, so the two unit embeddings are summed with equal weight and re-normalised.
+
+MMEB-V3, leaderboard snapshot of 16 Oct 2026:
+
+| Model | Overall | Overall-V3 (new) | Text | Audio | Agent | Image | Video | Visdoc |
+|---|---|---|---|---|---|---|---|---|
+| **Sentio** | **61.69** | **48.46** | **48.98** | 36.81 | **50.57** | 80.90 | **75.04** | 83.28 |
+| Wianmu-5B-v0 | 60.45 | 47.50 | 48.13 | 42.93 | 48.00 | 79.73 | 71.56 | 81.99 |
+| WeMM-Embedding-9B | 59.33 | 44.54 | 48.76 | — | 50.10 | 80.99 | 74.31 | 83.34 |
+| Wianmu-3B-v0 | 58.16 | 44.37 | 45.30 | 41.16 | 44.11 | 78.67 | 70.79 | 80.80 |
+
+MMEB-V2 (78 tasks) and the video block:
+
+| Model | Overall-V2 | Image (V2) | Video | V-CLS | V-QA | V-RET | V-MRET | Visdoc |
+|---|---|---|---|---|---|---|---|---|
+| **Sentio** | 80.68 | 81.77 | 75.04 | 87.54 | 77.89 | 67.80 | 61.53 | 83.28 |
+| WeMM-Embedding-9B | 80.58 | 81.87 | 74.31 | 87.36 | 77.69 | 67.38 | 58.46 | 83.34 |
+
+Audio tasks (hit@1):
+
+| Task | Sentio | Wianmu-3B-v0 | Wianmu-5B-v0 |
+|---|---|---|---|
+| NSynth | 76.10 | 74.10 | 60.90 |
+| ESC-50 | 56.15 | 79.20 | 86.60 |
+| UrbanSound8K | 55.92 | 64.05 | 66.63 |
+| SpeechCommands | 89.50 | 89.80 | 89.90 |
+| CREMA-D | 33.74 | 47.08 | 48.93 |
+| Clotho | 8.48 | 13.61 | 16.56 |
+| SoundDescs | 6.10 | 9.50 | 20.50 |
+| TUTSound | 61.76 | 52.96 | 46.89 |
+| TUTSound (hard) | 5.16 | 4.55 | 3.64 |
+| AVE | 9.20 | 16.42 | 18.66 |
+| SpeechCOCO | 2.80 | 1.50 | 13.00 |
+| **A-CLS / A-RET / Audio-Overall** | **62.28 / 15.58 / 36.81** | 70.85 / 16.42 / 41.16 | 70.59 / 19.88 / 42.93 |
+
+Sentio takes the top of the 190-task board on the strength of the 9B tower (text, agent, video and
+visual documents) and is the first of a series: the audio block is where it trails the smaller
+Wianmu models, and the next releases work on exactly that.
+
+## Known limitations
+
+- Audio is pooled into a single vector, so there is no fine timing inside a clip. Grounding and
+  moment-level tasks suffer from this the most.
+- Clips longer than 30 seconds get cut.
+- The audio stage saw seven datasets, dominated by NSynth instrument labels. Together with the
+  adaptable encoder this makes the model strong on instruments and sound events (NSynth, TUTSound)
+  and weaker than the smaller Wianmu models on environmental scenes (ESC-50, UrbanSound8K), emotion
+  (CREMA-D) and audio captions (Clotho, SoundDescs).
+- The tower is frozen during the audio stage, so whatever WeMM-9B can't tell apart in text, Sentio
+  can't either. Some label names also collide across datasets ("happy" is a spoken word in
+  SpeechCommands and an emotion in RAVDESS, and both share one target vector).
+- Most of the training targets are English.
+
+## Usage
+
+The repository holds the model card, `config.json` and `modeling_sentio.py`; the base checkpoints
+(WeMM-Embedding-9B, Qwen2.5-Omni-3B) are downloaded from their own repositories, and Sentio's
+trained parts (the vdr LoRA adapter, 43M, and `trainable.pt` with the audio head and the
+audio-tower LoRA, 109M) are not distributed here.
+
+```python
+from modeling_sentio import Sentio
+
+model = Sentio.from_pretrained(
+    ".",
+    wemm_path="path/to/WeMM-Embedding-9B",
+    wemm_adapter_path="path/to/sentio/vdr_lora_adapter",
+    qwen_omni_path="path/to/Qwen2.5-Omni-3B",
+    trainable_path="path/to/sentio/trainable.pt",
+)
+audio = model.encode_audio(["dog.wav", "piano.wav"])          # (2, 4096)
+text  = model.encode_text(["dog barking", "piano"])            # (2, 4096)
+image = model.encode_image(["photo.jpg"])                      # (1, 4096)
+scores = audio @ text.T
+```
+
+Preprocessing (128-bin mel features, 30-second head crop, mean pooling over valid frames) mirrors
+the MMEB-V3 evaluation that produced the scores above.
+
+## License
+
+Sentio's own code, model card and trained parameters are released under Apache-2.0. The two base
+models keep their own terms: WeMM-Embedding-9B is licensed by Tencent under Apache-2.0; the audio
+encoder is the audio tower of Qwen2.5-Omni-3B, distributed by Alibaba Cloud under the
+[Qwen Research License](https://huggingface.co/Qwen/Qwen2.5-Omni-3B/blob/main/LICENSE), whose terms
+apply to that component.
+
+## Credits
+Credits to [Wianmu-5B](https://huggingface.co/eduardo-terres2/Wianmu-5B), [Tianmu-Emb-Uni](https://huggingface.co/TianmuLab/Tianmu-Emb-Uni/tree/main) and the [WeMM-Embedding family](https://github.com/Tencent/WeMM-Embedding).

config.json

@@ -0,0 +1,18 @@
+{
+  "model_name": "Sentio",
+  "wemm_path": "path/to/WeMM-Embedding-9B",
+  "wemm_adapter_path": "path/to/sentio/vdr_lora_adapter",
+  "qwen_omni_path": "path/to/Qwen2.5-Omni-3B",
+  "trainable_path": "path/to/sentio/trainable.pt",
+  "audio_tower_dim": 2048,
+  "embedding_dim": 4096,
+  "audio_sample_rate": 16000,
+  "audio_max_seconds": 30,
+  "audio_min_seconds": 0.025,
+  "audio_lora": {
+    "r": 16,
+    "alpha": 32,
+    "dropout": 0.05,
+    "targets": ["q_proj", "k_proj", "v_proj", "out_proj", "fc1", "fc2"]
+  }
+}

modeling_sentio.py

@@ -0,0 +1,276 @@
+"""
+Sentio: WeMM-Embedding-9B (+ a LoRA for video moment retrieval) with an audio branch, all
+embedding into one 4096-d space.
+
+  text / image / video / page -> WeMM-Embedding-9B with the vdr LoRA merged in
+  audio                       -> Qwen2.5-Omni-3B audio tower + LoRA -> mean pool -> trained head
+                                 (connector -> modality adapter fusion -> RMSNorm)
+
+config.json points at the two base checkpoints, the vdr LoRA adapter (PEFT directory) and
+trainable.pt (audio head + audio-tower LoRA). Preprocessing mirrors the MMEB-V3 evaluation that
+produced the reported scores. Adapted from Wianmu-5B's modeling_wianmu.py.
+"""
+import json
+import os
+from typing import List, Optional, Sequence, Union
+
+import numpy as np
+import soundfile as sf
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchaudio
+from safetensors import safe_open
+from transformers import AutoConfig, AutoModel, AutoProcessor, WhisperFeatureExtractor
+from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import Qwen2_5OmniAudioEncoder
+
+AudioInput = Union[str, np.ndarray, torch.Tensor]
+
+TARGET_MODALITIES = ("text", "audio", "video", "label")
+WEMM_IMAGE_KWARGS = {"min_pixels": 64 * 32 * 32, "max_pixels": 8192 * 32 * 32}
+WEMM_TEXT_MAX_LENGTH = 8192
+
+
+# ============================================================================
+# Trained audio head
+# ============================================================================
+class AudioConnector(nn.Module):
+    def __init__(self, input_dim: int, output_dim: int):
+        super().__init__()
+        h1, h2 = 2 * input_dim, 4 * input_dim
+        self.mlp = nn.Sequential(
+            nn.Linear(input_dim, h1, bias=False), nn.LayerNorm(h1), nn.GELU(), nn.Dropout(0.1),
+            nn.Linear(h1, h2, bias=False), nn.LayerNorm(h2), nn.GELU(), nn.Dropout(0.1),
+            nn.Linear(h2, output_dim, bias=False),
+        )
+
+    def forward(self, x):
+        return self.mlp(x)
+
+
+class LightweightAdapter(nn.Module):
+    def __init__(self, dim: int, reduction: int = 8):
+        super().__init__()
+        self.down = nn.Linear(dim, dim // reduction, bias=False)
+        self.act = nn.GELU()
+        self.drop = nn.Dropout(0.1)
+        self.up = nn.Linear(dim // reduction, dim, bias=False)
+
+    def forward(self, x):
+        return x + self.up(self.drop(self.act(self.down(x))))
+
+
+class ModalityAdapterFusion(nn.Module):
+    """One adapter per target modality, mixed by a softmax gate computed from the input."""
+
+    def __init__(self, dim: int):
+        super().__init__()
+        self.adapters = nn.ModuleDict({m: LightweightAdapter(dim) for m in TARGET_MODALITIES})
+        self.gate = nn.Sequential(
+            nn.Linear(dim, dim // 4, bias=False), nn.LayerNorm(dim // 4), nn.GELU(),
+            nn.Linear(dim // 4, len(TARGET_MODALITIES)),
+        )
+
+    def forward(self, x):
+        weights = F.softmax(self.gate(x), dim=-1)
+        stacked = torch.stack([self.adapters[m](x) for m in TARGET_MODALITIES], dim=1)
+        return (weights.unsqueeze(-1) * stacked).sum(dim=1)
+
+
+class ProjectionHead(nn.Module):
+    def __init__(self, dim: int):
+        super().__init__()
+        self.norm = nn.RMSNorm(dim, eps=1e-6)
+
+    def forward(self, x):
+        return self.norm(x)
+
+
+class SentioAudioHead(nn.Module):
+    """Pooled audio-tower output (B, audio_dim) -> L2-normalized (B, dim)."""
+
+    def __init__(self, audio_dim: int, dim: int):
+        super().__init__()
+        self.audio_connector = AudioConnector(audio_dim, dim)
+        self.adapter_fusion = ModalityAdapterFusion(dim)
+        self.projection = ProjectionHead(dim)
+
+    def forward(self, x):
+        x = self.projection(self.adapter_fusion(self.audio_connector(x)))
+        return F.normalize(x, p=2, dim=-1, eps=1e-6)
+
+
+# ============================================================================
+# Qwen2.5-Omni-3B audio tower (frozen base + LoRA)
+# ============================================================================
+def load_omni_audio_tower(qwen_omni_path: str, prefix: str = "thinker.audio_tower.") -> Qwen2_5OmniAudioEncoder:
+    """Reads only the audio-tower tensors out of the full Qwen2.5-Omni checkpoint."""
+    audio_config = AutoConfig.from_pretrained(qwen_omni_path).thinker_config.audio_config
+    tower = Qwen2_5OmniAudioEncoder(audio_config)
+    with open(os.path.join(qwen_omni_path, "model.safetensors.index.json")) as f:
+        weight_map = json.load(f)["weight_map"]
+    state = {}
+    for shard in sorted({s for k, s in weight_map.items() if k.startswith(prefix)}):
+        with safe_open(os.path.join(qwen_omni_path, shard), framework="pt") as f:
+            for key in f.keys():
+                if key.startswith(prefix):
+                    state[key[len(prefix):]] = f.get_tensor(key)
+    tower.load_state_dict(state, strict=True)
+    return tower
+
+
+def inject_audio_lora(tower: nn.Module, lora: dict) -> nn.Module:
+    """Adds the (untrained) LoRA modules Sentio was trained with; weights come from trainable.pt."""
+    from peft import LoraConfig, inject_adapter_in_model
+    cfg = LoraConfig(r=lora["r"], lora_alpha=lora["alpha"], lora_dropout=lora["dropout"],
+                     target_modules=list(lora["targets"]), bias="none")
+    return inject_adapter_in_model(cfg, tower, adapter_name="default")
+
+
+def load_trainable(trainable_path: str, audio_tower: nn.Module, audio_head: nn.Module):
+    """trainable.pt holds the head blocks and the audio-tower LoRA tensors saved during training."""
+    state = torch.load(trainable_path, map_location="cpu", weights_only=True)
+    for block in ("audio_connector", "adapter_fusion", "projection"):
+        missing, unexpected = getattr(audio_head, block).load_state_dict(state[block], strict=False)
+        if missing or unexpected:
+            raise RuntimeError(f"{block}: missing {missing}, unexpected {unexpected}")
+    lora_state = {k[len("encoder."):] if k.startswith("encoder.") else k: v
+                  for k, v in state["audio_encoder"].items() if "lora_" in k}
+    missing, unexpected = audio_tower.load_state_dict(lora_state, strict=False)
+    if unexpected or any("lora_" in k for k in missing):
+        raise RuntimeError(f"audio-tower LoRA: unexpected {unexpected}, missing {[k for k in missing if 'lora_' in k]}")
+
+
+@torch.no_grad()
+def pool_audio_tower(tower: Qwen2_5OmniAudioEncoder, features: torch.Tensor, num_frames: int) -> torch.Tensor:
+    """(128, T) mel features of one clip -> mean of the tower's output states over its valid frames."""
+    if num_frames < 4:
+        return features.new_zeros(tower.config.output_dim)
+    features = features[:, :num_frames].contiguous()
+    feature_lens = torch.tensor([num_frames], dtype=torch.long, device=features.device)
+    aftercnn_lens = ((feature_lens - 1) // 2 + 1).clamp(min=1)
+    out = tower(features, feature_lens=feature_lens, aftercnn_lens=aftercnn_lens)
+    out = out[0] if isinstance(out, tuple) else out.last_hidden_state
+    return out.mean(dim=0) if out.dim() == 2 else out.mean(dim=1).squeeze(0)
+
+
+# ============================================================================
+# Sentio
+# ============================================================================
+class Sentio(nn.Module):
+    def __init__(self, config: dict, audio_tower, audio_head, feature_extractor, wemm, wemm_processor):
+        super().__init__()
+        self.config = config
+        self.audio_tower = audio_tower
+        self.audio_head = audio_head
+        self.feature_extractor = feature_extractor
+        self.wemm = wemm
+        self.wemm_processor = wemm_processor
+
+    @classmethod
+    def from_pretrained(cls, path: str, device: Optional[str] = None, dtype: torch.dtype = torch.bfloat16,
+                        **config_overrides):
+        """config_overrides replace config.json entries, e.g. wemm_path=..., trainable_path=..."""
+        device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        with open(os.path.join(path, "config.json")) as f:
+            config = {**json.load(f), **config_overrides}
+
+        audio_tower = inject_audio_lora(load_omni_audio_tower(config["qwen_omni_path"]), config["audio_lora"])
+        audio_head = SentioAudioHead(config["audio_tower_dim"], config["embedding_dim"])
+        load_trainable(config["trainable_path"], audio_tower, audio_head)
+        feature_extractor = WhisperFeatureExtractor.from_pretrained(config["qwen_omni_path"])
+
+        wemm = AutoModel.from_pretrained(
+            config["wemm_path"], torch_dtype=dtype, trust_remote_code=True, attn_implementation="sdpa",
+        )
+        if config.get("wemm_adapter_path"):
+            from peft import PeftModel
+            wemm = PeftModel.from_pretrained(wemm, config["wemm_adapter_path"]).merge_and_unload()
+        wemm_processor = AutoProcessor.from_pretrained(config["wemm_path"], trust_remote_code=True)
+
+        model = cls(config, audio_tower, audio_head, feature_extractor, wemm, wemm_processor)
+        return model.to(device=device, dtype=dtype).eval()
+
+    @property
+    def device(self) -> torch.device:
+        return next(self.audio_head.parameters()).device
+
+    def _autocast(self):
+        return torch.autocast(device_type=self.device.type, dtype=torch.bfloat16,
+                              enabled=self.device.type == "cuda")
+
+    # ---------------------------------------------------------------- audio
+    def load_audio(self, audio: AudioInput, sampling_rate: Optional[int] = None) -> np.ndarray:
+        """Path or waveform -> mono float32 at 16 kHz, head-cropped to audio_max_seconds.
+        Clips shorter than audio_min_seconds come back as 10 ms of silence (a zero embedding)."""
+        target_sr = self.config["audio_sample_rate"]
+        if isinstance(audio, str):
+            wav, sampling_rate = sf.read(audio, dtype="float32")
+            wav = torch.from_numpy(wav)
+        else:
+            if sampling_rate is None:
+                raise ValueError("sampling_rate is required when passing a waveform")
+            wav = torch.as_tensor(audio, dtype=torch.float32)
+        if wav.dim() > 1:
+            wav = wav.mean(dim=1)
+        if sampling_rate != target_sr and wav.numel() > 0:
+            wav = torchaudio.functional.resample(wav, sampling_rate, target_sr)
+        if wav.numel() < int(self.config["audio_min_seconds"] * target_sr):
+            return np.zeros(target_sr // 100, dtype=np.float32)
+        return wav[: int(self.config["audio_max_seconds"] * target_sr)].numpy()
+
+    @torch.no_grad()
+    def encode_audio(self, audios: Sequence[AudioInput], sampling_rate: Optional[int] = None) -> torch.Tensor:
+        """Audio file paths or 1-D/2-D (samples, channels) waveforms -> (N, 4096)."""
+        pooled = []
+        for audio in audios:
+            wav = self.load_audio(audio, sampling_rate)
+            feats = self.feature_extractor(wav, sampling_rate=self.config["audio_sample_rate"],
+                                           return_attention_mask=True, return_tensors="pt")
+            features = feats["input_features"][0].to(self.device, dtype=self.audio_tower.dtype)
+            num_frames = min(int(feats["attention_mask"][0].sum()), features.shape[-1])
+            with self._autocast():
+                pooled.append(pool_audio_tower(self.audio_tower, features, num_frames))
+        with self._autocast():
+            return self.audio_head(torch.stack(pooled)).float()
+
+    # ----------------------------------------------------------------- WeMM
+    @torch.no_grad()
+    def _encode_wemm(self, conversations: List[list], has_visual: bool) -> torch.Tensor:
+        chat_texts = [
+            self.wemm_processor.apply_chat_template(conv, tokenize=False, add_generation_prompt=False)
+            for conv in conversations
+        ]
+        kwargs = dict(text=chat_texts, padding=True, return_tensors="pt")
+        if has_visual:
+            from qwen_vl_utils import process_vision_info
+            images, _, _ = process_vision_info(conversations, image_patch_size=16,
+                                               return_video_kwargs=True, return_video_metadata=True)
+            kwargs.update(images=images, truncation=False)
+        else:
+            kwargs.update(truncation=True, max_length=WEMM_TEXT_MAX_LENGTH)
+        inputs = self.wemm_processor(**kwargs).to(self.device)
+        inner = getattr(self.wemm, "model", None)
+        if inner is not None and hasattr(inner, "rope_deltas"):
+            inner.rope_deltas = None
+        with self._autocast():
+            emb = self.wemm.embedding(**inputs)
+        return F.normalize(emb.float(), p=2, dim=-1)
+
+    def encode_text(self, texts: Sequence[str]) -> torch.Tensor:
+        """Texts -> (N, 4096). Audio<->text was trained and evaluated on raw labels and
+        captions (e.g. "dog", "chirping_birds"), with no instruction prefix."""
+        conversations = [[{"role": "user", "content": [{"type": "text", "text": t.strip()}]}] for t in texts]
+        return self._encode_wemm(conversations, has_visual=False)
+
+    def encode_image(self, images: Sequence, texts: Optional[Sequence[str]] = None) -> torch.Tensor:
+        """Image paths/URLs/PIL images, each with optional accompanying text -> (N, 4096).
+        This is WeMM-Embedding-9B with the vdr LoRA merged; for video, call self.wemm as in WeMM's model card."""
+        texts = texts or [""] * len(images)
+        conversations = []
+        for image, text in zip(images, texts):
+            content = [{"type": "image", "image": image, **WEMM_IMAGE_KWARGS}]
+            if text.strip():
+                content.append({"type": "text", "text": text.strip()})
+            conversations.append([{"role": "user", "content": content}])
+        return self._encode_wemm(conversations, has_visual=True)