From d0ecde7779417768216c134d017c488a84f8ef50 Mon Sep 17 00:00:00 2001 From: vdruts <20264016+vdruts@users.noreply.github.com> Date: Mon, 21 Sep 2026 09:21:44 -0400 Subject: [PATCH] YuE2: load safetensors semantic heads / NAR adapters (v5, v8, v9) and key the latent cache by head --- .../audio_models/yue2/src/model.py | 31 ++++++++++++++++--- .../audio_models/yue2/src/tokenizer.py | 14 +++++++-- .../audio_models/yue2/yue2_model.py | 9 ++++++ 3 files changed, 48 insertions(+), 6 deletions(-) diff --git a/extensions_built_in/audio_models/yue2/src/model.py b/extensions_built_in/audio_models/yue2/src/model.py index fb80541..0cc1207 100644 --- a/extensions_built_in/audio_models/yue2/src/model.py +++ b/extensions_built_in/audio_models/yue2/src/model.py @@ -402,13 +402,36 @@ class YuE2Model(nn.Module, OstrisModelMixin): @torch.no_grad() +def _nar_lora_from_safetensors(path: str) -> dict: + """The ``.safetensors`` releases of the community NAR adapter name every tensor + (``layers.N.nar_self_attn.q_proj.lora_A`` ..., ``vae2llm.*``, ``llm2vae.*``); rebuild the + ``.pt`` layout (``lora``: layer-major A/B list, ``io``: full projection weights) that the merge walks. + The ``_comfyui`` files use ComfyUI's fused layout and are not handled here.""" + from safetensors.torch import load_file + + tensors = load_file(path, device="cpu") + if not any(".nar_self_attn." in k and k.endswith(".lora_A") for k in tensors): + raise ValueError(f"{path}: not an unmerged NAR LoRA release (the _comfyui layout is not supported for merging)") + modules = [("nar_self_attn", n) for n in ("q_proj", "k_proj", "v_proj", "o_proj")] + modules += [("nar_mlp", n) for n in ("gate_proj", "up_proj", "down_proj")] + lora = [] + for layer in sorted({int(k.split(".")[1]) for k in tensors if k.startswith("layers.")}): + for block, proj in modules: + lora += [tensors[f"layers.{layer}.{block}.{proj}.lora_A"], tensors[f"layers.{layer}.{block}.{proj}.lora_B"]] + io = {m: {k.split(".", 1)[1]: v for k, v in tensors.items() if k.startswith(m + ".")} for m in ("vae2llm", "llm2vae")} + return {"lora": lora, "io": {m: sd for m, sd in io.items() if sd}} + + def merge_nar_lora(model: YuE2Model, ckpt_path: str, scale: float = 1.0): """Fold the community NAR adapter (unmerged q/k/v + gate/up LoRA pairs, plus full vae2llm/llm2vae weights) into the merged-projection base.""" - try: - ck = torch.load(ckpt_path, map_location="cpu", weights_only=True) - except Exception: - ck = torch.load(ckpt_path, map_location="cpu", weights_only=False) + if str(ckpt_path).endswith(".safetensors"): + ck = _nar_lora_from_safetensors(ckpt_path) + else: + try: + ck = torch.load(ckpt_path, map_location="cpu", weights_only=True) + except Exception: + ck = torch.load(ckpt_path, map_location="cpu", weights_only=False) if getattr(model, "aitk_is_quantized", False): raise ValueError("merge_nar_lora needs the bf16 checkpoint; the int8 convrot repack cannot take merged weights") tensors = iter(ck["lora"]) diff --git a/extensions_built_in/audio_models/yue2/src/tokenizer.py b/extensions_built_in/audio_models/yue2/src/tokenizer.py index c0954ae..40fc2f3 100644 --- a/extensions_built_in/audio_models/yue2/src/tokenizer.py +++ b/extensions_built_in/audio_models/yue2/src/tokenizer.py @@ -167,6 +167,17 @@ class SheetSage2Transcriber: return abc +def load_head_state_dict(head_path: str) -> dict: + """Head weights from either release format of the community tokenizer repo: the original ``.pt`` + (``{"model": state_dict, "cfg": ...}``) or the ``.safetensors`` files, which hold the same state dict + flat. The v5 / v8 / v9 joint heads (same architecture as v4) are published as safetensors only.""" + if str(head_path).endswith(".safetensors"): + from safetensors.torch import load_file + + return load_file(head_path, device="cpu") + return torch.load(head_path, map_location="cpu", weights_only=False)["model"] + + class SemanticTokenizer(nn.Module): """Waveform -> per-frame codec ids (0..32767) at 25 Hz.""" @@ -180,8 +191,7 @@ class SemanticTokenizer(nn.Module): self.mert.requires_grad_(False) _rebuild_rotary(self.mert) self.head = TokenizerHead() - ck = torch.load(head_path, map_location="cpu", weights_only=False) - self.head.load_state_dict(ck["model"]) + self.head.load_state_dict(load_head_state_dict(head_path)) self.head.eval().requires_grad_(False) @property diff --git a/extensions_built_in/audio_models/yue2/yue2_model.py b/extensions_built_in/audio_models/yue2/yue2_model.py index acb2c34..3b1c557 100644 --- a/extensions_built_in/audio_models/yue2/yue2_model.py +++ b/extensions_built_in/audio_models/yue2/yue2_model.py @@ -7,6 +7,10 @@ The official audio -> semantic-token tokenizer is unreleased; training condition from the community head (MERT-v2-FullSong + classifier) in ``Mothersuperior/yue2-mothersuperior-realaudio-tokenizer-v4``. Its companion NAR adapter can be merged into the base NAR on load (``model_kwargs.merge_nar_lora``, default false). +The default is the v4 pair; the repo's later pairs (v5 / v8 / v9, safetensors only) load through +``model_kwargs.semantic_head_path`` / ``nar_lora_path``, e.g. +``Mothersuperior/yue2-mothersuperior-realaudio-tokenizer-v4/tokenizer_head_joint_v9.safetensors`` +(or a local file). A non-default head gets its own latent cache key, since the cache holds its tokens. Training: per song the latent cache carries VAE latents plus the head's codec tokens; the "text embedding" is the AR prompt prefix (instruction + tags + lyrics) as token embeddings. @@ -303,6 +307,11 @@ class YuE2AudioModel(BaseAudioModel): def get_latent_space_version(self): version = super().get_latent_space_version() + # the cached codec tokens come from the semantic head: a non-default head needs its own cache + # (the default keeps the plain key so existing caches stay valid) + if self.semantic_head_path != f"{HEAD_REPO}/{HEAD_FILE}": + head_name = os.path.splitext(os.path.basename(self.semantic_head_path))[0] + version = f"{version}_{re.sub(r'[^A-Za-z0-9]+', '-', head_name)}" # stem tokens/sheets only ride the latent cache in separation mode return f"{version}_sep" if self.do_separation else version