Add lm lane (byte-identical mirror of the upstream release)

lm/LICENSE

@@ -0,0 +1,429 @@
+YuE2 model-weight license
+
+The YuE2-3B, YuE2-Vae and YuE2-Vae-legacy checkpoint weights are
+licensed under Creative Commons Attribution-NonCommercial 4.0 International
+(CC BY-NC 4.0): https://creativecommons.org/licenses/by-nc/4.0/
+
+Scope: the YuE2 checkpoint weights in model.safetensors, or the corresponding
+models/<model-name>/model.safetensors files in the inference kit.
+For attribution, identify YuE2, the model name and its source repository:
+https://huggingface.co/m-a-p/YuE2-3B
+https://huggingface.co/m-a-p/YuE2-Vae
+https://huggingface.co/m-a-p/YuE2-Vae-legacy
+
+This weight license does not replace separately applicable licenses for code,
+text tokenization files, evaluation assets or other bundled material.
+Third-party source notices and their original license texts are preserved in
+THIRD_PARTY_NOTICES.md and licenses/.
+
+The official license text follows without modification.
+Source: https://creativecommons.org/licenses/by-nc/4.0/legalcode.txt
+
+Attribution-NonCommercial 4.0 International
+
+=======================================================================
+
+Creative Commons Corporation ("Creative Commons") is not a law firm and
+does not provide legal services or legal advice. Distribution of
+Creative Commons public licenses does not create a lawyer-client or
+other relationship. Creative Commons makes its licenses and related
+information available on an "as-is" basis. Creative Commons gives no
+warranties regarding its licenses, any material licensed under their
+terms and conditions, or any related information. Creative Commons
+disclaims all liability for damages resulting from their use to the
+fullest extent possible.
+
+Using Creative Commons Public Licenses
+
+Creative Commons public licenses provide a standard set of terms and
+conditions that creators and other rights holders may use to share
+original works of authorship and other material subject to copyright
+and certain other rights specified in the public license below. The
+following considerations are for informational purposes only, are not
+exhaustive, and do not form part of our licenses.
+
+     Considerations for licensors: Our public licenses are
+     intended for use by those authorized to give the public
+     permission to use material in ways otherwise restricted by
+     copyright and certain other rights. Our licenses are
+     irrevocable. Licensors should read and understand the terms
+     and conditions of the license they choose before applying it.
+     Licensors should also secure all rights necessary before
+     applying our licenses so that the public can reuse the
+     material as expected. Licensors should clearly mark any
+     material not subject to the license. This includes other CC-
+     licensed material, or material used under an exception or
+     limitation to copyright. More considerations for licensors:
+    wiki.creativecommons.org/Considerations_for_licensors
+
+     Considerations for the public: By using one of our public
+     licenses, a licensor grants the public permission to use the
+     licensed material under specified terms and conditions. If
+     the licensor's permission is not necessary for any reason--for
+     example, because of any applicable exception or limitation to
+     copyright--then that use is not regulated by the license. Our
+     licenses grant only permissions under copyright and certain
+     other rights that a licensor has authority to grant. Use of
+     the licensed material may still be restricted for other
+     reasons, including because others have copyright or other
+     rights in the material. A licensor may make special requests,
+     such as asking that all changes be marked or described.
+     Although not required by our licenses, you are encouraged to
+     respect those requests where reasonable. More considerations
+     for the public:
+    wiki.creativecommons.org/Considerations_for_licensees
+
+=======================================================================
+
+Creative Commons Attribution-NonCommercial 4.0 International Public
+License
+
+By exercising the Licensed Rights (defined below), You accept and agree
+to be bound by the terms and conditions of this Creative Commons
+Attribution-NonCommercial 4.0 International Public License ("Public
+License"). To the extent this Public License may be interpreted as a
+contract, You are granted the Licensed Rights in consideration of Your
+acceptance of these terms and conditions, and the Licensor grants You
+such rights in consideration of benefits the Licensor receives from
+making the Licensed Material available under these terms and
+conditions.
+
+
+Section 1 -- Definitions.
+
+  a. Adapted Material means material subject to Copyright and Similar
+     Rights that is derived from or based upon the Licensed Material
+     and in which the Licensed Material is translated, altered,
+     arranged, transformed, or otherwise modified in a manner requiring
+     permission under the Copyright and Similar Rights held by the
+     Licensor. For purposes of this Public License, where the Licensed
+     Material is a musical work, performance, or sound recording,
+     Adapted Material is always produced where the Licensed Material is
+     synched in timed relation with a moving image.
+
+  b. Adapter's License means the license You apply to Your Copyright
+     and Similar Rights in Your contributions to Adapted Material in
+     accordance with the terms and conditions of this Public License.
+
+  c. Copyright and Similar Rights means copyright and/or similar rights
+     closely related to copyright including, without limitation,
+     performance, broadcast, sound recording, and Sui Generis Database
+     Rights, without regard to how the rights are labeled or
+     categorized. For purposes of this Public License, the rights
+     specified in Section 2(b)(1)-(2) are not Copyright and Similar
+     Rights.
+  d. Effective Technological Measures means those measures that, in the
+     absence of proper authority, may not be circumvented under laws
+     fulfilling obligations under Article 11 of the WIPO Copyright
+     Treaty adopted on December 20, 1996, and/or similar international
+     agreements.
+
+  e. Exceptions and Limitations means fair use, fair dealing, and/or
+     any other exception or limitation to Copyright and Similar Rights
+     that applies to Your use of the Licensed Material.
+
+  f. Licensed Material means the artistic or literary work, database,
+     or other material to which the Licensor applied this Public
+     License.
+
+  g. Licensed Rights means the rights granted to You subject to the
+     terms and conditions of this Public License, which are limited to
+     all Copyright and Similar Rights that apply to Your use of the
+     Licensed Material and that the Licensor has authority to license.
+
+  h. Licensor means the individual(s) or entity(ies) granting rights
+     under this Public License.
+
+  i. NonCommercial means not primarily intended for or directed towards
+     commercial advantage or monetary compensation. For purposes of
+     this Public License, the exchange of the Licensed Material for
+     other material subject to Copyright and Similar Rights by digital
+     file-sharing or similar means is NonCommercial provided there is
+     no payment of monetary compensation in connection with the
+     exchange.
+
+  j. Share means to provide material to the public by any means or
+     process that requires permission under the Licensed Rights, such
+     as reproduction, public display, public performance, distribution,
+     dissemination, communication, or importation, and to make material
+     available to the public including in ways that members of the
+     public may access the material from a place and at a time
+     individually chosen by them.
+
+  k. Sui Generis Database Rights means rights other than copyright
+     resulting from Directive 96/9/EC of the European Parliament and of
+     the Council of 11 March 1996 on the legal protection of databases,
+     as amended and/or succeeded, as well as other essentially
+     equivalent rights anywhere in the world.
+
+  l. You means the individual or entity exercising the Licensed Rights
+     under this Public License. Your has a corresponding meaning.
+
+
+Section 2 -- Scope.
+
+  a. License grant.
+
+       1. Subject to the terms and conditions of this Public License,
+          the Licensor hereby grants You a worldwide, royalty-free,
+          non-sublicensable, non-exclusive, irrevocable license to
+          exercise the Licensed Rights in the Licensed Material to:
+
+            a. reproduce and Share the Licensed Material, in whole or
+               in part, for NonCommercial purposes only; and
+
+            b. produce, reproduce, and Share Adapted Material for
+               NonCommercial purposes only.
+
+       2. Exceptions and Limitations. For the avoidance of doubt, where
+          Exceptions and Limitations apply to Your use, this Public
+          License does not apply, and You do not need to comply with
+          its terms and conditions.
+
+       3. Term. The term of this Public License is specified in Section
+          6(a).
+
+       4. Media and formats; technical modifications allowed. The
+          Licensor authorizes You to exercise the Licensed Rights in
+          all media and formats whether now known or hereafter created,
+          and to make technical modifications necessary to do so. The
+          Licensor waives and/or agrees not to assert any right or
+          authority to forbid You from making technical modifications
+          necessary to exercise the Licensed Rights, including
+          technical modifications necessary to circumvent Effective
+          Technological Measures. For purposes of this Public License,
+          simply making modifications authorized by this Section 2(a)
+          (4) never produces Adapted Material.
+
+       5. Downstream recipients.
+
+            a. Offer from the Licensor -- Licensed Material. Every
+               recipient of the Licensed Material automatically
+               receives an offer from the Licensor to exercise the
+               Licensed Rights under the terms and conditions of this
+               Public License.
+
+            b. No downstream restrictions. You may not offer or impose
+               any additional or different terms or conditions on, or
+               apply any Effective Technological Measures to, the
+               Licensed Material if doing so restricts exercise of the
+               Licensed Rights by any recipient of the Licensed
+               Material.
+
+       6. No endorsement. Nothing in this Public License constitutes or
+          may be construed as permission to assert or imply that You
+          are, or that Your use of the Licensed Material is, connected
+          with, or sponsored, endorsed, or granted official status by,
+          the Licensor or others designated to receive attribution as
+          provided in Section 3(a)(1)(A)(i).
+
+  b. Other rights.
+
+       1. Moral rights, such as the right of integrity, are not
+          licensed under this Public License, nor are publicity,
+          privacy, and/or other similar personality rights; however, to
+          the extent possible, the Licensor waives and/or agrees not to
+          assert any such rights held by the Licensor to the limited
+          extent necessary to allow You to exercise the Licensed
+          Rights, but not otherwise.
+
+       2. Patent and trademark rights are not licensed under this
+          Public License.
+
+       3. To the extent possible, the Licensor waives any right to
+          collect royalties from You for the exercise of the Licensed
+          Rights, whether directly or through a collecting society
+          under any voluntary or waivable statutory or compulsory
+          licensing scheme. In all other cases the Licensor expressly
+          reserves any right to collect such royalties, including when
+          the Licensed Material is used other than for NonCommercial
+          purposes.
+
+
+Section 3 -- License Conditions.
+
+Your exercise of the Licensed Rights is expressly made subject to the
+following conditions.
+
+  a. Attribution.
+
+       1. If You Share the Licensed Material (including in modified
+          form), You must:
+
+            a. retain the following if it is supplied by the Licensor
+               with the Licensed Material:
+
+                 i. identification of the creator(s) of the Licensed
+                    Material and any others designated to receive
+                    attribution, in any reasonable manner requested by
+                    the Licensor (including by pseudonym if
+                    designated);
+
+                ii. a copyright notice;
+
+               iii. a notice that refers to this Public License;
+
+                iv. a notice that refers to the disclaimer of
+                    warranties;
+
+                 v. a URI or hyperlink to the Licensed Material to the
+                    extent reasonably practicable;
+
+            b. indicate if You modified the Licensed Material and
+               retain an indication of any previous modifications; and
+
+            c. indicate the Licensed Material is licensed under this
+               Public License, and include the text of, or the URI or
+               hyperlink to, this Public License.
+
+       2. You may satisfy the conditions in Section 3(a)(1) in any
+          reasonable manner based on the medium, means, and context in
+          which You Share the Licensed Material. For example, it may be
+          reasonable to satisfy the conditions by providing a URI or
+          hyperlink to a resource that includes the required
+          information.
+
+       3. If requested by the Licensor, You must remove any of the
+          information required by Section 3(a)(1)(A) to the extent
+          reasonably practicable.
+
+       4. If You Share Adapted Material You produce, the Adapter's
+          License You apply must not prevent recipients of the Adapted
+          Material from complying with this Public License.
+
+
+Section 4 -- Sui Generis Database Rights.
+
+Where the Licensed Rights include Sui Generis Database Rights that
+apply to Your use of the Licensed Material:
+
+  a. for the avoidance of doubt, Section 2(a)(1) grants You the right
+     to extract, reuse, reproduce, and Share all or a substantial
+     portion of the contents of the database for NonCommercial purposes
+     only;
+
+  b. if You include all or a substantial portion of the database
+     contents in a database in which You have Sui Generis Database
+     Rights, then the database in which You have Sui Generis Database
+     Rights (but not its individual contents) is Adapted Material; and
+
+  c. You must comply with the conditions in Section 3(a) if You Share
+     all or a substantial portion of the contents of the database.
+
+For the avoidance of doubt, this Section 4 supplements and does not
+replace Your obligations under this Public License where the Licensed
+Rights include other Copyright and Similar Rights.
+
+
+Section 5 -- Disclaimer of Warranties and Limitation of Liability.
+
+  a. UNLESS OTHERWISE SEPARATELY UNDERTAKEN BY THE LICENSOR, TO THE
+     EXTENT POSSIBLE, THE LICENSOR OFFERS THE LICENSED MATERIAL AS-IS
+     AND AS-AVAILABLE, AND MAKES NO REPRESENTATIONS OR WARRANTIES OF
+     ANY KIND CONCERNING THE LICENSED MATERIAL, WHETHER EXPRESS,
+     IMPLIED, STATUTORY, OR OTHER. THIS INCLUDES, WITHOUT LIMITATION,
+     WARRANTIES OF TITLE, MERCHANTABILITY, FITNESS FOR A PARTICULAR
+     PURPOSE, NON-INFRINGEMENT, ABSENCE OF LATENT OR OTHER DEFECTS,
+     ACCURACY, OR THE PRESENCE OR ABSENCE OF ERRORS, WHETHER OR NOT
+     KNOWN OR DISCOVERABLE. WHERE DISCLAIMERS OF WARRANTIES ARE NOT
+     ALLOWED IN FULL OR IN PART, THIS DISCLAIMER MAY NOT APPLY TO YOU.
+
+  b. TO THE EXTENT POSSIBLE, IN NO EVENT WILL THE LICENSOR BE LIABLE
+     TO YOU ON ANY LEGAL THEORY (INCLUDING, WITHOUT LIMITATION,
+     NEGLIGENCE) OR OTHERWISE FOR ANY DIRECT, SPECIAL, INDIRECT,
+     INCIDENTAL, CONSEQUENTIAL, PUNITIVE, EXEMPLARY, OR OTHER LOSSES,
+     COSTS, EXPENSES, OR DAMAGES ARISING OUT OF THIS PUBLIC LICENSE OR
+     USE OF THE LICENSED MATERIAL, EVEN IF THE LICENSOR HAS BEEN
+     ADVISED OF THE POSSIBILITY OF SUCH LOSSES, COSTS, EXPENSES, OR
+     DAMAGES. WHERE A LIMITATION OF LIABILITY IS NOT ALLOWED IN FULL OR
+     IN PART, THIS LIMITATION MAY NOT APPLY TO YOU.
+
+  c. The disclaimer of warranties and limitation of liability provided
+     above shall be interpreted in a manner that, to the extent
+     possible, most closely approximates an absolute disclaimer and
+     waiver of all liability.
+
+
+Section 6 -- Term and Termination.
+
+  a. This Public License applies for the term of the Copyright and
+     Similar Rights licensed here. However, if You fail to comply with
+     this Public License, then Your rights under this Public License
+     terminate automatically.
+
+  b. Where Your right to use the Licensed Material has terminated under
+     Section 6(a), it reinstates:
+
+       1. automatically as of the date the violation is cured, provided
+          it is cured within 30 days of Your discovery of the
+          violation; or
+
+       2. upon express reinstatement by the Licensor.
+
+     For the avoidance of doubt, this Section 6(b) does not affect any
+     right the Licensor may have to seek remedies for Your violations
+     of this Public License.
+
+  c. For the avoidance of doubt, the Licensor may also offer the
+     Licensed Material under separate terms or conditions or stop
+     distributing the Licensed Material at any time; however, doing so
+     will not terminate this Public License.
+
+  d. Sections 1, 5, 6, 7, and 8 survive termination of this Public
+     License.
+
+
+Section 7 -- Other Terms and Conditions.
+
+  a. The Licensor shall not be bound by any additional or different
+     terms or conditions communicated by You unless expressly agreed.
+
+  b. Any arrangements, understandings, or agreements regarding the
+     Licensed Material not stated herein are separate from and
+     independent of the terms and conditions of this Public License.
+
+
+Section 8 -- Interpretation.
+
+  a. For the avoidance of doubt, this Public License does not, and
+     shall not be interpreted to, reduce, limit, restrict, or impose
+     conditions on any use of the Licensed Material that could lawfully
+     be made without permission under this Public License.
+
+  b. To the extent possible, if any provision of this Public License is
+     deemed unenforceable, it shall be automatically reformed to the
+     minimum extent necessary to make it enforceable. If the provision
+     cannot be reformed, it shall be severed from this Public License
+     without affecting the enforceability of the remaining terms and
+     conditions.
+
+  c. No term or condition of this Public License will be waived and no
+     failure to comply consented to unless expressly agreed to by the
+     Licensor.
+
+  d. Nothing in this Public License constitutes or may be interpreted
+     as a limitation upon, or waiver of, any privileges and immunities
+     that apply to the Licensor or You, including from the legal
+     processes of any jurisdiction or authority.
+
+=======================================================================
+
+Creative Commons is not a party to its public
+licenses. Notwithstanding, Creative Commons may elect to apply one of
+its public licenses to material it publishes and in those instances
+will be considered the “Licensor.” The text of the Creative Commons
+public licenses is dedicated to the public domain under the CC0 Public
+Domain Dedication. Except for the limited purpose of indicating that
+material is shared under a Creative Commons public license or as
+otherwise permitted by the Creative Commons policies published at
+creativecommons.org/policies, Creative Commons does not authorize the
+use of the trademark "Creative Commons" or any other trademark or logo
+of Creative Commons without its prior written consent including,
+without limitation, in connection with any unauthorized modifications
+to any of its public licenses or any other arrangements,
+understandings, or agreements concerning use of licensed material. For
+the avoidance of doubt, this paragraph does not form part of the
+public licenses.
+
+Creative Commons may be contacted at creativecommons.org.
+

lm/THIRD_PARTY_NOTICES.md

@@ -0,0 +1,15 @@
+# Third-party code notices
+
+The Oobleck VAE and SnakeBeta implementation in `modeling_vae.py` is derived
+from stable-audio-tools commit `a6ae0cdf8b2eb1567a4b42ceadddec3712d99d45`.
+The module hierarchy, weight normalization and activation equations preserve
+the checkpoint's original inference implementation.
+
+- Oobleck / stable-audio-tools: Copyright (c) 2023 Stability AI, MIT.
+  Full text: `licenses/stable-audio-tools-MIT.txt`.
+- SnakeBeta / BigVGAN: Copyright (c) 2022 NVIDIA CORPORATION, MIT.
+  Full text: `licenses/SnakeBeta-NVIDIA-MIT.txt`.
+
+These notices cover the identified source code and retain its original licenses.
+The YuE2 model checkpoint weights are separately licensed under CC BY-NC 4.0;
+see LICENSE for the scope and full terms. This does not relicense third-party code.

lm/config.json

@@ -0,0 +1,38 @@
+{
+  "return_dict": true,
+  "output_hidden_states": false,
+  "torchscript": false,
+  "dtype": "bfloat16",
+  "tie_word_embeddings": false,
+  "is_encoder_decoder": false,
+  "is_decoder": false,
+  "add_cross_attention": false,
+  "architectures": [
+    "YuE2ForCausalLM"
+  ],
+  "bos_token_id": null,
+  "pad_token_id": null,
+  "eos_token_id": null,
+  "decoder_start_token_id": null,
+  "transformers_version": "4.57.6",
+  "auto_map": {
+    "AutoConfig": "modeling_yue2.YuE2Config",
+    "AutoModelForCausalLM": "modeling_yue2.YuE2ForCausalLM"
+  },
+  "model_type": "yue2",
+  "hidden_size": 2048,
+  "num_hidden_layers": 28,
+  "num_attention_heads": 16,
+  "num_key_value_heads": 8,
+  "head_dim": 128,
+  "intermediate_size": 6144,
+  "vocab_size": 184704,
+  "rms_norm_eps": 1e-06,
+  "rope_theta": 1000000,
+  "max_position_embeddings": 24576,
+  "latent_type": "vae",
+  "latent_dim": 64,
+  "max_latent_frames": 24576,
+  "timestep_shift": 1.0,
+  "output_attentions": false
+}

lm/generation_config.json

@@ -0,0 +1,11 @@
+{
+  "bos_token_id": 151643,
+  "eos_token_id": 151852,
+  "pad_token_id": 151643,
+  "do_sample": true,
+  "temperature": 1.0,
+  "top_p": 0.95,
+  "top_k": 100,
+  "max_new_tokens": 9000,
+  "use_cache": true
+}

lm/licenses/SnakeBeta-NVIDIA-MIT.txt

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2022 NVIDIA CORPORATION.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software. 
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

lm/licenses/stable-audio-tools-MIT.txt

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 Stability AI
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

lm/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1d55c42c1a9875c34f5d736e15078449992b044e807ce2a138e6cf289a1e59e9
+size 7261441640

lm/modeling_yue2.py

@@ -0,0 +1,705 @@
+"""YuE2 AR–NAR Mixture-of-Transformers, with checkpoint-compatible names.
+
+This module is self contained for Transformers ``trust_remote_code`` loading.
+It imports no CUDA extension and implements the released model architecture.
+``generate`` returns token IDs; the package pipeline supplies song generation.
+"""
+from __future__ import annotations
+
+import math
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import GenerationMixin, PretrainedConfig, PreTrainedModel
+from transformers.cache_utils import DynamicCache
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+
+def sdpa(query, key, value, *, attn_mask=None, is_causal=False):
+    """Use native grouped-query attention, including a portable MPS fallback."""
+    grouped = query.shape[1] != key.shape[1]
+    if grouped and query.device.type == "mps":
+        # PyTorch's MPS attention does not implement enable_gqa on every release.
+        groups = query.shape[1] // key.shape[1]
+        key = key.repeat_interleave(groups, dim=1)
+        value = value.repeat_interleave(groups, dim=1)
+        grouped = False
+    return F.scaled_dot_product_attention(
+        query, key, value, attn_mask=attn_mask, is_causal=is_causal,
+        enable_gqa=grouped,
+    )
+
+
+def _causal_mask(attention_mask, cache_position, key_length, batch_size):
+    """Physical cache slots are causal; RoPE positions may exclude padding."""
+    device = cache_position.device
+    visible = torch.arange(key_length, device=device)[None, :] <= cache_position[:, None]
+    visible = visible[None, None].expand(batch_size, 1, -1, -1)
+    if attention_mask is None:
+        return visible
+    mask = attention_mask.to(device=device)
+    if mask.ndim == 2:
+        if mask.shape[0] != batch_size or mask.shape[1] > key_length:
+            raise ValueError("attention_mask must cover the batch and used cache slots")
+        # Static cache has unused capacity after the supplied 2D padding mask.
+        if mask.shape[1] < key_length:
+            mask = F.pad(mask, (0, key_length - mask.shape[1]), value=0)
+        return visible & mask[:, None, None, :].bool()
+    if mask.ndim != 4 or mask.shape[-2:] != visible.shape[-2:]:
+        raise ValueError("Expected a 2D padding mask or a matching 4D attention mask")
+    if mask.dtype == torch.bool:
+        return visible & mask
+    return mask.masked_fill(~visible, float("-inf"))
+
+# ══════════════════════════════════════════════════════════════════════════════
+# Config
+# ══════════════════════════════════════════════════════════════════════════════
+
+
+class YuE2Config(PretrainedConfig):
+    model_type = "yue2"
+
+    _hf_fields = frozenset({
+        "model_type", "architectures", "auto_map", "transformers_version",
+        "dtype", "torch_dtype", "return_dict", "output_hidden_states",
+        "output_attentions", "use_cache", "tie_word_embeddings", "torchscript",
+        "is_decoder", "is_encoder_decoder", "add_cross_attention",
+        "bos_token_id", "eos_token_id", "pad_token_id", "decoder_start_token_id",
+        "attn_implementation",
+    })
+
+    def to_dict(self):
+        return {key: value for key, value in super().to_dict().items()
+                if key in self._hf_fields or key in self._inference_fields}
+
+    _inference_fields = frozenset(['hidden_size', 'num_hidden_layers', 'num_attention_heads', 'num_key_value_heads', 'head_dim', 'intermediate_size', 'vocab_size', 'rms_norm_eps', 'rope_theta', 'max_position_embeddings', 'tie_word_embeddings', 'latent_type', 'latent_dim', 'max_latent_frames', 'timestep_shift'])
+
+    def __init__(
+        self,
+        hidden_size: int = 2048,
+        num_hidden_layers: int = 28,
+        num_attention_heads: int = 16,
+        num_key_value_heads: int = 8,
+        head_dim: int = 128,
+        intermediate_size: int = 6144,
+        vocab_size: int = 184704,
+        rms_norm_eps: float = 1e-6,
+        rope_theta: float = 1000000.0,
+        max_position_embeddings: int = 24576,
+        tie_word_embeddings: bool = False,
+        # Acoustic inference architecture
+        latent_type: str = "vae",
+        latent_dim: int = 64,
+        max_latent_frames: int = 24576,
+        timestep_shift: float = 1.0,
+        **kwargs,
+    ):
+        if latent_type != "vae":
+            raise ValueError("YuE2 inference supports only latent_type='vae'")
+        # Serialize only the documented model and Transformers configuration.
+        kwargs = {key: value for key, value in kwargs.items() if key in self._hf_fields}
+        super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.head_dim = head_dim
+        self.intermediate_size = intermediate_size
+        self.vocab_size = vocab_size
+        self.rms_norm_eps = rms_norm_eps
+        self.rope_theta = rope_theta
+        self.max_position_embeddings = max_position_embeddings
+        self.latent_type = latent_type
+        self.latent_dim = latent_dim
+        self.max_latent_frames = max_latent_frames
+        self.timestep_shift = timestep_shift
+
+
+# ══════════════════════════════════════════════════════════════════════════════
+# Building blocks
+# ══════════════════════════════════════════════════════════════════════════════
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps).to(x.dtype) * self.weight
+
+
+class RotaryEmbedding(nn.Module):
+    def __init__(self, head_dim: int, base: float = 1000000.0):
+        super().__init__()
+        self.head_dim = head_dim
+        self.base = base
+        self._inv_freq: Optional[torch.Tensor] = None
+
+    def forward(self, position_ids: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        if self._inv_freq is None or self._inv_freq.device != position_ids.device:
+            self._inv_freq = 1.0 / (self.base ** (
+                torch.arange(0, self.head_dim, 2, dtype=torch.float32, device=position_ids.device) / self.head_dim
+            ))
+        pos = position_ids.float().unsqueeze(-1)
+        angles = pos * self._inv_freq
+        return angles.cos(), angles.sin()
+
+
+def _apply_rotary(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    half = x.shape[-1] // 2
+    x1, x2 = x[..., :half], x[..., half:]
+    cos, sin = cos.to(x.dtype), sin.to(x.dtype)
+    return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1)
+
+
+class Attention(nn.Module):
+    def __init__(self, config: YuE2Config):
+        super().__init__()
+        self.num_heads = config.num_attention_heads
+        self.num_kv_heads = config.num_key_value_heads
+        self.head_dim = config.head_dim
+        self.num_kv_groups = self.num_heads // self.num_kv_heads
+
+        self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=False)
+        self.q_norm = RMSNorm(self.head_dim, config.rms_norm_eps)
+        self.k_norm = RMSNorm(self.head_dim, config.rms_norm_eps)
+
+    def project_qkv(
+        self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Project, normalize, and apply RoPE. No SDPA, no KV cache, no O proj.
+
+        Returns Q [B,T,num_heads,hd], K [B,T,num_kv_heads,hd], V [B,T,num_kv_heads,hd].
+        """
+        B, T, _ = x.shape
+        q = self.q_proj(x).view(B, T, self.num_heads, self.head_dim)
+        k = self.k_proj(x).view(B, T, self.num_kv_heads, self.head_dim)
+        v = self.v_proj(x).view(B, T, self.num_kv_heads, self.head_dim)
+        q, k = self.q_norm(q), self.k_norm(k)
+        rc, rs = cos.unsqueeze(2), sin.unsqueeze(2)
+        q = _apply_rotary(q, rc, rs)
+        k = _apply_rotary(k, rc, rs)
+        return q, k, v
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        past_key_value: Optional[DynamicCache] = None,
+        layer_idx: int = 0,
+        attention_mask: Optional[torch.Tensor] = None,
+        cache_position: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        B, T, _ = x.shape
+        q, k, v = self.project_qkv(x, cos, sin)
+
+        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
+
+        if past_key_value is not None:
+            k, v = past_key_value.update(k, v, layer_idx, {"cache_position": cache_position})
+
+        if attention_mask is not None:
+            out = sdpa(q, k, v, attn_mask=attention_mask[..., :k.shape[2]])
+        else:
+            out = sdpa(q, k, v, is_causal=(T > 1 and k.shape[2] == T))
+        return self.o_proj(out.transpose(1, 2).reshape(B, T, -1))
+
+
+class MLP(nn.Module):
+    def __init__(self, config: YuE2Config):
+        super().__init__()
+        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class DecoderLayer(nn.Module):
+    """Transformer layer with full MoT: dual attention projections + dual MLP."""
+
+    def __init__(self, config: YuE2Config):
+        super().__init__()
+        # AR attention path
+        self.input_layernorm = RMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.self_attn = Attention(config)
+        # NAR attention path (separate Q/K/V/O + layernorms)
+        self.nar_input_layernorm = RMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.nar_self_attn = Attention(config)
+        # AR MLP path
+        self.post_attention_layernorm = RMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.mlp = MLP(config)
+        # NAR MLP path
+        self.nar_pre_mlp_layernorm = RMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.nar_mlp = MLP(config)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        past_key_value: Optional[DynamicCache] = None,
+        layer_idx: int = 0,
+        attention_mask: Optional[torch.Tensor] = None,
+        ar_mask: Optional[torch.Tensor] = None,
+        cache_position: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if ar_mask is not None:
+            mask_3d = ar_mask.unsqueeze(-1)  # [B, S, 1]
+            mask_4d = ar_mask.unsqueeze(-1).unsqueeze(-1)  # [B, S, 1, 1]
+
+            # Per-type input layernorm
+            ln_ar = self.input_layernorm(x)
+            ln_nar = self.nar_input_layernorm(x)
+
+            # Per-type QKV projection (both process all tokens)
+            q_ar, k_ar, v_ar = self.self_attn.project_qkv(ln_ar, cos, sin)
+            q_nar, k_nar, v_nar = self.nar_self_attn.project_qkv(ln_nar, cos, sin)
+
+            # Merge Q/K/V per-position: AR positions use AR projections, NAR use NAR
+            query = torch.where(mask_4d, q_ar, q_nar)    # [B, S, num_heads, hd]
+            # K/V have num_kv_heads (fewer), same mask broadcast works
+            key = torch.where(mask_4d, k_ar, k_nar)      # [B, S, num_kv_heads, hd]
+            value = torch.where(mask_4d, v_ar, v_nar)
+
+            # Transpose to [B, H, S, D] for SDPA
+            B, S = x.shape[:2]
+            query = query.transpose(1, 2)
+            key = key.transpose(1, 2)
+            value = value.transpose(1, 2)
+
+            # Shared attention with hybrid mask
+            if attention_mask is not None and attention_mask.dtype != torch.bool:
+                attention_mask = attention_mask.to(query.dtype)
+            core_out = sdpa(query, key, value, attn_mask=attention_mask)
+            core_out = core_out.transpose(1, 2).reshape(B, S, -1)
+
+            # Per-type O projection, merge by mask
+            o_ar = self.self_attn.o_proj(core_out)
+            o_nar = self.nar_self_attn.o_proj(core_out)
+            h = torch.where(mask_3d, o_ar, o_nar)
+            x = x + h
+
+            # Per-type MLP
+            ar_out = self.mlp(self.post_attention_layernorm(x))
+            nar_out = self.nar_mlp(self.nar_pre_mlp_layernorm(x))
+            mlp_out = torch.where(mask_3d, ar_out, nar_out)
+        else:
+            # AR-only mode (generation): use AR path only
+            h = self.self_attn(self.input_layernorm(x), cos, sin, past_key_value, layer_idx,
+                               attention_mask, cache_position)
+            x = x + h
+            mlp_out = self.mlp(self.post_attention_layernorm(x))
+
+        x = x + mlp_out
+        return x
+
+
+# ══════════════════════════════════════════════════════════════════════════════
+# NAR auxiliary modules
+# ══════════════════════════════════════════════════════════════════════════════
+
+
+class TimestepEmbedder(nn.Module):
+    """Sinusoidal timestep → MLP → hidden_size (same as modules.py)."""
+
+    def __init__(self, hidden_size: int, frequency_embedding_size: int = 256):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(frequency_embedding_size, hidden_size),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+
+    def forward(self, t):
+        half = self.frequency_embedding_size // 2
+        freqs = torch.exp(
+            -math.log(10000) * torch.arange(half, device=t.device, dtype=torch.float32) / half
+        )
+        args = t.float().unsqueeze(-1) * freqs.unsqueeze(0)
+        emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        return self.mlp(emb.to(next(self.parameters()).dtype))
+
+
+class AudioPositionEmbedding(nn.Module):
+    """Non-learnable 1D sinusoidal PE for audio latent frames."""
+
+    def __init__(self, max_frames: int, hidden_size: int):
+        super().__init__()
+        pe = torch.zeros(max_frames, hidden_size)
+        position = torch.arange(0, max_frames, dtype=torch.float32).unsqueeze(1)
+        div_term = torch.exp(
+            torch.arange(0, hidden_size, 2, dtype=torch.float32) * (-math.log(10000.0) / hidden_size)
+        )
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        self.register_buffer("pe", pe)
+
+    def forward(self, position_ids):
+        return self.pe[position_ids]
+
+
+# ══════════════════════════════════════════════════════════════════════════════
+# Static KV Cache
+# ══════════════════════════════════════════════════════════════════════════════
+
+
+class StaticKVCache:
+    """Bounded, append-only cache for the explicit single-request AR loop.
+
+    Returns views of the used prefix and never reallocates/copies its history.
+    Standard HF ``generate`` also supports Transformers' own StaticCache.
+    """
+
+    def __init__(
+        self, num_layers: int, batch_size: int, num_kv_heads: int,
+        max_seq_len: int, head_dim: int, dtype: torch.dtype, device: torch.device,
+    ):
+        self.num_layers = num_layers
+        self.max_seq_len = max_seq_len
+        self._seen_tokens = 0
+        self.key_cache: List[torch.Tensor] = [
+            torch.zeros(batch_size, num_kv_heads, max_seq_len, head_dim, dtype=dtype, device=device)
+            for _ in range(num_layers)
+        ]
+        self.value_cache: List[torch.Tensor] = [
+            torch.zeros(batch_size, num_kv_heads, max_seq_len, head_dim, dtype=dtype, device=device)
+            for _ in range(num_layers)
+        ]
+
+    def get_seq_length(self, layer_idx=0) -> int:
+        return self._seen_tokens
+
+    def update(self, key_states, value_states, layer_idx, cache_kwargs=None):
+        T = key_states.shape[2]
+        pos = self._seen_tokens
+        end = pos + T
+        if end > self.max_seq_len:
+            raise ValueError(f"KV cache capacity {self.max_seq_len} exceeded by {end}; generation was not shortened")
+        self.key_cache[layer_idx][:, :, pos:end] = key_states
+        self.value_cache[layer_idx][:, :, pos:end] = value_states
+        if layer_idx == self.num_layers - 1:
+            self._seen_tokens = end
+        return self.key_cache[layer_idx][:, :, :end], self.value_cache[layer_idx][:, :, :end]
+
+    def reset(self):
+        self._seen_tokens = 0
+
+    def reorder_cache(self, beam_idx):
+        self.key_cache = [v.index_select(0, beam_idx.to(v.device)) for v in self.key_cache]
+        self.value_cache = [v.index_select(0, beam_idx.to(v.device)) for v in self.value_cache]
+
+
+# ══════════════════════════════════════════════════════════════════════════════
+# Model
+# ══════════════════════════════════════════════════════════════════════════════
+
+
+class Backbone(nn.Module):
+    """Transformer backbone with MoT dual MLP."""
+
+    def __init__(self, config: YuE2Config):
+        super().__init__()
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.layers = nn.ModuleList([DecoderLayer(config) for _ in range(config.num_hidden_layers)])
+        self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.rotary_emb = RotaryEmbedding(config.head_dim, config.rope_theta)
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values=None,
+        use_cache: bool = True,
+        attention_mask: Optional[torch.Tensor] = None,
+        ar_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        cache_position: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, ...]:
+        if inputs_embeds is not None:
+            x = inputs_embeds
+        else:
+            x = self.embed_tokens(input_ids)
+        cos, sin = self.rotary_emb(position_ids)
+
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache()
+
+        for i, layer in enumerate(self.layers):
+            x = layer(x, cos, sin, past_key_values if use_cache else None,
+                      layer_idx=i, attention_mask=attention_mask, ar_mask=ar_mask,
+                      cache_position=cache_position)
+
+        return self.norm(x), past_key_values
+
+
+class YuE2PreTrainedModel(PreTrainedModel):
+    config_class = YuE2Config
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["DecoderLayer"]
+    _supports_sdpa = True
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, std=0.01)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, std=0.01)
+
+
+class YuE2ForCausalLM(YuE2PreTrainedModel, GenerationMixin):
+    """YuE2 model: AR causal LM (generate) + NAR flow matching (ODE)."""
+
+    def __init__(self, config: YuE2Config):
+        super().__init__(config)
+        self.model = Backbone(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        # NAR auxiliary
+        self.llm2vae = nn.Linear(config.hidden_size, config.latent_dim)
+        self.vae2llm = nn.Linear(config.latent_dim, config.hidden_size)
+        self.time_embedder = TimestepEmbedder(config.hidden_size)
+        self.latent_pos_embed = AudioPositionEmbedding(config.max_latent_frames, config.hidden_size)
+
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    # ── AR forward (standard causal LM, KV cached) ───────────────────
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values=None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        logits_to_keep: Union[int, torch.Tensor] = 0,
+        return_dict: Optional[bool] = None,
+        **kwargs,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        if (input_ids is None) == (inputs_embeds is None):
+            raise ValueError("Supply exactly one of input_ids or inputs_embeds")
+        use_cache = use_cache if use_cache is not None else getattr(self.config, "use_cache", True)
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        tensor = input_ids if input_ids is not None else inputs_embeds
+        batch_size, seq_len = tensor.shape[:2]
+        if not seq_len:
+            raise ValueError("Input must contain at least one token")
+        device = tensor.device
+        past_len = past_key_values.get_seq_length() if past_key_values is not None and use_cache else 0
+        if cache_position is None:
+            cache_position = torch.arange(past_len, past_len + seq_len, device=device)
+        else:
+            cache_position = cache_position.to(device=device, dtype=torch.long)
+        if cache_position.ndim != 1 or cache_position.numel() != seq_len:
+            raise ValueError("cache_position must identify each current token's physical cache slot")
+        if position_ids is None:
+            if attention_mask is not None and attention_mask.ndim == 2:
+                position_ids = attention_mask.long().cumsum(-1) - 1
+                position_ids.masked_fill_(attention_mask == 0, 0)
+                position_ids = position_ids[:, -seq_len:].to(device)
+            else:
+                position_ids = cache_position[None]
+        else:
+            position_ids = position_ids.to(device=device, dtype=torch.long)
+        if position_ids.shape[-1] != seq_len:
+            raise ValueError("position_ids must cover the current input tokens")
+
+        key_length = past_len + seq_len
+        if use_cache and past_key_values is not None and hasattr(past_key_values, "get_max_cache_shape"):
+            capacity = past_key_values.get_max_cache_shape()
+            if capacity is not None and capacity > 0:
+                key_length = capacity
+        # No explicit mask is needed for unpadded prefill or single-token dynamic
+        # decode. Chunked prefill needs bottom-right causal alignment; a full
+        # static cache additionally needs to hide all unfilled slots.
+        needs_mask = attention_mask is not None or key_length != past_len + seq_len or (past_len > 0 and seq_len > 1)
+        causal_mask = _causal_mask(attention_mask, cache_position, key_length, batch_size) if needs_mask else None
+        hidden_states, past_key_values = self.model(
+            input_ids=input_ids, position_ids=position_ids,
+            past_key_values=past_key_values, use_cache=use_cache,
+            attention_mask=causal_mask, inputs_embeds=inputs_embeds,
+            cache_position=cache_position,
+        )
+
+        if isinstance(logits_to_keep, int):
+            if logits_to_keep < 0:
+                raise ValueError("logits_to_keep must be nonnegative")
+            selected = hidden_states[:, -logits_to_keep:, :] if logits_to_keep else hidden_states
+        else:
+            selected = hidden_states[:, logits_to_keep.to(device), :]
+        if labels is not None and selected.shape[1] != hidden_states.shape[1]:
+            raise ValueError("Loss computation requires logits_to_keep=0")
+        logits = self.lm_head(selected)
+
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss = F.cross_entropy(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
+
+        if not return_dict:
+            output = (logits, past_key_values) if use_cache else (logits,)
+            return ((loss,) + output) if loss is not None else output
+        return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=past_key_values if use_cache else None)
+
+    def prepare_inputs_for_generation(
+        self, input_ids, past_key_values=None, attention_mask=None,
+        inputs_embeds=None, cache_position=None, position_ids=None, **kwargs,
+    ):
+        """Keep physical cache slots separate from padding-aware RoPE positions."""
+        past_len = past_key_values.get_seq_length() if past_key_values is not None else 0
+        if cache_position is None:
+            total = inputs_embeds.shape[1] if inputs_embeds is not None and past_len == 0 else input_ids.shape[1]
+            count = max(total - past_len, 1) if past_len else total
+            cache_position = torch.arange(past_len, past_len + count, device=input_ids.device)
+        count = cache_position.numel()
+        use_embeds = inputs_embeds is not None and past_len == 0
+        if use_embeds:
+            current_ids, current_embeds = None, inputs_embeds[:, -count:]
+        else:
+            current_ids, current_embeds = input_ids[:, -count:].contiguous(), None
+        if position_ids is None and attention_mask is not None and attention_mask.ndim == 2:
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 0)
+        if position_ids is not None:
+            position_ids = position_ids[:, -count:].contiguous()
+        return {
+            "input_ids": current_ids, "inputs_embeds": current_embeds,
+            "past_key_values": past_key_values, "attention_mask": attention_mask,
+            "position_ids": position_ids, "cache_position": cache_position,
+            "use_cache": kwargs.get("use_cache", True),
+            "logits_to_keep": kwargs.get("logits_to_keep", 1),
+        }
+
+    # ── NAR velocity (flow matching, no KV cache) ────────────────────
+
+    def _shift_t_value(self, t_value: float, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
+        t_sig = torch.sigmoid(torch.tensor(t_value, dtype=dtype, device=device))
+        shift = self.config.timestep_shift
+        return shift * t_sig / (1 + (shift - 1) * t_sig)
+
+    @torch.no_grad()
+    def nar_velocity(
+        self,
+        tokens: torch.LongTensor,
+        ar_mask: torch.BoolTensor,
+        nar_mask: torch.BoolTensor,
+        nar_content_mask: torch.BoolTensor,
+        x_t: torch.Tensor,
+        t_value: float,
+        nar_cond_end: int = 0,
+    ) -> torch.Tensor:
+        """Compute v_theta(x_t, t) — flow-matching velocity field.
+
+        Args:
+            tokens: [1, S] full sequence (AR + NAR tokens)
+            ar_mask: [1, S] True for AR positions
+            nar_mask: [1, S] True for NAR positions
+            nar_content_mask: [1, S] True for actual latent positions (not LATENT_START/END)
+            x_t: [T_lat, D] current ODE state
+            t_value: raw timestep (will be sigmoid-shifted)
+            nar_cond_end: if > 0, NAR only sees positions < nar_cond_end (text-only mode)
+        Returns:
+            v_pred: [T_lat, D] predicted velocity
+        """
+        device = tokens.device
+        dtype = next(self.parameters()).dtype
+        B, S = tokens.shape
+
+        # 1. Token embeddings
+        token_emb = self.model.embed_tokens(tokens)  # [B, S, H]
+
+        # 2. Build latent hidden for ALL NAR positions (START + content + END)
+        # Training injects vae2llm(x_t) + time_emb + pos_emb at ALL NAR positions,
+        # including LATENT_START (clean=0) and LATENT_END (clean=0).
+        # NAR position IDs via cumsum: START=0, content=[1..T_lat], END=T_lat+1.
+        t_shifted = self._shift_t_value(t_value, device, dtype)
+        T_lat = x_t.shape[0]
+
+        nar_indices = nar_mask[0].nonzero(as_tuple=True)[0]  # all NAR positions
+        content_indices = nar_content_mask[0].nonzero(as_tuple=True)[0]
+        N_nar = nar_indices.shape[0]  # START + T_lat + END
+
+        # Build x_t for all NAR positions: zeros for START/END, actual x_t for content
+        x_nar = torch.zeros(N_nar, x_t.shape[1], device=device, dtype=dtype)
+        x_nar[1:1 + T_lat] = x_t.to(dtype)  # content frames at positions [1, T_lat]
+
+        latent_hidden_nar = self.vae2llm(x_nar.unsqueeze(0))  # [1, N_nar, H]
+
+        # Timestep embedding (same t for all NAR positions)
+        time_emb = self.time_embedder(t_shifted.expand(N_nar)).unsqueeze(0)
+        latent_hidden_nar = latent_hidden_nar + time_emb
+
+        # Position embedding: cumsum-style [0, 1, 2, ..., N_nar-1]
+        pos_ids = torch.arange(N_nar, device=device).clamp(max=self.config.max_latent_frames - 1)
+        pos_emb = self.latent_pos_embed(pos_ids).unsqueeze(0)
+        latent_hidden_nar = latent_hidden_nar + pos_emb
+
+        # Inject at ALL NAR positions (matching training's torch.where)
+        token_emb[0, nar_indices] = latent_hidden_nar[0]
+
+        # 3. Build hybrid attention mask [B, 1, S, S]
+        # AR→AR: causal, NAR→AR: full, NAR→NAR: bidirectional, AR→NAR: blocked
+        ar_q = ar_mask.unsqueeze(2).float()   # [B, S, 1]
+        ar_k = ar_mask.unsqueeze(1).float()   # [B, 1, S]
+        nar_q = nar_mask.unsqueeze(2).float()
+        nar_k = nar_mask.unsqueeze(1).float()
+        causal = torch.tril(torch.ones(S, S, device=device))
+
+        if nar_cond_end > 0:
+            # Codec dropout: NAR only sees positions < nar_cond_end (text) + NAR
+            text_k = torch.zeros(1, 1, S, device=device)
+            text_k[0, 0, :nar_cond_end] = 1.0
+            mask = (ar_q * ar_k * causal) + (nar_q * text_k) + (nar_q * nar_k)
+        else:
+            mask = (ar_q * ar_k * causal) + (nar_q * ar_k) + (nar_q * nar_k)
+        # Convert to additive: 0 → attend, -inf → block
+        attn_mask = mask.unsqueeze(1)  # [B, 1, S, S]
+        attn_mask = attn_mask.masked_fill(attn_mask == 0, float("-inf")).masked_fill(attn_mask > 0, 0.0)
+
+        # 4. Position IDs + RoPE
+        position_ids = torch.arange(S, device=device).unsqueeze(0)
+
+        # 5. Forward through decoder (with MoT routing)
+        ar_mask_bt = ar_mask  # [B, S] bool for MoT routing
+        hidden_states, _ = self.model(
+            inputs_embeds=token_emb, position_ids=position_ids,
+            use_cache=False, attention_mask=attn_mask, ar_mask=ar_mask_bt,
+        )
+
+        # 6. NAR head at content positions
+        nar_pred = self.llm2vae(hidden_states)  # [B, S, D]
+        v_pred = nar_pred[0, content_indices]  # [T_lat, D]
+        return v_pred
+
+# Keep custom code + auto_map when a local user calls save_pretrained as well
+# as when the release builder creates a Hub repository.
+YuE2Config.register_for_auto_class()
+YuE2ForCausalLM.register_for_auto_class("AutoModelForCausalLM")

lm/weights_manifest.json

@@ -0,0 +1,9 @@
+{
+  "schema": 1,
+  "files": {
+    "model.safetensors": {
+      "bytes": 7261441640,
+      "sha256": "1d55c42c1a9875c34f5d736e15078449992b044e807ce2a138e6cf289a1e59e9"
+    }
+  }
+}

lm/yue2_generation_config.json

@@ -0,0 +1,24 @@
+{
+  "abc": {
+    "temperature": 0.7,
+    "top_p": 0.9,
+    "top_k": 30,
+    "repetition_penalty": 1.005,
+    "penalty_window": 100,
+    "min_tokens": 32,
+    "max_tokens": 4096
+  },
+  "semantic": {
+    "temperature": 1.0,
+    "top_p": 0.95,
+    "top_k": 100,
+    "repetition_penalty": 1.2,
+    "penalty_window": 50,
+    "min_tokens": 200,
+    "max_tokens": 9000
+  },
+  "ode_steps": 32,
+  "ode_method": "midpoint",
+  "context": 24576,
+  "version": "yue2-native-v1"
+}