tsumugi-mrl

Pretrained music audio encoder from Tsumugi-MRL, with utilities for LoRA fine-tuning. The encoder reads stereo audio at 22,050 Hz and returns 512-dimensional frame features at 25 Hz.

Install

pip install tsumugi-mrl

The package requires PyTorch 2.7 or later. For GPU support, install the CUDA build of torch and torchaudio for your platform from pytorch.org first.

Extract features

import torch
from tsumugi_mrl import TsumugiMRLModel

model = TsumugiMRLModel.from_pretrained("anime-song/tsumugi-mrl").eval()
audio = torch.randn(1, 2, 22_050 * 4)  # [batch, 2, samples], stereo 22,050 Hz
with torch.no_grad():
    hidden = model(audio)  # [batch, frames, 512]
    clip = model.mean_pool(hidden)  # [batch, 512]

Convert mono audio by repeating the channel. No pretraining objective trains audio_projection, so use the frame features or their mean rather than encode_embedding() unless you fine-tune that head.

Fine-tune with LoRA

import torch
from tsumugi_mrl import TsumugiMRLModel, apply_lora, load_lora, save_lora

encoder = TsumugiMRLModel.from_pretrained("anime-song/tsumugi-mrl")
apply_lora(encoder, rank=8, alpha=16)  # q/k/v in all layers; other weights frozen
head = torch.nn.Linear(encoder.config.d_model, 10)
parameters = [p for p in encoder.parameters() if p.requires_grad] + list(head.parameters())
optimizer = torch.optim.AdamW(parameters, lr=1e-4)

# ... train ...
save_lora(encoder, "lora")  # lora_config.json + lora.safetensors

restored = load_lora(TsumugiMRLModel.from_pretrained("anime-song/tsumugi-mrl"), "lora")

LoRAConfig selects other targets (to_out.0, to_gates, feed-forward net.1/net.4), a subset of layers, per-pattern ranks, and extra trainable parameters. merge_lora(model, unload=True) folds the adapters into ordinary nn.Linear layers. See the repository README for details.

License

The code is released under the MIT License. The pretrained weights on the Hugging Face Hub are released separately under CC BY-NC 4.0.