| """ |
| FIGMA inference model (inference-only). |
| |
| Loads the packaged FIGMA checkpoint (frozen MuQ audio encoder + frozen E5 text |
| encoder + two Transformer projection heads) and produces the GLOBAL audio/text |
| embeddings used for retrieval. Retrieval score = cosine similarity between the |
| global audio and text embeddings. |
| |
| Note: retrieval uses only the global embeddings (mean-pooled audio frames + E5 |
| [CLS]). The frame/token projection path is used during training only. |
| |
| License: this model bundles MuQ weights (CC BY-NC 4.0), so the composite model is |
| released for NON-COMMERCIAL research use (CC BY-NC 4.0). MuQ and E5 are attributed. |
| """ |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from torch.nn import TransformerEncoder, TransformerEncoderLayer |
| from transformers import AutoModel, AutoTokenizer |
| from muq import MuQ |
|
|
| TEXT_ENCODER = "intfloat/multilingual-e5-large-instruct" |
| AUDIO_ENCODER = "OpenMuQ/MuQ-large-msd-iter" |
| SAMPLE_RATE = 24000 |
| CLIP_SECONDS = 10.0 |
|
|
|
|
| class ProjectionHead(nn.Module): |
| def __init__(self, input_dim, hidden_dim=512, output_dim=512): |
| super().__init__() |
| layer = TransformerEncoderLayer( |
| d_model=input_dim, nhead=8, dim_feedforward=hidden_dim, |
| dropout=0.1, batch_first=True, |
| ) |
| self.transformer = TransformerEncoder(layer, num_layers=2) |
| self.output_proj = nn.Linear(input_dim, output_dim) |
|
|
| def forward(self, x): |
| if x.dim() == 2: |
| x = x.unsqueeze(1) |
| x = self.transformer(x) |
| x = x.squeeze(1) |
| else: |
| x = self.transformer(x) |
| return self.output_proj(x) |
|
|
|
|
| class Figma(nn.Module): |
| """Submodule names (muq / text_encoder / audio_proj / text_proj) match the |
| training checkpoint, so its state_dict loads directly.""" |
|
|
| def __init__(self, audio_feat_dim=1024): |
| super().__init__() |
| self.muq = MuQ.from_pretrained(AUDIO_ENCODER).eval().requires_grad_(False) |
| self.text_encoder = AutoModel.from_pretrained(TEXT_ENCODER) |
| for p in self.text_encoder.parameters(): |
| p.requires_grad = False |
| self.audio_proj = ProjectionHead(input_dim=audio_feat_dim) |
| self.text_proj = ProjectionHead(input_dim=self.text_encoder.config.hidden_size) |
|
|
| @torch.no_grad() |
| def encode_audio(self, wavs): |
| """wavs: [B, 1, samples] float tensor at 24 kHz -> [B, 512] L2-normalized.""" |
| wavs = torch.nan_to_num(wavs, nan=0.0, posinf=0.0, neginf=0.0) |
| with torch.amp.autocast("cuda", enabled=False): |
| a_seq = self.muq(wavs).last_hidden_state |
| a_clip = a_seq.mean(dim=1) |
| return F.normalize(self.audio_proj(a_clip), dim=-1, eps=1e-8) |
|
|
| @torch.no_grad() |
| def encode_text(self, text_inputs): |
| """text_inputs: dict of tokenized tensors -> [B, 512] L2-normalized.""" |
| t_seq = self.text_encoder(**text_inputs).last_hidden_state |
| t_cls = t_seq[:, 0, :] |
| return F.normalize(self.text_proj(t_cls), dim=-1, eps=1e-8) |
|
|
| @classmethod |
| def from_checkpoint(cls, ckpt_path, device="cpu"): |
| model = cls().to(device).eval() |
| ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False) |
| sd = ckpt.get("state_dict", ckpt) |
| missing, _ = model.load_state_dict(sd, strict=False) |
| bad = [k for k in missing if k.startswith(("audio_proj", "text_proj"))] |
| if bad: |
| raise RuntimeError(f"Projection-head weights missing from checkpoint: {bad}") |
| return model |
|
|
|
|
| def get_tokenizer(): |
| return AutoTokenizer.from_pretrained(TEXT_ENCODER) |
|
|
