YuE2: load safetensors semantic heads / NAR adapters (v5, v8, v9) and key the latent cache by head

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"])

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

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