File size: 5,600 Bytes
90884df | 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 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 | """Fail-fast native-token encoding API for the released Music 3 checkpoint.
The public checkpoint does not contain the RVQ quantizer/codebooks required by
this operation. This module exists so callers get a precise, reusable failure
instead of accidentally treating continuous DAV latents as discrete tokens.
"""
from __future__ import annotations
import argparse
import json
import os
from pathlib import Path
from typing import Any
import torch
from inspect_dav import inspect_checkpoint
DAV_PATH_ENV = "MINIMAX_DAV_PATH"
class NativeTokenizerUnavailableError(RuntimeError):
"""The supplied release checkpoint cannot produce native Music 3 tokens."""
def __init__(self, message: str, *, report: dict[str, Any] | None = None) -> None:
super().__init__(message)
self.report = report
def resolve_dav_path(dav_path: str | os.PathLike[str] | None = None) -> Path:
value = dav_path if dav_path is not None else os.environ.get(DAV_PATH_ENV)
if value is None or not str(value).strip():
raise NativeTokenizerUnavailableError(
f"no DAV checkpoint supplied; pass dav_path or set {DAV_PATH_ENV}"
)
return Path(value).expanduser()
def require_native_tokenizer(
dav_path: str | os.PathLike[str] | None = None,
) -> dict[str, Any]:
"""Inspect ``dav.pth`` and return its report only if a native tokenizer exists."""
resolved = resolve_dav_path(dav_path)
report = inspect_checkpoint(resolved)
capabilities = report["capabilities"]
if not capabilities["can_encode_native_music3_tokens"]:
missing = [
label
for label, available in (
("waveform/tokenizer encoder weights", capabilities["waveform_analysis_encoder_weights"]),
("RVQ/VQ quantizer weights", capabilities["rvq_or_vq_quantizer_weights"]),
("codebook embedding weights", capabilities["codebook_embedding_weights"]),
("serialized tokenizer architecture/config", capabilities["serialized_architecture_config"]),
)
if not available
]
if missing:
reason = "released dav.pth cannot encode native Music 3 tokens: missing " + ", ".join(missing) + "."
else:
reason = (
"released dav.pth has tokenizer-like candidate weights/config, but no exact compatible "
"executable Music 3 tokenizer architecture/API has been implemented and verified."
)
continuous_evidence = (
capabilities["waveform_analysis_encoder_weights"]
and capabilities["continuous_gaussian_posterior_heads"]
and capabilities["continuous_flow_weights"]
)
if continuous_evidence:
reason += (
" encoder.* plus mean_proj.* and logs_proj.* form a continuous "
"Flow-VAE analysis path; they do not emit integer RVQ codes."
)
raise NativeTokenizerUnavailableError(reason, report=report)
return report
def encode_audio(
audio_path: str | os.PathLike[str],
*,
dav_path: str | os.PathLike[str] | None = None,
) -> torch.Tensor:
"""Return native tokens as ``[frames, 8]`` or fail before reading the WAV.
The release inspected for this project always takes the failure path. The
return annotation records the intended contract without manufacturing token
IDs or using a third-party codec.
"""
# Capability inspection intentionally precedes even checking the input path.
require_native_tokenizer(dav_path)
raise NativeTokenizerUnavailableError(
"checkpoint advertises tokenizer-like weights, but no released Music 3 "
"tokenizer architecture/API is available to execute them safely"
)
def _blocked_payload(error: NativeTokenizerUnavailableError, audio_path: str) -> dict[str, Any]:
payload: dict[str, Any] = {
"status": "BLOCKED",
"operation": "encode_audio",
"audio_path": audio_path,
"audio_was_read": False,
"reason": str(error),
}
if error.report is not None:
payload["capabilities"] = error.report["capabilities"]
payload["state_dict"] = error.report["state_dict"]
return payload
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("audio", help="Input WAV path (not opened when capability is blocked)")
parser.add_argument("--dav", help=f"Path to dav.pth; defaults to {DAV_PATH_ENV}")
parser.add_argument("--json", action="store_true", help="Emit machine-readable JSON")
return parser
def main(argv: list[str] | None = None) -> int:
args = build_parser().parse_args(argv)
try:
tokens = encode_audio(args.audio, dav_path=args.dav)
except (NativeTokenizerUnavailableError, FileNotFoundError, RuntimeError, ValueError) as error:
if isinstance(error, NativeTokenizerUnavailableError):
native_error = error
else:
native_error = NativeTokenizerUnavailableError(str(error))
payload = _blocked_payload(native_error, args.audio)
print(json.dumps(payload, indent=2) if args.json else f"BLOCKED: {payload['reason']}")
return 2
# Kept for API completeness if an official tokenizer is released later.
payload = {"status": "OK", "shape": list(tokens.shape), "dtype": str(tokens.dtype)}
print(json.dumps(payload, indent=2) if args.json else f"tokens: {tuple(tokens.shape)}")
return 0
if __name__ == "__main__":
raise SystemExit(main())
|