File size: 2,893 Bytes
c0aa76e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 | #!/usr/bin/env python3
"""
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())
|