initial commit

.gitignore

@@ -0,0 +1 @@
+__pycache__/

LICENSE

@@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

@@ -0,0 +1,65 @@
+---
+language:
+- en
+- zh
+license: apache-2.0
+base_model:
+  - m-a-p/YuE2-3B
+  - m-a-p/YuE2-Vae
+tags:
+- music-generation
+- yue2
+- modular-diffusers
+pipeline_tag: text-to-audio
+library_name: diffusers
+---
+# YuE2 Modular Diffusers
+
+<audio controls src="https://huggingface.co/datasets/OzzyGT/diffusers-examples/resolve/main/yue2/city_lights.mp3"></audio>
+
+*A 69-second song generated with the code below, seed `831001`.*
+
+Modular Diffusers blocks for [m-a-p/YuE2-3B](https://huggingface.co/m-a-p/YuE2-3B): full songs from a style description and lyrics, planned as an ABC score first, as 48 kHz stereo audio. The model code and weights are in [OzzyGT/YuE2-3B-Diffusers](https://huggingface.co/OzzyGT/YuE2-3B-Diffusers).
+
+Note: This model requires `tiktoken` and the example uses `soundfile`, install them with `pip install tiktoken soundfile`.
+
+## Licensing
+
+The code is Apache-2.0, adapted from [multimodal-art-projection/YuE](https://github.com/multimodal-art-projection/YuE). The weights are CC BY-NC 4.0 (non-commercial); see the [model repository](https://huggingface.co/OzzyGT/YuE2-3B-Diffusers).
+
+## Sample song
+
+The song above was generated with the following code:
+
+```python
+import soundfile as sf
+import torch
+from diffusers import ModularPipelineBlocks
+
+
+blocks = ModularPipelineBlocks.from_pretrained(  # load the blocks first to avoid warnings
+    "OzzyGT/YuE2-Modular",
+    trust_remote_code=True,
+    components_repo="OzzyGT/YuE2-3B-Diffusers",
+    trust_components_code=True,
+)
+pipe = blocks.init_pipeline()
+pipe.load_components(dtype={"transformer": torch.bfloat16, "vae": torch.float32})
+pipe.to("cuda")
+
+style = "English, warm piano pop, expressive female voice, acoustic piano, rounded bass and light drums, lyrical memorable melody, unhurried phrasing, 88 BPM"
+lyrics = """[Verse]
+Neon fades along the lane
+Footsteps keep the time of rain
+Fold the night and leave it here
+Morning has a sky to clear
+
+[Chorus]
+Let the day come into view
+Every road begins with you
+Hold a little room for light
+We will sing beyond the night"""
+
+state = pipe(style=style, lyrics=lyrics, seed=831001, use_cuda_graph=True)
+sf.write("song.flac", state.get("audios")[0].T.numpy(), state.get("sample_rate"), subtype="PCM_24")
+```

blocks.py

@@ -0,0 +1,516 @@
+# Adapted for diffusers from multimodal-art-projection/YuE at commit ef1936f2ee39fe8de486a0f47a481c95f8d4da87.
+# Licensed under Apache-2.0; see LICENSE.
+from dataclasses import replace
+
+import torch
+from diffusers import ModelMixin
+from diffusers.configuration_utils import FrozenDict
+from diffusers.modular_pipelines import (
+    ComponentSpec,
+    InputParam,
+    LoopSequentialPipelineBlocks,
+    ModularPipelineBlocks,
+    OutputParam,
+    SequentialPipelineBlocks,
+)
+from diffusers.utils.dynamic_modules_utils import get_class_from_dynamic_module
+from transformers import PreTrainedTokenizer
+
+from .cuda_graph import GraphAR
+from .guidance import YuE2SemanticGuider
+from .nar import YuE2PrefixKVCache, solve_midpoint, song_chunks
+from .protocol import (
+    ABC_SAMPLING,
+    CODEC_OFFSET,
+    CONTEXT,
+    SEMANTIC_SAMPLING,
+    SongRequest,
+    negative_prefix,
+    resolve_sampling,
+    token_prefixes,
+)
+from .sampling import generate_tokens
+
+
+def callback_inputs():
+    return [
+        InputParam(
+            "cancelled", type_hint=object, default=None, description="Optional callable returning whether to cancel."
+        ),
+        InputParam(
+            "on_token", type_hint=object, default=None, description="Optional callback receiving stage and token ID."
+        ),
+    ]
+
+
+def check_cancelled(callback):
+    if callback is not None and callback():
+        raise InterruptedError("Generation cancelled")
+
+
+class YuE2PrepareInputsStep(ModularPipelineBlocks):
+    model_name = "yue2"
+
+    @property
+    def description(self):
+        return "Validates the song request."
+
+    @property
+    def inputs(self):
+        return [
+            InputParam("style", type_hint=str, required=True, description="Music style description."),
+            InputParam("lyrics", type_hint=str, required=True, description="Lyrics with song section tags."),
+            InputParam("cot", type_hint=str, default="full", description="Score mode: full, melody, or off."),
+            InputParam(
+                "seed",
+                type_hint=int,
+                default=831001,
+                description="Request seed. Each stage reseeds from it, as the original release does.",
+            ),
+            InputParam("abc", type_hint=str, default=None, description="Optional supplied ABC score."),
+        ]
+
+    @property
+    def intermediate_outputs(self):
+        return [OutputParam("request", type_hint=SongRequest, description="Validated song request.")]
+
+    def __call__(self, components, state):
+        block_state = self.get_block_state(state)
+        block_state.request = SongRequest(
+            style=block_state.style,
+            lyrics=block_state.lyrics,
+            cot=block_state.cot,
+            seed=block_state.seed,
+            abc=block_state.abc,
+        )
+        self.set_block_state(state, block_state)
+        return components, state
+
+
+class YuE2PlanStep(ModularPipelineBlocks):
+    model_name = "yue2"
+    _requirements = {"tiktoken": ">=0.12.0"}
+
+    @property
+    def description(self):
+        return "Plans the song as an ABC score (or uses the supplied one) and builds the semantic-stage prompt."
+
+    @property
+    def expected_components(self):
+        return [ComponentSpec("transformer", ModelMixin), ComponentSpec("tokenizer", PreTrainedTokenizer)]
+
+    @property
+    def inputs(self):
+        return [
+            InputParam("request", type_hint=SongRequest, required=True, description="Validated song request."),
+            InputParam(
+                "abc_sampling", type_hint=dict, default=None, description="Overrides for the ABC sampling defaults."
+            ),
+            InputParam("use_cuda_graph", type_hint=bool, default=False, description="Decode tokens with CUDA graphs."),
+        ] + callback_inputs()
+
+    @property
+    def intermediate_outputs(self):
+        return [
+            OutputParam("score", type_hint=str, description="Generated or supplied ABC score; None with cot='off'."),
+            OutputParam("abc_ids", type_hint=list, description="Exact score token IDs."),
+            OutputParam("prefix", type_hint=list, description="Semantic-stage prompt token IDs."),
+            OutputParam("abc_truncated", type_hint=bool, description="Whether planning hit its token limit."),
+            OutputParam("abc_timing", type_hint=dict, description="Planning timings."),
+        ]
+
+    @torch.inference_mode()
+    def __call__(self, components, state):
+        block_state = self.get_block_state(state)
+        check_cancelled(block_state.cancelled)
+        request = block_state.request
+        score, abc_ids, timing, truncated = None, [], {}, False
+        if request.cot != "off":
+            if request.abc is not None:
+                score = request.abc
+                abc_ids = components.tokenizer.encode(score)
+            else:
+                abc_ids, timing, truncated = generate_tokens(
+                    components.transformer,
+                    token_prefixes(request, components.tokenizer),
+                    resolve_sampling(block_state.abc_sampling, ABC_SAMPLING),
+                    request.seed,
+                    "abc",
+                    components._execution_device,
+                    cancelled=block_state.cancelled,
+                    on_token=block_state.on_token,
+                    graph_decoder=GraphAR if block_state.use_cuda_graph else None,
+                )
+                score = components.tokenizer.decode(abc_ids)
+        block_state.score = score
+        block_state.abc_ids = abc_ids
+        block_state.prefix = token_prefixes(request, components.tokenizer, abc_ids)
+        block_state.abc_truncated = truncated
+        block_state.abc_timing = timing
+        self.set_block_state(state, block_state)
+        return components, state
+
+
+class YuE2SemanticStep(ModularPipelineBlocks):
+    model_name = "yue2"
+    _requirements = {"tiktoken": ">=0.12.0"}
+
+    @property
+    def description(self):
+        return "Generates the song's codec (semantic) tokens, with guidance from the `guider` component."
+
+    @property
+    def expected_components(self):
+        return [
+            ComponentSpec("transformer", ModelMixin),
+            ComponentSpec("tokenizer", PreTrainedTokenizer),
+            ComponentSpec(
+                "guider", YuE2SemanticGuider, config=FrozenDict({"scale": None}), default_creation_method="from_config"
+            ),
+        ]
+
+    @property
+    def inputs(self):
+        return [
+            InputParam("request", type_hint=SongRequest, required=True, description="Validated song request."),
+            InputParam("abc_ids", type_hint=list, required=True, description="Exact score token IDs."),
+            InputParam("prefix", type_hint=list, required=True, description="Semantic-stage prompt token IDs."),
+            InputParam(
+                "semantic_sampling",
+                type_hint=dict,
+                default=None,
+                description="Overrides for the semantic sampling defaults.",
+            ),
+            InputParam("use_cuda_graph", type_hint=bool, default=False, description="Decode tokens with CUDA graphs."),
+        ] + callback_inputs()
+
+    @property
+    def intermediate_outputs(self):
+        return [
+            OutputParam("semantic_tokens", type_hint=list, description="Codec token IDs, 25 per second of audio."),
+            OutputParam("semantic_truncated", type_hint=bool, description="Whether generation hit its token limit."),
+            OutputParam("semantic_timing", type_hint=dict, description="Semantic-stage timings."),
+        ]
+
+    @torch.inference_mode()
+    def __call__(self, components, state):
+        block_state = self.get_block_state(state)
+        check_cancelled(block_state.cancelled)
+        request = block_state.request
+        if token_prefixes(request, components.tokenizer, block_state.abc_ids) != block_state.prefix:
+            raise ValueError("The semantic prompt disagrees with the request and score IDs")
+        guider = components.guider
+        negative = None
+        if guider.scale_for(request.cot) != 1:
+            negative = negative_prefix(request, components.tokenizer, block_state.abc_ids)
+        tokens, timing, truncated = generate_tokens(
+            components.transformer,
+            block_state.prefix,
+            resolve_sampling(block_state.semantic_sampling, SEMANTIC_SAMPLING),
+            request.seed,
+            "semantic",
+            components._execution_device,
+            negative=negative,
+            combine_logits=lambda conditional, unconditional: guider.combine(conditional, unconditional, request.cot),
+            legacy_off=request.cot == "off",
+            cancelled=block_state.cancelled,
+            on_token=block_state.on_token,
+            graph_decoder=GraphAR if block_state.use_cuda_graph else None,
+        )
+        block_state.semantic_tokens = [int(t) - CODEC_OFFSET for t in tokens]
+        block_state.semantic_truncated = truncated
+        block_state.semantic_timing = timing
+        self.set_block_state(state, block_state)
+        return components, state
+
+
+class YuE2PrepareChunksStep(ModularPipelineBlocks):
+    model_name = "yue2"
+
+    @property
+    def description(self):
+        return "Splits the song into acoustic chunks and draws their seeded noise."
+
+    @property
+    def inputs(self):
+        return [
+            InputParam("request", type_hint=SongRequest, required=True, description="Validated song request."),
+            InputParam("prefix", type_hint=list, required=True, description="Semantic-stage prompt token IDs."),
+            InputParam("semantic_tokens", type_hint=list, required=True, description="Codec token IDs."),
+            InputParam(
+                "acoustic_context",
+                type_hint=int,
+                default=CONTEXT,
+                description="Context limit that sets the acoustic chunk size.",
+            ),
+        ]
+
+    @property
+    def intermediate_outputs(self):
+        return [OutputParam("chunks", type_hint=list, description="Acoustic chunk tokens and CPU FP32 noise.")]
+
+    def __call__(self, components, state):
+        block_state = self.get_block_state(state)
+        block_state.chunks = song_chunks(
+            block_state.prefix,
+            block_state.semantic_tokens,
+            block_state.request.seed,
+            block_state.acoustic_context,
+        )
+        self.set_block_state(state, block_state)
+        return components, state
+
+
+class YuE2ChunkConditionStep(ModularPipelineBlocks):
+    model_name = "yue2"
+
+    @property
+    def description(self):
+        return "Runs the chunk's tokens through the token stream once and keeps their keys and values."
+
+    @property
+    def expected_components(self):
+        return [ComponentSpec("transformer", ModelMixin)]
+
+    @property
+    def inputs(self):
+        return [
+            InputParam("chunks", type_hint=list, required=True, description="Prepared acoustic chunks."),
+            InputParam("cancelled", type_hint=object, default=None, description="Cancellation callable."),
+        ]
+
+    @torch.inference_mode()
+    def __call__(self, components, block_state, k):
+        check_cancelled(block_state.cancelled)
+        block_state.chunk = block_state.chunks[k]
+        block_state.kv_cache = YuE2PrefixKVCache()
+        ids = torch.tensor([block_state.chunk.ar_tokens], device=components._execution_device)
+        components.transformer(ids, kv_cache=block_state.kv_cache, logits_to_keep=1)
+        return components, block_state
+
+
+class YuE2ChunkSynthesizeStep(ModularPipelineBlocks):
+    model_name = "yue2"
+
+    @property
+    def description(self):
+        return "Solves the chunk's flow-matching ODE with the midpoint method."
+
+    @property
+    def expected_components(self):
+        return [ComponentSpec("transformer", ModelMixin)]
+
+    @property
+    def inputs(self):
+        return [
+            InputParam("chunks", type_hint=list, required=True, description="Prepared acoustic chunks."),
+            InputParam("ode_steps", type_hint=int, default=32, description="Midpoint steps per chunk."),
+            InputParam("cancelled", type_hint=object, default=None, description="Cancellation callable."),
+            InputParam(
+                "on_progress", type_hint=object, default=None, description="Callback(stage, completed, total)."
+            ),
+        ]
+
+    @torch.inference_mode()
+    def __call__(self, components, block_state, k):
+        report = None
+        if block_state.on_progress is not None:
+
+            def report(done, total):
+                block_state.on_progress("synthesis", k * total + done, len(block_state.chunks) * total)
+
+        block_state.chunk_latents = solve_midpoint(
+            components.transformer,
+            block_state.kv_cache,
+            block_state.chunk.noise,
+            components._execution_device,
+            block_state.ode_steps,
+            block_state.cancelled,
+            report,
+        )
+        return components, block_state
+
+
+class YuE2CollectChunkStep(ModularPipelineBlocks):
+    model_name = "yue2"
+
+    @property
+    def description(self):
+        return "Appends the chunk's latents."
+
+    def __call__(self, components, block_state, k):
+        block_state.latent_chunks.append(block_state.chunk_latents)
+        return components, block_state
+
+
+class YuE2AcousticChunkLoop(LoopSequentialPipelineBlocks):
+    model_name = "yue2"
+    block_classes = [YuE2ChunkConditionStep, YuE2ChunkSynthesizeStep, YuE2CollectChunkStep]
+    block_names = ["condition", "synthesize", "collect"]
+
+    @property
+    def description(self):
+        return "Synthesizes acoustic latents chunk by chunk."
+
+    @property
+    def loop_inputs(self):
+        return [InputParam("chunks", type_hint=list, required=True, description="Prepared acoustic chunks.")]
+
+    @property
+    def loop_intermediate_outputs(self):
+        return [OutputParam("latents", type_hint=torch.Tensor, description="Acoustic latents [frames, 64], CPU FP32.")]
+
+    @torch.inference_mode()
+    def __call__(self, components, state):
+        block_state = self.get_block_state(state)
+        if not block_state.chunks:
+            raise ValueError("At least one acoustic chunk is required")
+        block_state.latent_chunks = []
+        for k in range(len(block_state.chunks)):
+            try:
+                components, block_state = self.loop_step(components, block_state, k=k)
+            finally:
+                block_state.kv_cache = None
+        block_state.latents = torch.cat(block_state.latent_chunks, dim=0)
+        self.set_block_state(state, block_state)
+        return components, state
+
+
+class YuE2SynthesizeStep(SequentialPipelineBlocks):
+    model_name = "yue2"
+    block_classes = [YuE2PrepareChunksStep, YuE2AcousticChunkLoop]
+    block_names = ["prepare_chunks", "chunk_loop"]
+
+    @property
+    def description(self):
+        return "Turns codec tokens into acoustic latents."
+
+
+class YuE2DecodeStep(ModularPipelineBlocks):
+    model_name = "yue2"
+
+    @property
+    def description(self):
+        return (
+            "Decodes latents to 48 kHz stereo audio in overlapping tiles. The VAE runs in FP32 with TF32 and cuDNN "
+            "autotuning disabled, as in the original release."
+        )
+
+    @property
+    def expected_components(self):
+        return [ComponentSpec("vae", ModelMixin)]
+
+    @property
+    def inputs(self):
+        return [
+            InputParam(
+                "latents",
+                type_hint=torch.Tensor,
+                required=True,
+                description="Acoustic latents, [frames, 64] or [1, 64, frames].",
+            ),
+            InputParam("vae_tile_frames", type_hint=int, default=1024, description="Latent frames decoded per tile."),
+            InputParam(
+                "vae_tile_overlap_frames",
+                type_hint=int,
+                default=16,
+                description="Context frames decoded on each side of a tile and cropped away.",
+            ),
+            InputParam("cancelled", type_hint=object, default=None, description="Cancellation callable."),
+            InputParam(
+                "on_progress", type_hint=object, default=None, description="Callback(stage, completed, total)."
+            ),
+        ]
+
+    @property
+    def intermediate_outputs(self):
+        return [
+            OutputParam("audios", type_hint=torch.Tensor, description="Stereo audio [1, 2, samples], CPU FP32."),
+            OutputParam("sample_rate", type_hint=int, description="Sample rate in Hz."),
+        ]
+
+    @torch.inference_mode()
+    def __call__(self, components, state):
+        block_state = self.get_block_state(state)
+        check_cancelled(block_state.cancelled)
+        vae = components.vae
+        device = components._execution_device
+        latents = torch.as_tensor(block_state.latents, dtype=torch.float32)
+        if latents.ndim == 2:
+            latents = latents.T.unsqueeze(0)
+        frames = latents.shape[-1]
+        tile, overlap = block_state.vae_tile_frames, block_state.vae_tile_overlap_frames
+        if tile < 1 or overlap < 0:
+            raise ValueError("vae_tile_frames must be positive and vae_tile_overlap_frames nonnegative")
+        pieces = []
+        starts = range(0, frames, tile)
+        with torch.backends.cudnn.flags(enabled=True, benchmark=False, deterministic=True, allow_tf32=False):
+            for index, start in enumerate(starts):
+                check_cancelled(block_state.cancelled)
+                end = min(frames, start + tile)
+                left, right = max(0, start - overlap), min(frames, end + overlap)
+                decoded = vae.decode(latents[..., left:right].to(device=device, dtype=torch.float32)).sample
+                crop = (start - left) * vae.hop_length
+                if end < frames:
+                    decoded = decoded[..., crop : crop + (end - start) * vae.hop_length]
+                else:
+                    decoded = decoded[..., crop:]
+                pieces.append(decoded.cpu())
+                if block_state.on_progress is not None:
+                    block_state.on_progress("decode", index + 1, len(starts))
+        audio = torch.cat(pieces, dim=-1)
+        if not torch.isfinite(audio).all():
+            raise FloatingPointError("Decoded audio is non-finite")
+        block_state.audios = audio.clamp(-1, 1)
+        block_state.sample_rate = vae.config.sampling_rate
+        self.set_block_state(state, block_state)
+        return components, state
+
+
+class YuE2Blocks(SequentialPipelineBlocks):
+    model_name = "yue2"
+    block_names = ["prepare", "plan", "semantic", "synthesize", "decode"]
+    block_classes = [YuE2PrepareInputsStep, YuE2PlanStep, YuE2SemanticStep, YuE2SynthesizeStep, YuE2DecodeStep]
+
+    def __init__(self, components_repo=None, components_revision=None, trust_components_code=False):
+        super().__init__()
+        self.components_repo = components_repo
+        self.components_revision = components_revision
+        self.component_types = {}
+        if components_repo is not None:
+            if not trust_components_code:
+                raise ValueError("Set trust_components_code=True to load code from the model repository")
+            for name, class_name in (
+                ("transformer", "YuE2TransformerModel"),
+                ("vae", "YuE2VAE"),
+                ("tokenizer", "YuE2Tokenizer"),
+            ):
+                self.component_types[name] = get_class_from_dynamic_module(
+                    components_repo,
+                    module_file=f"{name}.py",
+                    class_name=class_name,
+                    revision=components_revision,
+                    trust_remote_code=True,
+                )
+
+    @property
+    def expected_components(self):
+        specs = super().expected_components
+        if self.components_repo is None:
+            return specs
+        return [
+            replace(
+                spec,
+                type_hint=self.component_types[spec.name],
+                pretrained_model_name_or_path=str(self.components_repo),
+                subfolder=spec.name,
+                revision=self.components_revision,
+            )
+            if spec.default_creation_method == "from_pretrained"
+            else spec
+            for spec in specs
+        ]
+
+    @property
+    def description(self):
+        return "YuE2 score planning, semantic generation, midpoint synthesis, and stereo decoding."

cuda_graph.py

@@ -0,0 +1,195 @@
+# Adapted for diffusers from multimodal-art-projection/YuE at commit ef1936f2ee39fe8de486a0f47a481c95f8d4da87.
+# Licensed under Apache-2.0; see LICENSE.
+"""CUDA-graph token decoding for one request or its two guidance branches.
+
+The graph replays the transformer's token stream for one new token per branch. The caller combines the branch logits,
+samples once and passes that token to `step`. Each branch keeps its own RoPE positions and cache slots.
+"""
+
+from __future__ import annotations
+
+import torch
+
+
+def apply_rotary_emb(hidden_states, cos, sin):
+    # Same arithmetic as `apply_rotary_emb` in the model repository's transformer.py.
+    half = hidden_states.shape[-1] // 2
+    x1, x2 = hidden_states[..., :half], hidden_states[..., half:]
+    cos, sin = cos.to(hidden_states.dtype), sin.to(hidden_states.dtype)
+    return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1)
+
+
+class _BranchPrefillCache:
+    """Writes one branch's prompt keys and values into the graph's sequence-major buffers."""
+
+    def __init__(self, keys, values, branch):
+        self.keys = [buffer[branch : branch + 1] for buffer in keys]
+        self.values = [buffer[branch : branch + 1] for buffer in values]
+        self.seen_tokens = 0
+
+    def get_seq_length(self):
+        return self.seen_tokens
+
+    def update(self, key, value, layer_idx):
+        end = self.seen_tokens + key.shape[1]
+        if end > self.keys[layer_idx].shape[1]:
+            raise ValueError("Prefix exceeds preallocated KV capacity")
+        self.keys[layer_idx][:, self.seen_tokens : end].copy_(key)
+        self.values[layer_idx][:, self.seen_tokens : end].copy_(value)
+        if layer_idx == len(self.keys) - 1:
+            self.seen_tokens = end
+        return self.keys[layer_idx][:, :end], self.values[layer_idx][:, :end]
+
+
+class GraphAR:
+    """Fixed-capacity decoding. `prefill()` returns the first logits; up to `max_tokens - 1` `step(token)` calls follow.
+
+    Returned logits share output storage and stay valid until the next step.
+    """
+
+    def __init__(self, transformer, prefixes, max_tokens, device):
+        if transformer.training:
+            raise ValueError("GraphAR requires transformer.eval()")
+        if not 1 <= len(prefixes) <= 2:
+            raise ValueError("GraphAR supports one request or exactly two guidance branches")
+        if device.type != "cuda":
+            raise ValueError("CUDA graphs require a CUDA device")
+        try:
+            from diffusers.hooks.group_offloading import _get_group_onload_device
+
+            _get_group_onload_device(transformer)
+        except ValueError:
+            pass
+        else:
+            raise ValueError(
+                "CUDA graphs cannot capture a group-offloaded transformer: streamed weights change address between "
+                "steps. Pass use_cuda_graph=False, or keep the transformer resident on the GPU."
+            )
+        config = transformer.config
+        prefixes = [[int(token) for token in prefix] for prefix in prefixes]
+        for prefix in prefixes:
+            if not prefix or not all(0 <= token < config.vocab_size for token in prefix):
+                raise ValueError("Prefixes require valid token IDs")
+            if len(prefix) + max_tokens > config.max_position_embeddings:
+                raise ValueError("Prefix plus generation budget exceeds model context; no length was shortened")
+        self.transformer, self.device, self.dtype = transformer, device, transformer.dtype
+        self.prefixes, self.max_tokens, self.branches = prefixes, int(max_tokens), len(prefixes)
+        self.capacity = max(map(len, prefixes)) + self.max_tokens
+        self.graph = self.output = None
+        # PyTorch's variable-length FlashAttention takes per-branch key lengths on the GPU, so no mask is built.
+        flash = (
+            self.dtype in {torch.bfloat16, torch.float16}
+            and config.attention_head_dim % 8 == 0
+            and config.attention_head_dim <= 256
+            and "seqused_k" in str(torch.ops.aten._flash_attention_forward.default._schema)
+        )
+        if not flash:
+            raise ValueError(
+                "CUDA graphs need PyTorch's variable-length FlashAttention with a BF16/FP16 transformer; "
+                "pass use_cuda_graph=False."
+            )
+        self.ready, self.closed, self.steps = False, False, 0
+        # Sequence-major buffers make the packed FlashAttention view contiguous without copying the cache each step.
+        shape = (self.branches, self.capacity, config.num_key_value_heads, config.attention_head_dim)
+        self.keys = [torch.zeros(shape, device=device, dtype=self.dtype) for _ in range(config.num_layers)]
+        self.values = [torch.zeros(shape, device=device, dtype=self.dtype) for _ in range(config.num_layers)]
+        self.positions = torch.tensor([len(prefix) for prefix in prefixes], dtype=torch.long, device=device)
+        self.initial_positions = self.positions.clone()
+        self.cu_q = torch.arange(self.branches + 1, dtype=torch.int32, device=device)
+        self.cu_k = self.cu_q * self.capacity
+        self.tokens = torch.tensor([[prefix[-1]] for prefix in prefixes], dtype=torch.long, device=device)
+
+    @torch.inference_mode()
+    def _decode(self):
+        transformer = self.transformer
+        config = transformer.config
+        heads, kv_heads, head_dim = config.num_attention_heads, config.num_key_value_heads, config.attention_head_dim
+        cos, sin = transformer.rotary_emb(self.positions[:, None])
+        cos, sin = cos.unsqueeze(2), sin.unsqueeze(2)
+        hidden_states = transformer.embed_tokens(self.tokens)
+        used_lengths = (self.positions + 1).to(torch.int32)
+        slots = self.positions[:, None, None, None].expand(self.branches, 1, kv_heads, head_dim)
+        for block, keys, values in zip(transformer.transformer_blocks, self.keys, self.values):
+            attn = block.attn
+            normalized = block.norm1(hidden_states)
+            query = attn.norm_q(attn.to_q(normalized).view(self.branches, 1, heads, head_dim))
+            key = attn.norm_k(attn.to_k(normalized).view(self.branches, 1, kv_heads, head_dim))
+            value = attn.to_v(normalized).view(self.branches, 1, kv_heads, head_dim)
+            query = apply_rotary_emb(query, cos, sin)
+            key = apply_rotary_emb(key, cos, sin)
+            keys.scatter_(1, slots, key)
+            values.scatter_(1, slots, value)
+            # Every slot is allocated, but only each branch's filled prefix and current token are visible.
+            # The packed variable-length entry point respects seqused_k; the 4D fixed-batch one ignores it.
+            attn_output = torch.ops.aten._flash_attention_forward(
+                query[:, 0],
+                keys.view(-1, kv_heads, head_dim),
+                values.view(-1, kv_heads, head_dim),
+                self.cu_q,
+                self.cu_k,
+                1,
+                self.capacity,
+                0.0,
+                False,
+                False,
+                seqused_k=used_lengths,
+            )[0][:, None]
+            hidden_states = hidden_states + attn.to_out[0](attn_output.reshape(self.branches, 1, -1))
+            hidden_states = hidden_states + block.ff(block.norm2(hidden_states))
+        output = transformer.lm_head(transformer.norm_out(hidden_states))[:, 0]
+        self.positions.add_(1)
+        return output
+
+    @torch.inference_mode()
+    def _capture(self):
+        with torch.cuda.device(self.device):
+            current = torch.cuda.current_stream(self.device)
+            warmup = torch.cuda.Stream(device=self.device)
+            warmup.wait_stream(current)
+            with torch.cuda.stream(warmup):
+                for _ in range(3):
+                    self.positions.copy_(self.initial_positions)
+                    self._decode()
+                self.positions.copy_(self.initial_positions)
+            current.wait_stream(warmup)
+            torch.cuda.synchronize(self.device)
+            self.graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(self.graph):
+                self.output = self._decode()
+            # Warmup and capture wrote one future slot; the first real step overwrites it before it becomes visible.
+            self.positions.copy_(self.initial_positions)
+
+    @torch.inference_mode()
+    def prefill(self):
+        if self.closed or self.ready:
+            raise RuntimeError("prefill must be called exactly once on an open GraphAR")
+        logits = []
+        for branch, prefix in enumerate(self.prefixes):
+            cache = _BranchPrefillCache(self.keys, self.values, branch)
+            ids = torch.tensor([prefix], dtype=torch.long, device=self.device)
+            logits.append(self.transformer(ids, kv_cache=cache, logits_to_keep=1).logits[:, -1])
+        if any(param.device != self.device for param in self.transformer.parameters()):
+            raise ValueError("CUDA graphs need the whole transformer resident on the GPU; pass use_cuda_graph=False.")
+        if self.max_tokens > 1:
+            self._capture()
+        self.ready = True
+        return torch.cat(logits, dim=0)
+
+    @torch.inference_mode()
+    def step(self, token):
+        if self.closed or not self.ready:
+            raise RuntimeError("Call prefill before step and do not use a closed GraphAR")
+        if self.steps >= self.max_tokens - 1:
+            raise ValueError("Requested generation budget is exhausted")
+        self.tokens.copy_(token.reshape(1, 1).expand(self.branches, 1))
+        self.graph.replay()
+        self.steps += 1
+        return self.output
+
+    def close(self):
+        self.graph = self.output = None
+        self.keys.clear()
+        self.values.clear()
+        self.tokens = self.positions = self.initial_positions = None
+        self.cu_q = self.cu_k = None
+        self.closed = True

guidance.py

@@ -0,0 +1,19 @@
+import math
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+
+
+class YuE2SemanticGuider(ConfigMixin):
+    config_name = "guider_config.json"
+
+    @register_to_config
+    def __init__(self, scale=None):
+        if scale is not None and (not math.isfinite(scale) or not 0 <= scale <= 20):
+            raise ValueError("Semantic guidance must be finite and in [0,20]")
+
+    def scale_for(self, cot):
+        return (1.01 if cot == "off" else 1.0) if self.config.scale is None else self.config.scale
+
+    def combine(self, conditional, unconditional, cot):
+        scale = self.scale_for(cot)
+        return conditional if scale == 1.0 else unconditional + scale * (conditional - unconditional)

modular_config.json

@@ -0,0 +1,10 @@
+{
+  "_class_name": "YuE2Blocks",
+  "_diffusers_version": "0.41.0.dev0",
+  "auto_map": {
+    "ModularPipelineBlocks": "blocks.YuE2Blocks"
+  },
+  "requirements": {
+    "tiktoken": ">=0.12.0"
+  }
+}

modular_model_index.json

@@ -0,0 +1,4 @@
+{
+  "_class_name": "ModularPipeline",
+  "_blocks_class_name": "YuE2Blocks"
+}

nar.py

@@ -0,0 +1,96 @@
+# Adapted for diffusers from multimodal-art-projection/YuE at commit ef1936f2ee39fe8de486a0f47a481c95f8d4da87.
+# Licensed under Apache-2.0; see LICENSE.
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from numbers import Integral
+
+import torch
+
+from .protocol import CODEC_OFFSET, CODEC_SIZE, CONTEXT, MUSIC_END, chunk_ranges
+
+
+@dataclass
+class Chunk:
+    ar_tokens: list[int]
+    noise: torch.Tensor
+
+
+class YuE2PrefixKVCache:
+    """Holds a chunk's token-prefix keys and values; later acoustic calls attend to the prefix without extending it."""
+
+    def __init__(self):
+        self.keys, self.values = [], []
+
+    def get_seq_length(self):
+        return self.keys[0].shape[1] if self.keys else 0
+
+    def update(self, key, value, layer_idx):
+        if layer_idx == len(self.keys):
+            self.keys.append(key)
+            self.values.append(value)
+            return key, value
+        return torch.cat((self.keys[layer_idx], key), dim=1), torch.cat((self.values[layer_idx], value), dim=1)
+
+
+def _integers(values, name):
+    result = list(values)
+    if not result or any(isinstance(v, bool) or not isinstance(v, Integral) for v in result):
+        raise ValueError(f"{name} must be a nonempty sequence of integer token IDs")
+    return [int(v) for v in result]
+
+
+def song_chunks(prefix, codec, seed, context=CONTEXT):
+    """Draw the whole song's CPU FP32 noise once, then split it at the release's chunk boundaries."""
+    prefix = _integers(prefix, "prefix")
+    codec = _integers(codec, "codec")
+    if min(prefix) < 0 or min(codec) < 0 or max(codec) >= CODEC_SIZE:
+        raise ValueError("Token IDs are outside their allowed vocabulary")
+    if isinstance(context, bool) or not isinstance(context, Integral) or not 1 <= context <= CONTEXT:
+        raise ValueError(f"context must be an integer in 1..{CONTEXT}")
+    ranges = chunk_ranges(len(codec), len(prefix), int(context))
+    generator = torch.Generator(device="cpu").manual_seed(int(seed))
+    noise = torch.randn((len(codec), 64), dtype=torch.float32, device="cpu", generator=generator)
+    return [
+        Chunk(prefix + [value + CODEC_OFFSET for value in codec[a:b]] + [MUSIC_END], noise[a:b]) for a, b in ranges
+    ]
+
+
+@torch.inference_mode()
+def solve_midpoint(
+    transformer,
+    kv_cache,
+    noise,
+    device,
+    steps=32,
+    cancelled: Callable[[], bool] | None = None,
+    on_progress: Callable[[int, int], None] | None = None,
+):
+    """Integrate the flow from t=1 (noise) to t=0 with the midpoint method; returns CPU FP32 [frames, 64] latents."""
+    if isinstance(steps, bool) or not isinstance(steps, Integral) or steps < 1:
+        raise ValueError("steps must be a positive integer")
+    if not torch.isfinite(noise).all():
+        raise ValueError("Acoustic noise contains non-finite values")
+
+    def velocity(state, t):
+        # The model takes time in logit space; logit(1) is clamped to 20.
+        raw = torch.logit(torch.tensor(t, dtype=torch.float64, device="cpu")).clamp(-20, 20).item()
+        return transformer(latents=state[None], timestep=raw, kv_cache=kv_cache).sample[0]
+
+    state = noise.to(device=device, dtype=transformer.dtype)
+    dt = 1.0 / steps
+    for step in range(steps):
+        if cancelled is not None and cancelled():
+            raise InterruptedError("Cancelled during acoustic flow matching")
+        t = 1.0 - step * dt
+        mid = state - velocity(state, t) * (dt / 2)
+        if cancelled is not None and cancelled():
+            raise InterruptedError("Cancelled during acoustic flow matching")
+        state = state - velocity(mid, t - dt / 2) * dt
+        if on_progress is not None:
+            on_progress(step + 1, int(steps))
+    result = state.float().cpu()
+    if not torch.isfinite(result).all():
+        raise FloatingPointError("Acoustic flow matching produced non-finite latents")
+    return result

protocol.py

@@ -0,0 +1,108 @@
+# Adapted for diffusers from multimodal-art-projection/YuE at commit ef1936f2ee39fe8de486a0f47a481c95f8d4da87.
+# Licensed under Apache-2.0; see LICENSE.
+from __future__ import annotations
+
+import math
+from dataclasses import asdict, dataclass
+
+EOD = 151643
+ABC_START, ABC_END = (151847, 151848)
+MUSIC_START, MUSIC_END = (151851, 151852)
+CODEC_OFFSET, CODEC_SIZE = (151853, 32768)
+CONTEXT = 24576
+INSTRUCTIONS = {
+    "off": "Generate music with codec tokens from the given conditions.",
+    "melody": "Generate a melody-only ABC transcription without chord symbols, then generate music with codec tokens from the given conditions.",
+    "full": "Generate a chord-annotated ABC transcription, then generate music with codec tokens from the given conditions.",
+}
+
+
+@dataclass(frozen=True)
+class Sampling:
+    temperature: float = 1.0
+    top_p: float = 0.95
+    top_k: int = 100
+    repetition_penalty: float = 1.2
+    penalty_window: int = 50
+    min_tokens: int = 200
+    max_tokens: int = 9000
+
+    def __post_init__(self):
+        if any(type(x) is not int for x in (self.top_k, self.penalty_window, self.min_tokens, self.max_tokens)):
+            raise ValueError("Sampling counts must be integers")
+        if not all(math.isfinite(x) for x in (self.temperature, self.top_p, self.repetition_penalty)):
+            raise ValueError("Sampling numbers must be finite")
+        if not 0 <= self.temperature <= 5 or not 0 < self.top_p <= 1 or self.top_k < 1:
+            raise ValueError("Invalid sampling temperature/top_p/top_k")
+        if self.repetition_penalty <= 0 or not 1 <= self.penalty_window <= 100:
+            raise ValueError("Invalid repetition penalty/window")
+        if not 0 <= self.min_tokens <= self.max_tokens or self.max_tokens < 1:
+            raise ValueError("Require 0 <= min_tokens <= max_tokens")
+
+
+ABC_SAMPLING = Sampling(0.7, 0.9, 30, 1.005, 100, 32, 4096)
+SEMANTIC_SAMPLING = Sampling()
+
+
+def resolve_sampling(overrides, default):
+    if overrides is None:
+        return default
+    if not isinstance(overrides, dict):
+        raise TypeError("Sampling overrides must be a dictionary")
+    return Sampling(**{**asdict(default), **overrides})
+
+
+@dataclass(frozen=True)
+class SongRequest:
+    style: str
+    lyrics: str
+    cot: str = "full"
+    seed: int = 831001
+    abc: str | None = None
+
+    def __post_init__(self):
+        if self.cot not in INSTRUCTIONS:
+            raise ValueError("cot must be off, melody or full")
+        if not isinstance(self.style, str) or not isinstance(self.lyrics, str):
+            raise TypeError("style and lyrics must be strings")
+        if type(self.seed) is not int or not 0 <= self.seed < 2**63:
+            raise ValueError("seed must be an integer in [0, 2**63)")
+        if self.abc is not None and (self.cot == "off" or not isinstance(self.abc, str) or not self.abc.strip()):
+            raise ValueError("External ABC requires nonempty text and cot=melody/full")
+
+    def text(self):
+        return f"{INSTRUCTIONS[self.cot]}\n[Tags]\n{self.style}\n[Lyrics]\n{self.lyrics}\n"
+
+
+def _check_abc_ids(abc_ids):
+    abc_ids = list(abc_ids)
+    if any(type(token) is not int or not 0 <= token < EOD for token in abc_ids):
+        raise ValueError("ABC IDs must remain inside the ordinary text vocabulary")
+    return abc_ids
+
+
+def token_prefixes(request, tokenizer, abc_ids=None):
+    base = [EOD] + tokenizer.encode(request.text())
+    if request.cot == "off":
+        return base + [ABC_START, ABC_END, MUSIC_START]
+    if abc_ids is None:
+        if request.abc is None:
+            return base + [ABC_START]
+        abc_ids = tokenizer.encode(request.abc)
+    return base + [ABC_START] + _check_abc_ids(abc_ids) + [ABC_END, MUSIC_START]
+
+
+def negative_prefix(request, tokenizer, abc_ids=None):
+    base = [EOD] + tokenizer.encode(INSTRUCTIONS[request.cot])
+    if request.cot == "off":
+        return base + [MUSIC_START]
+    if abc_ids is None:
+        raise ValueError("Symbolic CFG must retain the exact positive-branch ABC IDs")
+    return base + [ABC_START] + _check_abc_ids(abc_ids) + [ABC_END, MUSIC_START]
+
+
+def chunk_ranges(frames, prefix_tokens, context=CONTEXT):
+    size = min((context - prefix_tokens - 3) // 2, CONTEXT)
+    if frames < 1 or size < 1:
+        raise ValueError("Empty codec or prefix leaves no acoustic context")
+    return [(a, min(a + size, frames)) for a in range(0, frames, size)]

sampling.py

@@ -0,0 +1,200 @@
+# Adapted for diffusers from multimodal-art-projection/YuE at commit ef1936f2ee39fe8de486a0f47a481c95f8d4da87.
+# Licensed under Apache-2.0; see LICENSE.
+from __future__ import annotations
+
+import time
+
+import torch
+
+from .protocol import ABC_END, CODEC_OFFSET, CODEC_SIZE, CONTEXT, EOD, MUSIC_END
+
+
+class YuE2StaticKVCache:
+    """Preallocated token cache for eager decoding; returns views of the filled prefix without copying history."""
+
+    def __init__(self, num_layers, batch_size, num_kv_heads, max_seq_len, head_dim, dtype, device):
+        self.num_layers = num_layers
+        self.max_seq_len = max_seq_len
+        self.seen_tokens = 0
+        shape = (batch_size, num_kv_heads, max_seq_len, head_dim)
+        self.keys = [torch.zeros(shape, dtype=dtype, device=device) for _ in range(num_layers)]
+        self.values = [torch.zeros(shape, dtype=dtype, device=device) for _ in range(num_layers)]
+
+    def get_seq_length(self):
+        return self.seen_tokens
+
+    def update(self, key, value, layer_idx):
+        start, end = self.seen_tokens, self.seen_tokens + key.shape[1]
+        if end > self.max_seq_len:
+            raise ValueError(f"KV cache capacity {self.max_seq_len} exceeded by {end}; generation was not shortened")
+        self.keys[layer_idx][:, :, start:end] = key.transpose(1, 2)
+        self.values[layer_idx][:, :, start:end] = value.transpose(1, 2)
+        if layer_idx == self.num_layers - 1:
+            self.seen_tokens = end
+        return self.keys[layer_idx][:, :, :end].transpose(1, 2), self.values[layer_idx][:, :, :end].transpose(1, 2)
+
+
+def synchronize(device):
+    if device.type == "cuda":
+        torch.cuda.synchronize(device)
+    elif device.type == "mps":
+        torch.mps.synchronize()
+
+
+def window_penalty(logits, recent_ids, penalty):
+    if penalty == 1.0 or len(recent_ids) == 0:
+        return logits
+    recent = torch.as_tensor(recent_ids, dtype=torch.long, device=logits.device).reshape(1, -1)
+    freq = torch.zeros_like(logits)
+    freq.scatter_add_(-1, recent, torch.ones_like(recent, dtype=logits.dtype))
+    alpha = penalty**freq
+    return torch.where(logits < 0, logits * alpha, logits / alpha)
+
+
+def distribution(logits, sampling, history, step, phase, legacy_off=False):
+    # Planning-off requests keep the release's BF16 logits and top-3 floor; other modes sample from FP32 logits.
+    scores = logits.clone() if legacy_off else logits.float().clone()
+    end = ABC_END if phase == "abc" else MUSIC_END
+    allowed = torch.full_like(scores, float("-inf"))
+    if phase == "abc":
+        allowed[..., :EOD] = 0
+    else:
+        allowed[..., CODEC_OFFSET : CODEC_OFFSET + CODEC_SIZE] = 0
+    allowed[..., end] = 0
+    scores = scores + allowed
+    if step < sampling.min_tokens:
+        scores[..., end] = -torch.inf
+    scores = window_penalty(scores, history[-sampling.penalty_window :], sampling.repetition_penalty)
+    if sampling.temperature == 0:
+        return scores
+    if sampling.temperature != 1:
+        scores = scores / sampling.temperature
+    threshold = scores.topk(min(sampling.top_k, scores.shape[-1])).values[..., -1, None]
+    scores = scores.masked_fill(scores < threshold, -torch.inf)
+    if sampling.top_p < 1:
+        values, indices = scores.sort(descending=True)
+        probabilities = values.softmax(-1)
+        removed = probabilities.cumsum(-1) - probabilities > sampling.top_p
+        removed[..., : 3 if legacy_off else 1] = False
+        values = values.masked_fill(removed, -torch.inf)
+        scores = values.scatter(-1, indices, values)
+    return scores
+
+
+@torch.inference_mode()
+def generate_tokens(
+    transformer,
+    prefix,
+    sampling,
+    seed,
+    phase,
+    device,
+    negative=None,
+    combine_logits=None,
+    legacy_off=False,
+    cancelled=None,
+    on_token=None,
+    graph_decoder=None,
+):
+    """Sample one stage's tokens.
+
+    With `negative`, `combine_logits(conditional, unconditional)` applies guidance. `graph_decoder` (the `GraphAR`
+    class) decodes with CUDA graphs on CUDA devices; elsewhere decoding stays eager.
+    """
+    if len(prefix) + sampling.max_tokens > CONTEXT:
+        raise ValueError("Prefix + requested generation budget exceeds 24576; no implicit truncation")
+    if negative is not None and (combine_logits is None or len(negative) + sampling.max_tokens > CONTEXT):
+        raise ValueError("Guidance needs `combine_logits` and a negative prefix that fits the context")
+    if cancelled is not None and cancelled():
+        raise InterruptedError("Cancelled before prefill")
+    # Both stages reset the request seed, as the release does.
+    rng_device = device if device.type in {"cpu", "cuda"} else torch.device("cpu")
+    generator = torch.Generator(device=rng_device).manual_seed(seed)
+    config = transformer.config
+
+    def prefill(ids):
+        cache = YuE2StaticKVCache(
+            num_layers=config.num_layers,
+            batch_size=1,
+            num_kv_heads=config.num_key_value_heads,
+            max_seq_len=len(ids) + sampling.max_tokens,
+            head_dim=config.attention_head_dim,
+            dtype=transformer.dtype,
+            device=device,
+        )
+        logits = transformer(torch.tensor([ids], device=device), kv_cache=cache, logits_to_keep=1).logits
+        return logits[:, -1, :], cache
+
+    graph = None
+    positive_cache = negative_cache = None
+    synchronize(device)
+    start = time.perf_counter()
+    try:
+        if graph_decoder is not None and device.type == "cuda":
+            graph = graph_decoder(
+                transformer, [prefix] if negative is None else [prefix, negative], sampling.max_tokens, device
+            )
+            logits = graph.prefill()
+            conditional = logits[:1]
+            unconditional = logits[1:] if negative is not None else None
+        else:
+            conditional, positive_cache = prefill(prefix)
+            unconditional = None
+            if negative is not None:
+                unconditional, negative_cache = prefill(negative)
+        synchronize(device)
+        prefill_seconds = time.perf_counter() - start
+        history, first, eos = [], None, False
+        end = ABC_END if phase == "abc" else MUSIC_END
+        for step in range(sampling.max_tokens):
+            if cancelled is not None and cancelled():
+                raise InterruptedError(f"Cancelled during {phase}")
+            logits = conditional if negative is None else combine_logits(conditional, unconditional)
+            scores = distribution(logits, sampling, history, step, phase, legacy_off)
+            if sampling.temperature == 0:
+                next_id = scores.argmax(-1, keepdim=True)
+            else:
+                probabilities = scores.softmax(-1)
+                if device.type == "mps":
+                    next_id = torch.multinomial(probabilities.cpu(), 1, generator=generator).to(device)
+                else:
+                    next_id = torch.multinomial(probabilities, 1, generator=generator)
+            token = int(next_id.item())
+            if first is None:
+                first = time.perf_counter() - start
+            if on_token is not None:
+                on_token(phase, token)
+            if token == end:
+                eos = True
+                break
+            history.append(token)
+            if step + 1 < sampling.max_tokens:
+                if graph is not None:
+                    branch_logits = graph.step(next_id)
+                    conditional = branch_logits[:1]
+                    unconditional = branch_logits[1:] if negative is not None else None
+                else:
+                    conditional = transformer(next_id, kv_cache=positive_cache, logits_to_keep=1).logits[:, -1, :]
+                    if negative_cache is not None:
+                        unconditional = transformer(next_id, kv_cache=negative_cache, logits_to_keep=1).logits[
+                            :, -1, :
+                        ]
+        synchronize(device)
+        seconds = time.perf_counter() - start
+        count = len(history) + int(eos)
+        timing = {
+            "seconds": seconds,
+            "prefill_seconds": prefill_seconds,
+            "ttft_seconds": first,
+            "output_tokens": count,
+            "content_tokens": len(history),
+            "output_tps": count / seconds,
+            "prefix_tokens": len(prefix),
+            "cfg_branches": 1 if negative is None else 2,
+            "execution": "cuda_graph" if graph is not None else "eager",
+        }
+        return history, timing, not eos
+    finally:
+        if graph is not None:
+            graph.close()
+        positive_cache = negative_cache = None