| |
| """ |
| rec7_encode.py — real audio -> M3 states, codes, renderer condition. |
| |
| python rec7_encode.py song.mp3 --m3 /path/to/minimax_music3 --out song_rec7.pt |
| |
| Writes a .pt with: |
| h [T, 4096] float16 the states (what the LM would have produced) |
| codes [T, 8] int16 c0 (16384-way) + c1..c7 (1024-way each) |
| cond [T, 32768] float16 optional (--cond): the renderer condition |
| z [128, L] float16 optional (--z): the Flow-VAE latents |
| 25 frames per second. Works on the whole file (any length). |
| """ |
| import argparse |
| import sys |
| from pathlib import Path |
|
|
| import torch |
|
|
| HERE = Path(__file__).resolve().parent |
| sys.path.insert(0, str(HERE)) |
|
|
|
|
| def main() -> int: |
| ap = argparse.ArgumentParser() |
| ap.add_argument("audio", type=Path) |
| ap.add_argument("--m3", type=Path, required=True, |
| help="MiniMax-Music3 checkpoint directory") |
| ap.add_argument("--weights", type=Path, default=HERE / "weights") |
| ap.add_argument("--out", type=Path, default=None) |
| ap.add_argument("--seconds", type=float, default=0, |
| help="0 = whole file") |
| ap.add_argument("--cond", action="store_true", |
| help="also write the 32768-d renderer condition") |
| ap.add_argument("--z", action="store_true", help="also write latents") |
| a = ap.parse_args() |
|
|
| import torchaudio |
| from rec7_model import (load_rec7, load_dav, lm_semantic_tables, |
| encode_audio, read_states, states_to_streams) |
| dev = "cuda" |
| wav, sr = torchaudio.load(str(a.audio)) |
| if sr != 44100: |
| wav = torchaudio.transforms.Resample(sr, 44100)(wav) |
| if wav.shape[0] == 1: |
| wav = wav.repeat(2, 1) |
| wav = wav[:2].float() |
| if a.seconds > 0: |
| wav = wav[:, : int(a.seconds * 44100)] |
|
|
| dav = load_dav(a.m3, dev) |
| z = encode_audio(dav, wav, dev) |
| del dav |
| torch.cuda.empty_cache() |
|
|
| trunk, head = load_rec7(a.weights, dev) |
| h = read_states(trunk, head, z, dev) |
| del trunk, head |
| torch.cuda.empty_cache() |
|
|
| from diffusers.models.transformers.minimax_music3_rvq_depth_decoder \ |
| import MiniMaxMusic3RVQDepthDecoder |
| depth = MiniMaxMusic3RVQDepthDecoder.from_pretrained( |
| str(a.m3 / "rvq_depth_decoder"), torch_dtype=torch.bfloat16 |
| ).to(dev).eval() |
| W_head, W_embed = lm_semantic_tables(a.m3, dev) |
| codes, cond = states_to_streams(h, depth, W_head, W_embed, dev) |
|
|
| out = {"h": h.cpu().to(torch.float16), "codes": codes.to(torch.int16), |
| "frame_rate": 25, "source": str(a.audio)} |
| if a.cond: |
| out["cond"] = cond.cpu().to(torch.float16) |
| if a.z: |
| out["z"] = z.to(torch.float16) |
| p = a.out or a.audio.with_suffix(".rec7.pt") |
| torch.save(out, p) |
| print(f"{p}: {h.shape[0]} frames ({h.shape[0] / 25:.1f}s)") |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
