74faf62ef1
The AI hub serves local models on the LAN and now yields GPU memory to the UI (shrinks the chat context, then unloads; grows back after; never loads into the UI reserve), groups models by type, and reports each model's memory split. LLM decode on Vulkan with Splash kernels, one shared GPU device per process, vision on chat turns, an on-device text detector, music models and stems. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
194 lines
8.7 KiB
Python
194 lines
8.7 KiB
Python
# yue2_oracle_dump.py - reference-stage dumps for the native YuE2 port.
|
|
#
|
|
# Dev tool only (like music3_oracle_dump.py): it runs the official yue2
|
|
# package once, in an isolated venv on the CUDA dev node, and writes raw
|
|
# little-endian tensors + meta.json that the native port compares against
|
|
# (`yue2_render --check <dir>`). Nothing in the product imports or runs it.
|
|
#
|
|
# Usage (node 165):
|
|
# ~/ref/yue2-venv/bin/python yue2_oracle_dump.py \
|
|
# --model-dir ~/.makepad/weights/music/YuE2-3B \
|
|
# --vae-dir ~/.makepad/weights/music/YuE2-Vae \
|
|
# --fixture ~/.makepad/weights/music/YuE2-3B/examples/tonight-awake.json \
|
|
# --out-dir ~/ref/yue2_oracle/tonight-awake
|
|
#
|
|
# Files (all little-endian, shapes in meta.json):
|
|
# prefix_abc.i32 ABC-stage prefix ids (tokenizer check)
|
|
# tok_<name>.i32 tokenizer fixtures (style, lyrics, abc text)
|
|
# abc_prefill_logits.f32 [vocab] last-position logits of the ABC prefix
|
|
# greedy_abc.i32 temperature-0 ABC ids (eager, no CUDA graph)
|
|
# abc.i32 / abc.txt sampled ABC ids + decoded score (release path)
|
|
# semantic.i32 sampled codec ids (0..32768, offset removed)
|
|
# tf_abc_logits.f32 [n, 151649] teacher-forced logits, ABC phase rows 0..ABC_END
|
|
# tf_sem_logits.f32 [n, 32769] teacher-forced logits, rows MUSIC_END..codec end
|
|
# nar_noise.f32 / nar_v0.f32 / nar_vmid.f32 chunk-0 noise, velocity at t=1, at the first midpoint
|
|
# latents.f32 [frames, 64] final ODE state (release path)
|
|
# audio.f32 / audio.wav [samples, 2] VAE decode of latents.f32
|
|
import argparse
|
|
import json
|
|
import os
|
|
import struct
|
|
import sys
|
|
import time
|
|
import wave
|
|
|
|
os.environ.setdefault("HF_HUB_OFFLINE", "1")
|
|
os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")
|
|
os.environ.setdefault("HF_HUB_DISABLE_TELEMETRY", "1")
|
|
|
|
TF_ABC_ROWS = 64
|
|
TF_SEM_ROWS = 256
|
|
|
|
|
|
def save(out_dir, name, array, meta):
|
|
import numpy as np
|
|
|
|
array = np.ascontiguousarray(array)
|
|
if array.dtype.kind == "f":
|
|
array = array.astype("<f4")
|
|
else:
|
|
array = array.astype("<i4")
|
|
array.tofile(os.path.join(out_dir, name))
|
|
meta["files"][name] = {"shape": list(array.shape), "dtype": str(array.dtype)}
|
|
|
|
|
|
def write_wav_i16(path, audio, sample_rate):
|
|
import numpy as np
|
|
|
|
pcm = (np.clip(audio, -1.0, 1.0) * 32767.0).astype("<i2")
|
|
with wave.open(path, "wb") as handle:
|
|
handle.setnchannels(pcm.shape[1])
|
|
handle.setsampwidth(2)
|
|
handle.setframerate(sample_rate)
|
|
handle.writeframes(pcm.tobytes())
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--model-dir", required=True)
|
|
parser.add_argument("--vae-dir", required=True)
|
|
parser.add_argument("--fixture", required=True, help="json with style, lyrics, seed, cot")
|
|
parser.add_argument("--out-dir", required=True)
|
|
parser.add_argument("--max-semantic", type=int, default=None,
|
|
help="cap semantic tokens (short oracle); default = release 9000")
|
|
args = parser.parse_args()
|
|
|
|
import numpy as np
|
|
import torch
|
|
from yue2 import YuE2Pipeline
|
|
from yue2.protocol import (ABC_END, CODEC_OFFSET, CODEC_SIZE, EOD, MUSIC_END,
|
|
MUSIC_START, Sampling, token_prefixes)
|
|
from yue2.sampling import generate_tokens
|
|
from yue2.nar import CachedNAR, song_chunks
|
|
from yue2.pipeline import SemanticResult
|
|
|
|
os.makedirs(args.out_dir, exist_ok=True)
|
|
fixture = json.load(open(args.fixture, encoding="utf-8"))
|
|
meta = {"fixture": fixture, "files": {}, "timing": {},
|
|
"torch": torch.__version__, "cuda": torch.version.cuda,
|
|
"device": torch.cuda.get_device_name(0)}
|
|
|
|
pipe = YuE2Pipeline(args.model_dir, args.vae_dir, device="cuda", progress=False)
|
|
tok = pipe.tokenizer
|
|
request = pipe._request(fixture["style"], fixture["lyrics"],
|
|
cot=fixture.get("cot", "full"), seed=int(fixture["seed"]))
|
|
|
|
# Tokenizer fixtures.
|
|
prefix = token_prefixes(request, tok)
|
|
save(args.out_dir, "prefix_abc.i32", np.array(prefix), meta)
|
|
save(args.out_dir, "tok_style.i32", np.array(tok.encode(fixture["style"])), meta)
|
|
save(args.out_dir, "tok_lyrics.i32", np.array(tok.encode(fixture["lyrics"])), meta)
|
|
|
|
model = pipe._load_model()
|
|
|
|
# ABC prefill logits (eager forward, last position).
|
|
with torch.inference_mode():
|
|
out = model(torch.tensor([prefix], device="cuda"), use_cache=False, logits_to_keep=1)
|
|
save(args.out_dir, "abc_prefill_logits.f32", out.logits[0, -1].float().cpu().numpy(), meta)
|
|
|
|
# Greedy ABC (eager, no graph): exact-token gate independent of RNG.
|
|
greedy = Sampling(0.0, 1.0, 1, 1.0, 1, 0, 128)
|
|
ids, _, _ = generate_tokens(model, prefix, greedy, request.seed, "abc", use_cuda_graph=False)
|
|
save(args.out_dir, "greedy_abc.i32", np.array(ids), meta)
|
|
|
|
# Release path: plan + semantic with the shipped sampling (CUDA graph).
|
|
torch.cuda.synchronize()
|
|
t0 = time.perf_counter()
|
|
plan = pipe.plan(request=request)
|
|
torch.cuda.synchronize()
|
|
meta["timing"]["abc_seconds"] = time.perf_counter() - t0
|
|
meta["timing"]["abc"] = plan.timing
|
|
save(args.out_dir, "abc.i32", np.array(plan.abc_ids), meta)
|
|
with open(os.path.join(args.out_dir, "abc.txt"), "w", encoding="utf-8") as handle:
|
|
handle.write(plan.abc or "")
|
|
save(args.out_dir, "tok_abc.i32", np.array(tok.encode(plan.abc or "")), meta)
|
|
save(args.out_dir, "prefix_sem.i32", np.array(plan.prefix), meta)
|
|
|
|
sampling = None
|
|
if args.max_semantic:
|
|
from dataclasses import replace
|
|
sampling = replace(pipe.generation_config.semantic, max_tokens=args.max_semantic,
|
|
min_tokens=min(200, args.max_semantic))
|
|
t0 = time.perf_counter()
|
|
semantic = pipe.generate_semantic(plan, sampling=sampling)
|
|
torch.cuda.synchronize()
|
|
meta["timing"]["semantic_seconds"] = time.perf_counter() - t0
|
|
meta["timing"]["semantic"] = semantic.timing
|
|
meta["semantic_truncated"] = semantic.truncated
|
|
save(args.out_dir, "semantic.i32", np.array(semantic.tokens), meta)
|
|
|
|
# Teacher-forced logits over the sampled sequence (eager, one forward).
|
|
with torch.inference_mode():
|
|
seq = list(prefix) + list(plan.abc_ids[:TF_ABC_ROWS])
|
|
logits = model(torch.tensor([seq], device="cuda"), use_cache=False).logits[0]
|
|
rows = logits[len(prefix) - 1: len(prefix) - 1 + min(TF_ABC_ROWS, len(plan.abc_ids)), : ABC_END + 1]
|
|
save(args.out_dir, "tf_abc_logits.f32", rows.float().cpu().numpy(), meta)
|
|
sem_ids = [t + CODEC_OFFSET for t in semantic.tokens[:TF_SEM_ROWS]]
|
|
seq = list(plan.prefix) + sem_ids
|
|
logits = model(torch.tensor([seq], device="cuda"), use_cache=False).logits[0]
|
|
start = len(plan.prefix) - 1
|
|
rows = logits[start: start + len(sem_ids), MUSIC_END: CODEC_OFFSET + CODEC_SIZE]
|
|
save(args.out_dir, "tf_sem_logits.f32", rows.float().cpu().numpy(), meta)
|
|
del logits, rows
|
|
|
|
# NAR probes on chunk 0, then the release synthesis.
|
|
chunks = song_chunks(plan.prefix, semantic.tokens, request.seed)
|
|
meta["nar_chunks"] = [[len(c.ar_tokens), int(c.noise.shape[0])] for c in chunks]
|
|
engine = CachedNAR(model, chunks[0])
|
|
with torch.inference_mode():
|
|
state = chunks[0].noise.to(device="cuda", dtype=torch.bfloat16)
|
|
steps = pipe.generation_config.ode_steps
|
|
dt = 1.0 / steps
|
|
raw = torch.logit(torch.tensor(1.0, dtype=torch.float64)).clamp(-20, 20).item()
|
|
v0 = engine.velocity(state, raw)
|
|
mid = state - v0 * (dt / 2)
|
|
raw_mid = torch.logit(torch.tensor(1.0 - dt / 2, dtype=torch.float64)).clamp(-20, 20).item()
|
|
vmid = engine.velocity(mid, raw_mid)
|
|
save(args.out_dir, "nar_noise.f32", chunks[0].noise.numpy(), meta)
|
|
save(args.out_dir, "nar_v0.f32", v0.float().cpu().numpy(), meta)
|
|
save(args.out_dir, "nar_vmid.f32", vmid.float().cpu().numpy(), meta)
|
|
meta["nar_raw_t"] = [raw, raw_mid]
|
|
engine.close()
|
|
del engine
|
|
|
|
t0 = time.perf_counter()
|
|
latents = pipe.synthesize(semantic)
|
|
torch.cuda.synchronize()
|
|
meta["timing"]["nar_seconds"] = time.perf_counter() - t0
|
|
save(args.out_dir, "latents.f32", latents, meta)
|
|
|
|
t0 = time.perf_counter()
|
|
audio = pipe.decode(latents)
|
|
meta["timing"]["vae_seconds"] = time.perf_counter() - t0
|
|
save(args.out_dir, "audio.f32", audio, meta)
|
|
write_wav_i16(os.path.join(args.out_dir, "audio.wav"), audio, 48000)
|
|
meta["audio_seconds"] = audio.shape[0] / 48000.0
|
|
meta["peak_vram_gib"] = torch.cuda.max_memory_allocated() / 2**30
|
|
|
|
with open(os.path.join(args.out_dir, "meta.json"), "w", encoding="utf-8") as handle:
|
|
json.dump(meta, handle, indent=1, ensure_ascii=False)
|
|
print("ORACLE-DONE", args.out_dir, json.dumps(meta["timing"], default=str)[:400])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|