| |
| """ |
| FIGMA inference: encode music and text into a shared space and retrieve. |
| |
| Examples |
| -------- |
| # Text -> Audio: rank a folder of .wav clips by a text query |
| python inference.py --checkpoint figma.ckpt \ |
| --audio_dir ./clips --query "a song in F minor at 120 BPM in 4/4 time" --topk 5 |
| |
| # Dump audio embeddings for a folder of clips |
| python inference.py --checkpoint figma.ckpt --audio_dir ./clips --save_embeddings emb.pt |
| """ |
|
|
| import os |
| import glob |
| import argparse |
|
|
| import torch |
| import librosa |
|
|
| from figma_model import Figma, get_tokenizer, SAMPLE_RATE, CLIP_SECONDS |
|
|
|
|
| def load_audio_batch(paths, clip_seconds, sr=SAMPLE_RATE): |
| fixed = int(sr * clip_seconds) |
| out = torch.zeros(len(paths), 1, fixed) |
| for i, p in enumerate(paths): |
| wav, _ = librosa.load(p, sr=sr) |
| wav = torch.tensor(wav, dtype=torch.float32) |
| take = min(wav.numel(), fixed) |
| out[i, 0, :take] = wav[:take] |
| return out |
|
|
|
|
| @torch.no_grad() |
| def embed_audios(model, paths, clip_seconds, device, batch_size=32): |
| embs = [] |
| for i in range(0, len(paths), batch_size): |
| batch = load_audio_batch(paths[i:i + batch_size], clip_seconds).to(device) |
| embs.append(model.encode_audio(batch).cpu()) |
| return torch.cat(embs) if embs else torch.empty(0) |
|
|
|
|
| @torch.no_grad() |
| def embed_texts(model, tokenizer, texts, device, max_len=128, batch_size=64): |
| embs = [] |
| for i in range(0, len(texts), batch_size): |
| tok = tokenizer(texts[i:i + batch_size], return_tensors="pt", |
| padding="max_length", truncation=True, max_length=max_len) |
| tok = {k: v.to(device) for k, v in tok.items()} |
| embs.append(model.encode_text(tok).cpu()) |
| return torch.cat(embs) if embs else torch.empty(0) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--checkpoint", required=True, help="Path to the FIGMA checkpoint (.ckpt)") |
| ap.add_argument("--audio_dir", required=True, help="Folder of audio files") |
| ap.add_argument("--audio_glob", default="*.wav") |
| ap.add_argument("--query", help="Text query for text->audio retrieval") |
| ap.add_argument("--topk", type=int, default=5) |
| ap.add_argument("--save_embeddings", help="Path to save {paths, audio_emb} .pt") |
| ap.add_argument("--clip_seconds", type=float, default=CLIP_SECONDS) |
| ap.add_argument("--gpu", type=int, default=0) |
| args = ap.parse_args() |
|
|
| device = torch.device(f"cuda:{args.gpu}" if torch.cuda.is_available() else "cpu") |
| model = Figma.from_checkpoint(args.checkpoint, device=device) |
| tokenizer = get_tokenizer() |
|
|
| paths = sorted(glob.glob(os.path.join(args.audio_dir, args.audio_glob))) |
| if not paths: |
| raise SystemExit(f"No audio matching {args.audio_glob} in {args.audio_dir}") |
| print(f"Encoding {len(paths)} audio clips at {args.clip_seconds}s ...") |
| audio_emb = embed_audios(model, paths, args.clip_seconds, device) |
|
|
| if args.save_embeddings: |
| torch.save({"paths": paths, "audio_emb": audio_emb}, args.save_embeddings) |
| print(f"Saved embeddings -> {args.save_embeddings}") |
|
|
| if args.query: |
| q = embed_texts(model, tokenizer, [args.query], device)[0] |
| scores = audio_emb @ q |
| topk = scores.topk(min(args.topk, len(paths))) |
| print(f"\nTop {topk.indices.numel()} matches for: {args.query!r}") |
| for rank, (idx, sc) in enumerate(zip(topk.indices.tolist(), topk.values.tolist()), 1): |
| print(f" {rank:>2}. {sc:.4f} {os.path.basename(paths[idx])}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
