Launch complete Breeze TTS 2 ZeroGPU studio

.gitattributes

@@ -1,35 +1,7 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
 *.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
 *.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
 *.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.wav filter=lfs diff=lfs merge=lfs -text
+*.mp3 filter=lfs diff=lfs merge=lfs -text
+*.flac filter=lfs diff=lfs merge=lfs -text

MODEL_LICENSE

@@ -0,0 +1,427 @@
+BREEZEBLUE RESEARCH AND NON-COMMERCIAL LICENSE AGREEMENT
+
+Version 1.0
+Last Updated: August 25, 2026
+
+This BreezeBlue Research and Non-Commercial License Agreement (the
+"Agreement") is entered into between RESONIA, INC., doing business as
+BreezeBlue ("BreezeBlue," "Licensor," "we," or "us"), and the individual or
+entity accessing, using, modifying, or distributing any Model Materials
+("You," "Your," or "Licensee").
+
+By accessing, downloading, using, modifying, or distributing any portion of
+the Model Materials or a Derivative Model, You acknowledge that You have read,
+understood, and agreed to this Agreement. If You act on behalf of an entity,
+You represent that You have authority to bind that entity, and "You" includes
+both You and that entity.
+
+This Agreement permits research and non-commercial use of the Model Materials
+free of charge. It does not grant any commercial rights and is not an open
+source license. Commercial use requires a separate written license from
+BreezeBlue.
+
+
+1. DEFINITIONS
+
+1.1 "Model" means the Breeze TTS 2 model.
+
+1.2 "Model Materials" means the Model weights, checkpoints, parameters,
+adapters, model-specific tokenizer or codec weights, embeddings, and other
+machine-learning artifacts that BreezeBlue expressly makes available under
+this Agreement, together with any accompanying documentation expressly
+designated as part of the Model Materials.
+
+Model Materials do not include source code expressly licensed under the
+Apache License, Version 2.0, or any third-party component made available under
+its own license.
+
+1.3 "Derivative Model" means any modification, adaptation, fine-tune, LoRA,
+merge, quantization, distillation, or other model that is based on or derived
+from the Model Materials, including a model created using the weights,
+parameters, activations, embeddings, or Outputs of the Model Materials to
+replicate or materially transfer their speech or audio generation
+capabilities. Outputs themselves are not Derivative Models.
+
+1.4 "Output" means audio, speech, text, metadata, or other content generated
+by the Model Materials or a Derivative Model in response to an input.
+
+1.5 "Research Purpose" means academic or scientific research that is not
+primarily intended for commercial advantage or monetary compensation,
+including peer-reviewed research, reproducibility studies, benchmarking,
+evaluation, security testing, and red-teaming.
+
+1.6 "Non-Commercial Purpose" means personal, educational, hobbyist, or other
+use that is not primarily intended for commercial advantage or monetary
+compensation, including limited evaluation and testing.
+
+1.7 "Commercial Purpose" means any use primarily intended for or directed
+toward commercial advantage, business benefit, or monetary compensation,
+whether direct or indirect. Commercial Purpose includes, without limitation:
+
+(a) production use or use in a product, service, or business operation;
+
+(b) hosting or making the Model Materials, a Derivative Model, or their
+functionality available through an API, SaaS offering, application, plug-in,
+website, or other service;
+
+(c) creating or providing content, products, or services for a fee, for a
+client, in connection with advertising, sponsorship, subscription revenue,
+or any other monetized activity;
+
+(d) using the Model Materials, a Derivative Model, or Outputs for an
+organization's internal operations beyond limited evaluation; or
+
+(e) distributing, licensing, selling, or otherwise commercializing the Model
+Materials or a Derivative Model.
+
+The status of a Licensee as a for-profit entity does not, by itself, make
+limited internal research, benchmarking, evaluation, security testing, or
+red-teaming a Commercial Purpose, provided that the use is not deployed in
+production, made available to customers or other external end users, used to
+generate revenue, or used to train or improve a non-BreezeBlue generative
+model.
+
+1.8 "Affiliate" means an entity that directly or indirectly controls, is
+controlled by, or is under common control with a party, where "control" means
+ownership or control of more than fifty percent (50%) of the voting interests
+of that entity.
+
+
+2. RESEARCH AND NON-COMMERCIAL LICENSE
+
+Subject to Your continued compliance with this Agreement, BreezeBlue grants
+You a limited, non-exclusive, worldwide, non-transferable, non-sublicensable,
+royalty-free license under BreezeBlue's intellectual property rights in the
+Model Materials to:
+
+(a) access, use, reproduce, and modify the Model Materials;
+
+(b) create Derivative Models; and
+
+(c) distribute the Model Materials and Derivative Models as expressly
+permitted by Section 4,
+
+in each case solely for a Research Purpose or Non-Commercial Purpose.
+
+No rights are granted except as expressly stated in this Agreement.
+
+
+3. COMMERCIAL USE
+
+No Commercial Purpose is permitted under this Agreement. Any Commercial
+Purpose involving the Model Materials, a Derivative Model, or an Output
+requires a separate written commercial license from BreezeBlue.
+
+For clarity, this Agreement does not include a creator, content monetization,
+small-business, revenue-threshold, or other implied commercial-use exception.
+
+This Agreement governs only the open-weight Model Materials and their
+self-hosted use. Outputs generated through BreezeBlue's hosted products or
+APIs are governed by the applicable Terms of Service, plan terms, and any
+separate written agreement, rather than this Agreement.
+
+For commercial licensing inquiries, contact:
+
+Website: https://breezeblue.ai
+Email: contact@breeze.blue
+
+
+4. DISTRIBUTION AND ATTRIBUTION
+
+You may distribute the Model Materials or a Derivative Model solely for a
+Research Purpose or Non-Commercial Purpose and only if You:
+
+(a) provide every recipient with a complete copy of this Agreement and any
+NOTICE file accompanying the Model Materials;
+
+(b) retain all copyright, attribution, provenance, and proprietary notices;
+
+(c) include the following notice in a NOTICE file distributed with the Model
+Materials or Derivative Model:
+
+"Breeze TTS 2 is licensed under the BreezeBlue Research and Non-Commercial
+License Agreement. Copyright (c) 2026 RESONIA, INC. All Rights Reserved.";
+
+(d) clearly identify any modifications You made and, for a Derivative Model,
+state prominently in its model card, repository, and documentation:
+
+"Derived from Breeze TTS 2 by BreezeBlue and licensed for research and
+non-commercial use only.";
+
+(e) license the distributed Model Materials or Derivative Model under this
+Agreement without adding terms that conflict with or reduce the protections
+of this Agreement; and
+
+(f) ensure that each recipient independently accepts this Agreement before
+accessing or using the Model Materials or Derivative Model.
+
+You may not use "BreezeBlue," "Breeze TTS 2," or any confusingly similar name
+as the primary name of a Derivative Model, product, or service, or in any way
+that suggests endorsement, affiliation, official status, or successor status.
+Reasonable descriptive attribution required by this Section is permitted.
+
+Publication of non-commercial research results, benchmark results, and
+illustrative samples in papers, reports, conference presentations, model
+cards, repositories, or non-monetized demonstrations is permitted, provided
+that You reasonably acknowledge Breeze TTS 2 and comply with this Agreement.
+
+
+5. USE RESTRICTIONS
+
+You will not, and will not permit any third party to:
+
+(a) use the Model Materials, a Derivative Model, or any Output for a
+Commercial Purpose without a separate written commercial license;
+
+(b) use the Model Materials, a Derivative Model, or any Output to create,
+train, fine-tune, distill, evaluate for training selection, or otherwise
+improve any non-BreezeBlue foundational or generative speech, audio,
+language, or multimodal model, except for creating a Derivative Model solely
+for a Research Purpose or Non-Commercial Purpose under this Agreement;
+
+(c) use the Model Materials, a Derivative Model, or any Output to clone,
+simulate, imitate, or impersonate a real person's voice, identity, likeness,
+or persona without that person's explicit and legally sufficient consent;
+
+(d) submit or use reference audio, speaker embeddings, voice prompts,
+recordings, performances, or other materials unless You have all rights,
+licenses, permissions, and consents necessary for that use;
+
+(e) use the Model Materials, a Derivative Model, or any Output to defraud,
+deceive, harass, stalk, threaten, defame, exploit, or cause harm to any person;
+
+(f) generate or distribute sexual, exploitative, or otherwise harmful content
+involving a minor;
+
+(g) engage in deceptive political persuasion, election interference, voter
+suppression, synthetic robocalls, or false impersonation of a candidate,
+public official, or other person in a political context;
+
+(h) identify, authenticate, track, surveil, or profile a person using voice or
+other biometric characteristics without a lawful basis and all consent
+required by applicable law;
+
+(i) distribute synthetic audio in circumstances where a reasonable person
+would be materially misled about its origin or authenticity, or fail to make
+any disclosure required by applicable law;
+
+(j) remove, obscure, alter, or circumvent any copyright notice, attribution,
+watermark, provenance signal, safety mechanism, or other proprietary notice
+included in the Model Materials or Outputs;
+
+(k) attempt to extract or reconstruct non-public training data, source data,
+trade secrets, or proprietary components from the Model Materials, except to
+the extent such a restriction is prohibited by applicable law; or
+
+(l) use, export, re-export, transfer, or make available the Model Materials or
+a Derivative Model in violation of any applicable law, regulation, export
+control, or trade sanction.
+
+Nothing in Section 5(k) prohibits ordinary inspection, loading, conversion,
+quantization, modification, or fine-tuning of the released Model Materials as
+expressly permitted by this Agreement.
+
+
+6. LEGAL COMPLIANCE AND THIRD-PARTY RIGHTS
+
+You are solely responsible for ensuring that Your use of the Model Materials,
+Derivative Models, inputs, and Outputs complies with all applicable laws and
+regulations and does not infringe or violate any third-party right, including
+intellectual property, privacy, publicity, biometric privacy,
+recording-consent, consumer-protection, election, and other applicable rights.
+
+This Agreement grants no right to use, reproduce, clone, simulate, imitate,
+or otherwise exploit any third-party voice, identity, likeness, persona,
+performance, recording, copyrighted work, trademark, or other protected
+material. You are solely responsible for obtaining all necessary rights,
+licenses, permissions, notices, and consents.
+
+BreezeBlue does not represent or warrant that any Model Material or Output is
+cleared for Your intended use, unique, accurate, non-infringing, or suitable
+for any particular purpose.
+
+
+7. INTELLECTUAL PROPERTY
+
+7.1 Model Materials. BreezeBlue and its licensors retain all right, title, and
+interest in and to the Model Materials. Except for the limited license granted
+in Section 2, no right or interest in the Model Materials is transferred to
+You.
+
+7.2 Derivative Models. As between You and BreezeBlue, You own the original
+modifications You contribute to a Derivative Model, to the extent permitted by
+applicable law. That ownership is subject to BreezeBlue's ownership of the
+underlying Model Materials and to every restriction in this Agreement. A
+Derivative Model remains governed by this Agreement.
+
+7.3 Outputs. As between You and BreezeBlue, and to the extent permitted by
+applicable law, You own Outputs You lawfully generate. Output ownership does
+not remove or limit the non-commercial restrictions in this Agreement, grant
+rights in third-party material, or authorize any use prohibited by this
+Agreement.
+
+7.4 Trademarks. No trademark license is granted except the limited right to
+make accurate descriptive references required for attribution under Section
+4. All goodwill arising from permitted use of BreezeBlue's marks belongs to
+BreezeBlue.
+
+7.5 Feedback. If You voluntarily provide suggestions, comments, or other
+feedback about the Model Materials, You grant BreezeBlue a perpetual,
+irrevocable, worldwide, royalty-free, fully paid, transferable, sublicensable,
+non-exclusive license to use and exploit that feedback without restriction or
+compensation.
+
+7.6 Intellectual Property Claims. If You or Your Affiliate brings a claim or
+legal proceeding alleging that the Model Materials, a Derivative Model, or an
+Output infringes intellectual property rights owned or licensable by You or
+Your Affiliate, the licenses granted to You under this Agreement terminate
+automatically when the claim or proceeding is filed.
+
+
+8. REPRESENTATIONS
+
+You represent and warrant that:
+
+(a) You have legal capacity and authority to enter into this Agreement;
+
+(b) You will use and distribute the Model Materials, Derivative Models, and
+Outputs only as permitted by this Agreement;
+
+(c) You have all necessary rights and consents for Your inputs and intended
+uses; and
+
+(d) You are not prohibited from receiving or using the Model Materials under
+applicable export-control, sanctions, or other laws.
+
+
+9. DISCLAIMER OF WARRANTIES
+
+TO THE MAXIMUM EXTENT PERMITTED BY APPLICABLE LAW, THE MODEL MATERIALS,
+DERIVATIVE MODELS, OUTPUTS, AND ALL RELATED MATERIALS ARE PROVIDED "AS IS" AND
+"AS AVAILABLE," WITH ALL FAULTS AND WITHOUT WARRANTIES OF ANY KIND, WHETHER
+EXPRESS, IMPLIED, STATUTORY, OR OTHERWISE.
+
+BREEZEBLUE DISCLAIMS ALL WARRANTIES, INCLUDING WARRANTIES OF TITLE,
+NON-INFRINGEMENT, MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE,
+ACCURACY, RELIABILITY, SECURITY, AND ANY WARRANTY THAT THE MODEL MATERIALS OR
+OUTPUTS WILL BE ERROR-FREE, UNINTERRUPTED, UNIQUE, OR SUITABLE FOR YOUR USE.
+
+YOU ARE SOLELY RESPONSIBLE FOR DETERMINING THE APPROPRIATENESS AND LAWFULNESS
+OF YOUR USE AND ASSUME ALL RISKS ARISING FROM THAT USE.
+
+
+10. LIMITATION OF LIABILITY
+
+TO THE MAXIMUM EXTENT PERMITTED BY APPLICABLE LAW, BREEZEBLUE AND ITS
+AFFILIATES, LICENSORS, OFFICERS, DIRECTORS, EMPLOYEES, CONTRACTORS, AND AGENTS
+WILL NOT BE LIABLE UNDER ANY THEORY OF LIABILITY FOR ANY INDIRECT, INCIDENTAL,
+SPECIAL, CONSEQUENTIAL, EXEMPLARY, OR PUNITIVE DAMAGES, OR FOR ANY LOSS OF
+PROFITS, REVENUE, BUSINESS, GOODWILL, DATA, OR USE, ARISING OUT OF OR RELATING
+TO THIS AGREEMENT, THE MODEL MATERIALS, ANY DERIVATIVE MODEL, OR ANY OUTPUT,
+EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGES.
+
+TO THE MAXIMUM EXTENT PERMITTED BY APPLICABLE LAW, BREEZEBLUE'S TOTAL
+AGGREGATE LIABILITY ARISING OUT OF OR RELATING TO THIS AGREEMENT WILL NOT
+EXCEED ONE HUNDRED U.S. DOLLARS (US$100).
+
+
+11. INDEMNIFICATION
+
+You agree to defend, indemnify, and hold harmless BreezeBlue and its
+Affiliates, licensors, officers, directors, employees, contractors, and agents
+from and against any third-party claim, demand, action, proceeding, damage,
+loss, liability, judgment, settlement, cost, or expense, including reasonable
+attorneys' fees, arising out of or relating to:
+
+(a) Your access to, use, modification, or distribution of the Model Materials
+or a Derivative Model;
+
+(b) Your inputs or Outputs;
+
+(c) Your violation of this Agreement or applicable law; or
+
+(d) Your infringement, misappropriation, or violation of any third-party
+right.
+
+
+12. TERM AND TERMINATION
+
+This Agreement begins when You first access or use the Model Materials and
+continues until terminated.
+
+BreezeBlue may terminate this Agreement if You materially breach it. For a
+material breach that is reasonably capable of cure, BreezeBlue will provide
+written notice and at least thirty (30) days to cure before termination.
+BreezeBlue may terminate immediately for a breach of Section 3, Section 5, or
+applicable law, or where delay would reasonably create legal, security, fraud,
+or safety risk.
+
+Upon termination, You must immediately cease using and distributing the Model
+Materials and Derivative Models, delete all copies under Your control, and,
+upon reasonable request, confirm deletion in writing. Termination does not
+require deletion of an Output lawfully created before termination, but no
+Output may be used in violation of this Agreement or applicable law.
+
+BreezeBlue may stop distributing the Model Materials prospectively. Such
+discontinuation alone does not terminate rights already granted to a Licensee
+that remains in compliance with this Agreement.
+
+Sections 3, 5, 6, 7, 9, 10, 11, 12, and 13 survive termination to the extent
+necessary to give them effect.
+
+
+13. GOVERNING LAW AND JURISDICTION
+
+This Agreement is governed by and construed in accordance with the laws of the
+State of Delaware, without regard to its conflict-of-laws principles. The
+United Nations Convention on Contracts for the International Sale of Goods
+does not apply.
+
+To the extent permitted by applicable law, any dispute arising out of or
+relating to this Agreement shall be brought exclusively in the state or
+federal courts located in the State of Delaware, and each party consents to
+the personal jurisdiction and venue of those courts.
+
+
+14. MISCELLANEOUS
+
+14.1 Entire Agreement. This Agreement and any NOTICE file accompanying the
+Model Materials constitute the entire agreement concerning the Model
+Materials and supersede prior or contemporaneous understandings concerning
+the same subject matter.
+
+14.2 Separate Software Licenses. Source code and third-party components may be
+made available under separate licenses. Those separate licenses govern only
+the materials expressly identified as subject to them and do not grant any
+right to use the Model Materials.
+
+14.3 Assignment. You may not assign or transfer this Agreement or any right
+under it without BreezeBlue's prior written consent. BreezeBlue may assign
+this Agreement in connection with a merger, acquisition, reorganization, sale
+of assets, or operation of law.
+
+14.4 Severability. If any provision is held invalid or unenforceable, it will
+be enforced to the maximum extent permitted, and the remaining provisions
+will remain in effect.
+
+14.5 No Waiver. Failure to enforce a provision is not a waiver of the right to
+enforce it later.
+
+14.6 No Endorsement. Nothing in this Agreement creates an agency,
+partnership, joint venture, employment, fiduciary, sponsorship, or endorsement
+relationship between the parties.
+
+14.7 Updates. BreezeBlue may publish new versions of this Agreement. A new
+version applies only if You expressly accept it or access Model Materials
+released under that new version. Your continued use of a previously obtained
+version of the Model Materials remains governed by the version You accepted,
+unless applicable law requires otherwise.
+
+
+15. CONTACT
+
+Questions and commercial licensing requests may be sent to:
+
+RESONIA, INC., doing business as BreezeBlue
+Website: https://breezeblue.ai
+Email: contact@breeze.blue

README.md

@@ -1,13 +1,59 @@
 ---
 title: Breeze TTS 2 Studio
-emoji: 👁
-colorFrom: red
-colorTo: purple
+emoji: "🌬️"
+colorFrom: indigo
+colorTo: blue
 sdk: gradio
-sdk_version: 6.26.0
-python_version: '3.12'
 app_file: app.py
 pinned: false
+license: other
+license_name: breezeblue-research-and-non-commercial-license
+license_link: https://huggingface.co/BreezeBlue/Breeze-TTS-2/blob/main/LICENSE
+models:
+  - BreezeBlue/Breeze-TTS-2
+short_description: Design, clone, and direct expressive English and Chinese.
+python_version: 3.10.13
+startup_duration_timeout: 1h
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Breeze TTS 2 Studio
+
+An independent, full-capability research demo for
+[`BreezeBlue/Breeze-TTS-2`](https://huggingface.co/BreezeBlue/Breeze-TTS-2).
+
+## What the Space demonstrates
+
+- **Voice design:** create a new voice from a natural-language description without reference audio.
+- **Voice cloning:** preserve a speaker from clean reference audio and its exact transcript.
+- **Voice direction:** retain speaker identity while steering tone, emotion, pace, energy, and delivery.
+- **Vocal events:** inline English events such as `(laugh)`, `(cough)`, `(clears throat)`, and `(sigh)`, plus Chinese events such as `[笑]`, `[咳嗽]`, `[清嗓子]`, and `[叹气]`.
+- **Bilingual synthesis:** English and Chinese generation from one model.
+- **Streaming:** audio chunks are played as the model produces them; a complete 24 kHz WAV is provided afterward.
+- **Reproducibility:** deterministic seeds and pinned upstream model/code revisions.
+
+## ZeroGPU design
+
+The model is loaded on CUDA at module scope as required by ZeroGPU, while the generation
+function uses `@spaces.GPU(size="large")`. The app intentionally uses the upstream eager
+inference path because `torch.compile` is unsupported on ZeroGPU and reliability matters more
+than advertising fast-path numbers that shared hardware may not reproduce.
+
+## Source provenance
+
+- Model revision: `c1c8ca18b70b30822735633991d9ebf4898e47d4`
+- Inference-code revision: `ca632ce6c4d05f7985da4eab29b1a5d445b43f7b`
+- Upstream inference repository: <https://github.com/breezeblue-ai/breeze-tts>
+
+The `models/`, `breeze_infer/`, and `configs/` folders are vendored from that pinned inference
+revision so that the Space does not depend on mutable code at startup.
+
+## License and safety
+
+The upstream inference code is Apache-2.0. The model weights, derivative models, and
+self-hosted outputs are governed by the **BreezeBlue Research and Non-Commercial License**.
+Commercial use requires written authorization from RESONIA, INC.
+
+Only use reference voices with appropriate rights and consent. Do not use generated speech
+for impersonation, deception, fraud, harassment, or authentication bypass.
+
+This Space is unofficial and is not affiliated with or endorsed by BreezeBlue.

UPSTREAM_CODE_LICENSE

@@ -0,0 +1,175 @@
+
+                                 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.

app.py

@@ -0,0 +1,807 @@
+"""Breeze TTS 2 Studio — a polished ZeroGPU showcase."""
+
+from __future__ import annotations
+
+# ZeroGPU must patch PyTorch before any package imports torch.
+import spaces
+
+import html
+import math
+import random
+import tempfile
+import time
+import uuid
+from pathlib import Path
+
+import gradio as gr
+import numpy as np
+import soundfile as sf
+from huggingface_hub import snapshot_download
+
+from breeze_infer.runtime import (
+    load_runtime,
+    set_all_seeds,
+    update_generation_config_for_breeze,
+)
+from breeze_infer.templates import get_template, prepare_inputs
+from models.fast_streaming import FastBreezeStreamingRuntime, FastStreamingConfig
+
+
+MODEL_ID = "BreezeBlue/Breeze-TTS-2"
+MODEL_REVISION = "c1c8ca18b70b30822735633991d9ebf4898e47d4"
+CODE_REVISION = "ca632ce6c4d05f7985da4eab29b1a5d445b43f7b"
+MAX_TEXT_CHARS = 900
+MAX_INSTRUCTION_CHARS = 560
+MAX_REFERENCE_SECONDS = 60
+MAX_NEW_TOKENS = 1500
+MAX_SEQ_LEN = 2048
+
+
+MODEL_PATH = Path(
+    snapshot_download(repo_id=MODEL_ID, revision=MODEL_REVISION)
+)
+TOKENIZER, MODEL, AUDIO_TOKENIZER = load_runtime(
+    MODEL_PATH,
+    device="cuda",
+    attn_implementation="eager",
+)
+update_generation_config_for_breeze(MODEL)
+RUNTIME = FastBreezeStreamingRuntime(
+    MODEL,
+    AUDIO_TOKENIZER,
+    FastStreamingConfig(
+        max_new_tokens=MAX_NEW_TOKENS,
+        max_seq_len=MAX_SEQ_LEN,
+        collect_timing=True,
+        fast_all=False,
+        repetition_penalty=1.1,
+    ),
+    tokenizer=TOKENIZER,
+)
+
+
+MODE_HELP = {
+    "Voice design": (
+        "Describe a voice in plain language. No recording needed. Best for creating "
+        "a new character, narrator, host, or product voice."
+    ),
+    "Voice clone": (
+        "Upload a clean voice reference and its exact transcript. The model carries "
+        "the speaker's identity and natural delivery into new words."
+    ),
+    "Voice direction": (
+        "Start from a reference voice, then direct the performance—emotion, pace, "
+        "energy, intimacy, authority, or scene."
+    ),
+}
+
+DESIGN_EXAMPLES = [
+    [
+        "Voice design",
+        "English",
+        "(sigh) Welcome aboard. Your journey begins now.",
+        "A warm, thoughtful young woman with a clear voice and a calm, reflective delivery.",
+        4.0,
+    ],
+    [
+        "Voice design",
+        "English",
+        "The storm is close now. Keep the lantern low, and stay where I can see you.",
+        "An older coastal storyteller with a weathered, resonant voice; intimate, measured, and quietly suspenseful.",
+        4.0,
+    ],
+    [
+        "Voice design",
+        "中文",
+        "[笑] 歡迎來到今晚的故事時間，讓我們一起開始吧。",
+        "一位溫柔自信的年輕女性，聲音清晰，語氣親切，表達輕快而富有感染力。",
+        4.0,
+    ],
+    [
+        "Voice design",
+        "中文",
+        "請稍等一下，我想我已經找到問題的答案了。",
+        "一位沉穩理性的中年研究員，音色溫暖，語速稍慢，思考時帶著自然的停頓。",
+        4.0,
+    ],
+]
+
+CSS = r"""
+:root {
+  --ink: #172036;
+  --muted: #667085;
+  --line: #e8eaf0;
+  --paper: #ffffff;
+  --wash: #f6f7fb;
+  --accent: #5565e8;
+  --accent-dark: #3948c7;
+  --blue-soft: #eef0ff;
+  --mint: #dff7ed;
+}
+
+.gradio-container {
+  max-width: 1180px !important;
+  margin: 0 auto !important;
+  padding: 24px 22px 70px !important;
+  color: var(--ink) !important;
+  background:
+    radial-gradient(circle at 82% 4%, rgba(85,101,232,.10), transparent 23rem),
+    radial-gradient(circle at 9% 28%, rgba(72,187,156,.07), transparent 20rem),
+    #fbfbfd !important;
+}
+
+footer { display: none !important; }
+.prose { max-width: none !important; }
+
+#hero {
+  position: relative;
+  overflow: hidden;
+  padding: 44px 44px 38px;
+  border: 1px solid rgba(23,32,54,.08);
+  border-radius: 28px;
+  background: rgba(255,255,255,.86);
+  box-shadow: 0 22px 65px rgba(31,42,82,.08);
+  margin-bottom: 18px;
+}
+#hero:after {
+  content: "";
+  position: absolute;
+  width: 290px;
+  height: 290px;
+  right: -110px;
+  top: -135px;
+  border-radius: 50%;
+  border: 58px solid rgba(85,101,232,.08);
+}
+.hero-eyebrow {
+  display: flex;
+  align-items: center;
+  gap: 9px;
+  color: #4c5bd4;
+  font-size: 12px;
+  font-weight: 750;
+  letter-spacing: .13em;
+  text-transform: uppercase;
+}
+.hero-eyebrow i {
+  width: 8px;
+  height: 8px;
+  display: inline-block;
+  border-radius: 50%;
+  background: #35b784;
+  box-shadow: 0 0 0 5px rgba(53,183,132,.12);
+}
+#hero h1 {
+  max-width: 760px;
+  margin: 18px 0 12px;
+  font-size: clamp(38px, 6vw, 67px);
+  line-height: .98;
+  letter-spacing: -.052em;
+  color: #131b2e;
+}
+#hero .lede {
+  max-width: 720px;
+  margin: 0;
+  font-size: 18px;
+  line-height: 1.58;
+  color: #5b6477;
+}
+.cap-row {
+  display: flex;
+  flex-wrap: wrap;
+  gap: 8px;
+  margin-top: 24px;
+}
+.cap-row span {
+  display: inline-flex;
+  align-items: center;
+  padding: 8px 12px;
+  border: 1px solid #e4e6ef;
+  border-radius: 999px;
+  color: #434d63;
+  background: #fff;
+  font-size: 12px;
+  font-weight: 650;
+}
+
+.studio-shell {
+  padding: 18px !important;
+  border: 1px solid rgba(23,32,54,.08) !important;
+  border-radius: 26px !important;
+  background: rgba(255,255,255,.92) !important;
+  box-shadow: 0 20px 52px rgba(31,42,82,.07) !important;
+}
+.panel {
+  padding: 20px !important;
+  border: 1px solid var(--line) !important;
+  border-radius: 20px !important;
+  background: var(--paper) !important;
+}
+.output-panel {
+  min-height: 100%;
+  background: #f8f9fc !important;
+}
+.section-kicker {
+  margin: 0 0 2px !important;
+  color: #7a8396 !important;
+  font-size: 11px !important;
+  font-weight: 750 !important;
+  letter-spacing: .12em !important;
+  text-transform: uppercase !important;
+}
+.section-title h2, .section-title h3 {
+  margin: 0 0 14px !important;
+  color: var(--ink) !important;
+  letter-spacing: -.025em !important;
+}
+.mode-note {
+  padding: 11px 13px;
+  border-left: 3px solid #7c88ef;
+  border-radius: 0 10px 10px 0;
+  background: #f5f6ff;
+  color: #505a70;
+  font-size: 13px;
+  line-height: 1.5;
+}
+.event-label {
+  margin: 3px 0 7px;
+  color: #71798a;
+  font-size: 12px;
+}
+.event-chip {
+  min-width: 0 !important;
+  border: 1px solid #e1e4ec !important;
+  border-radius: 999px !important;
+  background: #fff !important;
+  color: #465067 !important;
+  font-size: 12px !important;
+  box-shadow: none !important;
+}
+.event-chip:hover {
+  border-color: #aeb6f5 !important;
+  color: #3948c7 !important;
+  background: #f7f7ff !important;
+}
+#generate-button {
+  min-height: 50px;
+  border: none !important;
+  border-radius: 14px !important;
+  background: var(--accent) !important;
+  color: white !important;
+  font-weight: 720 !important;
+  box-shadow: 0 10px 24px rgba(85,101,232,.22) !important;
+}
+#generate-button:hover { background: var(--accent-dark) !important; }
+.secondary-button {
+  min-height: 50px;
+  border-radius: 14px !important;
+  box-shadow: none !important;
+}
+#live-audio {
+  border: 1px solid #e0e3eb !important;
+  border-radius: 16px !important;
+  background: #fff !important;
+}
+.status-card {
+  min-height: 76px;
+  padding: 13px 15px;
+  border: 1px solid #e2e5ee;
+  border-radius: 14px;
+  background: white;
+  color: #596277;
+  font-size: 13px;
+  line-height: 1.55;
+}
+.status-card strong { color: #202a41; }
+.model-facts {
+  display: grid;
+  grid-template-columns: repeat(2, minmax(0,1fr));
+  gap: 9px;
+  margin-top: 16px;
+}
+.model-facts div {
+  padding: 13px;
+  border: 1px solid #e5e7ee;
+  border-radius: 13px;
+  background: rgba(255,255,255,.76);
+}
+.model-facts b {
+  display: block;
+  margin-bottom: 3px;
+  font-size: 16px;
+  color: #202a41;
+}
+.model-facts span {
+  color: #737b8c;
+  font-size: 11px;
+}
+.recipe-grid {
+  display: grid;
+  grid-template-columns: repeat(3, minmax(0,1fr));
+  gap: 12px;
+  margin: 18px 0;
+}
+.recipe {
+  padding: 18px;
+  border: 1px solid #e3e6ed;
+  border-radius: 17px;
+  background: rgba(255,255,255,.84);
+}
+.recipe b { display: block; margin-bottom: 7px; color: #29334a; }
+.recipe p { margin: 0; color: #6c7485; font-size: 13px; line-height: 1.5; }
+.fine-print {
+  color: #70798a;
+  font-size: 12px;
+  line-height: 1.55;
+}
+
+@media (max-width: 760px) {
+  .gradio-container { padding: 12px 10px 48px !important; }
+  #hero { padding: 30px 24px; border-radius: 22px; }
+  #hero h1 { font-size: 44px; }
+  #hero .lede { font-size: 16px; }
+  .studio-shell { padding: 10px !important; }
+  .panel { padding: 14px !important; }
+  .recipe-grid { grid-template-columns: 1fr; }
+}
+"""
+
+
+def mode_ui(mode: str):
+    needs_reference = mode in {"Voice clone", "Voice direction"}
+    needs_instruction = mode in {"Voice design", "Voice direction"}
+    return (
+        f'<div class="mode-note">{html.escape(MODE_HELP[mode])}</div>',
+        gr.update(visible=needs_reference),
+        gr.update(visible=needs_instruction),
+        gr.update(
+            visible=needs_instruction,
+            value=4.0 if needs_instruction else 1.0,
+            interactive=needs_instruction,
+        ),
+    )
+
+
+def insert_event(current_text: str, tag: str) -> str:
+    current_text = (current_text or "").lstrip()
+    return f"{tag} {current_text}".strip()
+
+
+def reset_output():
+    return None, None, (
+        '<div class="status-card"><strong>Ready when you are.</strong><br>'
+        "Your generated audio and performance details will appear here.</div>"
+    )
+
+
+def _clean_text(value: str | None, name: str, limit: int) -> str:
+    cleaned = (value or "").strip()
+    if not cleaned:
+        raise gr.Error(f"{name} is required.")
+    if len(cleaned) > limit:
+        raise gr.Error(f"{name} is too long ({len(cleaned)}/{limit} characters).")
+    return cleaned
+
+
+def _validate_reference(path: str | None, transcript: str | None, consent: bool):
+    if not consent:
+        raise gr.Error(
+            "Please confirm that you have permission to use the reference voice."
+        )
+    if not path or not Path(path).is_file():
+        raise gr.Error("Upload or record a clean reference voice first.")
+    ref_text = _clean_text(transcript, "Exact reference transcript", MAX_TEXT_CHARS)
+    try:
+        info = sf.info(path)
+    except Exception as exc:
+        raise gr.Error("The reference audio could not be read. Try WAV, FLAC, or MP3.") from exc
+    if info.duration < 1.0:
+        raise gr.Error("Use at least one second of reference speech.")
+    if info.duration > MAX_REFERENCE_SECONDS:
+        raise gr.Error(
+            f"Keep reference audio under {MAX_REFERENCE_SECONDS} seconds for this demo."
+        )
+    return str(Path(path)), ref_text, float(info.duration)
+
+
+def gpu_duration(
+    mode,
+    language,
+    text,
+    instruction,
+    reference_audio,
+    reference_text,
+    consent,
+    cfg_scale,
+    seed,
+    randomize_seed,
+):
+    del language, instruction, reference_audio, reference_text, consent, cfg_scale, seed, randomize_seed
+    estimate = 42 + len(text or "") // 7 + (12 if mode != "Voice design" else 0)
+    return max(45, min(120, estimate))
+
+
+@spaces.GPU(duration=gpu_duration, size="large")
+def generate(
+    mode: str,
+    language: str,
+    text: str,
+    instruction: str,
+    reference_audio: str | None,
+    reference_text: str,
+    consent: bool,
+    cfg_scale: float,
+    seed: float,
+    randomize_seed: bool,
+):
+    target_text = _clean_text(text, "Script", MAX_TEXT_CHARS)
+    if mode not in MODE_HELP:
+        raise gr.Error("Choose a generation mode.")
+
+    if randomize_seed:
+        final_seed = random.SystemRandom().randint(0, 2_147_483_647)
+    else:
+        final_seed = int(seed)
+        if final_seed < 0 or final_seed > 2_147_483_647:
+            raise gr.Error("Seed must be between 0 and 2,147,483,647.")
+
+    has_reference = mode in {"Voice clone", "Voice direction"}
+    ref_duration = None
+    ref_path = None
+    ref_transcript = None
+    if has_reference:
+        ref_path, ref_transcript, ref_duration = _validate_reference(
+            reference_audio, reference_text, consent
+        )
+
+    if mode == "Voice clone":
+        direction = "自然清晰地說話。" if language == "中文" else "Speak clearly and naturally."
+        guidance = 1.0
+    else:
+        direction = _clean_text(
+            instruction, "Voice description or direction", MAX_INSTRUCTION_CHARS
+        )
+        guidance = float(cfg_scale)
+        if not math.isfinite(guidance) or not 1.0 <= guidance <= 6.0:
+            raise gr.Error("Guidance strength must be between 1 and 6.")
+
+    request_id = f"space-{uuid.uuid4().hex}"
+    request = {
+        "id": request_id,
+        "text": target_text,
+        "instruction": direction,
+        "speaker": "S0",
+    }
+    template_name = "tts_instruction"
+    if has_reference:
+        request["ref_audio_path"] = ref_path
+        request["ref_text"] = ref_transcript
+        template_name = "ref_edit_tata"
+
+    set_all_seeds(final_seed)
+    inputs = prepare_inputs(
+        TOKENIZER,
+        AUDIO_TOKENIZER,
+        MODEL,
+        [request],
+        get_template(template_name),
+        guidance_scale=guidance,
+        guidance_scale_ref=None,
+        guidance_scale_ins=None,
+    )
+
+    output_path = Path(tempfile.gettempdir()) / f"breeze-tts-2-{uuid.uuid4().hex}.wav"
+    started = time.perf_counter()
+    first_audio_at = None
+    sample_count = 0
+    chunk_count = 0
+
+    with sf.SoundFile(
+        output_path,
+        mode="w",
+        samplerate=RUNTIME.sample_rate,
+        channels=1,
+        subtype="PCM_16",
+    ) as wav_file:
+        for chunk in RUNTIME.iter_audio_chunks(inputs, request_id=request_id):
+            audio = np.asarray(chunk.audio, dtype=np.float32).reshape(-1)
+            if audio.size == 0:
+                continue
+            audio = np.clip(audio, -1.0, 1.0)
+            wav_file.write(audio)
+            sample_count += audio.size
+            chunk_count += 1
+            if first_audio_at is None:
+                first_audio_at = time.perf_counter()
+            elapsed = time.perf_counter() - started
+            generated_seconds = sample_count / RUNTIME.sample_rate
+            live_status = (
+                '<div class="status-card"><strong>Streaming your voice…</strong><br>'
+                f"{generated_seconds:.1f}s rendered · {elapsed:.1f}s elapsed · seed {final_seed:,}</div>"
+            )
+            yield (chunk.sample_rate, audio), gr.skip(), live_status
+
+    if sample_count == 0:
+        raise gr.Error("The model completed without producing audio. Try a shorter script.")
+
+    elapsed = time.perf_counter() - started
+    audio_seconds = sample_count / RUNTIME.sample_rate
+    rtf = elapsed / audio_seconds if audio_seconds else 0.0
+    ttfa = (first_audio_at - started) * 1000 if first_audio_at else 0.0
+    reference_note = f" · reference {ref_duration:.1f}s" if ref_duration else ""
+    final_status = (
+        '<div class="status-card"><strong>Complete.</strong><br>'
+        f"{audio_seconds:.1f}s audio · {elapsed:.1f}s generation · {rtf:.2f} RTF · "
+        f"{ttfa:.0f}ms first chunk · {chunk_count} chunks · seed {final_seed:,}{reference_note}</div>"
+    )
+    yield gr.skip(), str(output_path), final_status
+
+
+theme = gr.themes.Base(
+    primary_hue=gr.themes.colors.indigo,
+    neutral_hue=gr.themes.colors.slate,
+    font=["Inter", "ui-sans-serif", "system-ui", "sans-serif"],
+    font_mono=["IBM Plex Mono", "ui-monospace", "monospace"],
+).set(
+    body_background_fill="#fbfbfd",
+    block_background_fill="#ffffff",
+    block_border_width="0px",
+    input_background_fill="#ffffff",
+    input_border_color="#dfe3ec",
+    input_border_color_focus="#7c88ef",
+    button_primary_background_fill="#5565e8",
+    button_primary_background_fill_hover="#3948c7",
+)
+
+
+with gr.Blocks(title="Breeze TTS 2 Studio") as demo:
+    gr.HTML(
+        """
+        <section id="hero">
+          <div class="hero-eyebrow"><i></i> Unofficial research space · Breeze TTS 2</div>
+          <h1>Give a voice<br>to intent.</h1>
+          <p class="lede">Design a speaker from words, carry a real voice into a new script,
+          or direct a performance with natural language. Built to make every capability clear,
+          creative, and immediately usable.</p>
+          <div class="cap-row">
+            <span>Voice design</span><span>Voice cloning</span><span>Voice direction</span>
+            <span>English + 中文</span><span>Live chunk playback</span><span>24 kHz WAV</span>
+          </div>
+        </section>
+        """
+    )
+
+    with gr.Row(elem_classes="studio-shell", equal_height=False):
+        with gr.Column(scale=7, min_width=360, elem_classes="panel"):
+            gr.Markdown("CREATE", elem_classes="section-kicker")
+            gr.Markdown("## Shape the performance", elem_classes="section-title")
+
+            mode = gr.Radio(
+                ["Voice design", "Voice clone", "Voice direction"],
+                value="Voice design",
+                label="Mode",
+            )
+            mode_note = gr.HTML(
+                f'<div class="mode-note">{MODE_HELP["Voice design"]}</div>'
+            )
+            language = gr.Radio(["English", "中文"], value="English", label="Language")
+            text = gr.Textbox(
+                label="Script",
+                placeholder="Write the words you want to hear…",
+                lines=6,
+                max_lines=10,
+                max_length=MAX_TEXT_CHARS,
+                buttons=["copy"],
+                info="You can add vocal events directly in the script.",
+            )
+
+            gr.HTML('<div class="event-label">Vocal events — click to insert</div>')
+            with gr.Row():
+                event_buttons = []
+                for label, tag in [
+                    ("Laugh", "(laugh)"),
+                    ("Sigh", "(sigh)"),
+                    ("Cough", "(cough)"),
+                    ("Clear throat", "(clears throat)"),
+                    ("笑", "[笑]"),
+                    ("嘆氣", "[叹气]"),
+                    ("咳嗽", "[咳嗽]"),
+                    ("清嗓", "[清嗓子]"),
+                ]:
+                    event_buttons.append(
+                        (gr.Button(label, size="sm", elem_classes="event-chip"), tag)
+                    )
+
+            with gr.Group(visible=True) as instruction_group:
+                instruction = gr.Textbox(
+                    label="Voice description",
+                    placeholder=(
+                        "Describe age, timbre, accent, pace, energy, personality, "
+                        "emotion, and scene…"
+                    ),
+                    lines=4,
+                    max_lines=7,
+                    max_length=MAX_INSTRUCTION_CHARS,
+                    info="Write this in the same language as your script.",
+                )
+
+            with gr.Group(visible=False) as reference_group:
+                reference_audio = gr.Audio(
+                    sources=["upload", "microphone"],
+                    type="filepath",
+                    label="Reference voice",
+                    format="wav",
+                )
+                reference_text = gr.Textbox(
+                    label="Exact reference transcript",
+                    placeholder="Type exactly what is spoken in the recording…",
+                    lines=3,
+                    max_lines=6,
+                    max_length=MAX_TEXT_CHARS,
+                    info="Matching punctuation and wording helps preserve the voice.",
+                )
+                consent = gr.Checkbox(
+                    label="I have permission to use this voice and will not use it deceptively.",
+                    value=False,
+                )
+
+            with gr.Accordion("Generation settings", open=False):
+                with gr.Group(visible=True) as guidance_group:
+                    cfg_scale = gr.Slider(
+                        1.0,
+                        6.0,
+                        value=4.0,
+                        step=0.25,
+                        label="Direction strength",
+                        info="Higher follows the description more strongly; 4 is the official recommendation.",
+                    )
+                with gr.Row():
+                    seed = gr.Number(
+                        value=42,
+                        precision=0,
+                        label="Seed",
+                        minimum=0,
+                        maximum=2_147_483_647,
+                    )
+                    randomize_seed = gr.Checkbox(
+                        value=True,
+                        label="New variation each time",
+                    )
+
+            with gr.Row():
+                generate_button = gr.Button(
+                    "Generate voice",
+                    variant="primary",
+                    elem_id="generate-button",
+                    scale=4,
+                )
+                clear_button = gr.Button(
+                    "Clear output",
+                    variant="secondary",
+                    elem_classes="secondary-button",
+                    scale=1,
+                )
+
+        with gr.Column(scale=5, min_width=320, elem_classes=["panel", "output-panel"]):
+            gr.Markdown("LISTEN", elem_classes="section-kicker")
+            gr.Markdown("## Your rendered voice", elem_classes="section-title")
+            live_audio = gr.Audio(
+                label="Live playback",
+                streaming=True,
+                autoplay=True,
+                interactive=False,
+                format="wav",
+                elem_id="live-audio",
+            )
+            download = gr.File(label="Download the complete WAV", interactive=False)
+            status = gr.HTML(
+                '<div class="status-card"><strong>Ready when you are.</strong><br>'
+                "Your generated audio and performance details will appear here.</div>"
+            )
+            gr.HTML(
+                """
+                <div class="model-facts">
+                  <div><b>~3B</b><span>MODEL PARAMETERS</span></div>
+                  <div><b>24 kHz</b><span>MONO OUTPUT</span></div>
+                  <div><b>3 modes</b><span>DESIGN · CLONE · DIRECT</span></div>
+                  <div><b>ZeroGPU</b><span>48 GB LARGE SLICE</span></div>
+                </div>
+                <p class="fine-print">This Space uses the reliable eager path. The published
+                &lt;40 ms TTFA and 0.32 RTF figures are official warmed-H100 fast-path results,
+                not a promise for shared ZeroGPU sessions.</p>
+                """
+            )
+
+    gr.HTML(
+        """
+        <div class="recipe-grid">
+          <div class="recipe"><b>Design with specificity</b><p>Combine vocal identity,
+          timbre, pace, energy, emotional state, and the scene. A coherent character works
+          better than a pile of adjectives.</p></div>
+          <div class="recipe"><b>Clone from clean speech</b><p>Use one speaker, little room
+          noise, and an exact transcript. Ten to thirty seconds is a practical starting point.</p></div>
+          <div class="recipe"><b>Direct the change</b><p>Describe the performance delta:
+          “keep the same identity, but speak more slowly with restrained urgency.”</p></div>
+        </div>
+        """
+    )
+
+    with gr.Accordion("Starter voices", open=False):
+        gr.Markdown(
+            "These examples demonstrate open-ended voice design in both supported languages."
+        )
+        gr.Examples(
+            examples=DESIGN_EXAMPLES,
+            inputs=[mode, language, text, instruction, cfg_scale],
+            label="Load a recipe",
+        )
+
+    with gr.Accordion("About this model, license, and responsible use", open=False):
+        gr.Markdown(
+            f"""
+**Breeze TTS 2** is an open-weight bilingual speech model from BreezeBlue. This independent
+Space exposes reference-free voice design, reference-guided cloning and direction, inline
+vocal events, and streaming generation.
+
+- Model: [{MODEL_ID}](https://huggingface.co/{MODEL_ID})
+- Inference code: [breezeblue-ai/breeze-tts](https://github.com/breezeblue-ai/breeze-tts)
+- Pinned model revision: `{MODEL_REVISION[:12]}`
+- Pinned code revision: `{CODE_REVISION[:12]}`
+
+The inference source is Apache-2.0. The model weights, derivative models, and self-hosted
+outputs are licensed for **research and non-commercial use only**. Commercial use requires
+written authorization from RESONIA, INC. Only upload voices you have the right and consent to
+use. Do not use generated speech for impersonation, deception, fraud, harassment, or
+authentication bypass.
+            """
+        )
+
+    mode.change(
+        mode_ui,
+        inputs=mode,
+        outputs=[mode_note, reference_group, instruction_group, guidance_group],
+        queue=False,
+    )
+    for button, tag in event_buttons:
+        button.click(
+            lambda current, event_tag=tag: insert_event(current, event_tag),
+            inputs=text,
+            outputs=text,
+            queue=False,
+        )
+
+    generation_event = generate_button.click(
+        generate,
+        inputs=[
+            mode,
+            language,
+            text,
+            instruction,
+            reference_audio,
+            reference_text,
+            consent,
+            cfg_scale,
+            seed,
+            randomize_seed,
+        ],
+        outputs=[live_audio, download, status],
+        api_name="generate",
+        api_description="Generate speech with Breeze TTS 2 using design, clone, or direction mode.",
+        show_progress="minimal",
+        concurrency_limit=1,
+        concurrency_id="breeze-tts-2",
+        stream_every=0.35,
+    )
+    clear_button.click(
+        reset_output,
+        outputs=[live_audio, download, status],
+        cancels=[generation_event],
+        queue=False,
+    )
+
+
+if __name__ == "__main__":
+    demo.queue(max_size=40, default_concurrency_limit=1).launch(
+        theme=theme,
+        css=CSS,
+    )

breeze_infer/__init__.py

@@ -0,0 +1 @@
+"""Breeze TTS inference runtime package."""

breeze_infer/api.py

@@ -0,0 +1,256 @@
+"""Thin streaming API over the PyTorch Breeze inference runtime."""
+
+from __future__ import annotations
+
+import argparse
+import tempfile
+import threading
+import uuid
+from collections.abc import AsyncIterator, Iterator
+from contextlib import asynccontextmanager
+from dataclasses import dataclass, replace
+from pathlib import Path
+
+import numpy as np
+from fastapi import FastAPI, File, Form, HTTPException, UploadFile
+from fastapi.responses import JSONResponse, StreamingResponse
+
+from breeze_infer.runtime import (
+    load_runtime,
+    resolve_device,
+    set_all_seeds,
+    update_generation_config_for_breeze,
+)
+from breeze_infer.templates import get_template, prepare_inputs
+from models.fast_streaming import FastBreezeStreamingRuntime, FastStreamingConfig
+from models.warmup_profile import load_warmup_profile
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+FAST_CONFIG = REPO_ROOT / "configs" / "fast.json"
+DEFAULT_CFG_SCALE = 1.0
+MAX_NEW_TOKENS = 1500
+MAX_SEQ_LEN = 2048
+REPETITION_PENALTY = 1.1
+OPTIONAL_AUDIO_FILE = File(None)
+
+
+@dataclass(frozen=True)
+class ApiSettings:
+    model: Path
+    fast_all: bool | None
+    fast_text_encoder: bool
+    fast_backbone_prefill: bool
+    fast_backbone_decode: bool
+    fast_depth_decoder: bool
+    fast_codec: bool
+
+
+_settings: ApiSettings | None = None
+_request_lock = threading.Lock()
+
+
+def _pcm16(audio: np.ndarray) -> bytes:
+    audio = np.asarray(audio, dtype=np.float32)
+    audio = np.clip(audio, -1.0, 1.0)
+    return (audio * 32767.0).astype("<i2", copy=False).tobytes()
+
+
+async def _save_upload(upload: UploadFile) -> Path:
+    suffix = Path(upload.filename or "reference.wav").suffix or ".wav"
+    with tempfile.NamedTemporaryFile(
+        prefix="breeze_ref_", suffix=suffix, delete=False
+    ) as temporary:
+        path = Path(temporary.name)
+        try:
+            payload = await upload.read()
+            if not payload:
+                raise HTTPException(status_code=400, detail="Reference audio is empty.")
+            temporary.write(payload)
+        except Exception:
+            path.unlink(missing_ok=True)
+            raise
+    return path
+
+
+def _load_app(app: FastAPI, settings: ApiSettings) -> None:
+    tokenizer, model, audio_tokenizer = load_runtime(
+        settings.model,
+        device=resolve_device(),
+        attn_implementation="eager",
+    )
+    update_generation_config_for_breeze(model)
+
+    config = FastStreamingConfig(
+        max_new_tokens=MAX_NEW_TOKENS,
+        max_seq_len=MAX_SEQ_LEN,
+        fast_all=settings.fast_all,
+        fast_text_encoder=settings.fast_text_encoder,
+        fast_backbone_prefill=settings.fast_backbone_prefill,
+        fast_backbone_decode=settings.fast_backbone_decode,
+        fast_depth_decoder=settings.fast_depth_decoder,
+        fast_codec=settings.fast_codec,
+        repetition_penalty=REPETITION_PENALTY,
+    )
+    runtime = FastBreezeStreamingRuntime(
+        model, audio_tokenizer, config, tokenizer=tokenizer
+    )
+    if runtime.fast_enabled:
+        profile = load_warmup_profile(FAST_CONFIG)
+        profile = replace(profile, codec_chunk_frames=runtime.codec_chunk_frames)
+        manifest = runtime.warmup_from_profile(profile)
+        print(f"fast warmup: {manifest['total_elapsed_ms']:.2f} ms", flush=True)
+
+    app.state.tokenizer = tokenizer
+    app.state.model = model
+    app.state.audio_tokenizer = audio_tokenizer
+    app.state.runtime = runtime
+
+
+@asynccontextmanager
+async def _lifespan(app: FastAPI) -> AsyncIterator[None]:
+    if _settings is None:
+        raise RuntimeError("API settings are not initialized")
+    _load_app(app, _settings)
+    yield
+
+
+app = FastAPI(title="Breeze TTS API", lifespan=_lifespan)
+
+
+@app.get("/health")
+def health() -> JSONResponse:
+    if not hasattr(app.state, "runtime"):
+        return JSONResponse({"status": "loading"}, status_code=503)
+    return JSONResponse({"status": "ok", "sample_rate": app.state.runtime.sample_rate})
+
+
+@app.post("/v1/audio/speech")
+async def speech(
+    text: str = Form(...),
+    instruction: str = Form("Speak clearly and naturally."),
+    cfg_scale: float = Form(DEFAULT_CFG_SCALE),
+    ref_audio: UploadFile | None = OPTIONAL_AUDIO_FILE,
+    ref_text: str = Form(""),
+    seed: int = Form(42),
+) -> StreamingResponse:
+    if not _request_lock.acquire(blocking=False):
+        raise HTTPException(
+            status_code=409, detail="An inference request is already running."
+        )
+
+    reference_path: Path | None = None
+    try:
+        if not np.isfinite(cfg_scale) or cfg_scale <= 0:
+            raise HTTPException(
+                status_code=400, detail="cfg_scale must be greater than 0."
+            )
+        ref_text = ref_text.strip()
+        has_reference = ref_audio is not None and bool(ref_audio.filename)
+        if has_reference != bool(ref_text):
+            raise HTTPException(
+                status_code=400,
+                detail="ref_audio and ref_text must be provided together or both omitted.",
+            )
+        if has_reference:
+            assert ref_audio is not None
+            reference_path = await _save_upload(ref_audio)
+
+        request_id = f"api-{uuid.uuid4().hex}"
+        request = {
+            "id": request_id,
+            "text": text,
+            "instruction": instruction,
+            "speaker": "S0",
+        }
+        template_name = "tts_instruction"
+        if reference_path is not None:
+            request["ref_audio_path"] = str(reference_path)
+            request["ref_text"] = ref_text
+            template_name = "ref_edit_tata"
+
+        set_all_seeds(seed)
+        inputs = prepare_inputs(
+            app.state.tokenizer,
+            app.state.audio_tokenizer,
+            app.state.model,
+            [request],
+            get_template(template_name),
+            guidance_scale=cfg_scale,
+            guidance_scale_ref=None,
+            guidance_scale_ins=None,
+        )
+    except Exception:
+        if reference_path is not None:
+            reference_path.unlink(missing_ok=True)
+        _request_lock.release()
+        raise
+
+    def body() -> Iterator[bytes]:
+        try:
+            for chunk in app.state.runtime.iter_audio_chunks(
+                inputs, request_id=request_id
+            ):
+                pcm = _pcm16(chunk.audio)
+                if pcm:
+                    yield pcm
+        finally:
+            if reference_path is not None:
+                reference_path.unlink(missing_ok=True)
+            _request_lock.release()
+
+    return StreamingResponse(
+        body(),
+        media_type="audio/pcm",
+        headers={
+            "X-Sample-Rate": str(app.state.runtime.sample_rate),
+            "X-Sample-Format": "s16le",
+            "Cache-Control": "no-store",
+        },
+    )
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="Serve Breeze TTS 2 streaming inference"
+    )
+    parser.add_argument("model", type=Path)
+    parser.add_argument("--host", default="127.0.0.1")
+    parser.add_argument("--port", type=int, default=7860)
+    parser.add_argument(
+        "--fast-all", action=argparse.BooleanOptionalAction, default=None
+    )
+    parser.add_argument(
+        "--fast-text-encoder", action=argparse.BooleanOptionalAction, default=False
+    )
+    parser.add_argument(
+        "--fast-backbone-prefill", action=argparse.BooleanOptionalAction, default=False
+    )
+    parser.add_argument(
+        "--fast-backbone-decode", action=argparse.BooleanOptionalAction, default=False
+    )
+    parser.add_argument(
+        "--fast-depth-decoder", action=argparse.BooleanOptionalAction, default=False
+    )
+    parser.add_argument(
+        "--fast-codec", action=argparse.BooleanOptionalAction, default=False
+    )
+    args = parser.parse_args()
+
+    global _settings
+    _settings = ApiSettings(
+        model=args.model,
+        fast_all=args.fast_all,
+        fast_text_encoder=args.fast_text_encoder,
+        fast_backbone_prefill=args.fast_backbone_prefill,
+        fast_backbone_decode=args.fast_backbone_decode,
+        fast_depth_decoder=args.fast_depth_decoder,
+        fast_codec=args.fast_codec,
+    )
+
+    import uvicorn
+
+    uvicorn.run(app, host=args.host, port=args.port, log_level="info")
+
+
+if __name__ == "__main__":
+    main()

breeze_infer/audio.py

@@ -0,0 +1,22 @@
+"""Reference-audio loading for Breeze inference."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import soundfile as sf
+import torch
+
+
+def encode_prompt_audio(audio_tokenizer: Any, audio_path: str | Path) -> torch.Tensor:
+    wav, sample_rate = sf.read(audio_path, always_2d=True, dtype="float32")
+    wav = np.mean(wav, axis=1)
+    encoded = audio_tokenizer.encode(wav, sr=sample_rate)
+    codes = torch.as_tensor(encoded["audio_codes"][0], dtype=torch.int16)
+    if codes.ndim != 2:
+        raise ValueError(
+            f"Expected 2D audio codes for '{audio_path}', got shape {tuple(codes.shape)}"
+        )
+    return codes.cpu().contiguous()

breeze_infer/runtime.py

@@ -0,0 +1,106 @@
+from __future__ import annotations
+
+import os
+import random
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+from transformers import AutoTokenizer
+
+from models.breeze import BreezeForConditionalGeneration
+
+
+def get_dist_info() -> tuple[int, int, int]:
+    rank = int(os.environ.get("RANK", "0"))
+    world_size = int(os.environ.get("WORLD_SIZE", "1"))
+    local_rank = int(os.environ.get("LOCAL_RANK", str(rank)))
+    return rank, world_size, local_rank
+
+
+def resolve_device(explicit_device: str | None = None) -> str:
+    if explicit_device:
+        return explicit_device
+
+    _, _, local_rank = get_dist_info()
+    if torch.cuda.is_available():
+        return f"cuda:{local_rank}"
+    return "cpu"
+
+
+def set_all_seeds(seed: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+def update_generation_config_for_breeze(
+    model: torch.nn.Module,
+    generation_config: dict[str, Any] | None = None,
+) -> None:
+    generation_config = generation_config or {
+        "depth_decoder_do_sample": True,
+        "depth_decoder_temperature": 0.9,
+        "depth_decoder_top_p": 1.0,
+        "depth_decoder_top_k": 50,
+        "do_sample": True,
+        "top_p": 1.0,
+        "top_k": 50,
+        "max_new_tokens": 750,
+        "temperature": 0.9,
+    }
+
+    prefix = "depth_decoder_"
+    depth_decoder_attrs = {
+        attr[len(prefix) :]: value
+        for attr, value in generation_config.items()
+        if attr.startswith(prefix)
+    }
+    vars(model.depth_decoder.generation_config).update(
+        {"_from_model_config": False, **depth_decoder_attrs}
+    )
+    vars(model.generation_config).update(generation_config)
+
+
+def load_runtime(
+    ckpt_dir: Path,
+    *,
+    device: str,
+    attn_implementation: str,
+) -> tuple[AutoTokenizer, BreezeForConditionalGeneration, Any]:
+
+    if device.startswith("cuda"):
+        try:
+            torch.cuda.set_device(device)
+        except Exception as exc:
+            rank, world_size, local_rank = get_dist_info()
+            raise RuntimeError(
+                "Failed to set CUDA device "
+                f"device={device} rank={rank} world_size={world_size} local_rank={local_rank} "
+                f"CUDA_VISIBLE_DEVICES={os.environ.get('CUDA_VISIBLE_DEVICES')} "
+                f"device_count={torch.cuda.device_count()}"
+            ) from exc
+    tokenizer = AutoTokenizer.from_pretrained(ckpt_dir)
+    model = BreezeForConditionalGeneration.from_pretrained(
+        ckpt_dir,
+        dtype=torch.bfloat16,
+        attn_implementation=attn_implementation,
+    )
+    model.to(device).eval()
+
+    from qwen_tts import Qwen3TTSTokenizer
+
+    bundled_audio_tokenizer = ckpt_dir / "audio_tokenizer"
+    if not bundled_audio_tokenizer.is_dir():
+        raise FileNotFoundError(
+            "Bundled audio tokenizer not found at "
+            f"{bundled_audio_tokenizer}. The Breeze model package must include "
+            "the audio_tokenizer directory."
+        )
+    audio_tokenizer = Qwen3TTSTokenizer.from_pretrained(
+        str(bundled_audio_tokenizer), device_map=device
+    )
+    return tokenizer, model, audio_tokenizer

breeze_infer/templates.py

@@ -0,0 +1,339 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+
+from .audio import encode_prompt_audio
+
+AUDIO_TAG = "<|AUDIO|>"
+AUDIO_EOS = "<|audio_eos|>"
+INSTRUCTION_BOS = "<ins_bos>"
+INSTRUCTION_EOS = "<ins_eos>"
+
+
+Segment = dict[str, Any]
+Request = dict[str, Any]
+
+
+@dataclass(frozen=True)
+class TemplateSpec:
+    name: str
+    required_fields: tuple[str, ...]
+    build_segments: Callable[[Request], list[Segment]]
+    build_negative_segments: Callable[[Request], list[Segment]] | None = None
+    build_dual_branches: Callable[[Request], dict[str, list[Segment]]] | None = None
+
+
+def _speaker_prefix(request: Request) -> str:
+    speaker = request.get("speaker", "S0")
+    if speaker in (None, ""):
+        return ""
+    if isinstance(speaker, str) and speaker.startswith("[") and speaker.endswith("]"):
+        return speaker
+    return f"[{speaker}]"
+
+
+def _tts_plain_segments(request: Request) -> list[Segment]:
+    return [{"type": "text", "text": f"{_speaker_prefix(request)}{request['text']}"}]
+
+
+def _tts_instruction_segments(request: Request) -> list[Segment]:
+    prefix = _speaker_prefix(request)
+    return [
+        {
+            "type": "text",
+            "text": f"{prefix}{INSTRUCTION_BOS}{request['instruction']}{INSTRUCTION_EOS}{request['text']}",
+        }
+    ]
+
+
+def _tts_instruction_negative_segments(request: Request) -> list[Segment]:
+    return _tts_plain_segments(request)
+
+
+def _ref_audio_segment(
+    request: Request,
+    *,
+    append_eos: bool = True,
+    drop_last_frame: bool = False,
+) -> Segment:
+    segment: Segment = {
+        "type": "audio",
+        "append_eos": append_eos,
+        "drop_last_frame": drop_last_frame,
+    }
+    if request.get("ref_audio_path"):
+        segment["audio_path"] = request["ref_audio_path"]
+    return segment
+
+
+def _ref_clone_tata_segments(request: Request) -> list[Segment]:
+    prefix = _speaker_prefix(request)
+    return [
+        {"type": "text", "text": f"{prefix}{request['ref_text']}"},
+        _ref_audio_segment(request),
+        {"type": "text", "text": f"{prefix}{request['text']}"},
+    ]
+
+
+def _ref_edit_tata_segments(request: Request) -> list[Segment]:
+    prefix = _speaker_prefix(request)
+    return [
+        {"type": "text", "text": f"{prefix}{request['ref_text']}"},
+        _ref_audio_segment(request),
+        {
+            "type": "text",
+            "text": f"{prefix}{INSTRUCTION_BOS}{request['instruction']}{INSTRUCTION_EOS}{request['text']}",
+        },
+    ]
+
+
+def _ref_edit_tata_negative_segments(request: Request) -> list[Segment]:
+    return _ref_clone_tata_segments(request)
+
+
+def _ref_edit_tata_dual_branches(request: Request) -> dict[str, list[Segment]]:
+    prefix = _speaker_prefix(request)
+    return {
+        "uncond": [{"type": "text", "text": f"{prefix}{request['text']}"}],
+        "ref": _ref_clone_tata_segments(request),
+        "ins": _tts_instruction_segments(request),
+    }
+
+
+TEMPLATES: dict[str, TemplateSpec] = {
+    "tts_instruction": TemplateSpec(
+        name="tts_instruction",
+        required_fields=("text", "instruction"),
+        build_segments=_tts_instruction_segments,
+        build_negative_segments=_tts_instruction_negative_segments,
+    ),
+    "ref_edit_tata": TemplateSpec(
+        name="ref_edit_tata",
+        required_fields=("text", "instruction", "ref_audio_path", "ref_text"),
+        build_segments=_ref_edit_tata_segments,
+        build_negative_segments=_ref_edit_tata_negative_segments,
+        build_dual_branches=_ref_edit_tata_dual_branches,
+    ),
+}
+
+
+def get_template(name: str) -> TemplateSpec:
+    try:
+        return TEMPLATES[name]
+    except KeyError as exc:
+        raise KeyError(
+            f"Unknown template '{name}'. Available: {sorted(TEMPLATES)}"
+        ) from exc
+
+
+def _encode_prompt_audio(audio_tokenizer: Any, audio_path: str) -> torch.Tensor:
+    return encode_prompt_audio(audio_tokenizer, audio_path)
+
+
+def _resolve_segment_audio_codes(
+    audio_tokenizer: Any, segment: Segment
+) -> torch.Tensor:
+    audio_path = segment.get("audio_path")
+    if not audio_path:
+        raise ValueError("Audio segment must include audio_path")
+    return _encode_prompt_audio(audio_tokenizer, audio_path)
+
+
+def _prepare_one(
+    tokenizer: Any,
+    audio_tokenizer: Any,
+    model_config: Any,
+    segments: list[Segment],
+) -> dict[str, torch.Tensor]:
+    rendered_segments: list[dict[str, str]] = []
+    audio_tokens_list: list[torch.Tensor] = []
+
+    for segment in segments:
+        segment_type = segment["type"]
+        if segment_type == "text":
+            encoded = tokenizer(segment["text"], add_special_tokens=True)
+            rendered = tokenizer.decode(encoded["input_ids"], skip_special_tokens=False)
+            rendered_segments.append({"type": "text", "value": rendered})
+            continue
+
+        if segment_type != "audio":
+            raise ValueError(f"Unknown segment type: {segment_type}")
+
+        codes = _resolve_segment_audio_codes(audio_tokenizer, segment)
+        if segment.get("drop_last_frame", False):
+            if codes.shape[0] <= 1:
+                raise ValueError(
+                    "Cannot drop the last frame from an audio segment with <= 1 frame"
+                )
+            codes = codes[:-1].contiguous()
+        audio_tokens_list.append(codes)
+
+        placeholders = AUDIO_TAG * codes.shape[0]
+        if segment.get("append_eos", False):
+            placeholders += AUDIO_EOS
+        rendered_segments.append({"type": "audio", "value": placeholders})
+
+    final_text = "".join(segment["value"] for segment in rendered_segments)
+    encoded = tokenizer(final_text, add_special_tokens=False, return_tensors="pt")
+
+    text_ids_mask: list[bool] = []
+    text_ids_len: list[int] = []
+    for segment in rendered_segments:
+        segment_len = len(
+            tokenizer(segment["value"], add_special_tokens=False)["input_ids"]
+        )
+        if segment["type"] == "text":
+            text_ids_mask.extend([True] * segment_len)
+            text_ids_len.append(segment_len)
+        else:
+            text_ids_mask.extend([False] * segment_len)
+
+    num_codebooks = getattr(model_config, "num_codebooks", 16)
+    if audio_tokens_list:
+        audio_tokens = torch.cat(audio_tokens_list, dim=0).unsqueeze(0)
+    else:
+        audio_tokens = torch.zeros((1, 0, num_codebooks), dtype=torch.int16)
+
+    encoded["audio_tokens"] = audio_tokens
+    encoded["text_ids_mask"] = torch.tensor([text_ids_mask], dtype=torch.bool)
+    encoded["text_ids_len"] = torch.tensor(text_ids_len, dtype=torch.long)
+    return encoded
+
+
+def _collate_inputs(
+    tokenizer: Any, inputs_list: list[dict[str, torch.Tensor]], device: str
+) -> dict[str, torch.Tensor | None]:
+    pad_token_id = tokenizer.pad_token_id
+    if pad_token_id is None:
+        pad_token_id = tokenizer.eos_token_id
+    if pad_token_id is None:
+        raise ValueError("Tokenizer has no pad_token_id or eos_token_id")
+
+    audio_tokens = torch.cat([item["audio_tokens"] for item in inputs_list], dim=1)
+    max_len = max(item["input_ids"].shape[1] for item in inputs_list)
+
+    input_ids_list = []
+    attention_mask_list = []
+    text_ids_mask_list = []
+    for item in inputs_list:
+        input_ids = item["input_ids"]
+        attention_mask = item["attention_mask"]
+        text_ids_mask = item["text_ids_mask"]
+        pad_len = max_len - input_ids.shape[1]
+        if pad_len > 0:
+            input_ids = F.pad(input_ids, (pad_len, 0), value=pad_token_id)
+            attention_mask = F.pad(attention_mask, (pad_len, 0), value=0)
+            text_ids_mask = F.pad(text_ids_mask, (pad_len, 0), value=False)
+        input_ids_list.append(input_ids)
+        attention_mask_list.append(attention_mask)
+        text_ids_mask_list.append(text_ids_mask)
+
+    input_values = audio_tokens.to(device) if audio_tokens.shape[1] > 0 else None
+    return {
+        "input_ids": torch.cat(input_ids_list, dim=0).to(device),
+        "attention_mask": torch.cat(attention_mask_list, dim=0).to(device),
+        "text_ids_mask": torch.cat(text_ids_mask_list, dim=0).to(device),
+        "text_ids_len": torch.cat(
+            [item["text_ids_len"] for item in inputs_list], dim=0
+        ).to(device),
+        "input_values": input_values,
+    }
+
+
+def _prepare_segment_batches(
+    tokenizer: Any,
+    audio_tokenizer: Any,
+    model_config: Any,
+    device: str,
+    segment_batches: list[list[Segment]],
+) -> dict[str, torch.Tensor | None]:
+    inputs_list = [
+        _prepare_one(tokenizer, audio_tokenizer, model_config, segments)
+        for segments in segment_batches
+    ]
+    return _collate_inputs(tokenizer, inputs_list, device)
+
+
+def prepare_inputs(
+    tokenizer: Any,
+    audio_tokenizer: Any,
+    model: Any,
+    requests: list[Request],
+    template: TemplateSpec,
+    *,
+    guidance_scale: float,
+    guidance_scale_ref: float | None,
+    guidance_scale_ins: float | None,
+) -> dict[str, torch.Tensor | None | float]:
+    for request in requests:
+        missing = []
+        for field in template.required_fields:
+            if not request.get(field):
+                missing.append(field)
+        if missing:
+            raise ValueError(
+                f"Request {request.get('id')} missing template fields: {missing}"
+            )
+
+    positive_segments = [template.build_segments(request) for request in requests]
+    inputs = _prepare_segment_batches(
+        tokenizer,
+        audio_tokenizer,
+        model.config,
+        model.device,
+        positive_segments,
+    )
+
+    use_dual_cfg = (
+        guidance_scale_ref is not None
+        and guidance_scale_ins is not None
+        and template.build_dual_branches is not None
+    )
+    if use_dual_cfg:
+        branch_batches = [template.build_dual_branches(request) for request in requests]
+        for branch_name, prefix in [
+            ("uncond", "cfg_uncond"),
+            ("ref", "cfg_ref"),
+            ("ins", "cfg_ins"),
+        ]:
+            branch_inputs = _prepare_segment_batches(
+                tokenizer,
+                audio_tokenizer,
+                model.config,
+                model.device,
+                [branches[branch_name] for branches in branch_batches],
+            )
+            inputs[f"{prefix}_prompt_ids"] = branch_inputs["input_ids"]
+            inputs[f"{prefix}_prompt_attention_mask"] = branch_inputs["attention_mask"]
+            inputs[f"{prefix}_text_ids_mask"] = branch_inputs["text_ids_mask"]
+            inputs[f"{prefix}_text_ids_len"] = branch_inputs["text_ids_len"]
+        inputs["cfg_scale_ref"] = guidance_scale_ref
+        inputs["cfg_scale_ins"] = guidance_scale_ins
+        return inputs
+
+    if guidance_scale != 1.0:
+        if template.build_negative_segments is None:
+            raise ValueError(
+                f"Template '{template.name}' does not define a negative prompt but cfg_scale={guidance_scale}"
+            )
+        negative_inputs = _prepare_segment_batches(
+            tokenizer,
+            audio_tokenizer,
+            model.config,
+            model.device,
+            [template.build_negative_segments(request) for request in requests],
+        )
+        inputs["cfg_negative_prompt_ids"] = negative_inputs["input_ids"]
+        inputs["cfg_negative_prompt_attention_mask"] = negative_inputs["attention_mask"]
+        inputs["cfg_negative_text_ids_mask"] = negative_inputs["text_ids_mask"]
+        inputs["cfg_negative_text_ids_len"] = negative_inputs["text_ids_len"]
+        if negative_inputs["input_values"] is not None:
+            inputs["cfg_negative_input_values"] = negative_inputs["input_values"]
+        inputs["cfg_scale"] = guidance_scale
+
+    return inputs

configs/fast.json

@@ -0,0 +1,249 @@
+{
+  "schema_version": 1,
+  "name": "fast",
+  "service": {
+    "concurrency": 1,
+    "freeze_after_warmup": true,
+    "cfg_scales": [
+      1.0,
+      4.0
+    ]
+  },
+  "stages": {
+    "text_encoder": {
+      "graphs": [
+        {
+          "batch_size": 1,
+          "token_length": 32
+        },
+        {
+          "batch_size": 1,
+          "token_length": 64
+        },
+        {
+          "batch_size": 1,
+          "token_length": 96
+        },
+        {
+          "batch_size": 1,
+          "token_length": 128
+        },
+        {
+          "batch_size": 1,
+          "token_length": 160
+        },
+        {
+          "batch_size": 1,
+          "token_length": 192
+        },
+        {
+          "batch_size": 1,
+          "token_length": 224
+        },
+        {
+          "batch_size": 1,
+          "token_length": 256
+        },
+        {
+          "batch_size": 2,
+          "token_length": 32
+        },
+        {
+          "batch_size": 2,
+          "token_length": 64
+        },
+        {
+          "batch_size": 2,
+          "token_length": 96
+        },
+        {
+          "batch_size": 2,
+          "token_length": 128
+        },
+        {
+          "batch_size": 2,
+          "token_length": 160
+        },
+        {
+          "batch_size": 2,
+          "token_length": 192
+        },
+        {
+          "batch_size": 2,
+          "token_length": 224
+        },
+        {
+          "batch_size": 2,
+          "token_length": 256
+        },
+        {
+          "batch_size": 2,
+          "token_length": 288
+        },
+        {
+          "batch_size": 2,
+          "token_length": 320
+        },
+        {
+          "batch_size": 2,
+          "token_length": 352
+        },
+        {
+          "batch_size": 2,
+          "token_length": 384
+        },
+        {
+          "batch_size": 2,
+          "token_length": 416
+        },
+        {
+          "batch_size": 2,
+          "token_length": 448
+        },
+        {
+          "batch_size": 2,
+          "token_length": 480
+        },
+        {
+          "batch_size": 2,
+          "token_length": 512
+        }
+      ]
+    },
+    "backbone_prefill": {
+      "graphs": [
+        {
+          "branch_batch_size": 1,
+          "sequence_length": 32
+        },
+        {
+          "branch_batch_size": 1,
+          "sequence_length": 64
+        },
+        {
+          "branch_batch_size": 1,
+          "sequence_length": 96
+        },
+        {
+          "branch_batch_size": 1,
+          "sequence_length": 128
+        },
+        {
+          "branch_batch_size": 1,
+          "sequence_length": 160
+        },
+        {
+          "branch_batch_size": 1,
+          "sequence_length": 192
+        },
+        {
+          "branch_batch_size": 1,
+          "sequence_length": 224
+        },
+        {
+          "branch_batch_size": 1,
+          "sequence_length": 256
+        },
+        {
+          "branch_batch_size": 2,
+          "sequence_length": 32
+        },
+        {
+          "branch_batch_size": 2,
+          "sequence_length": 64
+        },
+        {
+          "branch_batch_size": 2,
+          "sequence_length": 96
+        },
+        {
+          "branch_batch_size": 2,
+          "sequence_length": 128
+        },
+        {
+          "branch_batch_size": 2,
+          "sequence_length": 160
+        },
+        {
+          "branch_batch_size": 2,
+          "sequence_length": 192
+        },
+        {
+          "branch_batch_size": 2,
+          "sequence_length": 224
+        },
+        {
+          "branch_batch_size": 2,
+          "sequence_length": 256
+        },
+        {
+          "branch_batch_size": 2,
+          "sequence_length": 288
+        },
+        {
+          "branch_batch_size": 2,
+          "sequence_length": 320
+        },
+        {
+          "branch_batch_size": 2,
+          "sequence_length": 352
+        },
+        {
+          "branch_batch_size": 2,
+          "sequence_length": 384
+        },
+        {
+          "branch_batch_size": 2,
+          "sequence_length": 416
+        },
+        {
+          "branch_batch_size": 2,
+          "sequence_length": 448
+        },
+        {
+          "branch_batch_size": 2,
+          "sequence_length": 480
+        },
+        {
+          "branch_batch_size": 2,
+          "sequence_length": 512
+        }
+      ]
+    },
+    "backbone_decode": {
+      "graphs": [
+        {
+          "branch_batch_size": 1
+        },
+        {
+          "branch_batch_size": 2
+        }
+      ]
+    },
+    "depth_decoder": {
+      "graphs": [
+        {
+          "batch_size": 1
+        },
+        {
+          "batch_size": 2
+        }
+      ]
+    },
+    "codec": {
+      "graphs": [
+        {
+          "num_lanes": 1,
+          "chunk_frames": 1
+        }
+      ]
+    }
+  },
+  "warmup_request": {
+    "template": "tts_instruction",
+    "text": "The cat sat on the mat.",
+    "instruction": "Speak naturally and clearly.",
+    "speaker": "S0",
+    "seed": 42
+  }
+}

models/breeze.py

@@ -0,0 +1,1888 @@
+# Breeze model implementation adapted from Hugging Face Transformers.
+#
+# coding=utf-8
+# Copyright 2025 Sesame and The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+
+from __future__ import annotations
+
+from collections.abc import Callable
+
+from dataclasses import dataclass
+
+import torch
+from torch import nn
+from transformers.activations import ACT2FN
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.generation import GenerationMixin
+from transformers.integrations import use_kernel_forward_from_hub
+from transformers.masking_utils import create_causal_mask
+from transformers.modeling_layers import GradientCheckpointingLayer
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPast,
+    CausalLMOutputWithPast,
+)
+from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from transformers.models.auto import AutoModel
+from transformers.processing_utils import Unpack
+from transformers.utils import (
+    ModelOutput,
+    TransformersKwargs,
+    auto_docstring,
+    can_return_tuple,
+    logging,
+)
+from transformers.utils.deprecation import deprecate_kwarg
+
+from .breeze_config import BreezeConfig, BreezeDepthDecoderConfig
+from .generation_breeze import BreezeGenerationMixin
+
+logger = logging.get_logger(__name__)
+
+
+@dataclass
+@auto_docstring(
+    custom_intro="""
+    Base class for the model autoregressive outputs.
+    """
+)
+class BreezeOutputWithPast(ModelOutput):
+    r"""
+    loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+        Language modeling loss (for next-token prediction).
+    logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
+        Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
+    past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+        Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape
+        `(batch_size, num_heads, sequence_length, embed_size_per_head)`)
+
+        Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
+        `past_key_values` input) to speed up sequential decoding.
+    depth_decoder_loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+        Language modeling loss (for next-token prediction) of the depth decoder model.
+    depth_decoder_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
+        Prediction scores of the depth decoder (scores for each vocabulary token before SoftMax).
+    depth_decoder_past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+        Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape
+        `(batch_size, num_heads, sequence_length, embed_size_per_head)`)
+    depth_decoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
+        Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
+        one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
+
+        Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
+    depth_decoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
+        Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
+        sequence_length)`.
+    backbone_loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+        Language modeling loss (for next-token prediction) of the backbone model.
+    """
+
+    loss: torch.FloatTensor | None = None
+    logits: torch.FloatTensor = None
+    past_key_values: tuple[tuple[torch.FloatTensor]] | None = None
+    hidden_states: tuple[torch.FloatTensor, ...] | None = None
+    attentions: tuple[torch.FloatTensor, ...] | None = None
+    depth_decoder_loss: torch.FloatTensor | None = None
+    depth_decoder_logits: torch.FloatTensor = None
+    depth_decoder_past_key_values: tuple[tuple[torch.FloatTensor]] | None = None
+    depth_decoder_hidden_states: tuple[torch.FloatTensor, ...] | None = None
+    depth_decoder_attentions: tuple[torch.FloatTensor, ...] | None = None
+    backbone_loss: torch.FloatTensor | None = None
+
+
+@use_kernel_forward_from_hub("RMSNorm")
+class BreezeRMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        """
+        BreezeRMSNorm is equivalent to T5LayerNorm
+        """
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+    def extra_repr(self):
+        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+class BreezeRotaryEmbedding(nn.Module):
+    inv_freq: torch.Tensor  # fix linting for `register_buffer`
+
+    def __init__(self, config: BreezeConfig, device=None):
+        super().__init__()
+        # BC: "rope_type" was originally "type"
+        if hasattr(config, "rope_scaling") and isinstance(config.rope_scaling, dict):
+            self.rope_type = config.rope_scaling.get(
+                "rope_type", config.rope_scaling.get("type")
+            )
+        else:
+            self.rope_type = "default"
+        self.max_seq_len_cached = config.max_position_embeddings
+        self.original_max_seq_len = config.max_position_embeddings
+
+        self.config = config
+        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+
+        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.original_inv_freq = self.inv_freq
+
+    @torch.no_grad()
+    @dynamic_rope_update  # power user: used with advanced RoPE types (e.g. dynamic rope)
+    def forward(self, x, position_ids):
+        inv_freq_expanded = (
+            self.inv_freq[None, :, None]
+            .float()
+            .expand(position_ids.shape[0], -1, 1)
+            .to(x.device)
+        )
+        position_ids_expanded = position_ids[:, None, :].float()
+
+        device_type = (
+            x.device.type
+            if isinstance(x.device.type, str) and x.device.type != "mps"
+            else "cpu"
+        )
+        with torch.autocast(device_type=device_type, enabled=False):  # Force float32
+            freqs = (
+                inv_freq_expanded.float() @ position_ids_expanded.float()
+            ).transpose(1, 2)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            cos = emb.cos() * self.attention_scaling
+            sin = emb.sin() * self.attention_scaling
+
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+class BreezeMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.gate_proj = nn.Linear(
+            self.hidden_size, self.intermediate_size, bias=config.mlp_bias
+        )
+        self.up_proj = nn.Linear(
+            self.hidden_size, self.intermediate_size, bias=config.mlp_bias
+        )
+        self.down_proj = nn.Linear(
+            self.intermediate_size, self.hidden_size, bias=config.mlp_bias
+        )
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, x):
+        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+        return down_proj
+
+
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
+    """Applies Rotary Position Embedding to the query and key tensors.
+
+    Args:
+        q (`torch.Tensor`): The query tensor.
+        k (`torch.Tensor`): The key tensor.
+        cos (`torch.Tensor`): The cosine part of the rotary embedding.
+        sin (`torch.Tensor`): The sine part of the rotary embedding.
+        position_ids (`torch.Tensor`, *optional*):
+            Deprecated and unused.
+        unsqueeze_dim (`int`, *optional*, defaults to 1):
+            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+    Returns:
+        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+    """
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """
+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+    """
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(
+        batch, num_key_value_heads, n_rep, slen, head_dim
+    )
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+def eager_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: torch.Tensor | None,
+    scaling: float,
+    dropout: float = 0.0,
+    **kwargs: Unpack[TransformersKwargs],
+):
+    key_states = repeat_kv(key, module.num_key_value_groups)
+    value_states = repeat_kv(value, module.num_key_value_groups)
+
+    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+    if attention_mask is not None:
+        causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
+        attn_weights = attn_weights + causal_mask
+
+    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(
+        query.dtype
+    )
+    attn_weights = nn.functional.dropout(
+        attn_weights, p=dropout, training=module.training
+    )
+    attn_output = torch.matmul(attn_weights, value_states)
+    attn_output = attn_output.transpose(1, 2).contiguous()
+
+    return attn_output, attn_weights
+
+
+class BreezeAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config: BreezeConfig, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.head_dim = getattr(
+            config, "head_dim", config.hidden_size // config.num_attention_heads
+        )
+        self.num_key_value_groups = (
+            config.num_attention_heads // config.num_key_value_heads
+        )
+        self.scaling = self.head_dim**-0.5
+        self.attention_dropout = config.attention_dropout
+        self.is_causal = True
+
+        self.q_proj = nn.Linear(
+            config.hidden_size,
+            config.num_attention_heads * self.head_dim,
+            bias=config.attention_bias,
+        )
+        self.k_proj = nn.Linear(
+            config.hidden_size,
+            config.num_key_value_heads * self.head_dim,
+            bias=config.attention_bias,
+        )
+        self.v_proj = nn.Linear(
+            config.hidden_size,
+            config.num_key_value_heads * self.head_dim,
+            bias=config.attention_bias,
+        )
+        self.o_proj = nn.Linear(
+            config.num_attention_heads * self.head_dim,
+            config.hidden_size,
+            bias=config.attention_bias,
+        )
+
+    @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: torch.Tensor | None,
+        past_key_values: Cache | None = None,
+        cache_position: torch.LongTensor | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+        cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(
+            query_states, key_states, cos, sin
+        )
+
+        if past_key_values is not None:
+            # sin and cos are specific to RoPE models; cache_position needed for the static cache
+            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+            key_states, value_states = past_key_values.update(
+                key_states, value_states, self.layer_idx, cache_kwargs
+            )
+
+        attention_interface: Callable = eager_attention_forward
+        if self.config._attn_implementation != "eager":
+            attention_interface = ALL_ATTENTION_FUNCTIONS[
+                self.config._attn_implementation
+            ]
+
+        attn_output, attn_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            **kwargs,
+        )
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+
+
+class BreezeDecoderLayer(GradientCheckpointingLayer):
+    def __init__(self, config: BreezeConfig, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+
+        self.self_attn = BreezeAttention(config=config, layer_idx=layer_idx)
+
+        self.mlp = BreezeMLP(config)
+        self.input_layernorm = BreezeRMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+        self.post_attention_layernorm = BreezeRMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+
+    @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        use_cache: bool | None = False,
+        cache_position: torch.LongTensor | None = None,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,  # necessary, but kept here for BC
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        # Self Attention
+        hidden_states, _ = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            position_embeddings=position_embeddings,
+            **kwargs,
+        )
+        hidden_states = residual + hidden_states
+
+        # Fully Connected
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+        return hidden_states
+
+
+@auto_docstring(
+    custom_intro="""
+    The bare Breeze Model outputting raw hidden-states without any specific head on top.
+    """
+)
+@auto_docstring
+class BreezePreTrainedModel(PreTrainedModel):
+    config_class = BreezeConfig
+    config: BreezeConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["BreezeDecoderLayer"]
+    _skip_keys_device_placement = ["past_key_values"]
+    _supports_flash_attn = True
+    _supports_sdpa = True
+    # does not because of Mimi codec model
+    # _supports_flex_attn = True
+
+    _can_compile_fullgraph = True
+    _supports_attention_backend = True
+    _can_record_outputs = {
+        "hidden_states": BreezeDecoderLayer,
+        "attentions": BreezeAttention,
+    }
+
+    def _init_weights(self, module):
+        super()._init_weights(module)
+        if isinstance(module, BreezeCodebooksHead):
+            num_codebooks = module.num_codebooks
+            for i in range(num_codebooks - 1):
+                module.weight.data[i].normal_(
+                    mean=0.0, std=self.config.initializer_range
+                )
+
+
+@auto_docstring
+class BreezeDepthDecoderModel(BreezePreTrainedModel):
+    config: BreezeDepthDecoderConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+        if hasattr(config, "audio_embed_size") and config.audio_embed_size:
+            self.audio_embed_size = config.audio_embed_size
+        else:
+            self.audio_embed_size = config.backbone_hidden_size
+        self.embed_tokens = nn.Embedding(
+            (config.num_codebooks * config.vocab_size), self.audio_embed_size
+        )
+        self.layers = nn.ModuleList(
+            [
+                BreezeDecoderLayer(config, layer_idx)
+                for layer_idx in range(config.num_hidden_layers)
+            ]
+        )
+        self.norm = BreezeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = BreezeRotaryEmbedding(config=config)
+        self.gradient_checkpointing = False
+        self.inputs_embeds_projector = nn.Linear(
+            self.audio_embed_size, config.hidden_size, bias=False
+        )
+
+        if config.backbone_hidden_size != self.audio_embed_size:
+            self.backbone_hidden_state_projector = nn.Linear(
+                config.backbone_hidden_size, self.audio_embed_size, bias=False
+            )
+        else:
+            self.backbone_hidden_state_projector = None
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        backbone_last_hidden_state: torch.FloatTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        use_cache: bool | None = None,
+        cache_position: torch.LongTensor | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple | BaseModelOutputWithPast:
+        r"""
+        backbone_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, backbone_hidden_size)`, *optional*):
+            The last hidden state of the backbone model. Such input is required when the first codebook token (the one generated by the backbone model)
+            is provided in the `input_ids` argument.
+        """
+        if position_ids is not None and not torch.compiler.is_compiling():
+            logger.warning_once(
+                "Custom `position_ids` were provided but will be ignored. Breeze depth decoder automatically determines position_ids "
+                "from `cache_position` and as it requires them to be identical across the batch, the provided position_ids will be ignored."
+            )
+            position_ids = None
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError(
+                "You must specify exactly one of input_ids or inputs_embeds."
+            )
+
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache(config=self.config)
+
+        if cache_position is None:
+            past_seen_tokens = (
+                past_key_values.get_seq_length() if past_key_values is not None else 0
+            )
+            inputs_seq_length = (
+                inputs_embeds.shape[1]
+                if inputs_embeds is not None
+                else input_ids.shape[1]
+            )
+            device = (
+                inputs_embeds.device if inputs_embeds is not None else input_ids.device
+            )
+            cache_position = torch.arange(
+                past_seen_tokens, past_seen_tokens + inputs_seq_length, device=device
+            )
+
+        if inputs_embeds is None:
+            codebook_idxs = torch.clamp(cache_position - 1, min=0)
+            offset = codebook_idxs * self.vocab_size
+            inputs_embeds = self.embed_tokens(input_ids + offset)
+
+            input_ids_are_first_codebook = cache_position[0] == 0
+            if backbone_last_hidden_state is not None:
+                if self.backbone_hidden_state_projector is not None:
+                    inputs_embeds[:, 0] = self.backbone_hidden_state_projector(
+                        backbone_last_hidden_state
+                    )
+                else:
+                    inputs_embeds[:, 0] = backbone_last_hidden_state
+            else:
+                if not torch.compiler.is_compiling() and input_ids_are_first_codebook:
+                    logger.warning(
+                        "When the first codebook token is provided, `backbone_last_hidden_state` should also be provided for correct inference."
+                    )
+
+        inputs_embeds = self.inputs_embeds_projector(inputs_embeds)
+
+        causal_mask = create_causal_mask(
+            config=self.config,
+            input_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            cache_position=cache_position,
+            past_key_values=past_key_values,
+            position_ids=position_ids,
+        )
+
+        hidden_states = inputs_embeds
+
+        # create position embeddings to be shared across the decoder layers
+        position_ids = cache_position.unsqueeze(0)
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+            hidden_states = decoder_layer(
+                hidden_states,
+                attention_mask=causal_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+                cache_position=cache_position,
+                position_embeddings=position_embeddings,
+                **kwargs,
+            )
+
+        hidden_states = self.norm(hidden_states)
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values if use_cache else None,
+        )
+
+
+class BreezeCodebooksHead(nn.Module):
+    def __init__(self, hidden_size, num_codebooks, vocab_size):
+        super().__init__()
+        self.num_codebooks = num_codebooks
+        self.weight = nn.Parameter(
+            torch.empty(self.num_codebooks - 1, hidden_size, vocab_size)
+        )
+
+    def forward(self, hidden_states, cache_position=None):
+        if cache_position is None:
+            seq_length = hidden_states.shape[1]
+            codebook_weight = self.weight[torch.arange(seq_length)]
+        else:
+            codebook_idxs = cache_position - 1
+            codebook_weight = self.weight[codebook_idxs]
+
+        hidden_states = [
+            nn.functional.linear(
+                hidden_states[:, codebook_idx, :], codebook_weight[codebook_idx].T
+            )
+            for codebook_idx in range(codebook_weight.shape[0])
+        ]
+        hidden_states = torch.stack(hidden_states, dim=1)
+
+        return hidden_states
+
+    # 重写 extra_repr 方法以显示参数细节
+    def extra_repr(self):
+        return f"weight_shape={list(self.weight.shape)}, num_codebooks={self.num_codebooks}"
+
+
+@auto_docstring(
+    custom_intro="""
+    The BreezeDepthDecoder Model transformer, with a [`BreezeCodebooksHead`] on top,
+    which can be seen a position-specific language modeling head, allowing to use a different linear layer for each codebook
+    (e.g. position 0 is the first codebook and uses the first codebook head, etc.)
+    """
+)
+class BreezeDepthDecoderForCausalLM(BreezePreTrainedModel, GenerationMixin):
+    _tied_weights_keys = None
+    _tp_plan = None
+    _pp_plan = None
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = BreezeDepthDecoderModel(config)
+        self.vocab_size = config.vocab_size
+        self.codebooks_head = BreezeCodebooksHead(
+            config.hidden_size, config.num_codebooks, config.vocab_size
+        )
+
+        # Compute per-codebook loss weights
+        # The depth decoder predicts num_codebooks-1 codebooks (31 for default config)
+        # We distribute weights by prepending a virtual codebook to make division even
+        codebook_loss_weights_list = getattr(config, "codebook_loss_weights", [1.0])
+        num_depth_codebooks = config.num_codebooks - 1  # 31 for default 32 codebooks
+        num_weights = len(codebook_loss_weights_list)
+        num_codebooks_with_virtual = num_depth_codebooks + 1
+        codebooks_per_weight = num_codebooks_with_virtual // num_weights
+        weights_expanded = []
+        for weight in codebook_loss_weights_list:
+            weights_expanded.extend([weight] * codebooks_per_weight)
+        weights_expanded = weights_expanded[1:]
+        self.codebook_loss_weights = weights_expanded
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @can_return_tuple
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        backbone_last_hidden_state: torch.FloatTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | list[torch.FloatTensor] | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        cache_position: torch.LongTensor | None = None,
+        logits_to_keep: int | torch.Tensor = 0,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple | CausalLMOutputWithPast:
+        r"""
+        backbone_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, backbone_hidden_size)`, *optional*):
+            The last hidden state of the backbone model. Such input is required when the first codebook token (the one generated by the backbone model)
+            is provided in the `input_ids` argument.
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+        """
+        outputs = self.model(
+            input_ids=input_ids,
+            backbone_last_hidden_state=backbone_last_hidden_state,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            **kwargs,
+        )
+
+        hidden_states = outputs[0]
+        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+        if isinstance(logits_to_keep, int):
+            if logits_to_keep == 0:
+                # skip idx 0 logits since it's for the concatenated backbone last hidden state
+                slice_indices = slice(1, None)
+            else:
+                slice_indices = slice(-logits_to_keep, None)
+        else:
+            slice_indices = logits_to_keep
+
+        loss = None
+        logits = None
+
+        logits = self.codebooks_head(
+            hidden_states[:, slice_indices, :],
+            cache_position[slice_indices] if cache_position is not None else None,
+        )
+        logits = logits.contiguous()
+
+        if labels is not None:
+            shift_labels = labels[..., 1:].contiguous()
+            loss = self.loss_function(
+                logits=logits,
+                labels=None,
+                vocab_size=self.config.vocab_size,
+                shift_labels=shift_labels,
+                **kwargs,
+            )
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: torch.LongTensor,
+        past_key_values: Cache | None = None,
+        attention_mask: torch.LongTensor | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        cache_position: torch.LongTensor | None = None,
+        **kwargs,
+    ):
+        model_inputs = super().prepare_inputs_for_generation(
+            input_ids,
+            past_key_values,
+            attention_mask,
+            inputs_embeds,
+            cache_position,
+            **kwargs,
+        )
+
+        is_first_generation_step = model_inputs["cache_position"][0] == 0
+        if not is_first_generation_step:
+            model_inputs.pop("backbone_last_hidden_state")
+
+        # breeze depth decoder does not use position_ids
+        model_inputs.pop("position_ids")
+
+        return model_inputs
+
+
+class BreezeBackboneModelEmbeddings(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        if hasattr(config, "audio_embed_size") and config.audio_embed_size:
+            self.audio_embed_size = config.audio_embed_size
+        else:
+            self.audio_embed_size = self.hidden_size
+        self.embed_audio_tokens = nn.Embedding(
+            (config.num_codebooks * config.vocab_size), self.hidden_size
+        )
+        if self.audio_embed_size != self.hidden_size:
+            self.audio_embeds_projector = nn.Linear(
+                self.audio_embed_size, self.hidden_size, bias=False
+            )
+        else:
+            self.audio_embeds_projector = None
+        self.register_buffer(
+            "audio_tokens_offsets",
+            torch.arange(config.num_codebooks) * config.vocab_size,
+            persistent=False,
+        )
+
+    def forward(self, input_ids):
+        input_embeds = self.embed_audio_tokens(input_ids + self.audio_tokens_offsets)
+        if self.audio_embeds_projector:
+            input_embeds = self.audio_embeds_projector(input_embeds)
+        input_embeds = input_embeds.sum(dim=2)
+        return input_embeds
+
+
+@auto_docstring
+class BreezeBackboneModel(BreezePreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+        self.embed_tokens = BreezeBackboneModelEmbeddings(config)
+        self.layers = nn.ModuleList(
+            [
+                BreezeDecoderLayer(config, layer_idx)
+                for layer_idx in range(config.num_hidden_layers)
+            ]
+        )
+        self.norm = BreezeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = BreezeRotaryEmbedding(config=config)
+        self.gradient_checkpointing = False
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        cache_position: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> BaseModelOutputWithPast:
+        r"""
+        input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length, num_codebooks) or (batch_size, sequence_length)`):
+            1. (batch_size, sequence_length): corresponds to the input sequence prepared with the processor from the text prompt. Such input
+            requires `input_values` to be provided so that audio can be encoded in codebook tokens and then merged with the text tokens.
+
+            2. (batch_size, sequence_length, num_codebooks): codebook tokens generated during the autoregressive decoding. Such input is not meant to be used by end users.
+
+            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+            [`PreTrainedTokenizer.__call__`] for details.
+
+            [What are input IDs?](../glossary#input-ids)
+        """
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError(
+                "You must specify exactly one of input_ids or inputs_embeds"
+            )
+
+        if inputs_embeds is None:
+            inputs_embeds: torch.Tensor = self.embed_tokens(input_ids)
+
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache(config=self.config)
+
+        if cache_position is None:
+            past_seen_tokens = (
+                past_key_values.get_seq_length() if past_key_values is not None else 0
+            )
+            cache_position: torch.Tensor = torch.arange(
+                past_seen_tokens,
+                past_seen_tokens + inputs_embeds.shape[1],
+                device=inputs_embeds.device,
+            )
+
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+
+        causal_mask = create_causal_mask(
+            config=self.config,
+            input_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            cache_position=cache_position,
+            past_key_values=past_key_values,
+            position_ids=position_ids,
+        )
+
+        hidden_states = inputs_embeds
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+            hidden_states = decoder_layer(
+                hidden_states,
+                attention_mask=causal_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                cache_position=cache_position,
+                position_embeddings=position_embeddings,
+                **kwargs,
+            )
+
+        hidden_states = self.norm(hidden_states)
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values,
+        )
+
+
+from transformers import AutoConfig, PretrainedConfig
+
+
+@auto_docstring(
+    custom_intro="""
+    The Breeze model consists of two llama-like auto-regressive transformer models: a backbone model that predicts the first codebook token and a depth decoder that predicts the other codebook tokens.
+    """
+)
+class BreezeForConditionalGeneration(BreezePreTrainedModel, BreezeGenerationMixin):
+    _tied_weights_keys = [
+        "backbone_model.embed_tokens.embed_audio_tokens.weight",
+        "depth_decoder.model.embed_tokens.weight",
+    ]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.vocab_size = config.vocab_size
+        # Extra backbone EOS class at index vocab_size; audio codebook ids remain [0, vocab_size).
+        self.backbone_eos_token_id = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size + 1, bias=False)
+
+        # embed_text_tokens needs to include audio special tokens (audio_token_id, audio_eos_token_id)
+        # since they appear in input_ids before being replaced with audio embeddings
+        # text_vocab_size already includes these tokens after tokenizer extension
+        self.embed_text_tokens = nn.Embedding(
+            config.text_vocab_size, config.hidden_size
+        )
+
+        # Use factory to create backbone (supports breeze/qwen3/llama etc.)
+        from .breeze_backbone_factory import BreezeBackboneFactory
+
+        self.backbone_model = BreezeBackboneFactory.create_backbone(config)
+
+        self.depth_decoder = BreezeDepthDecoderForCausalLM._from_config(
+            config.depth_decoder_config
+        )
+
+        self.codec_model = AutoModel.from_config(config.codec_config)
+        # Codec model is always in eval mode (not trainable)
+        self.codec_model.eval()
+        for param in self.codec_model.parameters():
+            param.requires_grad = False
+
+        self.text_encoder_trainable = False
+        text_encoder_config = getattr(config, "text_encoder_config", None)
+        if text_encoder_config is None:
+            self.text_encoder = None
+            self.text_encoder_proj = None
+        elif isinstance(text_encoder_config, dict):
+            config.text_encoder_config = AutoConfig.for_model(**text_encoder_config)
+        elif isinstance(config.text_encoder_config, PretrainedConfig):
+            config.text_encoder_config = text_encoder_config
+        else:
+            raise ValueError(
+                "text_encoder_config must be None, dict or PretrainedConfig instance"
+            )
+        if getattr(config, "text_encoder_config", None) is not None:
+            self.text_encoder_layer_projs = None
+            self.text_encoder_dimfusion_layer_start_idx = getattr(
+                config, "text_encoder_dimfusion_layer_start_idx", 1
+            )  # 1 feature indicates from layer 1, feature 0 is embedding result
+            self.text_encoder_dimfusion_layer_end_idx = getattr(
+                config, "text_encoder_dimfusion_layer_end_idx", None
+            )  # None indicates up to last layer
+            self.text_encoder_dimfusion_fuse_first_layer = getattr(
+                config, "text_encoder_dimfusion_fuse_first_layer", False
+            )
+
+            self.text_encoder_feature_layer_idx = getattr(
+                config, "text_encoder_feature_layer_idx", -1
+            )
+            if isinstance(self.text_encoder_feature_layer_idx, int):
+                self.text_encoder_feature_layer_idx = (
+                    self.text_encoder_feature_layer_idx,
+                )
+            if self.text_encoder_feature_layer_idx != (-1,):
+                print(
+                    f"  - [BreezeForConditionalGeneration] Text encoder feature layer idx: {self.text_encoder_feature_layer_idx}"
+                )
+
+            text_encoder_attn_implementation = getattr(
+                config.text_encoder_config,
+                "preferred_attn_implementation",
+                None,
+            )
+            if text_encoder_attn_implementation is None:
+                text_encoder_attn_implementation = "flash_attention_2"
+            config.text_encoder_config._attn_implementation = (
+                text_encoder_attn_implementation
+            )
+            # Skip random weight initialization since we'll load pretrained weights later
+            from transformers.modeling_utils import no_init_weights
+
+            with no_init_weights():
+                self.text_encoder = AutoModel.from_config(
+                    config.text_encoder_config,
+                    attn_implementation=text_encoder_attn_implementation,
+                    dtype=torch.bfloat16,
+                )
+
+            text_encoder_proj_type = getattr(config, "text_encoder_proj_type", "linear")
+            if text_encoder_proj_type == "linear":
+                self.text_encoder_proj = nn.Linear(
+                    config.text_encoder_config.hidden_size,
+                    config.hidden_size,
+                    bias=False,
+                )
+            elif text_encoder_proj_type == "mlp":
+                self.text_encoder_proj = nn.Sequential(
+                    nn.Linear(
+                        config.text_encoder_config.hidden_size,
+                        config.hidden_size * 2,
+                        bias=False,
+                    ),
+                    nn.GELU(),
+                    nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False),
+                )
+            elif text_encoder_proj_type == "breeze_text_encoder_adapter":
+                from .t5_adapter import BreezeTextEncoderAdapter
+
+                self.text_encoder_proj = BreezeTextEncoderAdapter(config=config)
+            elif text_encoder_proj_type == "breeze_dimfusion":
+                proj_out_hidden_size = config.hidden_size
+                if self.text_encoder_dimfusion_fuse_first_layer:
+                    proj_out_hidden_size //= 2
+                self.text_encoder_proj = nn.Linear(
+                    config.text_encoder_config.hidden_size
+                    * len(self.text_encoder_feature_layer_idx),
+                    proj_out_hidden_size,
+                    bias=True,  # bias=True, follow DimFusion design： https://github.com/huggingface/diffusers/blob/d0ae34d3136f7cb119a66fe0b9d526202343674c/src/diffusers/models/transformers/transformer_bria_fibo.py#L477
+                )
+                self.text_encoder_layer_projs = nn.ModuleList(
+                    [
+                        nn.Linear(
+                            config.text_encoder_config.hidden_size,
+                            config.hidden_size // 2,
+                            bias=False,
+                        )
+                        for _ in range(self.config.num_hidden_layers)
+                    ]
+                )
+
+                dimfusion_kwargs = {
+                    "text_encoder_dimfusion_layer_start_idx": self.text_encoder_dimfusion_layer_start_idx,
+                    "text_encoder_dimfusion_layer_end_idx": self.text_encoder_dimfusion_layer_end_idx,
+                    "text_encoder_dimfusion_fuse_first_layer": self.text_encoder_dimfusion_fuse_first_layer,
+                }
+                print(
+                    "  - [BreezeForConditionalGeneration] Text encoder DimFusion layer projections created."
+                )
+                for k, v in dimfusion_kwargs.items():
+                    print(f"    - {k}: {v}")
+            else:
+                raise ValueError(
+                    f"Unsupported text_encoder_proj_type: {text_encoder_proj_type}"
+                )
+            print(
+                "  - [BreezeForConditionalGeneration] Text encoder projection type:",
+                text_encoder_proj_type,
+            )
+
+            requires_grad = getattr(config.text_encoder_config, "requires_grad", False)
+            self.text_encoder_trainable = bool(requires_grad)
+            if not self.text_encoder_trainable:
+                nonzero_dropout_fields = self._get_nonzero_text_encoder_dropout_fields(
+                    config.text_encoder_config
+                )
+                if nonzero_dropout_fields:
+                    formatted_fields = ", ".join(
+                        f"{name}={value}" for name, value in nonzero_dropout_fields
+                    )
+                    raise ValueError(
+                        "Frozen text_encoder must be deterministic, but its config has non-zero dropout fields: "
+                        f"{formatted_fields}. Set train_text_encoder=true or use a text encoder config with dropout=0."
+                    )
+                self.text_encoder.eval()
+            for param in self.text_encoder.parameters():
+                param.requires_grad = requires_grad
+            import os
+
+            if os.environ.get("RANK", "0") == "0":
+                logger.info(
+                    f"model.text_encoder initialized with requires_grad={requires_grad}"
+                )
+
+        # Disable num_items_in_batch to avoid incorrect loss normalization
+        # for depth decoder (which operates on different token counts than backbone)
+        self.accepts_loss_kwargs = False
+
+        # Temporarily remove text_encoder before post_init to avoid slow re-initialization
+        # (text_encoder weights will be loaded from pretrained checkpoint later)
+        _text_encoder = self.text_encoder
+        self.text_encoder = None
+        self.post_init()
+        self.text_encoder = _text_encoder
+
+    def get_input_embeddings(self):
+        return self.backbone_model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.backbone_model.embed_tokens = value
+
+    def _tie_weights(self):
+        if self.config.tie_codebooks_embeddings:
+            self._tie_or_clone_weights(
+                self.backbone_model.embed_tokens.embed_audio_tokens,
+                self.depth_decoder.model.embed_tokens,
+            )
+
+    @staticmethod
+    def _get_nonzero_text_encoder_dropout_fields(text_encoder_config):
+        dropout_field_names = (
+            "dropout",
+            "dropout_rate",
+            "attention_dropout",
+            "hidden_dropout",
+            "hidden_dropout_prob",
+            "activation_dropout",
+            "classifier_dropout",
+            "classifier_dropout_rate",
+            "embd_pdrop",
+            "resid_pdrop",
+        )
+        nonzero_fields = []
+        for field_name in dropout_field_names:
+            value = getattr(text_encoder_config, field_name, None)
+            if value is None:
+                continue
+            try:
+                numeric_value = float(value)
+            except (TypeError, ValueError):
+                continue
+            if numeric_value != 0.0:
+                nonzero_fields.append((field_name, value))
+        return nonzero_fields
+
+    def train(self, mode: bool = True):
+        """Override train() to keep non-trainable submodules in eval mode."""
+        super().train(mode)
+        # Always keep codec_model in eval mode
+        self.codec_model.eval()
+        if getattr(self, "text_encoder", None) is not None and not getattr(
+            self, "text_encoder_trainable", False
+        ):
+            self.text_encoder.eval()
+        return self
+
+    @classmethod
+    def from_pretrained(cls, *args, **kwargs):
+        if kwargs.get("output_loading_info", False):
+            model, loading_info = super().from_pretrained(*args, **kwargs)
+        else:
+            model = super().from_pretrained(*args, **kwargs)
+
+        # copy depth decoder generation conf attr to the depth decoder generation config
+        prefix = "depth_decoder_"
+        prefix_len = len(prefix)
+        depth_decoder_attrs = {
+            attr[prefix_len:]: value
+            for attr, value in vars(model.generation_config).items()
+            if attr.startswith(prefix)
+        }
+
+        vars(model.depth_decoder.generation_config).update(
+            {"_from_model_config": False, **depth_decoder_attrs}
+        )
+
+        # remove the depth decoder generation conf attr from the model generation config
+        for attr in depth_decoder_attrs:
+            delattr(model.generation_config, prefix + attr)
+
+        if "output_loading_info" in kwargs:
+            return model, loading_info
+        else:
+            return model
+
+    def save_pretrained(self, *args, **kwargs):
+        # copy the depth decoder generation config attributes to the model generation config
+        prefix = "depth_decoder_"
+        depth_decoder_attrs = self.depth_decoder.generation_config.to_diff_dict()
+        depth_decoder_attrs.pop("transformers_version", None)
+        for attr, value in depth_decoder_attrs.items():
+            setattr(self.generation_config, prefix + attr, value)
+
+        super().save_pretrained(*args, **kwargs)
+
+    def gradient_checkpointing_enable(self, *args, **kwargs):
+        gradient_checkpointing_kwargs = kwargs.pop("gradient_checkpointing_kwargs", {})
+        sub_modules_to_disable_default = [
+            # 'codec_model',
+            # # 'backbone_model',
+            # 'text_encoder',
+            # 'text_encoder_proj',
+            # 'depth_decoder'
+        ]
+
+        if gradient_checkpointing_kwargs is None:
+            gradient_checkpointing_kwargs = {}
+        sub_modules_to_disable = gradient_checkpointing_kwargs.pop(
+            "sub_modules_to_disable", sub_modules_to_disable_default
+        )
+
+        # call super first to enable gradient checkpointing
+        ret = super().gradient_checkpointing_enable(*args, **kwargs)
+
+        # override sub-modules to disable gradient checkpointing
+        sub_modules_to_disable = gradient_checkpointing_kwargs.get(
+            "sub_modules_to_disable", sub_modules_to_disable
+        )
+
+        # 1. disable gradient checkpointing for specified sub-modules
+        import os
+
+        rank = int(os.environ.get("RANK", "0"))
+        print0 = lambda *args, **kwargs: print(*args, **kwargs) if rank == 0 else None
+        for module_name in sub_modules_to_disable:
+            module = getattr(self, module_name, None)
+            module_layers = (
+                [
+                    (name, mod)
+                    for name, mod in module.named_modules()
+                    if isinstance(mod, GradientCheckpointingLayer)
+                ]
+                if module is not None
+                else []
+            )
+
+            if not module_layers:
+                print0(
+                    f"  - [BreezeForConditionalGeneration] gradient_checkpointing_disable skipped for sub-module: {module_name}"
+                )
+                continue
+
+            print0(
+                f"  - [BreezeForConditionalGeneration] gradient_checkpointing_disable applied for sub-module: {module_name}"
+            )
+            for layer_name, layer in module_layers:
+                layer.gradient_checkpointing = False
+                print0(
+                    f"    - Disabled gradient checkpointing for layer: {module_name}.{layer_name}"
+                )
+        return ret
+
+    def _batched_text_encoder_forward(self, segments, output_hidden_states=False):
+        """Run text encoder on a list of variable-length token ID tensors using padded batching.
+
+        Args:
+            segments: list of 1D tensors, each (seg_len,)
+
+        Returns:
+            hidden_states: list of 2D tensors, each (seg_len, hidden) with padding removed
+            layer_hidden_states: list of lists, per-segment layer hidden states (each (seg_len, hidden)),
+                or None when not requested.
+        """
+        if not segments:
+            return [], []
+
+        if (
+            getattr(self, "_fast_text_encoder_cudagraph", False)
+            and not output_hidden_states
+        ):
+            from .text_encoder_graph import TextEncoderGraphCache
+
+            cache = getattr(self, "_fast_text_encoder_graph_cache", None)
+            if cache is None:
+                cache = TextEncoderGraphCache(self.text_encoder, token_granularity=32)
+                self._fast_text_encoder_graph_cache = cache
+            return cache(segments)
+
+        device = segments[0].device
+        lengths = [s.shape[0] for s in segments]
+        feature_layer_idx = self.text_encoder_feature_layer_idx
+        if isinstance(feature_layer_idx, int):
+            feature_layer_idx = (feature_layer_idx,)
+        elif isinstance(feature_layer_idx, list):
+            feature_layer_idx = tuple(feature_layer_idx)
+
+        sorted_indices = sorted(range(len(segments)), key=lambda i: lengths[i])
+        buckets = []
+        current_bucket = []
+        current_min_len = None
+        for idx in sorted_indices:
+            length = lengths[idx]
+            if not current_bucket:
+                current_bucket = [idx]
+                current_min_len = length
+            elif length / max(current_min_len, 1) <= 2:
+                current_bucket.append(idx)
+            else:
+                buckets.append(current_bucket)
+                current_bucket = [idx]
+                current_min_len = length
+        if current_bucket:
+            buckets.append(current_bucket)
+
+        hidden_states = [None] * len(segments)
+        layer_hidden_states = [None] * len(segments) if output_hidden_states else None
+
+        for bucket_indices in buckets:
+            bucket_lengths = [lengths[i] for i in bucket_indices]
+            max_len = max(bucket_lengths)
+
+            padded_ids = torch.zeros(
+                len(bucket_indices), max_len, dtype=segments[0].dtype, device=device
+            )
+            attn_mask = torch.zeros(
+                len(bucket_indices), max_len, dtype=torch.long, device=device
+            )
+            pos_ids = torch.zeros(
+                len(bucket_indices), max_len, dtype=torch.long, device=device
+            )
+            for bucket_pos, idx in enumerate(bucket_indices):
+                length = lengths[idx]
+                padded_ids[bucket_pos, :length] = segments[idx]
+                attn_mask[bucket_pos, :length] = 1
+                pos_ids[bucket_pos, :length] = torch.arange(length, device=device)
+
+            output = self.text_encoder(
+                input_ids=padded_ids,
+                attention_mask=attn_mask,
+                position_ids=pos_ids,
+                output_hidden_states=output_hidden_states,
+            )
+
+            if feature_layer_idx == (-1,):
+                full_hs = output.last_hidden_state
+            elif isinstance(feature_layer_idx, tuple):
+                if output.hidden_states is None:
+                    raise ValueError(
+                        "output_hidden_states must be enabled when selecting non-final text encoder layers"
+                    )
+                full_hs = torch.concat(
+                    [output.hidden_states[li] for li in feature_layer_idx], dim=-1
+                )
+            else:
+                raise ValueError(
+                    f"Unsupported text_encoder_feature_layer_idx: {feature_layer_idx}"
+                )
+
+            for bucket_pos, idx in enumerate(bucket_indices):
+                hidden_states[idx] = full_hs[bucket_pos, : lengths[idx]]
+            if output.hidden_states is not None:
+                all_layer_hs = list(output.hidden_states)
+                if layer_hidden_states is None:
+                    layer_hidden_states = [None] * len(segments)
+                for bucket_pos, idx in enumerate(bucket_indices):
+                    layer_hidden_states[idx] = [
+                        lhs[bucket_pos, : lengths[idx]] for lhs in all_layer_hs
+                    ]
+
+        return hidden_states, layer_hidden_states
+
+    def convert_input_ids_to_embeds(
+        self,
+        input_ids: torch.Tensor | None = None,
+        text_ids_mask: torch.Tensor | None = None,
+        text_ids_len: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+    ):
+        assert self.text_encoder is not None, (
+            "text_encoder is None, cannot convert input_ids to embeds"
+        )
+        assert text_ids_mask is not None, (
+            "text_ids_mask is None, cannot convert input_ids to embeds"
+        )
+        assert text_ids_len is not None, (
+            "text_ids_len is None, cannot convert input_ids to embeds"
+        )
+
+        # Ensure text_ids_len matches mask-derived count (helps catch bad batching)
+        total_text_tokens = int(text_ids_mask.sum().item())
+        assert total_text_tokens == int(text_ids_len.sum().item()), (
+            f"text_ids_mask sum {total_text_tokens} does not match text_ids_len sum {int(text_ids_len.sum().item())}"
+        )
+
+        text_ids_len_list = [int(x) for x in text_ids_len.reshape(-1).tolist()]
+        text_ids_len_idx = 0
+
+        # == Phase 1: Parse all samples and collect segments ==
+        # Keep row-major order so the projected embeddings line up with
+        # boolean assignment into inputs_embeds[text_ids_mask].
+        all_segment_lengths = []
+        all_seg_token_ids = []
+        for batch_idx in range(input_ids.shape[0]):
+            mask_row = text_ids_mask[batch_idx]
+            sample_text_tokens = int(mask_row.sum().item())
+            if sample_text_tokens == 0:
+                continue
+
+            text_ids = input_ids[batch_idx][mask_row]  # (total_text_len,)
+            segment_lengths = []
+            running = 0
+            while running < sample_text_tokens:
+                if text_ids_len_idx >= len(text_ids_len_list):
+                    raise ValueError(
+                        "text_ids_len exhausted before covering all text tokens in the batch"
+                    )
+                seg_len = text_ids_len_list[text_ids_len_idx]
+                text_ids_len_idx += 1
+                if seg_len <= 0:
+                    continue
+                segment_lengths.append(seg_len)
+                running += seg_len
+            if running != sample_text_tokens:
+                raise ValueError(
+                    f"text_ids_len segments sum {running} does not match text_ids_mask sum {sample_text_tokens} for batch index {batch_idx}"
+                )
+
+            seg_token_ids = text_ids.split(segment_lengths, dim=0)  # list of 1D tensors
+            all_segment_lengths.extend(segment_lengths)
+            all_seg_token_ids.extend(seg_token_ids)
+
+        if text_ids_len_idx != len(text_ids_len_list):
+            raise ValueError(
+                f"Unused text_ids_len entries detected: consumed {text_ids_len_idx} of {len(text_ids_len_list)}"
+            )
+
+        # == Phase 2: Text encoder forward ==
+        # Batched independent encoding: each text segment is its own row in
+        # the padded text-encoder batch. This prevents cross-sample and
+        # cross-segment attention contamination under the flattening collator.
+        if all_seg_token_ids:
+            feature_layer_idx = self.text_encoder_feature_layer_idx
+            if isinstance(feature_layer_idx, int):
+                feature_layer_idx = (feature_layer_idx,)
+            elif isinstance(feature_layer_idx, list):
+                feature_layer_idx = tuple(feature_layer_idx)
+            needs_layer_hidden_states = (
+                self.text_encoder_layer_projs is not None or feature_layer_idx != (-1,)
+            )
+            seg_hidden_states, seg_layer_hidden_states = (
+                self._batched_text_encoder_forward(
+                    all_seg_token_ids,
+                    output_hidden_states=needs_layer_hidden_states,
+                )
+            )
+            seg_hidden_states = [hs.unsqueeze(0) for hs in seg_hidden_states]
+            text_embeds, text_encoder_layer_hidden_states = self._project_segments(
+                all_segment_lengths,
+                seg_hidden_states,
+                seg_layer_hidden_states,
+                is_separate=True,
+            )
+        else:
+            text_embeds = torch.empty(
+                (0, self.config.hidden_size),
+                device=input_ids.device,
+                dtype=self.embed_text_tokens.weight.dtype,
+            )
+            text_encoder_layer_hidden_states = None
+
+        # == Phase 3: Assemble final outputs ==
+        # Create inputs_embeds tensor
+        inputs_embeds = torch.zeros(
+            (input_ids.shape[0], input_ids.shape[1], self.config.hidden_size),
+            dtype=text_embeds.dtype,
+            device=text_embeds.device,
+        )
+        inputs_embeds[text_ids_mask] = text_embeds
+
+        return inputs_embeds, text_encoder_layer_hidden_states
+
+    def _project_segments(
+        self, segment_lengths, seg_hidden_states, seg_layer_hidden_states, is_separate
+    ):
+        """Project text encoder hidden states per segment and handle layer projections.
+
+        Args:
+            segment_lengths: list of int, length of each segment
+            seg_hidden_states: list of (1, seg_len, hidden) tensors
+            seg_layer_hidden_states: list of layer hidden state lists per segment.
+                - If is_separate=True: each entry is a list of (seg_len, hidden) tensors (one per layer).
+                - If is_separate=False: each entry is a list of (1, total_len, hidden) tensors (shared across segments).
+            is_separate: whether segments were encoded separately
+
+        Returns:
+            text_embeds: (total_text_len, hidden) projected embeddings
+            text_encoder_layer_hidden_states: list of (total_text_len, hidden) or None
+        """
+        projected_segments = []
+        for seg_idx, seg_hs in enumerate(seg_hidden_states):
+            proj = self.text_encoder_proj
+            if isinstance(proj, (nn.Linear, nn.Sequential)):
+                projected_segments.append(proj(seg_hs).squeeze(0))
+            else:
+                seg_pos = torch.arange(seg_hs.shape[1], device=seg_hs.device).unsqueeze(
+                    0
+                )
+                projected_segments.append(
+                    proj(inputs_embeds=seg_hs, position_ids=seg_pos).squeeze(0)
+                )
+        text_embeds = torch.cat(projected_segments, dim=0)
+
+        if self.text_encoder_layer_projs is not None:
+            num_layers = len(self.text_encoder_layer_projs)
+            projected_layer_hidden_states = []
+
+            for layer_idx in range(num_layers):
+                layer_parts = []
+                for seg_idx, seg_layer_hs_list in enumerate(seg_layer_hidden_states):
+                    selected = seg_layer_hs_list[
+                        self.text_encoder_dimfusion_layer_start_idx : self.text_encoder_dimfusion_layer_end_idx
+                    ]
+                    if len(selected) < num_layers:
+                        selected = selected + [selected[-1]] * (
+                            num_layers - len(selected)
+                        )
+                    elif len(selected) > num_layers:
+                        selected = selected[-num_layers:]
+
+                    seg_layer_hs = selected[layer_idx]
+                    if is_separate:
+                        # Already per-segment: (seg_len, hidden) -> (1, seg_len, hidden)
+                        seg_layer_hs = seg_layer_hs.unsqueeze(0)
+                    else:
+                        # Joint encoding: (1, total_len, hidden) -> split -> (1, seg_len, hidden)
+                        seg_layer_hs = (
+                            seg_layer_hs.squeeze(0)
+                            .split(segment_lengths, dim=0)[seg_idx]
+                            .unsqueeze(0)
+                        )
+
+                    layer_proj = self.text_encoder_layer_projs[layer_idx]
+                    layer_parts.append(layer_proj(seg_layer_hs).squeeze(0))
+                projected_layer_hidden_states.append(torch.cat(layer_parts, dim=0))
+
+            text_encoder_layer_hidden_states = projected_layer_hidden_states
+        else:
+            text_encoder_layer_hidden_states = None
+
+        if self.text_encoder_dimfusion_fuse_first_layer:
+            assert text_encoder_layer_hidden_states is not None, (
+                "text_encoder_layer_hidden_states is None, cannot fuse first layer hidden states"
+            )
+            layer0_hidden_states = text_encoder_layer_hidden_states[0]
+            assert (
+                layer0_hidden_states.shape[-1] + text_embeds.shape[-1]
+            ) == self.config.hidden_size, (
+                f"Dimension mismatch when fusing layer 0 hidden states with text embeddings, but got {layer0_hidden_states.shape[-1]} + {text_embeds.shape[-1]} != {self.config.hidden_size}"
+            )
+            text_embeds = torch.cat([text_embeds, layer0_hidden_states], dim=-1)
+
+        return text_embeds, text_encoder_layer_hidden_states
+
+    def _merge_input_ids_with_input_values(
+        self,
+        input_ids: torch.Tensor | None = None,
+        input_values: torch.Tensor | None = None,
+        labels: torch.Tensor | None = None,
+        text_ids_mask: torch.Tensor | None = None,
+        text_ids_len: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor | None:
+        """
+        Merges the input_ids and input_values to produce a single inputs_embeds tensor:
+        1 - Uses pre-encoded codebook tokens from input_values (no codec inference needed).
+        2 - Embeds codebook tokens and places them at the correct positions in the inputs_embeds tensor.
+        3 - If labels are provided, expands them to match codebook dimensions and position the target codebook tokens in the inputs_embeds tensor.
+
+        Args:
+            input_ids (`torch.Tensor` of shape `(batch_size, sequence_length)`):
+                The input ids to embed.
+            input_values (`torch.Tensor` of shape `(batch_size, audio_sequence_length, num_codebooks)`):
+                The pre-encoded audio tokens (RVQ codebook indices).
+            labels (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+                Labels for computing the loss.
+        """
+        text_encoder_layer_hidden_states = None
+        if self.text_encoder is not None:
+            inputs_embeds, text_encoder_layer_hidden_states = (
+                self.convert_input_ids_to_embeds(
+                    input_ids,
+                    text_ids_mask,
+                    text_ids_len,
+                    attention_mask=attention_mask,
+                )
+            )
+        else:
+            inputs_embeds = self.embed_text_tokens(input_ids)
+
+        if input_values is not None:
+            # input_values are already pre-encoded tokens: (batch_size, audio_seq_len, num_codebooks)
+            batched_audio_token_ids = input_values
+
+            # Get mask for audio token positions in input_ids
+            audio_token_id = self.config.audio_token_id
+            audio_token_mask = input_ids == audio_token_id
+
+            # Debug: check if audio tokens are found
+            num_audio_tokens = audio_token_mask.sum().item()
+            if num_audio_tokens == 0 and input_values is not None:
+                # This should not happen - we have audio data but no audio token markers
+                unique_ids = torch.unique(input_ids)
+                raise RuntimeError(
+                    f"Audio token mismatch: Expected audio_token_id={audio_token_id} in input_ids, "
+                    f"but found 0 matches. input_values shape: {input_values.shape}. "
+                    f"Unique input_ids: {unique_ids.tolist()[:20]}... (showing first 20)"
+                )
+
+            # Convert patches tokens to embeddings
+            audio_embeds = self.backbone_model.embed_tokens(batched_audio_token_ids)
+
+            # Flatten audio_embeds for assignment to masked positions
+            # audio_embeds shape: (batch_size, audio_seq_len, hidden_size)
+            # Need to flatten to match the flattened audio_token_mask
+            audio_embeds_flat = audio_embeds.reshape(-1, audio_embeds.shape[-1])
+            # Ensure dtype matches
+            audio_embeds_flat = audio_embeds_flat.to(inputs_embeds.dtype)
+            inputs_embeds[audio_token_mask] = audio_embeds_flat
+
+            # same for the audio eos token.
+            audio_eos_frame_ids = (
+                torch.ones(
+                    (1, 1, self.config.num_codebooks),
+                    device=input_ids.device,
+                    dtype=torch.long,
+                )
+                * self.config.codebook_eos_token_id
+            )
+            audio_eos_embeds = self.backbone_model.embed_tokens(
+                audio_eos_frame_ids
+            ).squeeze(1)
+            # Ensure dtype matches
+            audio_eos_embeds = audio_eos_embeds.to(inputs_embeds.dtype)
+
+            audio_eos_token_mask = input_ids == self.config.audio_eos_token_id
+            inputs_embeds[audio_eos_token_mask] = audio_eos_embeds.repeat(
+                audio_eos_token_mask.sum(), 1
+            )
+
+            # if the labels are provided, we need to expand the labels to (batch_size, seq_length, num_codebooks)
+            if labels is not None:
+                labels_expanded = labels.unsqueeze(-1).repeat(
+                    1, 1, self.config.num_codebooks
+                )
+                # Flatten batched_audio_token_ids for assignment to masked positions
+                batched_audio_token_ids_flat = batched_audio_token_ids.reshape(
+                    -1, batched_audio_token_ids.shape[-1]
+                )
+                labels_expanded[audio_token_mask] = batched_audio_token_ids_flat
+                labels_expanded[audio_eos_token_mask] = -100
+                eos_positions = audio_eos_token_mask.nonzero(as_tuple=True)
+                labels_expanded[eos_positions[0], eos_positions[1], 0] = (
+                    self.backbone_eos_token_id
+                )
+                # mask depth decoder
+                depth_decoder_ignore_frames_mask = labels == -101
+                depth_decoder_ignore_frames_mask &= ~audio_eos_token_mask
+                depth_decoder_ignore_frames_idxs = (
+                    depth_decoder_ignore_frames_mask.nonzero(as_tuple=True)
+                )
+                # set the first codebook patch to -100 to ignore
+                labels_expanded[
+                    depth_decoder_ignore_frames_idxs[0],
+                    depth_decoder_ignore_frames_idxs[1],
+                    1 : self.config.num_codebooks,
+                ] = -100
+                labels = labels_expanded
+
+        return {
+            "inputs_embeds": inputs_embeds,
+            "labels": labels,
+            "text_encoder_layer_hidden_states": text_encoder_layer_hidden_states,
+            "text_ids_mask": text_ids_mask,
+        }
+
+    def _get_audio_token_from_batch(self, audio_batch):
+        with torch.no_grad():
+            print("Encoding audio batch of shape:", audio_batch.shape)
+            codec_outputs = self.codec_model.encode(audio_batch.unsqueeze(0))
+            codebook_ids = codec_outputs.audio_codes.transpose(1, -1)[0]
+        return codebook_ids
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: torch.LongTensor,
+        past_key_values: Cache | None = None,
+        attention_mask: torch.LongTensor | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        cache_position: torch.LongTensor | None = None,
+        **kwargs,
+    ):
+        model_inputs = super().prepare_inputs_for_generation(
+            input_ids=input_ids,
+            past_key_values=past_key_values,
+            attention_mask=attention_mask,
+            inputs_embeds=inputs_embeds,
+            cache_position=cache_position,
+            **kwargs,
+        )
+
+        # if input_ids is not None and input_ids.ndim == 2 and model_inputs.get("inputs_embeds") is None:
+        #     print("------- Merging input_ids with input_values for generation...")
+        #     merged_inputs = self._merge_input_ids_with_input_values(
+        #         input_ids=input_ids,
+        #         input_values=kwargs.get("input_values"),
+        #         labels=kwargs.get("labels"),
+        #     )
+        #     model_inputs.update(
+        #         {"inputs_embeds": merged_inputs["inputs_embeds"], "labels": merged_inputs["labels"], "input_ids": None}
+        #     )
+
+        # cache_position = model_inputs["cache_position"]
+        # is_prefill = cache_position is not None and cache_position[0] == 0
+        # if is_prefill:
+        #     pass
+        # else:
+        #     # Remove text_encoder_layer_hidden_states to indicate no fusion needed during decoding
+        #     model_inputs.pop("text_encoder_layer_hidden_states", None)
+
+        return model_inputs
+
+    @can_return_tuple
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        input_values: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | list[torch.FloatTensor] | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        cache_position: torch.LongTensor | None = None,
+        logits_to_keep: int | torch.Tensor = 0,
+        text_ids_mask: torch.Tensor | None = None,
+        text_ids_len: torch.Tensor | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple | BreezeOutputWithPast:
+        r"""
+        input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length, num_codebooks) or (batch_size, sequence_length)`):
+            1. (batch_size, sequence_length): corresponds to the input sequence prepared with the processor from the text prompt. Such input
+            requires `input_values` to be provided so that audio can be encoded in codebook tokens and then merged with the text tokens.
+
+            2. (batch_size, sequence_length, num_codebooks): codebook tokens generated during the autoregressive decoding. Such input is not meant to be used by end users.
+
+            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+            [`PreTrainedTokenizer.__call__`] for details.
+
+            [What are input IDs?](../glossary#input-ids)
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should be in `[config.audio_token_id, -100, -101]`.
+            Requires targeted `input_values` to be provided as audio tokens will be inferred from it using the `codec_model`.
+            - `config.audio_token_id` indicates an audio frames (considering sequence length elements as frames)
+            - `-100` will be ignored in the loss computation
+            - `-101` indicates the audio frame will be used only for the backbone model (using the first codebook token as labels)
+
+            Such labels can be prepared using `output_labels=True` when calling [`the repository input preparation utilities`].
+        logits_to_keep (`int` or `torch.Tensor`, *optional*):
+            Kept for compatibility. Does not support another value than:
+            1. `0`, which is equivalent to keeping all logits, used in the training regime
+            2. `1`, which is equivalent to keeping only the last logit, used in the generation regime
+
+        text_ids_mask (`torch.BoolTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Marks text-token positions consumed by the text encoder.
+        text_ids_len (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+            Number of valid text tokens for each item in the batch.
+        """
+        text_encoder_layer_hidden_states = None
+        if input_ids is not None and input_ids.ndim == 2:
+            merged_inputs = self._merge_input_ids_with_input_values(
+                input_ids,
+                input_values,
+                labels,
+                text_ids_mask=text_ids_mask,
+                text_ids_len=text_ids_len,
+                attention_mask=attention_mask,
+            )
+            inputs_embeds = merged_inputs["inputs_embeds"]
+            labels = merged_inputs["labels"]
+            input_ids = None
+            text_encoder_layer_hidden_states = merged_inputs.get(
+                "text_encoder_layer_hidden_states", None
+            )
+
+        backbone_outputs = self.backbone_model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            text_encoder_layer_hidden_states=text_encoder_layer_hidden_states,
+            text_ids_mask=text_ids_mask,
+            **kwargs,
+        )
+
+        backbone_hidden_states = backbone_outputs[0]
+        backbone_hidden_states_for_depth_decoder = backbone_hidden_states
+
+        backbone_all_hidden_states = backbone_outputs.hidden_states
+
+        if backbone_all_hidden_states is None:
+            backbone_all_hidden_states = (backbone_hidden_states_for_depth_decoder,)
+        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+        slice_indices = (
+            slice(-logits_to_keep, None)
+            if isinstance(logits_to_keep, int)
+            else logits_to_keep
+        )
+
+        loss = None
+        backbone_loss = None
+        backbone_logits = None
+        depth_decoder_loss = None
+        depth_decoder_outputs = None
+
+        if labels is not None:
+            # Filter out num_items_in_batch from kwargs to ensure reduction='mean' is used
+            # This prevents incorrect loss normalization when backbone and depth decoder
+            # operate on different numbers of tokens
+            loss_kwargs = {k: v for k, v in kwargs.items() if k != "num_items_in_batch"}
+
+            # select first codebook as labels for the backbone model
+            backbone_labels = labels[:, :, 0]
+
+            backbone_logits = self.lm_head(backbone_hidden_states[:, slice_indices, :])
+            backbone_loss = self.loss_function(
+                logits=backbone_logits,
+                labels=backbone_labels,
+                vocab_size=self.lm_head.out_features,
+                **loss_kwargs,
+            )
+
+        # Compute logits only when actually needed
+        if backbone_logits is None and labels is None:
+            backbone_logits = self.lm_head(backbone_hidden_states[:, slice_indices, :])
+
+        if labels is not None:
+            # for the depth decoder, we need to select the frames to train on
+            # those are frames where the label is not uniformly `ignore_index` along the codebook dimension
+            train_mask = ~(labels[:, :, 1:] == -100).all(dim=-1)
+            depth_decoder_input_ids = labels[train_mask][
+                ..., : self.config.num_codebooks - 1
+            ]
+            # add place holder in position 0 that will be replaced by the backbone_last_hidden_state
+            depth_decoder_input_ids = nn.functional.pad(
+                depth_decoder_input_ids, (1, 0), value=0
+            )
+
+            train_idxs = train_mask.nonzero(as_tuple=True)
+            backbone_last_hidden_states = backbone_hidden_states_for_depth_decoder[
+                train_idxs[0], train_idxs[1] - 1, :
+            ]
+            depth_decoder_labels = labels[train_mask]
+
+            depth_decoder_outputs = self.depth_decoder(
+                input_ids=depth_decoder_input_ids,
+                backbone_last_hidden_state=backbone_last_hidden_states,
+                use_cache=use_cache,
+                return_dict=True,
+                labels=depth_decoder_labels,
+                **loss_kwargs,
+            )
+
+            depth_decoder_loss = depth_decoder_outputs.loss
+            if self.training:
+                loss = (
+                    backbone_loss
+                    + depth_decoder_loss * self.config.depth_header_loss_weight
+                )
+            else:
+                loss = backbone_loss + depth_decoder_loss
+        depth_decoder_hidden_states = None
+        if depth_decoder_outputs is not None:
+            depth_decoder_hidden_states = depth_decoder_outputs.hidden_states
+
+        return BreezeOutputWithPast(
+            loss=loss,
+            backbone_loss=backbone_loss,
+            depth_decoder_loss=depth_decoder_loss,
+            logits=backbone_logits,
+            past_key_values=backbone_outputs.past_key_values,
+            hidden_states=backbone_all_hidden_states,
+            attentions=backbone_outputs.attentions,
+            depth_decoder_logits=depth_decoder_outputs.logits
+            if depth_decoder_outputs is not None
+            else None,
+            depth_decoder_past_key_values=depth_decoder_outputs.past_key_values
+            if depth_decoder_outputs is not None
+            else None,
+            depth_decoder_hidden_states=depth_decoder_hidden_states,
+            depth_decoder_attentions=depth_decoder_outputs.attentions
+            if depth_decoder_outputs is not None
+            else None,
+        )
+
+
+__all__ = [
+    "BreezeBackboneModel",
+    "BreezeDepthDecoderForCausalLM",
+    "BreezeDepthDecoderModel",
+    "BreezeForConditionalGeneration",
+    "BreezePreTrainedModel",
+]
+
+AutoModel.register(BreezeConfig, BreezeForConditionalGeneration, exist_ok=True)

models/breeze_backbone_factory.py

@@ -0,0 +1,309 @@
+"""
+Breeze Backbone Factory and Adapter
+
+This module provides:
+1. BreezeBackboneFactory: Factory class to create different types of backbones
+2. BreezeBackboneAdapter: Adapter to use external pretrained LLMs as Breeze backbone
+
+Usage:
+    from models.breeze_backbone_factory import BreezeBackboneFactory
+
+    # In BreezeForConditionalGeneration.__init__:
+    self.backbone_model = BreezeBackboneFactory.create_backbone(config)
+"""
+
+
+import torch
+from torch import nn
+from transformers import AutoConfig
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.masking_utils import create_causal_mask
+from transformers.modeling_outputs import BaseModelOutputWithPast
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+
+class BreezeBackboneFactory:
+    """Factory class to create Breeze backbone models"""
+
+    SUPPORTED_MODELS = {
+        "breeze": "Native Breeze backbone",
+        "qwen3": "Qwen3ForCausalLM",
+        "llama3": "Llama3ForCausalLM",
+        # Add more models as needed
+    }
+
+    @staticmethod
+    def create_backbone(config):
+        """
+        Create backbone model based on config.backbone_model_type
+
+        Args:
+            config: BreezeConfig instance with backbone_model_type field
+
+        Returns:
+            Backbone model with interface compatible with BreezeBackboneModel
+        """
+        backbone_type = getattr(config, "backbone_model_type", "breeze")
+
+        if backbone_type == "breeze":
+            # Import here to avoid circular dependency
+            from .breeze import BreezeBackboneModel
+
+            return BreezeBackboneModel(config)
+        else:
+            # Use adapter for external LLMs
+            if backbone_type not in BreezeBackboneFactory.SUPPORTED_MODELS:
+                logger.warning(
+                    f"Backbone type '{backbone_type}' not in known types. "
+                    f"Attempting to create adapter anyway..."
+                )
+            return BreezeBackboneAdapter.create_from_config(config)
+
+
+class BreezeBackboneAdapter(nn.Module):
+    """
+    Adapter to use external pretrained LLM as Breeze Backbone
+
+    Key modifications from original LLM:
+    1. Replace LLM's embedding with BreezeBackboneModelEmbeddings (for audio tokens)
+    2. Keep LLM's transformer layers
+    3. Keep LLM's norm layer
+    4. Remove LLM's lm_head (Breeze has its own)
+
+    Interface: Compatible with BreezeBackboneModel's forward signature
+    """
+
+    def __init__(self, config, layers, norm, rotary_emb):
+        """
+        Args:
+            config: BreezeConfig (updated with LLM's architecture params)
+            layers: nn.ModuleList of transformer layers from LLM
+            norm: RMSNorm layer from LLM
+            rotary_emb: Rotary embedding layer from LLM
+        """
+        super().__init__()
+        self.config = config
+
+        # Ensure _attn_implementation is set on config
+        if not hasattr(config, "_attn_implementation"):
+            config._attn_implementation = "eager"
+
+        # Breeze's custom embedding for audio+text tokens
+        from .breeze import BreezeBackboneModelEmbeddings
+
+        self.embed_tokens = BreezeBackboneModelEmbeddings(config)
+
+        # Pretrained LLM's transformer components
+        self.layers = layers
+        self.norm = norm
+        self.rotary_emb = rotary_emb
+        self.gradient_checkpointing = False
+
+    @classmethod
+    def create_from_config(cls, config):
+        """
+        Create adapter from config (architecture only, no weight loading)
+
+        Weight loading is handled separately by breeze_model_init()
+
+        Args:
+            config: BreezeConfig with backbone_model_name_or_path set
+
+        Returns:
+            BreezeBackboneAdapter instance with initialized architecture
+        """
+        backbone_config = getattr(config, "backbone_config", None)
+        backbone_path = getattr(config, "backbone_model_name_or_path", None)
+        if backbone_config is None and backbone_path is None:
+            raise ValueError(
+                "config.backbone_config or config.backbone_model_name_or_path must be set"
+            )
+
+        backbone_type = config.backbone_model_type
+        if backbone_config is not None:
+            logger.info(
+                f"Creating backbone adapter for {backbone_type} from bundled config"
+            )
+            llm_config = AutoConfig.for_model(**backbone_config)
+        else:
+            logger.info(
+                f"Creating backbone adapter for {backbone_type} from {backbone_path}"
+            )
+            llm_config = AutoConfig.from_pretrained(backbone_path)
+
+        # Transfer attention implementation setting from BreezeConfig to LLM config
+        if (
+            hasattr(config, "_attn_implementation")
+            and config._attn_implementation is not None
+        ):
+            llm_config._attn_implementation = config._attn_implementation
+            logger.info(
+                f"Set attention implementation to: {config._attn_implementation}"
+            )
+        else:
+            # Default to eager if not specified
+            llm_config._attn_implementation = "eager"
+            logger.info("Attention implementation not specified, defaulting to: eager")
+
+        # Create layers based on backbone type
+        if backbone_type == "qwen3":
+            layers, norm, rotary_emb = cls._create_qwen3_layers(llm_config)
+        elif backbone_type == "llama3":
+            layers, norm, rotary_emb = cls._create_llama3_layers(llm_config)
+        else:
+            raise ValueError(
+                f"Unsupported backbone_model_type: {backbone_type}. "
+                f"Supported types: {list(BreezeBackboneFactory.SUPPORTED_MODELS.keys())}"
+            )
+
+        return cls(config, layers, norm, rotary_emb)
+
+    @staticmethod
+    def _create_qwen3_layers(llm_config):
+        """Create Qwen3 layer architecture"""
+        from transformers.models.qwen3.modeling_qwen3 import (
+            Qwen3DecoderLayer,
+            Qwen3RMSNorm,
+            Qwen3RotaryEmbedding,
+        )
+
+        layers = nn.ModuleList(
+            [
+                Qwen3DecoderLayer(llm_config, layer_idx)
+                for layer_idx in range(llm_config.num_hidden_layers)
+            ]
+        )
+        norm = Qwen3RMSNorm(llm_config.hidden_size, eps=llm_config.rms_norm_eps)
+        rotary_emb = Qwen3RotaryEmbedding(config=llm_config)
+
+        return layers, norm, rotary_emb
+
+    @staticmethod
+    def _create_llama3_layers(llm_config):
+        """Create Llama3 layer architecture"""
+        from transformers.models.llama.modeling_llama import (
+            LlamaDecoderLayer,
+            LlamaRMSNorm,
+            LlamaRotaryEmbedding,
+        )
+
+        layers = nn.ModuleList(
+            [
+                LlamaDecoderLayer(llm_config, layer_idx)
+                for layer_idx in range(llm_config.num_hidden_layers)
+            ]
+        )
+        norm = LlamaRMSNorm(llm_config.hidden_size, eps=llm_config.rms_norm_eps)
+        rotary_emb = LlamaRotaryEmbedding(config=llm_config)
+
+        return layers, norm, rotary_emb
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        cache_position: torch.LongTensor | None = None,
+        text_encoder_layer_hidden_states: torch.FloatTensor | None = None,
+        text_ids_mask: torch.BoolTensor | None = None,
+        use_cache: bool | None = None,
+        **kwargs,
+    ) -> BaseModelOutputWithPast:
+        """
+        Forward pass - interface compatible with BreezeBackboneModel
+
+        Args:
+            input_ids: (batch_size, seq_len, num_codebooks) or (batch_size, seq_len)
+            inputs_embeds: Pre-computed embeddings
+            attention_mask: Attention mask
+            position_ids: Position IDs
+            past_key_values: Cached key/values for generation
+            cache_position: Cache position tensor
+            use_cache: Whether to use KV cache
+
+        Returns:
+            BaseModelOutputWithPast with last_hidden_state
+        """
+        # Embed tokens if not provided
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        # Initialize cache
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache(config=self.config)
+
+        # Compute cache_position if not provided
+        if cache_position is None:
+            past_seen_tokens = (
+                past_key_values.get_seq_length() if past_key_values is not None else 0
+            )
+            cache_position = torch.arange(
+                past_seen_tokens,
+                past_seen_tokens + inputs_embeds.shape[1],
+                device=inputs_embeds.device,
+            )
+
+        # Compute position_ids if not provided
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+
+        # Create causal mask
+        causal_mask = create_causal_mask(
+            config=self.config,
+            input_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            cache_position=cache_position,
+            past_key_values=past_key_values,
+            position_ids=position_ids,
+        )
+
+        # Compute position embeddings (for RoPE)
+        hidden_states = inputs_embeds
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        # Pass through transformer layers
+        for layer_idx, decoder_layer in enumerate(self.layers):
+            if text_encoder_layer_hidden_states is not None and layer_idx > 0:
+                # skip for first layer fusion, cause it may has aleady been done in convert_input_ids_to_embeds() function
+                text_layer_hidden_states = text_encoder_layer_hidden_states[layer_idx]
+                # hidden_states[text_ids_mask][..., -self.config.hidden_size//2:] = text_layer_hidden_states
+
+                hidden_states_half_first, hidden_states_half_second = (
+                    hidden_states.split(self.config.hidden_size // 2, dim=-1)
+                )
+                if self.training:
+                    hidden_states_half_second = hidden_states_half_second.clone()
+                hidden_states_half_second[text_ids_mask] = text_layer_hidden_states
+                # hidden_states_half_second = torch.where(
+                #     text_ids_mask.unsqueeze(-1),
+                #     text_layer_hidden_states,
+                #     hidden_states_half_second,
+                # )
+                hidden_states = torch.cat(
+                    [hidden_states_half_first, hidden_states_half_second], dim=-1
+                )
+
+            hidden_states = decoder_layer(
+                hidden_states,
+                attention_mask=causal_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                cache_position=cache_position,
+                position_embeddings=position_embeddings,
+                **kwargs,
+            )
+
+        # Final norm
+        hidden_states = self.norm(hidden_states)
+
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values,
+        )
+
+
+__all__ = ["BreezeBackboneAdapter", "BreezeBackboneFactory"]

models/breeze_base_config.py

@@ -0,0 +1,454 @@
+# Copyright 2025 Sesame and The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.modeling_rope_utils import rope_config_validation
+from transformers.models.auto.configuration_auto import AutoConfig
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+
+class BreezeDepthDecoderConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`BreezeDepthDecoderModel`]. It is used to instantiate an Breeze depth decoder
+    model according to the specified arguments, defining the model architecture. Instantiating a configuration with the defaults will yield
+    a similar configuration to that of the breeze-1b.
+
+    e.g. a Breeze checkpoint
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+
+    Args:
+        num_codebooks (`int`, *optional*, defaults to 32):
+            Number of codebooks used in the underlying codec model responsible for tokenizing the audio.
+        backbone_hidden_size (`int`, *optional*, defaults to 2048):
+            Dimension of the hidden representations of the backbone model used with this depth decoder.
+        vocab_size (`int`, *optional*, defaults to 2051):
+            Vocabulary size of the BreezeDepthDecoder model. Defines the number of different audio tokens that can be represented by each codebook.
+        hidden_size (`int`, *optional*, defaults to 1024):
+            Dimension of the hidden representations.
+        intermediate_size (`int`, *optional*, defaults to 8192):
+            Dimension of the MLP representations.
+        num_hidden_layers (`int`, *optional*, defaults to 4):
+            Number of hidden layers in the Transformer decoder.
+        num_attention_heads (`int`, *optional*, defaults to 8):
+            Number of attention heads for each attention layer in the Transformer decoder.
+        num_key_value_heads (`int`, *optional*, defaults to 2):
+            This is the number of key_value heads that should be used to implement Grouped Query Attention. If
+            `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
+            `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When
+            converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
+            by meanpooling all the original heads within that group. For more details, check out [this
+            paper](https://huggingface.co/papers/2305.13245). If it is not specified, will default to
+            `num_attention_heads`.
+        hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
+            The non-linear activation function (function or string) in the decoder.
+        max_position_embeddings (`int`, *optional*, defaults to 33):
+            The maximum sequence length that this model might ever be used with.
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        rms_norm_eps (`float`, *optional*, defaults to 1e-05):
+            The epsilon used by the rms normalization layers.
+        use_cache (`bool`, *optional*, defaults to `True`):
+            Whether or not the model should return the last key/values attentions (not used by all models). Only
+            relevant if `config.is_decoder=True`.
+        pad_token_id (`int`, *optional*, defaults to 2050):
+            Padding token id.
+        bos_token_id (`int`, *optional*):
+            Beginning of stream token id.
+        eos_token_id (`int`, *optional*):
+            End of stream token id.
+        rope_theta (`float`, *optional*, defaults to 500000):
+            The base period of the RoPE embeddings.
+        rope_scaling (`Dict`, *optional*):
+            Dictionary containing the scaling configuration for the RoPE embeddings. NOTE: if you apply new rope type
+            and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value
+            accordingly.
+            Expected contents:
+                `rope_type` (`str`):
+                    The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope',
+                    'llama3'], with 'default' being the original RoPE implementation.
+                `factor` (`float`, *optional*):
+                    Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings. In
+                    most scaling types, a `factor` of x will enable the model to handle sequences of length x *
+                    original maximum pre-trained length.
+                `original_max_position_embeddings` (`int`, *optional*):
+                    Used with 'dynamic', 'longrope' and 'llama3'. The original max position embeddings used during
+                    pretraining.
+                `attention_factor` (`float`, *optional*):
+                    Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention
+                    computation. If unspecified, it defaults to value recommended by the implementation, using the
+                    `factor` field to infer the suggested value.
+                `beta_fast` (`float`, *optional*):
+                    Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear
+                    ramp function. If unspecified, it defaults to 32.
+                `beta_slow` (`float`, *optional*):
+                    Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear
+                    ramp function. If unspecified, it defaults to 1.
+                `short_factor` (`list[float]`, *optional*):
+                    Only used with 'longrope'. The scaling factor to be applied to short contexts (<
+                    `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
+                    size divided by the number of attention heads divided by 2
+                `long_factor` (`list[float]`, *optional*):
+                    Only used with 'longrope'. The scaling factor to be applied to long contexts (<
+                    `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
+                    size divided by the number of attention heads divided by 2
+                `low_freq_factor` (`float`, *optional*):
+                    Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE
+                `high_freq_factor` (`float`, *optional*):
+                    Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE
+        attention_bias (`bool`, *optional*, defaults to `False`):
+            Whether to use a bias in the query, key, value and output projection layers during self-attention.
+        attention_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+        mlp_bias (`bool`, *optional*, defaults to `False`):
+            Whether to use a bias in up_proj, down_proj and gate_proj layers in the MLP layers.
+        head_dim (`int`, *optional*):
+            The attention head dimension. If None, it will default to hidden_size // num_attention_heads
+
+    ```python
+    >>> from models.breeze import BreezeDepthDecoderModel
+    from models.breeze_config import BreezeDepthDecoderConfig
+
+    >>> # Initializing a BreezeDepthDecoder
+    >>> configuration = BreezeDepthDecoderConfig()
+    >>> model = BreezeDepthDecoderModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "breeze_depth_decoder_model"
+    base_config_key = "depth_decoder_config"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        num_codebooks=32,
+        backbone_hidden_size=2048,
+        vocab_size=2051,
+        hidden_size=1024,
+        intermediate_size=8192,
+        num_hidden_layers=4,
+        num_attention_heads=8,
+        num_key_value_heads=2,
+        hidden_act="silu",
+        max_position_embeddings=33,
+        initializer_range=0.02,
+        rms_norm_eps=1e-5,
+        use_cache=True,
+        pad_token_id=None,
+        bos_token_id=None,
+        eos_token_id=None,
+        rope_theta=500000,
+        rope_scaling=None,
+        attention_bias=False,
+        attention_dropout=0.0,
+        mlp_bias=False,
+        head_dim=None,
+        **kwargs,
+    ):
+        if kwargs.pop("tie_word_embeddings", False):
+            raise ValueError(
+                "`tie_word_embeddings=True` is not supported for BreezeDepthDecoderConfig"
+            )
+
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=False,
+            **kwargs,
+        )
+        self.num_codebooks = num_codebooks
+        self.vocab_size = vocab_size
+        self.backbone_hidden_size = backbone_hidden_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+
+        # for backward compatibility
+        if num_key_value_heads is None:
+            num_key_value_heads = num_attention_heads
+
+        self.num_key_value_heads = num_key_value_heads
+        self.hidden_act = hidden_act
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        self.rope_theta = rope_theta
+        self.rope_scaling = rope_scaling
+        self.attention_bias = attention_bias
+        self.attention_dropout = attention_dropout
+        self.mlp_bias = mlp_bias
+        self.head_dim = (
+            head_dim
+            if head_dim is not None
+            else self.hidden_size // self.num_attention_heads
+        )
+        # Validate the correctness of rotary position embeddings parameters
+        # BC: if there is a 'type' field, copy it it to 'rope_type'.
+        if self.rope_scaling is not None and "type" in self.rope_scaling:
+            self.rope_scaling["rope_type"] = self.rope_scaling["type"]
+        rope_config_validation(self)
+
+
+class BreezeConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`BreezeForConditionalGeneration`]. It is used to instantiate an Breeze
+    model according to the specified arguments, defining the model architecture. Instantiating a configuration
+    with the defaults will yield a similar configuration to that of the breeze-1b.
+
+    e.g. a Breeze checkpoint
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+    Args:
+        num_codebooks (`int`, *optional*, defaults to 32):
+            Number of codebooks used in the underlying codec model responsible for tokenizing the audio.
+        vocab_size (`int`, *optional*, defaults to 2051):
+            Vocabulary size of the Breeze model. Defines the number of different audio tokens that can be represented by each codebook.
+        text_vocab_size (`int`, *optional*, defaults to 128256):
+            Vocabulary size of the text input for the Breeze model. Defines the number of different text tokens that can be represented.
+        hidden_size (`int`, *optional*, defaults to 2048):
+            Dimension of the hidden representations of the backbone model.
+        intermediate_size (`int`, *optional*, defaults to 8192):
+            Dimension of the MLP representations of the backbone model.
+        num_hidden_layers (`int`, *optional*, defaults to 16):
+            Number of hidden layers in the backbone model Transformer decoder.
+        num_attention_heads (`int`, *optional*, defaults to 32):
+            Number of attention heads for each attention layer in the backbone model Transformer decoder.
+        num_key_value_heads (`int`, *optional*, defaults to 8):
+            This is the number of key_value heads that should be used to implement Grouped Query Attention. If
+            `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
+            `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When
+            converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
+            by meanpooling all the original heads within that group. For more details, check out [this
+            paper](https://huggingface.co/papers/2305.13245).
+        hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
+            The non-linear activation function (function or string) in the backbone model Transformer decoder.
+        max_position_embeddings (`int`, *optional*, defaults to 2048):
+            The maximum sequence length that this model might ever be used with.
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        rms_norm_eps (`float`, *optional*, defaults to 1e-05):
+            The epsilon used by the rms normalization layers.
+        use_cache (`bool`, *optional*, defaults to `True`):
+            Whether or not the model should return the last key/values attentions (not used by all models). Only
+            relevant if `config.is_decoder=True`.
+        pad_token_id (`int`, *optional*, defaults to 128002):
+            Padding token id.
+        codebook_pad_token_id (`int`, *optional*, defaults to 2050):
+            Padding token id for codebook tokens.
+        codebook_eos_token_id (`int`, *optional*, defaults to 0):
+            End of stream token id for codebook tokens.
+        bos_token_id (`int`, *optional*, defaults to 128000):
+            Beginning of stream token id.
+        eos_token_id (`int`, *optional*):
+            End of stream token id.
+        audio_token_id (`int`, *optional*, defaults to 128002):
+            Audio token id in the text input.
+        audio_eos_token_id (`int`, *optional*, defaults to 128003):
+            End of stream token id for audio in the text input.
+        rope_theta (`float`, *optional*, defaults to 500000):
+            The base period of the RoPE embeddings.
+        rope_scaling (`Dict`, *optional*, defaults to `{'factor': 32.0, 'high_freq_factor': 0.5, 'low_freq_factor': 0.125, 'original_max_position_embeddings': 1024, 'rope_type': 'llama3'}`):
+            Dictionary containing the scaling configuration for the RoPE embeddings. NOTE: if you apply new rope type
+            and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value
+            accordingly.
+            Expected contents:
+                `rope_type` (`str`):
+                    The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope',
+                    'llama3'], with 'default' being the original RoPE implementation.
+                `factor` (`float`, *optional*):
+                    Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings. In
+                    most scaling types, a `factor` of x will enable the model to handle sequences of length x *
+                    original maximum pre-trained length.
+                `original_max_position_embeddings` (`int`, *optional*):
+                    Used with 'dynamic', 'longrope' and 'llama3'. The original max position embeddings used during
+                    pretraining.
+                `attention_factor` (`float`, *optional*):
+                    Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention
+                    computation. If unspecified, it defaults to value recommended by the implementation, using the
+                    `factor` field to infer the suggested value.
+                `beta_fast` (`float`, *optional*):
+                    Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear
+                    ramp function. If unspecified, it defaults to 32.
+                `beta_slow` (`float`, *optional*):
+                    Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear
+                    ramp function. If unspecified, it defaults to 1.
+                `short_factor` (`list[float]`, *optional*):
+                    Only used with 'longrope'. The scaling factor to be applied to short contexts (<
+                    `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
+                    size divided by the number of attention heads divided by 2
+                `long_factor` (`list[float]`, *optional*):
+                    Only used with 'longrope'. The scaling factor to be applied to long contexts (<
+                    `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
+                    size divided by the number of attention heads divided by 2
+                `low_freq_factor` (`float`, *optional*):
+                    Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE
+                `high_freq_factor` (`float`, *optional*):
+                    Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE
+        attention_bias (`bool`, *optional*, defaults to `False`):
+            Whether to use a bias in the query, key, value and output projection layers during self-attention.
+        attention_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+        mlp_bias (`bool`, *optional*, defaults to `False`):
+            Whether to use a bias in up_proj, down_proj and gate_proj layers in the MLP layers.
+        head_dim (`int`, *optional*):
+            The attention head dimension. If None, it will default to hidden_size // num_attention_heads
+        tie_codebooks_embeddings (`bool`, *optional*, defaults to `True`):
+            Whether to tie the codebook tokens embeddings of the backbone model to the codebook tokens embeddings of the depth decoder.
+        depth_decoder_config (`BreezeDepthDecoderConfig`, *optional*):
+            Configuration for the depth decoder.
+        codec_config (`PretrainedConfig`, *optional*):
+            Configuration for the codec.
+
+    ```python
+    >>> from models.breeze import BreezeForConditionalGeneration
+    from models.breeze_config import BreezeConfig
+
+    >>> # Initializing a BreezeConfig
+    >>> configuration = BreezeConfig()
+
+    >>> # Initializing a model
+    >>> model = BreezeForConditionalGeneration(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "breeze"
+    base_config_key = "breeze_config"
+    keys_to_ignore_at_inference = ["past_key_values"]
+    sub_configs = {
+        "codec_config": AutoConfig,
+        "depth_decoder_config": BreezeDepthDecoderConfig,
+    }
+
+    def __init__(
+        self,
+        num_codebooks=32,
+        vocab_size=2051,
+        text_vocab_size=128256,
+        hidden_size=2048,
+        intermediate_size=8192,
+        num_hidden_layers=16,
+        num_attention_heads=32,
+        num_key_value_heads=8,
+        hidden_act="silu",
+        max_position_embeddings=2048,
+        initializer_range=0.02,
+        rms_norm_eps=1e-5,
+        use_cache=True,
+        pad_token_id=128002,
+        codebook_pad_token_id=2050,
+        codebook_eos_token_id=0,
+        bos_token_id=128000,
+        eos_token_id=None,
+        audio_token_id=128002,
+        audio_eos_token_id=128003,
+        rope_theta=500000,
+        rope_scaling=None,
+        attention_bias=False,
+        attention_dropout=0.0,
+        mlp_bias=False,
+        head_dim=None,
+        tie_codebooks_embeddings=True,
+        depth_decoder_config=None,
+        codec_config=None,
+        **kwargs,
+    ):
+        if kwargs.pop("tie_word_embeddings", False):
+            raise ValueError(
+                "`tie_word_embeddings=True` is not supported for BreezeConfig"
+            )
+
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=False,
+            **kwargs,
+        )
+
+        if depth_decoder_config is None:
+            self.depth_decoder_config = BreezeDepthDecoderConfig()
+            logger.info(
+                "depth_decoder_config is None, using default depth decoder config."
+            )
+        elif isinstance(depth_decoder_config, dict):
+            self.depth_decoder_config = BreezeDepthDecoderConfig(**depth_decoder_config)
+        elif isinstance(depth_decoder_config, BreezeDepthDecoderConfig):
+            self.depth_decoder_config = depth_decoder_config
+
+        if codec_config is None:
+            self.codec_config = AutoConfig.for_model("mimi")
+            logger.info("codec_config is None, using default audio encoder config.")
+        elif isinstance(codec_config, dict):
+            self.codec_config = AutoConfig.for_model(**codec_config)
+        elif isinstance(codec_config, PretrainedConfig):
+            self.codec_config = codec_config
+
+        self.text_vocab_size = text_vocab_size
+        self.num_codebooks = num_codebooks
+        self.audio_token_id = audio_token_id
+        self.audio_eos_token_id = audio_eos_token_id
+        self.codebook_pad_token_id = codebook_pad_token_id
+        self.codebook_eos_token_id = codebook_eos_token_id
+        self.tie_codebooks_embeddings = tie_codebooks_embeddings
+
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+
+        # for backward compatibility
+        if num_key_value_heads is None:
+            num_key_value_heads = num_attention_heads
+
+        self.num_key_value_heads = num_key_value_heads
+        self.hidden_act = hidden_act
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        self.rope_theta = rope_theta
+        self.rope_scaling = rope_scaling
+        self.attention_bias = attention_bias
+        self.attention_dropout = attention_dropout
+        self.mlp_bias = mlp_bias
+        self.head_dim = (
+            head_dim
+            if head_dim is not None
+            else self.hidden_size // self.num_attention_heads
+        )
+        # Validate the correctness of rotary position embeddings parameters
+        # BC: if there is a 'type' field, copy it it to 'rope_type'.
+        if self.rope_scaling is not None and "type" in self.rope_scaling:
+            self.rope_scaling["rope_type"] = self.rope_scaling["type"]
+        rope_config_validation(self)
+
+
+__all__ = [
+    "BreezeConfig",
+    "BreezeDepthDecoderConfig",
+]

models/breeze_config.py

@@ -0,0 +1,54 @@
+import os
+
+from transformers import AutoConfig, AutoModel
+from transformers.models.t5gemma.configuration_t5gemma import T5GemmaModuleConfig
+from transformers.models.t5gemma.modeling_t5gemma import T5GemmaEncoder
+
+from .breeze_base_config import BreezeConfig as BreezeConfig_transformers
+from .breeze_base_config import BreezeDepthDecoderConfig
+from .t5gemma2_compat import T5Gemma2TextConfig, T5Gemma2TextEncoder
+
+
+class T5GemmaEncoderWrapper(T5GemmaEncoder):
+    config_class = T5GemmaModuleConfig
+
+
+# register T5GemmaModuleConfig
+AutoConfig.register("t5_gemma_module", T5GemmaModuleConfig)
+AutoModel.register(T5GemmaModuleConfig, T5GemmaEncoderWrapper)
+AutoConfig.register("t5gemma2_text", T5Gemma2TextConfig)
+AutoModel.register(T5Gemma2TextConfig, T5Gemma2TextEncoder)
+
+
+# Update BreezeConfig to handle text_encoder_config properly
+class BreezeConfig(BreezeConfig_transformers):
+    model_type = "breeze"
+
+    sub_configs = {
+        "codec_config": AutoConfig,
+        "depth_decoder_config": BreezeDepthDecoderConfig,
+    }
+
+    def __init__(self, **kwargs):
+        text_encoder_config = kwargs.get("text_encoder_config", None)
+        super().__init__(**kwargs)
+
+        if text_encoder_config is None:
+            self.text_encoder_config = None
+        elif isinstance(text_encoder_config, str):
+            assert os.path.isdir(text_encoder_config), (
+                f"text_encoder_config as str must be a valid directory path: '{text_encoder_config}'"
+            )
+            self.text_encoder_config = AutoConfig.from_pretrained(text_encoder_config)
+        elif isinstance(text_encoder_config, dict):
+            self.text_encoder_config = AutoConfig.for_model(**text_encoder_config)
+        else:
+            raise ValueError(
+                f"text_encoder_config must be a str (path), dict, or AutoConfig instance. but got {type(text_encoder_config)}"
+            )
+
+
+AutoConfig.register(
+    "breeze_depth_decoder_model", BreezeDepthDecoderConfig, exist_ok=True
+)
+AutoConfig.register("breeze", BreezeConfig, exist_ok=True)

models/cudagraph/__init__.py

@@ -0,0 +1,7 @@
+"""CUDA Graph + StaticCache accelerated pieces for Breeze TTS."""
+
+from .backbone_graph import BackboneGraph
+from .depth_decoder_graph import DepthDecoderGraph
+from .sampling import sample_logits
+
+__all__ = ["BackboneGraph", "DepthDecoderGraph", "sample_logits"]

models/cudagraph/backbone_graph.py

@@ -0,0 +1,397 @@
+"""
+CUDA graph capture for the Breeze backbone's single-token decode step,
+using transformers StaticCache.
+
+The backbone (BreezeBackboneModel) predicts the first codebook token at each step.
+For CFG, we run batch=2 (cond + uncond) and apply the guidance formula inside
+the graph so it's captured as a single replay.
+
+Strategy:
+- Use transformers StaticCache for KV cache management
+- Use the model's forward method (handles mask, RoPE, attention internally)
+- Capture single-token decode + lm_head + CFG as a CUDA graph
+- Update cache_position buffer between replays
+- Batch size changes trigger automatic recapture via ensure_batch_size()
+"""
+
+import torch
+from transformers import StaticCache
+from transformers.masking_utils import create_causal_mask
+
+
+class BackboneGraph:
+    """
+    Captures the Breeze backbone's single-token decode step as a CUDA graph,
+    using the model's own forward with transformers StaticCache.
+
+    Supports batch=2 for CFG (cond=batch[0], uncond=batch[1]).
+
+    Batch size changes trigger automatic recapture:
+        graph.ensure_batch_size(4)
+
+    guidance_scale is runtime-mutable without recapture:
+        graph.guidance_scale.fill_(5.0)
+    """
+
+    def __init__(
+        self,
+        backbone_model,
+        lm_head,
+        embed_tokens,
+        config,
+        device="cuda:0",
+        dtype=torch.bfloat16,
+        max_seq_len=512,
+        guidance_scale=3.0,
+        debug: bool = False,
+        batch_size: int = 1,
+    ):
+        self.device = device
+        self.dtype = dtype
+        self.debug = debug
+        self.no_graph = (
+            False  # runtime flag: True = skip graph replay (for layer-diff hooks)
+        )
+        self.max_seq_len = max_seq_len
+        self.hidden_size = config.hidden_size
+        self.num_layers = config.num_hidden_layers
+        self.vocab_size = getattr(lm_head, "out_features", config.vocab_size)
+        self.num_codebooks = config.num_codebooks
+        self.config = config
+        self.batch_size = batch_size
+        self.half = self._real_batch_size(batch_size)
+
+        # Keep references to model components
+        self.model = backbone_model  # BreezeBackboneModel
+        self.lm_head = lm_head  # nn.Linear(hidden_size, vocab_size)
+        self.embed_tokens = embed_tokens  # BreezeBackboneModelEmbeddings
+
+        # Cast lm_head to fp32 for numerical stability (hooks still fire through module forward)
+        self.lm_head = self.lm_head.float()
+
+        # Transformers StaticCache — batch=2 for CFG
+        self.static_cache = StaticCache(
+            config=config, max_cache_len=max_seq_len, batch_size=self.batch_size
+        )
+
+        # Cache position buffer — [1] for KV cache slot (shared across batch)
+        self.cache_position = torch.zeros(1, dtype=torch.long, device=device)
+        # Per-batch position_ids for RoPE — [batch_size, 1]
+        self.position_ids = torch.zeros(
+            self.batch_size, 1, dtype=torch.long, device=device
+        )
+        # Base position for each batch (set by set_generation_state from attention_mask)
+        self._base_position = torch.zeros(
+            self.batch_size, dtype=torch.long, device=device
+        )
+
+        # Guidance scale buffer — can be updated at runtime via fill_()
+        self.guidance_scale = torch.tensor(
+            [guidance_scale], dtype=torch.float32, device=device
+        )
+
+        self._alloc_buffers()
+
+        self.graph = None
+        self.captured = False
+        self.attn_mask = None
+        # Per-batch left-padding lengths for vectorized mask construction
+        self._pad_lens = torch.zeros(self.batch_size, dtype=torch.long, device=device)
+        # Pre-allocated index range [0, 1, ..., max_seq_len-1] for broadcasting
+        self._kv_indices = torch.arange(max_seq_len, dtype=torch.long, device=device)
+
+    @staticmethod
+    def _real_batch_size(batch_size: int) -> int:
+        if batch_size <= 1:
+            return batch_size
+        return batch_size // 2
+
+    # ------------------------------------------------------------------
+    # Buffer allocation (called from __init__ and _rebuild_for_batch)
+    # ------------------------------------------------------------------
+
+    def _alloc_buffers(self):
+        """Allocate / re-allocate I/O buffers for current batch_size."""
+        # Input: codebook IDs [batch, 1, num_codebooks]
+        self.input_ids_buf = torch.zeros(
+            self.batch_size, 1, self.num_codebooks, dtype=torch.long, device=self.device
+        )
+        # Hidden states output [batch, 1, hidden_size]
+        self.hidden_buf = torch.zeros(
+            self.batch_size, 1, self.hidden_size, dtype=self.dtype, device=self.device
+        )
+        # Raw logits [batch, vocab_size]
+        self.logits_buf = torch.zeros(
+            self.batch_size, self.vocab_size, dtype=torch.float32, device=self.device
+        )
+        # CFG-applied logits [half, vocab_size]
+        self.cfg_logits_buf = torch.zeros(
+            self.half, self.vocab_size, dtype=torch.float32, device=self.device
+        )
+
+    # ------------------------------------------------------------------
+    # Batch size management — auto-recapture on change
+    # ------------------------------------------------------------------
+
+    @torch.inference_mode()
+    def ensure_batch_size(self, batch_size: int):
+        """No-op if batch_size matches current captured graph; else rebuild and recapture."""
+        if batch_size == self.batch_size:
+            return
+        self._rebuild_for_batch(batch_size)
+
+    @torch.inference_mode()
+    def _rebuild_for_batch(self, batch_size: int):
+        self.batch_size = batch_size
+        self.half = self._real_batch_size(batch_size)
+        self.static_cache = StaticCache(
+            config=self.config, max_cache_len=self.max_seq_len, batch_size=batch_size
+        )
+        # Re-allocate buffers for new batch_size
+        self.cache_position = torch.zeros(1, dtype=torch.long, device=self.device)
+        self.position_ids = torch.zeros(
+            batch_size, 1, dtype=torch.long, device=self.device
+        )
+        self._base_position = torch.zeros(
+            batch_size, dtype=torch.long, device=self.device
+        )
+        self._alloc_buffers()
+        self.captured = False
+        self.graph = None
+        self.attn_mask = None
+        self._pad_lens = torch.zeros(batch_size, dtype=torch.long, device=self.device)
+        self._kv_indices = torch.arange(
+            self.max_seq_len, dtype=torch.long, device=self.device
+        )
+        if self.no_graph:
+            self.prepare_eager()
+        else:
+            self.capture()
+
+    @torch.inference_mode()
+    def prepare_eager(self):
+        """Initialize fixed eager buffers without compile or CUDA Graph capture."""
+        self.no_graph = True
+        self._init_cache_layers()
+        self._build_initial_attention_mask()
+        self.captured = False
+        return self
+
+    def _init_cache_layers(self):
+        """Force lazy initialization of StaticCache layers before graph capture."""
+        config = self.model.config
+        num_kv_heads = getattr(
+            config, "num_key_value_heads", config.num_attention_heads
+        )
+        head_dim = getattr(
+            config, "head_dim", config.hidden_size // config.num_attention_heads
+        )
+        dummy_k = torch.zeros(
+            self.batch_size,
+            num_kv_heads,
+            1,
+            head_dim,
+            dtype=self.dtype,
+            device=self.device,
+        )
+        for layer in self.static_cache.layers:
+            if not layer.is_initialized:
+                layer.lazy_initialization(dummy_k)
+
+    def _build_initial_attention_mask(self):
+        """Build a default causal mask (no padding) for graph capture/warmup.
+
+        Uses create_causal_mask once at position 0 to determine the correct
+        shape and dtype, then allocates self.attn_mask as a static buffer.
+        """
+        dummy = torch.zeros(
+            self.batch_size, 1, self.hidden_size, dtype=self.dtype, device=self.device
+        )
+        pos = torch.tensor([0], device=self.device)
+        mask = create_causal_mask(
+            config=self.model.config,
+            input_embeds=dummy,
+            attention_mask=None,
+            cache_position=pos,
+            past_key_values=self.static_cache,
+        )
+        # mask shape: [batch, 1, 1, max_seq_len], dtype may be float or bool
+        if mask is None or mask.dtype == torch.bool:
+            # Some attn implementations return bool or None; allocate float mask
+            self.attn_mask = torch.zeros(
+                self.batch_size,
+                1,
+                1,
+                self.max_seq_len,
+                dtype=self.dtype,
+                device=self.device,
+            )
+        elif self.attn_mask is None:
+            self.attn_mask = mask.clone()
+        else:
+            self.attn_mask.copy_(mask)
+        print(
+            f"[BackboneGraph] Initial attention mask shape: {self.attn_mask.shape}, dtype: {self.attn_mask.dtype}"
+        )
+
+        self._mask_min_val = torch.finfo(self.attn_mask.dtype).min
+
+    def _set_attention_mask(self, position: int):
+        """Construct attention mask in-place from pad_lens and cache_position.
+
+        For each batch i, the valid KV range is [pad_lens[i], position].
+        Everything outside is masked with -inf. Fully vectorized, no loops.
+
+        attn_mask shape: [batch, 1, 1, max_seq_len]
+        """
+        # kv_indices: [max_seq_len], pad_lens: [batch, 1], position: scalar
+        kv = self._kv_indices  # [S]
+        lo = self._pad_lens.unsqueeze(1)  # [B, 1]
+        valid = (kv >= lo) & (kv <= position)  # [B, S]
+        self.attn_mask.fill_(self._mask_min_val)
+        self.attn_mask[:, 0, 0, :].masked_fill_(valid, 0.0)
+
+    def _decode_step(self):
+        """Single-token decode: embed → backbone forward → lm_head → CFG."""
+        # 1. Embed codebook tokens (sum of num_codebooks embeddings)
+        inputs_embeds = self.embed_tokens(self.input_ids_buf)  # [batch, 1, H]
+
+        # 2. Backbone forward with per-batch position_ids and shared cache_position
+        # position_ids: [batch_size, 1] — per-batch logical position for RoPE
+        # cache_position: [1] — shared KV cache slot for all batches
+        out = self.model(
+            inputs_embeds=inputs_embeds,
+            attention_mask=self.attn_mask,
+            past_key_values=self.static_cache,
+            position_ids=self.position_ids,
+            cache_position=self.cache_position,
+            use_cache=True,
+        )
+        self.hidden_buf.copy_(out.last_hidden_state)
+
+        # 3. lm_head → logits (lm_head is fp32, input cast to fp32)
+        logits = self.lm_head(self.hidden_buf[:, -1, :].float())  # [batch, vocab] fp32
+        self.logits_buf.copy_(logits)
+
+        # 4. CFG when paired cond/uncond rows are present; otherwise pass through.
+        if self.batch_size >= 2:
+            cond_logits = self.logits_buf[: self.half]
+            uncond_logits = self.logits_buf[self.half :]
+            cfg_result = uncond_logits + self.guidance_scale * (
+                cond_logits - uncond_logits
+            )
+        else:
+            cfg_result = self.logits_buf[: self.half]
+        self.cfg_logits_buf.copy_(cfg_result)
+
+    @torch.inference_mode()
+    def capture(self, prefill_len=100, num_warmup=3):
+        """Capture CUDA graph for single-token decode."""
+        print(f"Warming up backbone graph ({num_warmup} runs)...")
+
+        self._init_cache_layers()
+        self._build_initial_attention_mask()
+
+        # Initialize positions for warmup
+        self.cache_position[0] = prefill_len
+        self.position_ids.fill_(prefill_len)
+        self._pad_lens.zero_()
+        self._set_attention_mask(prefill_len)
+
+        for _ in range(num_warmup):
+            self._decode_step()
+        torch.cuda.synchronize(device=self.device)
+
+        print("Capturing CUDA graph for backbone decode...")
+        self.graph = torch.cuda.CUDAGraph()
+
+        s = torch.cuda.Stream(device=self.device)
+        s.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(s):
+            self._decode_step()
+            torch.cuda.synchronize(device=self.device)
+
+            with torch.cuda.graph(self.graph):
+                self._decode_step()
+
+        torch.cuda.current_stream().wait_stream(s)
+        torch.cuda.synchronize(device=self.device)
+        self.captured = True
+        print("Backbone CUDA graph captured!")
+
+    @torch.inference_mode()
+    def reset(self):
+        """Reset cache for new sequence."""
+        self.static_cache.reset()
+
+    @torch.inference_mode()
+    def finish_direct_prefill(self, seq_len: int) -> int:
+        self._prefill_len = int(seq_len)
+        return self._prefill_len
+
+    @torch.inference_mode()
+    def prefill_kv(self, past_key_values):
+        """
+        Copy HF DynamicCache from prefill into our StaticCache.
+        past_key_values: DynamicCache with num_layers layers of [batch, kv_heads, seq_len, head_dim]
+
+        For CFG: the DynamicCache should already have batch=2.
+        """
+        self.static_cache.reset()
+        seq_len = 0
+        for li in range(self.num_layers):
+            k, v = past_key_values[li]
+            seq_len = k.shape[2]
+            if seq_len > self.max_seq_len:
+                raise RuntimeError(
+                    f"Input too long: prefill has {seq_len} tokens but max_seq_len={self.max_seq_len}."
+                )
+            cache_pos = torch.arange(seq_len, device=self.device)
+            self.static_cache.update(k, v, li, {"cache_position": cache_pos})
+        self._prefill_len = seq_len  # remember for decode cache_position
+        return seq_len
+
+    @torch.inference_mode()
+    def set_generation_state(self, attention_mask=None):
+        """Initialize per-batch position_ids and pad_lens from attention_mask.
+
+        For left-padded inputs, pad_lens[i] = number of leading zeros in
+        attention_mask[i]. The mask is then constructed on-the-fly in
+        _set_attention_mask() using pad_lens — no 512-iter table rebuild.
+        """
+        if attention_mask is not None:
+            per_batch_pos = attention_mask.sum(dim=1).long()  # [batch_size]
+            self._base_position.copy_(per_batch_pos)
+            self.position_ids.copy_(per_batch_pos.unsqueeze(-1))
+            self.cache_position[0] = per_batch_pos.max().item()
+
+            # Compute per-batch left-padding length:
+            # pad_lens[i] = seq_len - num_valid_tokens[i]
+            seq_len = attention_mask.shape[1]
+            self._pad_lens.copy_(seq_len - per_batch_pos)
+        else:
+            self._pad_lens.zero_()
+
+    @torch.inference_mode()
+    def run(self, input_ids, step_idx):
+        """
+        Run one decode step.
+        input_ids: [batch, 1, num_codebooks] long tensor
+        step_idx: decode step index (0-based), added to per-batch base position
+        Returns: (hidden_states [batch, 1, H], cfg_logits [half, vocab])
+        """
+        self.input_ids_buf.copy_(input_ids)
+        # Update per-batch position_ids: base_position + step_idx
+        self.position_ids.copy_((self._base_position + step_idx).unsqueeze(-1))
+        # KV cache slot: prefill_len + step_idx (append after prefill KV, never overwrite)
+        self.cache_position[0] = self._prefill_len + step_idx
+        self._set_attention_mask(self.cache_position[0].item())
+        if self.no_graph:
+            self._decode_step()
+        else:
+            self.graph.replay()
+            if self.debug:
+                # Fake forward: trigger lm_head hooks using buffer data from replay.
+                # Result is discarded — only hook side-effects matter.
+                self.lm_head(self.hidden_buf[:, -1, :].float())
+        return self.hidden_buf, self.cfg_logits_buf

models/cudagraph/backbone_prefill_graph.py

@@ -0,0 +1,255 @@
+"""Static-shape CUDA Graph buckets for batch-1 backbone prefill."""
+
+from __future__ import annotations
+
+import threading
+from dataclasses import dataclass
+
+import torch
+
+
+@dataclass
+class BackbonePrefillOutput:
+    hidden_states: torch.Tensor
+    logits: torch.Tensor
+    attention_mask: torch.Tensor
+    prefill_len: int
+
+
+@dataclass
+class _PrefillRecord:
+    graph: torch.cuda.CUDAGraph
+    stream: torch.cuda.Stream
+    inputs_embeds: torch.Tensor
+    attention_mask: torch.Tensor
+    position_ids: torch.Tensor
+    causal_mask: torch.Tensor
+    cache_position: torch.Tensor
+    hidden_states: torch.Tensor
+    logits: torch.Tensor
+
+
+class BackbonePrefillGraphCache:
+    """Capture backbone prefill + lm_head while writing the decode StaticCache."""
+
+    def __init__(self, backbone_graph, *, token_granularity: int = 32) -> None:
+        if token_granularity <= 0:
+            raise ValueError("token_granularity must be > 0")
+        self.backbone_graph = backbone_graph
+        self.model = backbone_graph.model
+        self.lm_head = backbone_graph.lm_head
+        self.device = torch.device(backbone_graph.device)
+        self.dtype = backbone_graph.dtype
+        self.token_granularity = int(token_granularity)
+        self._records: dict[tuple[int, int], _PrefillRecord] = {}
+        self._graph_pool = torch.cuda.graph_pool_handle()
+        self._lock = threading.RLock()
+        self.captures = 0
+        self.replays = 0
+        self._frozen = False
+
+    def _bucket(self, length: int) -> int:
+        return (
+            (int(length) + self.token_granularity - 1) // self.token_granularity
+        ) * self.token_granularity
+
+    @staticmethod
+    def _copy_inputs(
+        inputs_embeds: torch.Tensor,
+        attention_mask: torch.Tensor,
+        static_embeds: torch.Tensor,
+        static_mask: torch.Tensor,
+        static_positions: torch.Tensor,
+    ) -> None:
+        seq_len = int(inputs_embeds.shape[1])
+        pad_len = int(static_embeds.shape[1]) - seq_len
+        static_embeds.zero_()
+        static_mask.zero_()
+        static_positions.fill_(1)
+        static_embeds[:, pad_len:].copy_(inputs_embeds)
+        static_mask[:, pad_len:].copy_(attention_mask)
+        positions = static_mask.long().cumsum(-1) - 1
+        positions.masked_fill_(static_mask == 0, 1)
+        static_positions.copy_(positions)
+
+    @staticmethod
+    def _update_causal_mask(
+        attention_mask: torch.Tensor, causal_mask: torch.Tensor
+    ) -> None:
+        batch_size, query_len = attention_mask.shape
+        key_len = causal_mask.shape[-1]
+        query_idx = torch.arange(query_len, device=attention_mask.device).view(
+            1, query_len, 1
+        )
+        key_idx = torch.arange(key_len, device=attention_mask.device).view(
+            1, 1, key_len
+        )
+        valid_keys = torch.zeros(
+            (batch_size, key_len), dtype=torch.bool, device=attention_mask.device
+        )
+        valid_keys[:, :query_len].copy_(attention_mask.to(torch.bool))
+        allowed = (key_idx <= query_idx) & valid_keys[:, None, :]
+        causal_mask.fill_(torch.finfo(causal_mask.dtype).min)
+        causal_mask[:, 0].masked_fill_(allowed, 0.0)
+
+    def _forward(
+        self,
+        inputs_embeds: torch.Tensor,
+        attention_mask: torch.Tensor,
+        position_ids: torch.Tensor,
+        cache_position: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        output = self.model(
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=self.backbone_graph.static_cache,
+            cache_position=cache_position,
+            use_cache=True,
+        )
+        hidden_states = output.last_hidden_state
+        logits = self.lm_head(hidden_states[:, -1, :].float()).float()
+        return hidden_states, logits
+
+    @torch.inference_mode()
+    def __call__(
+        self,
+        inputs_embeds: torch.Tensor,
+        attention_mask: torch.Tensor,
+    ) -> BackbonePrefillOutput:
+        batch_size, seq_len, hidden_size = inputs_embeds.shape
+        if batch_size != self.backbone_graph.batch_size:
+            raise ValueError(
+                f"prefill batch={batch_size} does not match decode graph "
+                f"batch={self.backbone_graph.batch_size}"
+            )
+        bucket_len = self._bucket(int(seq_len))
+        if bucket_len > self.backbone_graph.max_seq_len:
+            raise RuntimeError(
+                f"prefill bucket {bucket_len} exceeds max_seq_len "
+                f"{self.backbone_graph.max_seq_len}"
+            )
+        key = (int(batch_size), bucket_len)
+
+        with self._lock:
+            record = self._records.get(key)
+            if record is None:
+                if self._frozen:
+                    raise RuntimeError(
+                        f"backbone prefill CUDA graph {key} was not declared in the warmup profile"
+                    )
+                static_embeds = torch.zeros(
+                    (batch_size, bucket_len, hidden_size),
+                    dtype=inputs_embeds.dtype,
+                    device=inputs_embeds.device,
+                )
+                static_mask = torch.zeros(
+                    (batch_size, bucket_len),
+                    dtype=attention_mask.dtype,
+                    device=attention_mask.device,
+                )
+                static_positions = torch.ones_like(static_mask, dtype=torch.long)
+                static_causal_mask = torch.empty(
+                    (batch_size, 1, bucket_len, self.backbone_graph.max_seq_len),
+                    dtype=inputs_embeds.dtype,
+                    device=inputs_embeds.device,
+                )
+                cache_position = torch.arange(
+                    bucket_len, device=inputs_embeds.device, dtype=torch.long
+                )
+                self._copy_inputs(
+                    inputs_embeds,
+                    attention_mask,
+                    static_embeds,
+                    static_mask,
+                    static_positions,
+                )
+                self._update_causal_mask(static_mask, static_causal_mask)
+
+                capture_stream = torch.cuda.Stream(device=inputs_embeds.device)
+                capture_stream.wait_stream(
+                    torch.cuda.current_stream(inputs_embeds.device)
+                )
+                with torch.cuda.stream(capture_stream):
+                    for _ in range(3):
+                        hidden_states, logits = self._forward(
+                            static_embeds,
+                            static_causal_mask,
+                            static_positions,
+                            cache_position,
+                        )
+                capture_stream.synchronize()
+
+                graph = torch.cuda.CUDAGraph()
+                with torch.cuda.graph(
+                    graph,
+                    stream=capture_stream,
+                    pool=self._graph_pool,
+                    capture_error_mode="thread_local",
+                ):
+                    hidden_states, logits = self._forward(
+                        static_embeds,
+                        static_causal_mask,
+                        static_positions,
+                        cache_position,
+                    )
+                record = _PrefillRecord(
+                    graph=graph,
+                    stream=capture_stream,
+                    inputs_embeds=static_embeds,
+                    attention_mask=static_mask,
+                    position_ids=static_positions,
+                    causal_mask=static_causal_mask,
+                    cache_position=cache_position,
+                    hidden_states=hidden_states,
+                    logits=logits,
+                )
+                self._records[key] = record
+                self.captures += 1
+
+            self._copy_inputs(
+                inputs_embeds,
+                attention_mask,
+                record.inputs_embeds,
+                record.attention_mask,
+                record.position_ids,
+            )
+            self._update_causal_mask(record.attention_mask, record.causal_mask)
+            current_stream = torch.cuda.current_stream(inputs_embeds.device)
+            record.stream.wait_stream(current_stream)
+            record.graph.replay()
+            current_stream.wait_stream(record.stream)
+            self.backbone_graph.finish_direct_prefill(bucket_len)
+            self.replays += 1
+            return BackbonePrefillOutput(
+                hidden_states=record.hidden_states,
+                logits=record.logits,
+                attention_mask=record.attention_mask,
+                prefill_len=bucket_len,
+            )
+
+    @property
+    def graph_keys(self) -> tuple[tuple[int, int], ...]:
+        return tuple(sorted(self._records))
+
+    @torch.inference_mode()
+    def warmup_graph(self, *, branch_batch_size: int, sequence_length: int) -> None:
+        if self._frozen:
+            raise RuntimeError("backbone prefill CUDA graph cache is already frozen")
+        inputs_embeds = torch.zeros(
+            branch_batch_size,
+            sequence_length,
+            int(self.backbone_graph.hidden_size),
+            dtype=self.dtype,
+            device=self.device,
+        )
+        attention_mask = torch.ones(
+            branch_batch_size,
+            sequence_length,
+            dtype=torch.long,
+            device=self.device,
+        )
+        self(inputs_embeds, attention_mask)
+
+    def freeze(self) -> None:
+        self._frozen = True

models/cudagraph/depth_decoder_graph.py

@@ -0,0 +1,987 @@
+"""
+CUDA graph capture for the Breeze depth decoder's decode loop,
+using transformers StaticCache.
+
+The depth decoder generates num_codebooks-1 codebook tokens autoregressively
+(typically 15 for production models with 16 RVQ codebooks):
+- Step 0: prefill with 2 tokens (backbone_hidden placeholder + first_cb_token),
+  get logits for codebook[1]
+- Steps 1..num_codebooks-2: decode 1 token at a time using previous codebook token's embedding
+
+Strategy:
+- Use transformers StaticCache for KV cache management
+- Use the depth decoder's inner model forward
+- Unroll the full loop for deterministic shapes
+- Capture the entire loop as a single CUDA graph
+- CFG is baked in: batch=2 (cond + uncond), guidance applied per step
+- All tensors pre-allocated — no torch.tensor() calls inside the graph
+- Sampling params (temperature, top_k, top_p, do_sample, guidance_scale) are
+  per-sample tensor buffers [half, 1] or [half]: update via set_*() methods
+  at runtime without recapture. Accepts both scalar (broadcast) and tensor
+  (per-sample) values.
+- top_k uses a fixed MAX_K workspace; top_k_buf can be set to any value <= _max_k.
+- Multi-batch bucket support: pre-capture graphs for multiple batch sizes
+  (powers of 2), select the smallest fitting bucket at runtime. No recapture
+  needed when batch size changes within the pre-captured range.
+"""
+
+import logging
+from dataclasses import dataclass
+
+import torch
+import torch._dynamo
+import torch.nn.functional as F
+from transformers import StaticCache
+from transformers.masking_utils import create_causal_mask
+
+from ..logits_process import mask_invalid_codec_token_logits
+
+_log = logging.getLogger(__name__)
+
+
+@dataclass
+class _BucketState:
+    """Stores all per-bucket-size state needed for CUDA graph replay."""
+
+    batch_size: int
+    half: int
+    graph: torch.cuda.CUDAGraph
+    static_cache: StaticCache
+    # I/O buffers
+    backbone_hidden_buf: torch.Tensor
+    first_cb_token_buf: torch.Tensor
+    output_tokens: torch.Tensor
+    _tok_buf: torch.Tensor
+    prefill_input_ids: torch.Tensor
+    # Attention masks
+    prefill_attn: torch.Tensor
+    decode_attn: list
+    # Sampling param buffers
+    temperature_buf: torch.Tensor
+    top_k_buf: torch.Tensor
+    top_p_buf: torch.Tensor
+    do_sample_buf: torch.Tensor
+    guidance_scale: torch.Tensor
+    # Debug buffers (optional)
+    debug_logits: torch.Tensor | None = None
+    _debug_head_input: torch.Tensor | None = None
+    _debug_probs: torch.Tensor | None = None
+    _debug_probs_slot: torch.Tensor | None = None
+
+
+class DepthDecoderGraph:
+    """
+    Captures the full depth decoder loop as a CUDA graph,
+    using the model's forward with transformers StaticCache.
+
+    For CFG: batch=2 (cond=batch[0], uncond=batch[1]).
+
+    Sampling params are per-sample tensor buffers — no recapture needed:
+        # Scalar (broadcast to all samples):
+        graph.set_temperature(0.5)
+        graph.set_guidance_scale(3.0)
+
+        # Per-sample (tensor):
+        graph.set_temperature(torch.tensor([0.5, 1.2]))
+        graph.set_guidance_scale(torch.tensor([3.0, 1.0]))
+
+        # Or pass directly in run():
+        graph.run(backbone_h, first_cb, temperature=torch.tensor([0.5, 1.2]))
+
+    Batch size changes trigger automatic recapture:
+        graph.ensure_batch_size(4)
+    """
+
+    def __init__(
+        self,
+        depth_decoder,
+        config,
+        device="cuda:0",
+        dtype=torch.bfloat16,
+        do_sample=True,
+        top_k=50,
+        top_p=1.0,
+        temperature=0.9,
+        guidance_scale=3.0,
+        num_codebooks=None,
+        codec_codebook_size=None,
+        fast: bool = False,
+        debug: bool = False,
+        batch_size=2,
+        bucket_sizes: list[int] | None = None,
+    ):
+        self.device = device
+        self.dtype = dtype
+        self.config = config
+        self.debug = debug
+        self.no_graph = (
+            False  # runtime flag: True = skip graph replay (for layer-diff hooks)
+        )
+        self.num_layers = config.num_hidden_layers
+        self.hidden_size = config.hidden_size
+        self.backbone_hidden_size = config.backbone_hidden_size
+        self.vocab_size = config.vocab_size
+        self.codec_codebook_size = (
+            int(codec_codebook_size)
+            if codec_codebook_size is not None
+            else self.vocab_size
+        )
+        if not 0 < self.codec_codebook_size <= self.vocab_size:
+            raise ValueError(
+                f"codec_codebook_size must be in [1, {self.vocab_size}], "
+                f"got {self.codec_codebook_size}"
+            )
+        self.num_codebooks = (
+            num_codebooks if num_codebooks is not None else config.num_codebooks
+        )
+        self.num_decode_codebooks = self.num_codebooks - 1
+        self.max_seq = 2 + self.num_decode_codebooks
+        self.batch_size = batch_size
+        self.half = self._real_batch_size(self.batch_size)
+
+        # Multi-batch bucket support
+        self.bucket_sizes = sorted(bucket_sizes) if bucket_sizes else [batch_size]
+        self._bucket_graphs: dict[int, _BucketState] = {}
+
+        # Extract model components (references, not copies)
+        self.depth_model = depth_decoder.model  # BreezeDepthDecoderModel
+        self.codebooks_head = depth_decoder.codebooks_head  # BreezeCodebooksHead
+        self._orig_codebooks_head = (
+            depth_decoder.codebooks_head
+        )  # unwrapped ref for fake forward hooks
+        self.embed_tokens = self.depth_model.embed_tokens
+        self.inputs_embeds_projector = self.depth_model.inputs_embeds_projector
+        self.backbone_hidden_state_projector = getattr(
+            self.depth_model, "backbone_hidden_state_projector", None
+        )
+        self.audio_embed_size = self.depth_model.audio_embed_size
+
+        # Cast codebooks_head weight to fp32 for numerical stability
+        # (module forward still fires hooks, and F.linear works with fp32 input)
+        self.codebooks_head.weight.data = self.codebooks_head.weight.data.float()
+        self.fast = bool(fast)
+
+        # The fast path uses one maintained compile configuration before
+        # manual CUDA Graph capture. Compile modes are intentionally not public.
+        if self.fast:
+            compile_mode = "default"
+            _limit = self.num_layers * 4 + 16
+            torch._dynamo.config.cache_size_limit = max(
+                torch._dynamo.config.cache_size_limit, _limit
+            )
+            torch._dynamo.config.recompile_limit = max(
+                torch._dynamo.config.recompile_limit, _limit
+            )
+            already_compiled = hasattr(self.depth_model.layers[0], "_orig_mod")
+            if already_compiled:
+                _log.info("Depth decoder layers already compiled, skipping.")
+            else:
+                _log.info("Compiling depth decoder with the fast configuration.")
+                for index, layer in enumerate(self.depth_model.layers):
+                    self.depth_model.layers[index] = torch.compile(
+                        layer, mode=compile_mode, fullgraph=True
+                    )
+                self.depth_model.norm = torch.compile(
+                    self.depth_model.norm, mode=compile_mode, fullgraph=True
+                )
+            if not hasattr(self.codebooks_head, "_orig_mod"):
+                self.codebooks_head = torch.compile(
+                    self.codebooks_head, mode=compile_mode, fullgraph=True
+                )
+
+        # Sampling param defaults — saved for _alloc_buffers / _rebuild_for_batch
+        self._default_temperature = temperature
+        self._default_top_k = top_k
+        self._default_top_p = top_p
+        self._default_do_sample = do_sample
+        self._default_guidance_scale = guidance_scale
+
+        # Fixed MAX_K workspace for graph-safe top_k.
+        # top_k_buf can be set to any value in [1, _max_k] at runtime without recapture.
+        # Default 1024 so users can freely adjust top_k up to 1024.
+        self._max_k = min(1024, self.vocab_size)
+        self._topk_ranks = torch.arange(
+            self._max_k, device=device
+        )  # pre-alloc, never recreated
+        # top_p: always keep the highest-prob token (graph-safe alternative to remove[0]=False)
+        self._keep_first = torch.zeros(self.vocab_size, dtype=torch.bool, device=device)
+        self._keep_first[0] = True
+
+        # Pre-allocate ALL cache_position tensors (no torch.tensor() inside graph)
+        self.prefill_cache_pos = torch.arange(2, device=device)
+        self.decode_cache_positions = [
+            torch.tensor([2 + i], device=device)
+            for i in range(self.num_decode_codebooks - 1)
+        ]
+        self.head_prefill_pos = torch.tensor([1], device=device)
+        self.codebook_offsets = [
+            torch.tensor([i * self.vocab_size], dtype=torch.long, device=device)
+            for i in range(self.num_decode_codebooks)
+        ]
+
+        # Transformers StaticCache
+        self.static_cache = StaticCache(
+            config=config, max_cache_len=self.max_seq, batch_size=self.batch_size
+        )
+
+        self._alloc_buffers()
+
+        self.graph = None
+        self.captured = False
+        self.prefill_attn = None
+        self.decode_attn = None
+
+    @staticmethod
+    def _real_batch_size(batch_size: int) -> int:
+        if batch_size <= 1:
+            return batch_size
+        return batch_size // 2
+
+    # ------------------------------------------------------------------
+    # Runtime setters — no recapture needed
+    # ------------------------------------------------------------------
+
+    def set_temperature(self, v):
+        """v: scalar float (broadcast) or [half] / [half,1] tensor (per-sample)."""
+        if isinstance(v, (int, float)):
+            self.temperature_buf.fill_(v)
+        else:
+            self.temperature_buf.copy_(v.view(self.half, 1))
+
+    def set_top_k(self, v):
+        """v: scalar int (broadcast) or [half] / [half,1] tensor (per-sample)."""
+        if isinstance(v, (int, float)):
+            assert v <= self._max_k, (
+                f"top_k={v} exceeds _max_k={self._max_k}; rebuild graph to increase MAX_K"
+            )
+            self.top_k_buf.fill_(v)
+        else:
+            assert v.max().item() <= self._max_k, (
+                f"top_k max={v.max().item()} exceeds _max_k={self._max_k}"
+            )
+            self.top_k_buf.copy_(v.view(self.half, 1))
+
+    def set_top_p(self, v):
+        """v: scalar float (broadcast) or [half] / [half,1] tensor (per-sample)."""
+        if isinstance(v, (int, float)):
+            self.top_p_buf.fill_(v)
+        else:
+            self.top_p_buf.copy_(v.view(self.half, 1))
+
+    def set_do_sample(self, v):
+        """v: scalar bool (broadcast) or [half] long tensor (per-sample, 0/1)."""
+        if isinstance(v, (bool, int, float)):
+            self.do_sample_buf.fill_(1 if v else 0)
+        else:
+            self.do_sample_buf.copy_(v.view(self.half))
+
+    def set_guidance_scale(self, v):
+        """v: scalar float (broadcast) or [half] / [half,1] tensor (per-sample)."""
+        if isinstance(v, (int, float)):
+            self.guidance_scale.fill_(v)
+        else:
+            self.guidance_scale.copy_(v.view(self.half, 1))
+
+    # ------------------------------------------------------------------
+    # Buffer allocation (called from __init__ and _rebuild_for_batch)
+    # ------------------------------------------------------------------
+
+    def _alloc_buffers(self):
+        """Allocate / re-allocate I/O buffers for current batch_size.
+
+        Also (re-)creates per-sample sampling param buffers whose first
+        dimension is ``self.half``.  On first call the ``_default_*``
+        values are used; on subsequent calls (batch-size change) the
+        current first-sample value is preserved as the uniform default.
+        """
+        self.backbone_hidden_buf = torch.zeros(
+            self.batch_size,
+            self.backbone_hidden_size,
+            dtype=self.dtype,
+            device=self.device,
+        )
+        self.first_cb_token_buf = torch.zeros(
+            self.batch_size, dtype=torch.long, device=self.device
+        )
+        self.output_tokens = torch.zeros(
+            self.batch_size,
+            self.num_decode_codebooks,
+            dtype=torch.long,
+            device=self.device,
+        )
+        self._tok_buf = torch.zeros(
+            self.batch_size, dtype=torch.long, device=self.device
+        )
+        self.prefill_input_ids = torch.zeros(
+            self.batch_size, 2, dtype=torch.long, device=self.device
+        )
+
+        # --- Per-sample sampling param buffers [half, 1] or [half] ---
+        # Preserve current value when resizing, fall back to defaults on first init.
+        _temp = (
+            float(self.temperature_buf[0, 0])
+            if hasattr(self, "temperature_buf") and self.temperature_buf.numel()
+            else self._default_temperature
+        )
+        _top_k = (
+            int(self.top_k_buf[0, 0])
+            if hasattr(self, "top_k_buf") and self.top_k_buf.numel()
+            else self._default_top_k
+        )
+        _top_p = (
+            float(self.top_p_buf[0, 0])
+            if hasattr(self, "top_p_buf") and self.top_p_buf.numel()
+            else self._default_top_p
+        )
+        _do_sample = (
+            int(self.do_sample_buf[0])
+            if hasattr(self, "do_sample_buf") and self.do_sample_buf.numel()
+            else (1 if self._default_do_sample else 0)
+        )
+        _gs = (
+            float(self.guidance_scale[0, 0])
+            if hasattr(self, "guidance_scale") and self.guidance_scale.numel()
+            else self._default_guidance_scale
+        )
+
+        self.temperature_buf = torch.full(
+            (self.half, 1), _temp, dtype=torch.float32, device=self.device
+        )
+        self.top_k_buf = torch.full(
+            (self.half, 1), _top_k, dtype=torch.long, device=self.device
+        )
+        self.top_p_buf = torch.full(
+            (self.half, 1), _top_p, dtype=torch.float32, device=self.device
+        )
+        self.do_sample_buf = torch.full(
+            (self.half,), _do_sample, dtype=torch.long, device=self.device
+        )
+        self.guidance_scale = torch.full(
+            (self.half, 1), _gs, dtype=torch.float32, device=self.device
+        )
+
+        # Debug: per-codebook raw logits [num_decode_codebooks, batch_size, vocab_size]
+        if self.debug:
+            self.debug_logits = torch.zeros(
+                self.num_decode_codebooks,
+                self.batch_size,
+                self.vocab_size,
+                dtype=torch.float32,
+                device=self.device,
+            )
+            # Debug: hidden states before each codebooks_head call (for fake forward hooks)
+            self._debug_head_input = torch.zeros(
+                self.num_decode_codebooks,
+                self.batch_size,
+                1,
+                self.hidden_size,
+                dtype=torch.float32,
+                device=self.device,
+            )
+            # Debug: probs before torch.multinomial (for fake multinomial hooks)
+            self._debug_probs = torch.zeros(
+                self.num_decode_codebooks,
+                self.half,
+                self.vocab_size,
+                dtype=torch.float32,
+                device=self.device,
+            )
+            # Single-slot buffer written by _cfg_sample, copied per-cb in _full_loop
+            self._debug_probs_slot = torch.zeros(
+                self.half,
+                self.vocab_size,
+                dtype=torch.float32,
+                device=self.device,
+            )
+        else:
+            self.debug_logits = None
+            self._debug_head_input = None
+            self._debug_probs = None
+            self._debug_probs_slot = None
+
+    # ------------------------------------------------------------------
+    # Batch size management — auto-recapture on change
+    # ------------------------------------------------------------------
+
+    @torch.inference_mode()
+    def ensure_batch_size(self, batch_size: int):
+        """No-op if batch_size matches current captured graph; else rebuild and recapture."""
+        if batch_size == self.batch_size:
+            return
+        self._rebuild_for_batch(batch_size)
+
+    @torch.inference_mode()
+    def _rebuild_for_batch(self, batch_size: int):
+        self.batch_size = batch_size
+        self.half = self._real_batch_size(self.batch_size)
+        self.static_cache = StaticCache(
+            config=self.config, max_cache_len=self.max_seq, batch_size=batch_size
+        )
+        self._alloc_buffers()
+        self.captured = False
+        self.graph = None
+        self.prefill_attn = None
+        self.decode_attn = None
+        if self.no_graph:
+            self.prepare_eager()
+        else:
+            self.capture()
+
+    @torch.inference_mode()
+    def prepare_eager(self):
+        """Initialize the depth loop for native eager execution only."""
+        self.no_graph = True
+        self._init_cache_layers()
+        self._build_attention_masks()
+        self.captured = False
+        return self
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+
+    def _init_cache_layers(self):
+        """Force lazy initialization of StaticCache layers before graph capture."""
+        config = self.depth_model.config
+        num_kv_heads = getattr(
+            config, "num_key_value_heads", config.num_attention_heads
+        )
+        head_dim = getattr(
+            config, "head_dim", config.hidden_size // config.num_attention_heads
+        )
+        dummy_k = torch.zeros(
+            self.batch_size,
+            num_kv_heads,
+            1,
+            head_dim,
+            dtype=self.dtype,
+            device=self.device,
+        )
+        for layer in self.static_cache.layers:
+            if not layer.is_initialized:
+                layer.lazy_initialization(dummy_k)
+
+    def _make_attn_mask(self, input_embeds, cache_position):
+        return create_causal_mask(
+            config=self.depth_model.config,
+            input_embeds=input_embeds,
+            attention_mask=None,
+            cache_position=cache_position,
+            past_key_values=self.static_cache,
+        )
+
+    def _build_attention_masks(self):
+        dummy_prefill = torch.zeros(
+            self.batch_size, 2, self.hidden_size, dtype=self.dtype, device=self.device
+        )
+        dummy_decode = torch.zeros(
+            self.batch_size, 1, self.hidden_size, dtype=self.dtype, device=self.device
+        )
+        self.prefill_attn = self._make_attn_mask(dummy_prefill, self.prefill_cache_pos)
+        self.decode_attn = []
+        for pos in self.decode_cache_positions:
+            self.decode_attn.append(self._make_attn_mask(dummy_decode, pos))
+
+    def _cfg_sample(self, logits):
+        """Graph-safe CFG + sampling matching HF transformers order.
+
+        HF order: temperature -> top_k -> top_p -> softmax -> multinomial.
+        All filtering on raw logits (not log-space) for parity with
+        TemperatureLogitsWarper / TopKLogitsWarper / TopPLogitsWarper.
+
+        Writes sampled tokens into self._tok_buf [batch_size].
+        For batch_size=2*N: processes N samples, writes to both halves (cond+uncond).
+        """
+        # logits: [batch_size, 1, vocab]
+        if self.batch_size >= 2:
+            cond_logits = logits[: self.half, 0, :]
+            uncond_logits = logits[self.half :, 0, :]
+            cfg = uncond_logits + self.guidance_scale * (cond_logits - uncond_logits)
+        else:
+            cfg = logits[: self.half, 0, :]
+
+        mask_invalid_codec_token_logits(
+            cfg,
+            codebook_size=self.codec_codebook_size,
+            token_vocab_size=self.vocab_size,
+        )
+
+        # temperature scaling (on raw logits, same as TemperatureLogitsWarper)
+        scaled = cfg / self.temperature_buf  # [half, vocab]
+
+        # top_k on raw logits (graph-safe: fixed _max_k workspace)
+        effective_k = torch.where(
+            self.top_k_buf > 0,
+            self.top_k_buf,
+            torch.full_like(self.top_k_buf, self._max_k),
+        )
+        topk_vals, topk_idx = torch.topk(scaled, self._max_k)  # [half, _max_k]
+        topk_mask = self._topk_ranks < effective_k  # [_max_k] broadcasts
+        topk_vals = torch.where(
+            topk_mask, topk_vals, torch.full_like(topk_vals, float("-inf"))
+        )
+        scaled = torch.full_like(scaled, float("-inf")).scatter_(1, topk_idx, topk_vals)
+
+        # top_p on raw logits (HF-style: softmax -> cumsum -> mask -> apply to raw logits)
+        sorted_logits, sorted_idx = torch.sort(scaled, descending=True)  # [half, vocab]
+        cumprobs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+        remove = cumprobs > self.top_p_buf
+        remove = F.pad(
+            remove[..., :-1], (1, 0), value=False
+        )  # shift right, first always kept
+        sorted_logits = torch.where(
+            remove, torch.full_like(sorted_logits, float("-inf")), sorted_logits
+        )
+        scaled = torch.full_like(scaled, float("-inf")).scatter_(
+            1, sorted_idx, sorted_logits
+        )
+
+        # softmax -> multinomial (same as HF _sample)
+        probs = F.softmax(scaled, dim=-1)
+
+        # Debug: save probs before multinomial for fake forward hooks
+        if self._debug_probs_slot is not None:
+            self._debug_probs_slot.copy_(probs)
+
+        # do_sample: compute both paths, select via torch.where (no Python branch in graph)
+        greedy_toks = torch.argmax(cfg, dim=-1)  # [half]
+        sampled_toks = torch.multinomial(probs, 1).squeeze(-1)  # [half]
+        toks = torch.where(
+            self.do_sample_buf.bool(), sampled_toks, greedy_toks
+        )  # [half] per-sample
+
+        # Write to the active slots; duplicate for paired cond/uncond mode.
+        self._tok_buf[: self.half] = toks
+        if self.batch_size >= 2:
+            self._tok_buf[self.half :] = toks
+
+    def _full_loop(self):
+        """The full depth decoder loop on static buffers. Graph-safe."""
+        # Build [B, 2] input_ids: pos0=0 (placeholder), pos1=first_cb_token
+        self.prefill_input_ids[:, 0] = 0
+        self.prefill_input_ids[:, 1] = self.first_cb_token_buf
+
+        # Embed full 2-token sequence, then overwrite pos0 with backbone hidden
+        # (matches parity baseline: embed placeholder, then replace with backbone_h)
+        prefill_embeds = self.embed_tokens(
+            self.prefill_input_ids
+        )  # [B, 2, audio_embed_size]
+        # Project backbone hidden to audio_embed_size if dimensions differ
+        backbone_h = self.backbone_hidden_buf  # [B, backbone_hidden_size]
+        if self.backbone_hidden_state_projector is not None:
+            backbone_h = self.backbone_hidden_state_projector(
+                backbone_h
+            )  # [B, audio_embed_size]
+        prefill_embeds[:, 0] = backbone_h  # overwrite pos0
+        prefill_embeds = self.inputs_embeds_projector(prefill_embeds)  # [B, 2, hidden]
+
+        hidden_states = prefill_embeds
+        position_ids = self.prefill_cache_pos.unsqueeze(0)  # [1, 2]
+        position_embeddings = self.depth_model.rotary_emb(hidden_states, position_ids)
+
+        for decoder_layer in self.depth_model.layers:
+            hidden_states = decoder_layer(
+                hidden_states,
+                attention_mask=self.prefill_attn,
+                position_ids=position_ids,
+                past_key_values=self.static_cache,
+                use_cache=True,
+                cache_position=self.prefill_cache_pos,
+                position_embeddings=position_embeddings,
+            )
+
+        hidden_states = self.depth_model.norm(hidden_states)
+
+        # First codebook logits from position 1 (fp32: weight already cast, input cast here)
+        if self._debug_head_input is not None:
+            self._debug_head_input[0].copy_(hidden_states[:, 1:, :].float())
+        first_logits = self.codebooks_head(
+            hidden_states[:, 1:, :].float(),
+            cache_position=self.head_prefill_pos,
+        )  # [B, 1, vocab] fp32
+
+        if self.debug_logits is not None:
+            self.debug_logits[0].copy_(first_logits[:, 0, :])
+
+        self._cfg_sample(first_logits)
+        if self._debug_probs is not None:
+            self._debug_probs[0].copy_(self._debug_probs_slot)
+        self._tok_buf.clamp_(0, self.vocab_size - 1)
+        self.output_tokens[:, 0] = self._tok_buf
+
+        # Remaining codebooks: decode one at a time
+        for cb_idx in range(1, self.num_decode_codebooks):
+            offset_tok = self._tok_buf + self.codebook_offsets[cb_idx]
+            emb = self.embed_tokens(
+                offset_tok.unsqueeze(1).clamp_(
+                    0, self.num_codebooks * self.vocab_size - 1
+                )
+            )  # [B, 1, backbone_H]
+            emb = self.inputs_embeds_projector(emb)  # [B, 1, hidden]
+
+            cache_pos = self.decode_cache_positions[cb_idx - 1]
+            pos_ids = cache_pos.unsqueeze(0)
+            pos_emb = self.depth_model.rotary_emb(emb, pos_ids)
+
+            hidden_states = emb
+            for decoder_layer in self.depth_model.layers:
+                hidden_states = decoder_layer(
+                    hidden_states,
+                    attention_mask=self.decode_attn[cb_idx - 1],
+                    position_ids=pos_ids,
+                    past_key_values=self.static_cache,
+                    use_cache=True,
+                    cache_position=cache_pos,
+                    position_embeddings=pos_emb,
+                )
+
+            hidden_states = self.depth_model.norm(hidden_states)
+
+            if self._debug_head_input is not None:
+                self._debug_head_input[cb_idx].copy_(hidden_states.float())
+            logits = self.codebooks_head(
+                hidden_states.float(),
+                cache_position=cache_pos,
+            )  # [B, 1, vocab] fp32
+
+            if self.debug_logits is not None:
+                self.debug_logits[cb_idx].copy_(logits[:, 0, :])
+
+            self._cfg_sample(logits)
+            if self._debug_probs is not None:
+                self._debug_probs[cb_idx].copy_(self._debug_probs_slot)
+            self._tok_buf.clamp_(0, self.vocab_size - 1)
+            self.output_tokens[:, cb_idx] = self._tok_buf
+
+    # ------------------------------------------------------------------
+    # Snapshot / restore helpers for multi-bucket capture
+    # ------------------------------------------------------------------
+
+    def _snapshot_state(self) -> _BucketState:
+        """Snapshot current self.* buffers into a _BucketState."""
+        return _BucketState(
+            batch_size=self.batch_size,
+            half=self.half,
+            graph=self.graph,
+            static_cache=self.static_cache,
+            backbone_hidden_buf=self.backbone_hidden_buf,
+            first_cb_token_buf=self.first_cb_token_buf,
+            output_tokens=self.output_tokens,
+            _tok_buf=self._tok_buf,
+            prefill_input_ids=self.prefill_input_ids,
+            prefill_attn=self.prefill_attn,
+            decode_attn=self.decode_attn,
+            temperature_buf=self.temperature_buf,
+            top_k_buf=self.top_k_buf,
+            top_p_buf=self.top_p_buf,
+            do_sample_buf=self.do_sample_buf,
+            guidance_scale=self.guidance_scale,
+            debug_logits=self.debug_logits,
+            _debug_head_input=self._debug_head_input,
+            _debug_probs=self._debug_probs,
+            _debug_probs_slot=self._debug_probs_slot,
+        )
+
+    def _swap_to_bucket(self, state: _BucketState):
+        """Point self.* at the given bucket's buffers (for graph replay)."""
+        self.batch_size = state.batch_size
+        self.half = state.half
+        self.graph = state.graph
+        self.static_cache = state.static_cache
+        self.backbone_hidden_buf = state.backbone_hidden_buf
+        self.first_cb_token_buf = state.first_cb_token_buf
+        self.output_tokens = state.output_tokens
+        self._tok_buf = state._tok_buf
+        self.prefill_input_ids = state.prefill_input_ids
+        self.prefill_attn = state.prefill_attn
+        self.decode_attn = state.decode_attn
+        self.temperature_buf = state.temperature_buf
+        self.top_k_buf = state.top_k_buf
+        self.top_p_buf = state.top_p_buf
+        self.do_sample_buf = state.do_sample_buf
+        self.guidance_scale = state.guidance_scale
+        self.debug_logits = state.debug_logits
+        self._debug_head_input = state._debug_head_input
+        self._debug_probs = state._debug_probs
+        self._debug_probs_slot = state._debug_probs_slot
+
+    # ------------------------------------------------------------------
+    # Bucket selection
+    # ------------------------------------------------------------------
+
+    def _select_bucket(self, actual_batch: int) -> int | None:
+        """Return smallest bucket_size >= actual_batch, or None if none fits."""
+        for bsz in self.bucket_sizes:
+            if bsz >= actual_batch:
+                return bsz
+        return None
+
+    @staticmethod
+    def _copy_cfg_padded(dst, src, actual_batch: int, bucket_half: int):
+        """Copy [cond..., uncond...] rows into a padded CFG bucket.
+
+        ``_cfg_sample`` slices the static bucket as
+        ``[:bucket_half]`` for conditional rows and ``[bucket_half:]`` for
+        unconditional rows. When an actual batch is padded to a larger bucket,
+        real uncond rows must start at ``bucket_half`` rather than immediately
+        after the real cond rows.
+        """
+        actual_half = actual_batch // 2
+        bucket_batch = bucket_half * 2
+
+        dst[:actual_half].copy_(src[:actual_half])
+        dst[bucket_half : bucket_half + actual_half].copy_(
+            src[actual_half:actual_batch]
+        )
+
+        if actual_half < bucket_half:
+            dst[actual_half:bucket_half].zero_()
+            dst[bucket_half + actual_half : bucket_batch].zero_()
+
+    # ------------------------------------------------------------------
+    # Padded param setter (for bucket padding)
+    # ------------------------------------------------------------------
+
+    def _set_param_padded(self, buf, value, actual_n, bucket_n, is_1d=False):
+        """Set a sampling param buffer with padding for bucket size mismatch.
+
+        Args:
+            buf: target buffer, shape [bucket_n, 1] or [bucket_n]
+            value: scalar or tensor
+            actual_n: number of actual samples
+            bucket_n: bucket's half size
+            is_1d: True if buf is 1D (e.g. do_sample_buf)
+        """
+        if isinstance(value, (int, float, bool)):
+            buf.fill_(int(value) if is_1d else value)
+        else:
+            t = value.view(actual_n) if is_1d else value.view(actual_n, -1)
+            if is_1d:
+                buf[:actual_n].copy_(t)
+                if actual_n < bucket_n:
+                    buf[actual_n:].fill_(t[0].item())
+            else:
+                buf[:actual_n].copy_(t[:, : buf.shape[-1]])
+                if actual_n < bucket_n:
+                    buf[actual_n:].fill_(t[0, 0].item())
+
+    # ------------------------------------------------------------------
+    # Capture: single bucket + multi-bucket entry point
+    # ------------------------------------------------------------------
+
+    @torch.inference_mode()
+    def _capture_single_bucket(self, bsz: int, num_warmup: int):
+        """Capture a CUDA graph for a single batch size."""
+        # Set up state for this bucket size
+        self.batch_size = bsz
+        self.half = self._real_batch_size(bsz)
+        self.static_cache = StaticCache(
+            config=self.config, max_cache_len=self.max_seq, batch_size=bsz
+        )
+        self._alloc_buffers()
+        self._init_cache_layers()
+        self._build_attention_masks()
+
+        # Warmup
+        for _ in range(num_warmup):
+            self.static_cache.reset()
+            self._full_loop()
+        torch.cuda.synchronize(device=self.device)
+
+        # Capture
+        s = torch.cuda.Stream(device=self.device)
+        s.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(s):
+            self.graph = torch.cuda.CUDAGraph()
+            self.static_cache.reset()
+            self._full_loop()
+            torch.cuda.synchronize(device=self.device)
+
+            self.static_cache.reset()
+            with torch.cuda.graph(self.graph):
+                self._full_loop()
+
+        torch.cuda.current_stream().wait_stream(s)
+        torch.cuda.synchronize(device=self.device)
+
+        # Save state for this bucket
+        self._bucket_graphs[bsz] = self._snapshot_state()
+        print(f"  Depth decoder CUDA graph captured for bucket_size={bsz}")
+
+    @torch.inference_mode()
+    def capture(self, num_warmup=3):
+        """Warmup and capture CUDA graphs for all bucket sizes."""
+        if self.fast:
+            num_warmup = max(num_warmup, 8)
+        print(
+            f"Capturing depth decoder graphs for buckets={self.bucket_sizes} "
+            f"({num_warmup} warmup runs, fast={self.fast})..."
+        )
+
+        for bsz in self.bucket_sizes:
+            self._capture_single_bucket(bsz, num_warmup)
+
+        # Leave self.* pointing at the smallest bucket (default)
+        self._swap_to_bucket(self._bucket_graphs[self.bucket_sizes[0]])
+        self.captured = True
+        print(
+            f"Depth decoder CUDA graphs captured for all {len(self.bucket_sizes)} buckets!"
+        )
+        return self
+
+    # ------------------------------------------------------------------
+    # Eager fallback (for batch sizes exceeding all buckets)
+    # ------------------------------------------------------------------
+
+    @torch.inference_mode()
+    def _run_eager(
+        self,
+        backbone_hidden,
+        first_cb_token,
+        temperature=None,
+        top_k=None,
+        top_p=None,
+        do_sample=None,
+        guidance_scale=None,
+    ):
+        """Fallback: ensure_batch_size + _full_loop (no CUDA graph)."""
+        actual_batch = backbone_hidden.shape[0]
+        self.ensure_batch_size(actual_batch)
+
+        self.backbone_hidden_buf.copy_(backbone_hidden)
+        self.first_cb_token_buf.copy_(first_cb_token)
+
+        if temperature is not None:
+            self.set_temperature(temperature)
+        if top_k is not None:
+            self.set_top_k(top_k)
+        if top_p is not None:
+            self.set_top_p(top_p)
+        if do_sample is not None:
+            self.set_do_sample(do_sample)
+        if guidance_scale is not None:
+            self.set_guidance_scale(guidance_scale)
+
+        self.static_cache.reset()
+        self._full_loop()
+        return self.output_tokens[: self.half].clone()
+
+    # ------------------------------------------------------------------
+    # Run: bucket selection + graph replay
+    # ------------------------------------------------------------------
+
+    @torch.inference_mode()
+    def run(
+        self,
+        backbone_hidden,
+        first_cb_token,
+        temperature=None,
+        top_k=None,
+        top_p=None,
+        do_sample=None,
+        guidance_scale=None,
+    ):
+        """
+        Run the captured graph with automatic bucket selection.
+
+        backbone_hidden: [batch_size, backbone_hidden_size]
+        first_cb_token: [batch_size] long tensor
+
+        Optional per-sample sampling params (applied before graph replay):
+          temperature:    float or [half] / [half,1] tensor
+          top_k:          int   or [half] / [half,1] tensor
+          top_p:          float or [half] / [half,1] tensor
+          do_sample:      bool  or [half] long tensor (0/1)
+          guidance_scale: float or [half] / [half,1] tensor
+
+        Returns: [actual_half, num_decode_codebooks] long tensor of codebook tokens
+        """
+        actual_batch = backbone_hidden.shape[0]
+        actual_half = self._real_batch_size(actual_batch)
+
+        # Safety guard: single CFG expects paired cond+uncond rows. no-CFG uses
+        # batch=1 and is valid.
+        if actual_batch != 1 and actual_batch % 2 != 0:
+            _log.warning(
+                "DepthDecoderGraph.run: invalid CFG actual_batch=%d, "
+                "returning zeros to avoid CUDA graph mismatch",
+                actual_batch,
+            )
+            return torch.zeros(
+                1,
+                self.num_decode_codebooks,
+                dtype=torch.long,
+                device=self.device,
+            )
+
+        # Eager fallback: no_graph flag or no fitting bucket
+        bucket_bsz = self._select_bucket(actual_batch)
+        if self.no_graph or bucket_bsz is None:
+            return self._run_eager(
+                backbone_hidden,
+                first_cb_token,
+                temperature,
+                top_k,
+                top_p,
+                do_sample,
+                guidance_scale,
+            )
+
+        # Swap to the selected bucket's state
+        state = self._bucket_graphs[bucket_bsz]
+        self._swap_to_bucket(state)
+        bucket_half = self.half
+
+        if actual_batch == 1:
+            self.backbone_hidden_buf[:1].copy_(backbone_hidden[:1])
+            self.first_cb_token_buf[:1].copy_(first_cb_token[:1])
+        else:
+            # Preserve the CFG layout expected by _cfg_sample:
+            # [cond real][cond pad][uncond real][uncond pad].
+            self._copy_cfg_padded(
+                self.backbone_hidden_buf, backbone_hidden, actual_batch, bucket_half
+            )
+            self._copy_cfg_padded(
+                self.first_cb_token_buf, first_cb_token, actual_batch, bucket_half
+            )
+
+        # Set sampling params with padding
+        if temperature is not None:
+            self._set_param_padded(
+                self.temperature_buf, temperature, actual_half, bucket_half
+            )
+        if top_k is not None:
+            self._set_param_padded(self.top_k_buf, top_k, actual_half, bucket_half)
+        if top_p is not None:
+            self._set_param_padded(self.top_p_buf, top_p, actual_half, bucket_half)
+        if do_sample is not None:
+            self._set_param_padded(
+                self.do_sample_buf, do_sample, actual_half, bucket_half, is_1d=True
+            )
+        if guidance_scale is not None:
+            self._set_param_padded(
+                self.guidance_scale, guidance_scale, actual_half, bucket_half
+            )
+
+        # Replay graph
+        self.static_cache.reset()
+        state.graph.replay()
+
+        if self.debug and self._debug_head_input is not None:
+            head = self._orig_codebooks_head
+            for cb_idx in range(self.num_decode_codebooks):
+                if cb_idx == 0:
+                    head(
+                        self._debug_head_input[0], cache_position=self.head_prefill_pos
+                    )
+                else:
+                    head(
+                        self._debug_head_input[cb_idx],
+                        cache_position=self.decode_cache_positions[cb_idx - 1],
+                    )
+                torch.multinomial(
+                    self._debug_probs[cb_idx], 1, _hook_only_for_debug=True
+                )
+
+        # Slice output to actual_half (discard padding)
+        return self.output_tokens[:actual_half].clone()

models/cudagraph/sampling.py

@@ -0,0 +1,99 @@
+"""Shared sampling helpers for backbone and depth decoder generation."""
+
+from __future__ import annotations
+
+import os
+import random
+from collections.abc import Iterable
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+
+def apply_repetition_penalty(
+    logits: torch.Tensor,
+    token_history: torch.Tensor,
+    repetition_penalty: float,
+) -> torch.Tensor:
+    """Apply repetition penalty to logits in-place and return them.
+
+    Args:
+        logits: Tensor shaped [1, 1, vocab] or [1, vocab].
+        token_history: 1-D tensor of previously generated token ids.
+        repetition_penalty: HF-style repetition penalty (>1.0).
+    """
+    if repetition_penalty == 1.0 or token_history.numel() == 0:
+        return logits
+    unique_toks = token_history.unique()
+    tok_logits = logits[..., unique_toks]
+    logits[..., unique_toks] = torch.where(
+        tok_logits > 0, tok_logits / repetition_penalty, tok_logits * repetition_penalty
+    )
+    return logits
+
+
+def sample_logits(
+    logits: torch.Tensor,
+    *,
+    temperature: float,
+    top_k: int,
+    top_p: float,
+    do_sample: bool,
+    token_history: torch.Tensor | None = None,
+    repetition_penalty: float = 1.0,
+    suppress_mask: torch.Tensor | None = None,
+    suppress_tokens: Iterable[int] | None = None,
+) -> torch.Tensor:
+    """Sample a token from logits.
+
+    HF-compatible order: suppress -> temperature -> top_k -> top_p -> softmax -> sample.
+    Matches transformers logits_processor (TemperatureLogitsWarper, TopKLogitsWarper,
+    TopPLogitsWarper) followed by softmax + multinomial.
+    """
+    logits = logits.clone().float()
+    if token_history is not None:
+        apply_repetition_penalty(logits, token_history, repetition_penalty)
+    if suppress_mask is not None:
+        logits[..., suppress_mask] = float("-inf")
+    if suppress_tokens:
+        logits[..., list(suppress_tokens)] = float("-inf")
+    if not do_sample:
+        return torch.argmax(logits, dim=-1)
+    # temperature scaling (on raw logits, same as TemperatureLogitsWarper)
+    logits = logits / temperature
+    # top_k filtering (on raw logits, same as TopKLogitsWarper)
+    if top_k > 0:
+        k = min(top_k, logits.size(-1))
+        topk_vals, _ = torch.topk(logits, k)
+        threshold = topk_vals[..., -1:]
+        logits = logits.masked_fill(logits < threshold, float("-inf"))
+    # top_p filtering (on raw logits, same as TopPLogitsWarper)
+    if top_p < 1.0:
+        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+        cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+        sorted_indices_to_remove = cumulative_probs > top_p
+        # HF-style shift: keep the token that pushes cumsum over top_p
+        sorted_indices_to_remove = F.pad(
+            sorted_indices_to_remove[..., :-1], (1, 0), value=False
+        )
+        sorted_logits[sorted_indices_to_remove] = float("-inf")
+        logits = torch.full_like(logits, float("-inf")).scatter_(
+            -1, sorted_indices, sorted_logits
+        )
+    # softmax -> multinomial (same as HF _sample)
+    probs = F.softmax(logits, dim=-1)
+    return torch.multinomial(probs, 1).squeeze(-1)
+
+
+def set_deterministic(seed=42):
+    """Enable full deterministic mode and seed all RNGs for exact reproducibility."""
+    random.seed(seed)
+    os.environ["PYTHONHASHSEED"] = str(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
+    torch.use_deterministic_algorithms(True, warn_only=True)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False

models/fast_streaming.py

@@ -0,0 +1,943 @@
+from __future__ import annotations
+
+import json
+import time
+import uuid
+from collections.abc import Iterator
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Literal
+
+import numpy as np
+import torch
+
+from .cudagraph.backbone_graph import BackboneGraph
+from .cudagraph.backbone_prefill_graph import BackbonePrefillGraphCache
+from .cudagraph.depth_decoder_graph import DepthDecoderGraph
+from .cudagraph.sampling import sample_logits
+from .warmup_profile import FastStreamingWarmupProfile
+
+FastCfgMode = Literal["no_cfg", "single_cfg"]
+
+_DUAL_CFG_KEYS = (
+    "cfg_scale_ref",
+    "cfg_scale_ins",
+    "cfg_uncond_prompt_ids",
+    "cfg_ref_prompt_ids",
+    "cfg_ins_prompt_ids",
+)
+
+
+@dataclass(frozen=True)
+class FastStreamingConfig:
+    max_new_tokens: int = 750
+    max_seq_len: int = 1024
+    collect_timing: bool = False
+    fast_all: bool | None = None
+    fast_text_encoder: bool = False
+    fast_backbone_prefill: bool = False
+    fast_backbone_decode: bool = False
+    fast_depth_decoder: bool = False
+    fast_codec: bool = False
+    temperature: float | None = None
+    top_k: int | None = None
+    top_p: float | None = None
+    do_sample: bool | None = None
+    repetition_penalty: float = 1.1
+
+    def stage_fast(self, stage: str) -> bool:
+        """Resolve the master switch before the per-stage setting."""
+        if self.fast_all is not None:
+            return self.fast_all
+        field_name = f"fast_{stage}"
+        if field_name not in self.__dataclass_fields__ or field_name == "fast_all":
+            raise ValueError(f"unknown fast path stage: {stage!r}")
+        return bool(getattr(self, field_name))
+
+
+@dataclass(frozen=True)
+class FastStreamingChunk:
+    audio: np.ndarray
+    sample_rate: int
+    codec_frames: int
+    is_final: bool
+    timing: dict[str, float | int | bool] = field(default_factory=dict)
+
+
+@dataclass(frozen=True)
+class FastCfgSelection:
+    mode: FastCfgMode
+    guidance_scale: float
+    use_negative_as_main: bool
+
+
+@dataclass
+class _BranchBatch:
+    inputs_embeds: torch.Tensor
+    attention_mask: torch.Tensor
+    branch_batch_size: int
+    cfg: FastCfgSelection
+
+
+def reject_dual_cfg(inputs: dict[str, Any]) -> None:
+    present = [key for key in _DUAL_CFG_KEYS if inputs.get(key) is not None]
+    if present:
+        raise ValueError(
+            "fast streaming supports only no_cfg and single_cfg; "
+            f"dual CFG fields are not supported: {present}"
+        )
+
+
+def select_fast_cfg(inputs: dict[str, Any]) -> FastCfgSelection:
+    reject_dual_cfg(inputs)
+    cfg_scale = float(inputs.get("cfg_scale", 1.0))
+    has_negative = inputs.get("cfg_negative_prompt_ids") is not None
+    if cfg_scale == 0.0 and has_negative:
+        return FastCfgSelection(
+            mode="no_cfg", guidance_scale=1.0, use_negative_as_main=True
+        )
+    if cfg_scale != 1.0:
+        if not has_negative:
+            raise ValueError(
+                "single_cfg requires cfg_negative_prompt_ids when cfg_scale != 1.0"
+            )
+        return FastCfgSelection(
+            mode="single_cfg", guidance_scale=cfg_scale, use_negative_as_main=False
+        )
+    return FastCfgSelection(
+        mode="no_cfg", guidance_scale=1.0, use_negative_as_main=False
+    )
+
+
+def is_backbone_eos_token(token: torch.Tensor | int, config: Any) -> bool:
+    token_id = int(token.item() if isinstance(token, torch.Tensor) else token)
+    return token_id == int(config.vocab_size)
+
+
+def is_terminal_pad_frame(frame: torch.Tensor, config: Any) -> bool:
+    return bool((frame == int(config.codebook_pad_token_id)).all().item())
+
+
+def should_decode_codec_frame(frame: torch.Tensor, config: Any) -> bool:
+    return not is_terminal_pad_frame(frame, config)
+
+
+def _get_device(model: torch.nn.Module) -> torch.device:
+    try:
+        return next(model.parameters()).device
+    except StopIteration:
+        return torch.device("cpu")
+
+
+def _get_dtype(model: torch.nn.Module) -> torch.dtype:
+    try:
+        return next(model.parameters()).dtype
+    except StopIteration:
+        return torch.bfloat16
+
+
+def _left_pad_tensor(
+    tensor: torch.Tensor, target_len: int, value: float = 0
+) -> torch.Tensor:
+    pad_len = target_len - tensor.shape[1]
+    if pad_len <= 0:
+        return tensor
+    pad_shape = (tensor.shape[0], pad_len, *tensor.shape[2:])
+    pad = torch.full(pad_shape, value, dtype=tensor.dtype, device=tensor.device)
+    return torch.cat([pad, tensor], dim=1)
+
+
+def _extract_audio_np(audio: torch.Tensor) -> np.ndarray:
+    while audio.dim() > 1:
+        audio = audio[0]
+    return audio.detach().float().cpu().numpy()
+
+
+class FastBreezeStreamingRuntime:
+    def __init__(
+        self,
+        model: torch.nn.Module,
+        audio_tokenizer: Any,
+        config: FastStreamingConfig | None = None,
+        *,
+        tokenizer: Any | None = None,
+    ) -> None:
+        self.model = model
+        self.audio_tokenizer = audio_tokenizer
+        self.tokenizer = tokenizer
+        self.config = config or FastStreamingConfig()
+        self._fast_text_encoder = self.config.stage_fast("text_encoder")
+        self._fast_backbone_prefill = self.config.stage_fast("backbone_prefill")
+        self._fast_backbone_decode = self.config.stage_fast("backbone_decode")
+        self._fast_depth_decoder = self.config.stage_fast("depth_decoder")
+        self._fast_codec = self.config.stage_fast("codec")
+        self.model._fast_text_encoder_cudagraph = self._fast_text_encoder
+        self._codec_chunk_frames = 1 if self._fast_codec else 2
+        self.device = _get_device(model)
+        self.dtype = _get_dtype(model)
+        self._backbone_graph: BackboneGraph | None = None
+        self._backbone_graphs: dict[int, BackboneGraph] = {}
+        self._backbone_prefill_graph: BackbonePrefillGraphCache | None = None
+        self._backbone_prefill_graphs: dict[int, BackbonePrefillGraphCache] = {}
+        self._depth_decoder_graph: DepthDecoderGraph | None = None
+        self._codec_runtime = None
+        self._warmup_profile: FastStreamingWarmupProfile | None = None
+        self._warmup_manifest: dict[str, Any] | None = None
+        self._frozen_branch_batch_sizes: frozenset[int] | None = None
+        self._codec_codebook_size = int(self.model.config.codec_config.codebook_size)
+        self._reserved_codec_token_ids = tuple(
+            range(self._codec_codebook_size, int(self.model.config.vocab_size))
+        )
+
+        if self.device.type != "cuda":
+            raise RuntimeError("fast streaming requires a CUDA device")
+        if self.config.repetition_penalty <= 0:
+            raise ValueError("repetition_penalty must be > 0")
+
+    @property
+    def sample_rate(self) -> int:
+        return int(self.model.config.codec_config.sampling_rate)
+
+    def _sampling_params(self, generation_config: Any) -> dict[str, Any]:
+        def _value(name: str, override: Any, default: Any) -> Any:
+            if override is not None:
+                return override
+            value = getattr(generation_config, name, default)
+            return default if value is None else value
+
+        return {
+            "temperature": float(_value("temperature", self.config.temperature, 1.0)),
+            "top_k": int(_value("top_k", self.config.top_k, 0)),
+            "top_p": float(_value("top_p", self.config.top_p, 1.0)),
+            "do_sample": bool(_value("do_sample", self.config.do_sample, True)),
+        }
+
+    def _ensure_graphs(
+        self,
+        branch_batch_size: int,
+        guidance_scale: float,
+        *,
+        depth_bucket_sizes: list[int] | None = None,
+    ) -> None:
+        if (
+            self._frozen_branch_batch_sizes is not None
+            and branch_batch_size not in self._frozen_branch_batch_sizes
+        ):
+            raise RuntimeError(
+                "request branch batch size "
+                f"{branch_batch_size} is not declared by the frozen warmup config; "
+                f"expected one of {sorted(self._frozen_branch_batch_sizes)}"
+            )
+
+        backbone_graph = self._backbone_graphs.get(branch_batch_size)
+        if backbone_graph is None:
+            backbone_graph = BackboneGraph(
+                backbone_model=self.model.backbone_model,
+                lm_head=self.model.lm_head,
+                embed_tokens=self.model.backbone_model.embed_tokens,
+                config=self.model.config,
+                device=str(self.device),
+                dtype=self.dtype,
+                max_seq_len=self.config.max_seq_len,
+                guidance_scale=guidance_scale,
+                batch_size=branch_batch_size,
+            )
+            if self._fast_backbone_decode:
+                backbone_graph.capture()
+            else:
+                backbone_graph.prepare_eager()
+            self._backbone_graphs[branch_batch_size] = backbone_graph
+        else:
+            backbone_graph.guidance_scale.fill_(guidance_scale)
+        self._backbone_graph = backbone_graph
+        self._backbone_prefill_graph = self._backbone_prefill_graphs.get(
+            branch_batch_size
+        )
+
+        if self._depth_decoder_graph is None:
+            depth_gen = self.model.depth_decoder.generation_config
+            depth_params = self._sampling_params(depth_gen)
+            depth_bucket_sizes = depth_bucket_sizes or [1, 2]
+            self._depth_decoder_graph = DepthDecoderGraph(
+                depth_decoder=self.model.depth_decoder,
+                config=self.model.config.depth_decoder_config,
+                device=str(self.device),
+                dtype=self.dtype,
+                guidance_scale=guidance_scale,
+                num_codebooks=self.model.config.num_codebooks,
+                codec_codebook_size=self._codec_codebook_size,
+                fast=self._fast_depth_decoder,
+                batch_size=depth_bucket_sizes[0],
+                bucket_sizes=depth_bucket_sizes,
+                **depth_params,
+            )
+            if self._fast_depth_decoder:
+                self._depth_decoder_graph.capture()
+            else:
+                self._depth_decoder_graph.prepare_eager()
+        else:
+            self._depth_decoder_graph.set_guidance_scale(guidance_scale)
+
+    def _codec(self):
+        if self._codec_runtime is not None:
+            return self._codec_runtime
+        from .stream_runtime import MultiRequestStreamRuntime, QwenStreamRuntimeConfig
+
+        if getattr(self.audio_tokenizer, "model", None) is None:
+            raise RuntimeError(
+                "audio_tokenizer.model is required for fast streaming codec decode"
+            )
+        try:
+            codec_dtype = next(self.audio_tokenizer.model.parameters()).dtype
+        except StopIteration:
+            codec_dtype = torch.float32
+        runtime_config = QwenStreamRuntimeConfig(
+            chunk_frames=self._codec_chunk_frames,
+            num_lanes=1,
+            max_active_reqs=1,
+            fast=self._fast_codec,
+            device=self.device,
+            dtype=codec_dtype,
+        )
+        self._codec_runtime = MultiRequestStreamRuntime(
+            self.audio_tokenizer, runtime_config
+        )
+        return self._codec_runtime
+
+    def _merge_branch(
+        self,
+        *,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        text_ids_mask: torch.Tensor,
+        text_ids_len: torch.Tensor,
+        input_values: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        merged = self.model._merge_input_ids_with_input_values(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            text_ids_mask=text_ids_mask,
+            text_ids_len=text_ids_len,
+            input_values=input_values,
+        )
+        return merged["inputs_embeds"], attention_mask
+
+    def _merge_cfg_branches(
+        self, inputs: dict[str, Any]
+    ) -> tuple[torch.Tensor, torch.Tensor] | None:
+        """Merge cond/uncond together so their text segments use one batched graph."""
+        cond_values = inputs.get("input_values")
+        uncond_values = inputs.get("cfg_negative_input_values")
+        if (cond_values is None) != (uncond_values is None):
+            return None
+        if cond_values is not None and tuple(cond_values.shape[1:]) != tuple(
+            uncond_values.shape[1:]
+        ):
+            return None
+
+        cond_ids = inputs["input_ids"]
+        uncond_ids = inputs["cfg_negative_prompt_ids"]
+        if cond_ids.shape[0] != 1 or uncond_ids.shape[0] != 1:
+            return None
+        max_len = max(cond_ids.shape[1], uncond_ids.shape[1])
+        input_ids = torch.cat(
+            [
+                _left_pad_tensor(cond_ids, max_len, 0),
+                _left_pad_tensor(uncond_ids, max_len, 0),
+            ],
+            dim=0,
+        )
+        attention_mask = torch.cat(
+            [
+                _left_pad_tensor(inputs["attention_mask"], max_len, 0),
+                _left_pad_tensor(
+                    inputs["cfg_negative_prompt_attention_mask"], max_len, 0
+                ),
+            ],
+            dim=0,
+        )
+        text_ids_mask = torch.cat(
+            [
+                _left_pad_tensor(inputs["text_ids_mask"], max_len, False),
+                _left_pad_tensor(inputs["cfg_negative_text_ids_mask"], max_len, False),
+            ],
+            dim=0,
+        )
+        text_ids_len = torch.cat(
+            [inputs["text_ids_len"], inputs["cfg_negative_text_ids_len"]], dim=0
+        )
+        input_values = (
+            None
+            if cond_values is None
+            else torch.cat([cond_values, uncond_values], dim=0)
+        )
+        merged = self.model._merge_input_ids_with_input_values(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            text_ids_mask=text_ids_mask,
+            text_ids_len=text_ids_len,
+            input_values=input_values,
+        )
+        return merged["inputs_embeds"].contiguous(), attention_mask.contiguous()
+
+    def _build_branch_batch(self, inputs: dict[str, Any]) -> _BranchBatch:
+        cfg = select_fast_cfg(inputs)
+        if cfg.use_negative_as_main:
+            embeds, mask = self._merge_branch(
+                input_ids=inputs["cfg_negative_prompt_ids"],
+                attention_mask=inputs["cfg_negative_prompt_attention_mask"],
+                text_ids_mask=inputs["cfg_negative_text_ids_mask"],
+                text_ids_len=inputs["cfg_negative_text_ids_len"],
+                input_values=inputs.get("cfg_negative_input_values"),
+            )
+            return _BranchBatch(embeds.contiguous(), mask.contiguous(), 1, cfg)
+
+        if cfg.mode == "no_cfg":
+            cond_embeds, cond_mask = self._merge_branch(
+                input_ids=inputs["input_ids"],
+                attention_mask=inputs["attention_mask"],
+                text_ids_mask=inputs["text_ids_mask"],
+                text_ids_len=inputs["text_ids_len"],
+                input_values=inputs.get("input_values"),
+            )
+            return _BranchBatch(
+                cond_embeds.contiguous(), cond_mask.contiguous(), 1, cfg
+            )
+
+        joint = self._merge_cfg_branches(inputs) if self._fast_text_encoder else None
+        if joint is not None:
+            inputs_embeds, attention_mask = joint
+            return _BranchBatch(inputs_embeds, attention_mask, 2, cfg)
+
+        cond_embeds, cond_mask = self._merge_branch(
+            input_ids=inputs["input_ids"],
+            attention_mask=inputs["attention_mask"],
+            text_ids_mask=inputs["text_ids_mask"],
+            text_ids_len=inputs["text_ids_len"],
+            input_values=inputs.get("input_values"),
+        )
+        uncond_embeds, uncond_mask = self._merge_branch(
+            input_ids=inputs["cfg_negative_prompt_ids"],
+            attention_mask=inputs["cfg_negative_prompt_attention_mask"],
+            text_ids_mask=inputs["cfg_negative_text_ids_mask"],
+            text_ids_len=inputs["cfg_negative_text_ids_len"],
+            input_values=inputs.get("cfg_negative_input_values"),
+        )
+        max_len = max(cond_embeds.shape[1], uncond_embeds.shape[1])
+        inputs_embeds = torch.cat(
+            [
+                _left_pad_tensor(cond_embeds, max_len, 0),
+                _left_pad_tensor(uncond_embeds, max_len, 0),
+            ],
+            dim=0,
+        ).contiguous()
+        attention_mask = torch.cat(
+            [
+                _left_pad_tensor(cond_mask, max_len, 0),
+                _left_pad_tensor(uncond_mask, max_len, 0),
+            ],
+            dim=0,
+        ).contiguous()
+        return _BranchBatch(inputs_embeds, attention_mask, 2, cfg)
+
+    def _decode_codec_frames(
+        self,
+        *,
+        frames: list[torch.Tensor],
+        request_id: str,
+        reset: bool,
+        is_final: bool,
+        timing: dict[str, float | int | bool],
+    ) -> FastStreamingChunk:
+        frame_tensor = torch.stack(frames).to(self.device, dtype=torch.long)
+        codes = frame_tensor.transpose(0, 1).unsqueeze(0).contiguous()
+        codec_started = time.perf_counter()
+        audio = self._codec().decode_request_chunk(request_id, codes, reset=reset)
+        codec_launch_ms = (time.perf_counter() - codec_started) * 1000.0
+        d2h_started = time.perf_counter()
+        audio_np = _extract_audio_np(audio)
+        d2h_ms = (time.perf_counter() - d2h_started) * 1000.0
+        timing = {
+            **timing,
+            "codec_launch_ms": codec_launch_ms,
+            "audio_d2h_ms": d2h_ms,
+        }
+        return FastStreamingChunk(
+            audio=audio_np,
+            sample_rate=self.sample_rate,
+            codec_frames=len(frames),
+            is_final=is_final,
+            timing=timing,
+        )
+
+    @property
+    def fast_enabled(self) -> bool:
+        return any(
+            (
+                self._fast_text_encoder,
+                self._fast_backbone_prefill,
+                self._fast_backbone_decode,
+                self._fast_depth_decoder,
+                self._fast_codec,
+            )
+        )
+
+    @property
+    def codec_chunk_frames(self) -> int:
+        return self._codec_chunk_frames
+
+    @torch.inference_mode()
+    def warmup_from_profile(
+        self,
+        profile: FastStreamingWarmupProfile,
+        *,
+        manifest_path: str | Path | None = None,
+    ) -> dict[str, Any]:
+        """Deterministically create every graph declared by a serving profile."""
+        if self._warmup_manifest is not None:
+            raise RuntimeError("runtime has already completed profile warmup")
+        if self.tokenizer is None:
+            raise ValueError(
+                "profile warmup requires FastBreezeStreamingRuntime(tokenizer=...)"
+            )
+        if profile.codec_chunk_frames != self._codec_chunk_frames:
+            raise ValueError(
+                f"profile codec chunk_frames={profile.codec_chunk_frames} does not match "
+                f"runtime chunk_frames={self._codec_chunk_frames}"
+            )
+        if profile.codec_num_lanes != 1:
+            raise ValueError("fast streaming runtime supports exactly one codec lane")
+        if any(
+            graph.sequence_length > self.config.max_seq_len
+            for graph in profile.backbone_prefill_graphs
+        ):
+            raise ValueError(
+                "profile prefill sequence_length exceeds runtime max_seq_len"
+            )
+
+        total_started = time.perf_counter()
+        stages: dict[str, Any] = {}
+
+        graph_started = time.perf_counter()
+        cfg_scale_by_batch = {
+            1 if cfg_scale == 1.0 else 2: cfg_scale for cfg_scale in profile.cfg_scales
+        }
+        for branch_batch_size in profile.backbone_decode_branch_batch_sizes:
+            self._ensure_graphs(
+                branch_batch_size,
+                cfg_scale_by_batch[branch_batch_size],
+                depth_bucket_sizes=list(profile.depth_decoder_batch_sizes),
+            )
+        assert self._depth_decoder_graph is not None
+        torch.cuda.synchronize(self.device)
+        graph_elapsed_ms = (time.perf_counter() - graph_started) * 1000.0
+        stages["backbone_decode"] = {
+            "fast": self._fast_backbone_decode,
+            "graphs": [
+                {"branch_batch_size": batch_size}
+                for batch_size in sorted(self._backbone_graphs)
+            ],
+            "elapsed_ms_with_depth_decoder": graph_elapsed_ms,
+        }
+        stages["depth_decoder"] = {
+            "fast": self._fast_depth_decoder,
+            "graphs": [
+                {"batch_size": int(batch_size)}
+                for batch_size in sorted(self._depth_decoder_graph._bucket_graphs)
+            ],
+            "elapsed_ms_with_backbone_decode": graph_elapsed_ms,
+        }
+
+        sampling_started = time.perf_counter()
+        backbone_params = self._sampling_params(self.model.generation_config)
+        sampling_logits = torch.zeros(
+            1,
+            int(self.model.config.vocab_size) + 1,
+            dtype=torch.float32,
+            device=self.device,
+        )
+        for _ in range(2):
+            sample_logits(
+                sampling_logits,
+                suppress_tokens=self._reserved_codec_token_ids,
+                **backbone_params,
+            )
+        torch.cuda.synchronize(self.device)
+        stages["backbone_decode"]["sampling_warmup_ms"] = (
+            time.perf_counter() - sampling_started
+        ) * 1000.0
+
+        text_started = time.perf_counter()
+        text_graph_keys: tuple[tuple[int, int], ...] = ()
+        if self._fast_text_encoder:
+            from .text_encoder_graph import TextEncoderGraphCache
+
+            text_cache = getattr(self.model, "_fast_text_encoder_graph_cache", None)
+            if text_cache is None:
+                text_cache = TextEncoderGraphCache(
+                    self.model.text_encoder, token_granularity=32
+                )
+                self.model._fast_text_encoder_graph_cache = text_cache
+            for graph in profile.text_encoder_graphs:
+                text_cache.warmup_graph(
+                    batch_size=graph.batch_size,
+                    token_length=graph.token_length,
+                    device=self.device,
+                )
+            text_graph_keys = text_cache.graph_keys
+            if profile.freeze_after_warmup:
+                text_cache.freeze()
+        torch.cuda.synchronize(self.device)
+        stages["text_encoder"] = {
+            "fast": self._fast_text_encoder,
+            "graphs": [
+                {"batch_size": batch_size, "token_length": token_length}
+                for batch_size, token_length in text_graph_keys
+            ],
+            "elapsed_ms": (time.perf_counter() - text_started) * 1000.0,
+        }
+
+        prefill_started = time.perf_counter()
+        prefill_graph_keys: list[tuple[int, int]] = []
+        if self._fast_backbone_prefill:
+            for branch_batch_size in profile.backbone_decode_branch_batch_sizes:
+                backbone_graph = self._backbone_graphs[branch_batch_size]
+                prefill_cache = BackbonePrefillGraphCache(
+                    backbone_graph, token_granularity=32
+                )
+                for graph in profile.backbone_prefill_graphs:
+                    if graph.branch_batch_size == branch_batch_size:
+                        prefill_cache.warmup_graph(
+                            branch_batch_size=graph.branch_batch_size,
+                            sequence_length=graph.sequence_length,
+                        )
+                if profile.freeze_after_warmup:
+                    prefill_cache.freeze()
+                self._backbone_prefill_graphs[branch_batch_size] = prefill_cache
+                prefill_graph_keys.extend(prefill_cache.graph_keys)
+            if self._backbone_graph is not None:
+                self._backbone_prefill_graph = self._backbone_prefill_graphs.get(
+                    self._backbone_graph.batch_size
+                )
+        torch.cuda.synchronize(self.device)
+        stages["backbone_prefill"] = {
+            "fast": self._fast_backbone_prefill,
+            "graphs": [
+                {"branch_batch_size": batch_size, "sequence_length": sequence_length}
+                for batch_size, sequence_length in prefill_graph_keys
+            ],
+            "elapsed_ms": (time.perf_counter() - prefill_started) * 1000.0,
+        }
+
+        codec_started = time.perf_counter()
+        codec = self._codec()
+        codec_request_id = "profile-codec-warmup"
+        codec.open_request(codec_request_id, reset=True, is_first_decode=True)
+        codec_codes = torch.zeros(
+            1,
+            int(self.model.config.num_codebooks),
+            self._codec_chunk_frames,
+            dtype=torch.long,
+            device=self.device,
+        )
+        codec_audio = codec.decode_request_chunk(
+            codec_request_id,
+            codec_codes,
+            reset=True,
+        )
+        _extract_audio_np(codec_audio)
+        codec.close_request(codec_request_id)
+        torch.cuda.synchronize(self.device)
+        stages["codec"] = {
+            "fast": self._fast_codec,
+            "graphs": [
+                {
+                    "num_lanes": len(codec.lanes),
+                    "chunk_frames": self._codec_chunk_frames,
+                    "captured": all(lane.cuda_graph is not None for lane in codec.lanes)
+                    if self._fast_codec
+                    else False,
+                }
+            ],
+            "elapsed_ms": (time.perf_counter() - codec_started) * 1000.0,
+        }
+
+        if profile.freeze_after_warmup:
+            self._frozen_branch_batch_sizes = frozenset(
+                profile.backbone_decode_branch_batch_sizes
+            )
+        self._warmup_profile = profile
+
+        from breeze_infer.templates import get_template, prepare_inputs
+
+        request_spec = profile.warmup_request
+        synthetic_results: list[dict[str, Any]] = []
+        for cfg_scale in profile.cfg_scales:
+            synthetic_started = time.perf_counter()
+            request_id = f"profile-first-frame-cfg{cfg_scale:g}"
+            torch.manual_seed(request_spec.seed)
+            torch.cuda.manual_seed_all(request_spec.seed)
+            synthetic_inputs = prepare_inputs(
+                self.tokenizer,
+                self.audio_tokenizer,
+                self.model,
+                [
+                    {
+                        "id": request_id,
+                        "text": request_spec.text,
+                        "instruction": request_spec.instruction,
+                        "speaker": request_spec.speaker,
+                    }
+                ],
+                get_template(request_spec.template),
+                guidance_scale=cfg_scale,
+                guidance_scale_ref=None,
+                guidance_scale_ins=None,
+            )
+            synthetic_iterator = self.iter_audio_chunks(
+                synthetic_inputs, request_id=request_id
+            )
+            try:
+                try:
+                    synthetic_chunk = next(synthetic_iterator)
+                except StopIteration as exc:
+                    raise RuntimeError(
+                        "warmup request completed without producing an audio frame"
+                    ) from exc
+            finally:
+                synthetic_iterator.close()
+            torch.cuda.synchronize(self.device)
+            synthetic_results.append(
+                {
+                    "template": request_spec.template,
+                    "cfg_scale": cfg_scale,
+                    "cfg_mode": "no_cfg" if cfg_scale == 1.0 else "single_cfg",
+                    "audio_samples": int(synthetic_chunk.audio.size),
+                    "ttfa_internal_ms": synthetic_chunk.timing.get("ttfa_internal_ms"),
+                    "elapsed_ms": (time.perf_counter() - synthetic_started) * 1000.0,
+                }
+            )
+        stages["synthetic_first_frames"] = synthetic_results
+
+        manifest = {
+            "schema_version": 1,
+            "status": "ready",
+            "profile_source": profile.source,
+            "profile": profile.to_dict(),
+            "stages": stages,
+            "frozen": profile.freeze_after_warmup,
+            "total_elapsed_ms": (time.perf_counter() - total_started) * 1000.0,
+        }
+        self._warmup_manifest = manifest
+        if manifest_path is not None:
+            path = Path(manifest_path)
+            path.parent.mkdir(parents=True, exist_ok=True)
+            path.write_text(
+                json.dumps(manifest, ensure_ascii=False, indent=2) + "\n",
+                encoding="utf-8",
+            )
+        return manifest
+
+    @torch.inference_mode()
+    def iter_audio_chunks(
+        self,
+        inputs: dict[str, Any],
+        *,
+        request_id: str | None = None,
+    ) -> Iterator[FastStreamingChunk]:
+        cfg = select_fast_cfg(inputs)
+        branch_batch_size = 2 if cfg.mode == "single_cfg" else 1
+        self._ensure_graphs(branch_batch_size, cfg.guidance_scale)
+        assert self._backbone_graph is not None
+        assert self._depth_decoder_graph is not None
+
+        codec = self._codec()
+        branch = self._build_branch_batch(inputs)
+        request_id = request_id or f"local-{uuid.uuid4().hex}"
+        codec.open_request(request_id, reset=True, is_first_decode=True)
+
+        backbone_params = self._sampling_params(self.model.generation_config)
+        depth_params = self._sampling_params(self.model.depth_decoder.generation_config)
+        chunk_buffer: list[torch.Tensor] = []
+        chunk_index = 0
+        total_frames = 0
+        backbone_token_history = torch.empty(
+            self.config.max_new_tokens,
+            dtype=torch.long,
+            device=self.device,
+        )
+        first_decode = True
+        t_start = time.perf_counter()
+        t_chunk = t_start
+        prefill_start_event = None
+        prefill_end_event = None
+        if self.config.collect_timing:
+            prefill_start_event = torch.cuda.Event(enable_timing=True)
+            prefill_end_event = torch.cuda.Event(enable_timing=True)
+            prefill_start_event.record()
+
+        try:
+            attention_mask = branch.attention_mask
+            if self._fast_backbone_prefill:
+                if self._backbone_prefill_graph is None:
+                    self._backbone_prefill_graph = BackbonePrefillGraphCache(
+                        self._backbone_graph, token_granularity=32
+                    )
+                prefill_output = self._backbone_prefill_graph(
+                    branch.inputs_embeds, attention_mask
+                )
+                hidden = prefill_output.hidden_states
+                logits = prefill_output.logits
+                prefill_len = prefill_output.prefill_len
+                generation_attention_mask = prefill_output.attention_mask
+            else:
+                position_ids = attention_mask.long().cumsum(-1) - 1
+                position_ids.masked_fill_(attention_mask == 0, 1)
+                prefill_cache = None
+                cache_position = None
+                backbone_out = self.model.backbone_model(
+                    inputs_embeds=branch.inputs_embeds,
+                    attention_mask=attention_mask,
+                    position_ids=position_ids,
+                    past_key_values=prefill_cache,
+                    cache_position=cache_position,
+                    use_cache=True,
+                )
+                hidden = backbone_out.last_hidden_state
+                logits = self.model.lm_head(hidden[:, -1, :].float()).float()
+                prefill_len = self._backbone_graph.prefill_kv(
+                    backbone_out.past_key_values
+                )
+                generation_attention_mask = attention_mask
+
+            if branch.branch_batch_size == 2:
+                cond_logits = logits[:1]
+                uncond_logits = logits[1:]
+                logits = uncond_logits + branch.cfg.guidance_scale * (
+                    cond_logits - uncond_logits
+                )
+            else:
+                logits = logits[:1]
+            token = sample_logits(
+                logits,
+                suppress_tokens=self._reserved_codec_token_ids,
+                **backbone_params,
+            ).view(1)
+            self._backbone_graph.set_generation_state(generation_attention_mask)
+            if prefill_end_event is not None:
+                prefill_end_event.record()
+
+            for step_idx in range(self.config.max_new_tokens):
+                if is_backbone_eos_token(token, self.model.config):
+                    break
+
+                if prefill_len + step_idx >= self.config.max_seq_len - 1:
+                    break
+
+                if branch.branch_batch_size == 2:
+                    token_batch = token.repeat(2)
+                    depth_hidden = hidden[:, -1, :]
+                else:
+                    token_batch = token
+                    depth_hidden = hidden[:1, -1, :]
+
+                depth_tokens = self._depth_decoder_graph.run(
+                    depth_hidden,
+                    token_batch,
+                    guidance_scale=branch.cfg.guidance_scale,
+                    **depth_params,
+                )
+                frame = torch.cat([token.view(1), depth_tokens[0]], dim=0)
+                if should_decode_codec_frame(frame, self.model.config):
+                    chunk_buffer.append(frame.detach())
+
+                # A complete codec frame can be decoded immediately. Emit it
+                # before computing the next backbone token so that one full
+                # backbone decode step is no longer on the TTFA critical path.
+                reached_limit = step_idx == self.config.max_new_tokens - 1
+                chunk_ready = len(chunk_buffer) >= self._codec_chunk_frames
+                if chunk_ready or reached_limit:
+                    if chunk_buffer:
+                        decode_started = time.perf_counter()
+                        frames = chunk_buffer
+                        chunk_buffer = []
+                        total_frames += len(frames)
+                        timing: dict[str, float | int | bool] = {
+                            "chunk_index": chunk_index,
+                            "codec_frames": len(frames),
+                            "decode_launch_ms": (decode_started - t_chunk) * 1000.0,
+                            "total_frames": total_frames,
+                            "is_final": reached_limit,
+                        }
+                        chunk = self._decode_codec_frames(
+                            frames=frames,
+                            request_id=request_id,
+                            reset=first_decode,
+                            is_final=reached_limit,
+                            timing=timing,
+                        )
+                        if chunk_index == 0:
+                            first_timing = {
+                                **chunk.timing,
+                                "ttfa_internal_ms": (time.perf_counter() - t_start)
+                                * 1000.0,
+                            }
+                            if (
+                                prefill_start_event is not None
+                                and prefill_end_event is not None
+                            ):
+                                first_timing["prefill_gpu_ms"] = (
+                                    prefill_start_event.elapsed_time(prefill_end_event)
+                                )
+                            chunk = FastStreamingChunk(
+                                audio=chunk.audio,
+                                sample_rate=chunk.sample_rate,
+                                codec_frames=chunk.codec_frames,
+                                is_final=chunk.is_final,
+                                timing=first_timing,
+                            )
+                        yield chunk
+                        first_decode = False
+                        chunk_index += 1
+                        t_chunk = time.perf_counter()
+                    if reached_limit:
+                        break
+
+                frame_for_backbone = frame.view(1, 1, -1)
+                if branch.branch_batch_size == 2:
+                    frame_for_backbone = frame_for_backbone.repeat(2, 1, 1)
+
+                hidden, logits = self._backbone_graph.run(
+                    frame_for_backbone, step_idx=step_idx
+                )
+                logits = logits.float()
+                backbone_token_history[step_idx] = token[0]
+                token = sample_logits(
+                    logits,
+                    token_history=backbone_token_history[: step_idx + 1],
+                    repetition_penalty=self.config.repetition_penalty,
+                    suppress_tokens=self._reserved_codec_token_ids,
+                    **backbone_params,
+                ).view(1)
+
+                if is_backbone_eos_token(token, self.model.config):
+                    break
+
+            if chunk_buffer:
+                frames = chunk_buffer
+                total_frames += len(frames)
+                yield self._decode_codec_frames(
+                    frames=frames,
+                    request_id=request_id,
+                    reset=first_decode,
+                    is_final=True,
+                    timing={
+                        "chunk_index": chunk_index,
+                        "codec_frames": len(frames),
+                        "decode_launch_ms": (time.perf_counter() - t_chunk) * 1000.0,
+                        "total_frames": total_frames,
+                        "is_final": True,
+                    },
+                )
+        finally:
+            codec.close_request(request_id)

models/generation_breeze.py

@@ -0,0 +1,1352 @@
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+
+import os
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Optional
+
+import numpy as np
+import torch
+from torch import nn
+from transformers.generation import (
+    GenerateDecoderOnlyOutput,
+    GenerationConfig,
+    GenerationMixin,
+    GenerationMode,
+)
+from transformers.generation.logits_process import LogitsProcessorList
+from transformers.generation.stopping_criteria import (
+    MaxLengthCriteria,
+    StoppingCriteriaList,
+)
+from transformers.generation.utils import GenerateNonBeamOutput
+from transformers.utils import logging
+
+from .logits_process import mask_invalid_codec_token_logits
+
+if TYPE_CHECKING:
+    from transformers.generation.streamers import BaseStreamer
+
+
+logger = logging.get_logger(__name__)
+
+
+def _extract_decoded_audio_tensor(decoded_audio: Any) -> torch.Tensor:
+    if hasattr(decoded_audio, "audio_values"):
+        decoded_audio = decoded_audio.audio_values
+
+    while isinstance(decoded_audio, (list, tuple)):
+        decoded_audio = decoded_audio[0]
+
+    if isinstance(decoded_audio, np.ndarray):
+        decoded_audio = torch.as_tensor(decoded_audio)
+
+    if not isinstance(decoded_audio, torch.Tensor):
+        raise TypeError(f"Unsupported decoded audio type: {type(decoded_audio)!r}")
+
+    while decoded_audio.dim() > 1:
+        decoded_audio = decoded_audio[0]
+
+    return decoded_audio
+
+
+@dataclass
+class BreezeGenerateOutput(GenerateDecoderOnlyOutput):
+    """
+    Outputs of BreezeForConditionalGeneration.generate.
+
+    Args:
+        sequences (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+            The generated sequences. The second dimension (sequence_length) is either equal to `max_length` or shorter
+            if all batches finished early due to the `eos_token_id`.
+        scores (`tuple(torch.FloatTensor)` *optional*, returned when `output_scores=True`):
+            Processed prediction scores of the language modeling head (scores for each vocabulary token before SoftMax)
+            at each generation step. Tuple of `torch.FloatTensor` with up to `max_new_tokens` elements (one element for
+            each generated token), with each tensor of shape `(batch_size, config.vocab_size)`.
+        logits (`tuple(torch.FloatTensor)` *optional*, returned when `output_logits=True`):
+            Unprocessed prediction scores of the language modeling head (scores for each vocabulary token before SoftMax)
+            at each generation step. Tuple of `torch.FloatTensor` with up to `max_new_tokens` elements (one element for
+            each generated token), with each tensor of shape `(batch_size, config.vocab_size)`.
+        attentions (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `output_attentions=True`):
+            Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of
+            `torch.FloatTensor` of shape `(batch_size, num_heads, generated_length, sequence_length)`.
+        hidden_states (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `output_hidden_states=True`):
+            Tuple (one element for each generated token) of tuples (one element for each layer of the decoder) of
+            `torch.FloatTensor` of shape `(batch_size, generated_length, hidden_size)`.
+        past_key_values (`Cache`, *optional*, returned when `use_cache=True`):
+            Returns the model cache, used to speed up decoding. Different models have a different cache format, check
+        audio (`list(torch.FloatTensor)` of length `batch_size`):
+            The generated audio.
+    """
+
+    audio: list[torch.Tensor] | None = None
+
+
+class BreezeGenerationMixin(GenerationMixin):
+    # CFG kwargs that should bypass validation (Breeze requires custom CFG implementation
+    # because transformers' UnbatchedClassifierFreeGuidanceLogitsProcessor doesn't pass
+    # text_ids_mask and text_ids_len which Breeze model requires)
+    # Note: we use 'cfg_*' names to avoid triggering transformers' built-in CFG handling
+    _cfg_kwargs = {
+        "cfg_scale",
+        "cfg_negative_prompt_ids",
+        "cfg_negative_prompt_attention_mask",
+        "cfg_negative_text_ids_mask",
+        "cfg_negative_text_ids_len",
+        "cfg_negative_input_values",
+        # Dual CFG kwargs
+        "cfg_scale_ref",
+        "cfg_scale_ins",
+        "cfg_uncond_prompt_ids",
+        "cfg_uncond_prompt_attention_mask",
+        "cfg_uncond_text_ids_mask",
+        "cfg_uncond_text_ids_len",
+        "cfg_ref_prompt_ids",
+        "cfg_ref_prompt_attention_mask",
+        "cfg_ref_text_ids_mask",
+        "cfg_ref_text_ids_len",
+        "cfg_ins_prompt_ids",
+        "cfg_ins_prompt_attention_mask",
+        "cfg_ins_text_ids_mask",
+        "cfg_ins_text_ids_len",
+    }
+
+    def _mask_reserved_codec_logits(self, scores: torch.Tensor) -> torch.Tensor:
+        return mask_invalid_codec_token_logits(
+            scores,
+            codebook_size=int(self.config.codec_config.codebook_size),
+            token_vocab_size=int(self.config.vocab_size),
+        )
+
+    def _reserved_codec_token_ids(self) -> list[int]:
+        return list(
+            range(
+                int(self.config.codec_config.codebook_size),
+                int(self.config.vocab_size),
+            )
+        )
+
+    def _validate_model_kwargs(self, model_kwargs: dict[str, Any]):
+        """Override to allow CFG-related kwargs."""
+        cfg_kwargs = {
+            k: model_kwargs.pop(k)
+            for k in list(model_kwargs.keys())
+            if k in self._cfg_kwargs
+        }
+        super()._validate_model_kwargs(model_kwargs)
+        model_kwargs.update(cfg_kwargs)
+
+    def _depth_decoder_generate_with_cfg(
+        self,
+        depth_decoder_input_ids: torch.LongTensor,
+        cond_backbone_hidden_state: torch.FloatTensor,
+        uncond_backbone_hidden_state: torch.FloatTensor,
+        cfg_scale: float,
+    ) -> torch.LongTensor:
+        """
+        Custom depth decoder generation with CFG.
+
+        At each step, we run the depth decoder twice (cond and uncond) and apply CFG to the logits.
+        """
+        depth_decoder = self.depth_decoder
+        generation_config = depth_decoder.generation_config
+        num_codebooks = self.config.num_codebooks
+
+        # Get generation params from depth decoder config
+        do_sample = generation_config.do_sample
+        temperature = generation_config.temperature
+        top_p = generation_config.top_p
+        top_k = generation_config.top_k
+
+        batch_size = depth_decoder_input_ids.shape[0]
+        device = depth_decoder_input_ids.device
+
+        # Initialize sequences with input_ids (which contains placeholder + token0)
+        sequences = depth_decoder_input_ids  # [batch_size, 2]
+
+        # Generate tokens 1 to num_codebooks-1 (31 tokens for 32 codebooks)
+        for step in range(num_codebooks - 1):
+            # Cond forward pass
+            cond_outputs = depth_decoder(
+                input_ids=sequences,
+                backbone_last_hidden_state=cond_backbone_hidden_state,
+                use_cache=False,
+                return_dict=True,
+            )
+            cond_logits = cond_outputs.logits[:, -1, :].float()
+
+            # Uncond forward pass
+            uncond_outputs = depth_decoder(
+                input_ids=sequences,
+                backbone_last_hidden_state=uncond_backbone_hidden_state,
+                use_cache=False,
+                return_dict=True,
+            )
+            uncond_logits = uncond_outputs.logits[:, -1, :].float()
+
+            # Apply CFG
+            next_token_logits = uncond_logits + cfg_scale * (
+                cond_logits - uncond_logits
+            )
+            self._mask_reserved_codec_logits(next_token_logits)
+
+            # Apply temperature
+            if temperature is not None and temperature != 1.0:
+                next_token_logits = next_token_logits / temperature
+
+            # Sample or argmax
+            if do_sample:
+                probs = nn.functional.softmax(next_token_logits, dim=-1)
+
+                # Apply top_k
+                if top_k is not None and top_k > 0:
+                    top_k_probs, top_k_indices = torch.topk(
+                        probs, min(top_k, probs.size(-1))
+                    )
+                    probs = torch.zeros_like(probs).scatter_(
+                        -1, top_k_indices, top_k_probs
+                    )
+                    probs = probs / probs.sum(dim=-1, keepdim=True)
+
+                # Apply top_p
+                if top_p is not None and top_p < 1.0:
+                    sorted_probs, sorted_indices = torch.sort(probs, descending=True)
+                    cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
+                    sorted_indices_to_remove = cumulative_probs > top_p
+                    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[
+                        ..., :-1
+                    ].clone()
+                    sorted_indices_to_remove[..., 0] = 0
+                    indices_to_remove = sorted_indices_to_remove.scatter(
+                        -1, sorted_indices, sorted_indices_to_remove
+                    )
+                    probs = probs.masked_fill(indices_to_remove, 0.0)
+                    probs = probs / probs.sum(dim=-1, keepdim=True)
+
+                next_tokens = torch.multinomial(probs, num_samples=1)
+            else:
+                next_tokens = torch.argmax(next_token_logits, dim=-1, keepdim=True)
+
+            # Append to sequences
+            sequences = torch.cat([sequences, next_tokens], dim=-1)
+
+        return sequences
+
+    def _depth_decoder_generate_with_dual_cfg(
+        self,
+        depth_decoder_input_ids: torch.LongTensor,
+        uncond_backbone_hidden_state: torch.FloatTensor,
+        ref_backbone_hidden_state: torch.FloatTensor,
+        ins_backbone_hidden_state: torch.FloatTensor,
+        cfg_scale_ref: float,
+        cfg_scale_ins: float,
+    ) -> torch.LongTensor:
+        """
+        Custom depth decoder generation with dual CFG.
+
+        At each step, we run the depth decoder three times (uncond, ref, ins) and apply dual CFG to the logits:
+        logits = uncond + cfg_ref * (ref - uncond) + cfg_ins * (ins - uncond)
+        """
+        depth_decoder = self.depth_decoder
+        generation_config = depth_decoder.generation_config
+        num_codebooks = self.config.num_codebooks
+
+        do_sample = generation_config.do_sample
+        temperature = generation_config.temperature
+        top_p = generation_config.top_p
+        top_k = generation_config.top_k
+
+        batch_size = depth_decoder_input_ids.shape[0]
+        device = depth_decoder_input_ids.device
+
+        sequences = depth_decoder_input_ids  # [batch_size, 2]
+
+        for step in range(num_codebooks - 1):
+            # Uncond forward pass
+            uncond_outputs = depth_decoder(
+                input_ids=sequences,
+                backbone_last_hidden_state=uncond_backbone_hidden_state,
+                use_cache=False,
+                return_dict=True,
+            )
+            uncond_logits = uncond_outputs.logits[:, -1, :].float()
+
+            # Ref-only forward pass
+            ref_outputs = depth_decoder(
+                input_ids=sequences,
+                backbone_last_hidden_state=ref_backbone_hidden_state,
+                use_cache=False,
+                return_dict=True,
+            )
+            ref_logits = ref_outputs.logits[:, -1, :].float()
+
+            # Ins-only forward pass
+            ins_outputs = depth_decoder(
+                input_ids=sequences,
+                backbone_last_hidden_state=ins_backbone_hidden_state,
+                use_cache=False,
+                return_dict=True,
+            )
+            ins_logits = ins_outputs.logits[:, -1, :].float()
+
+            # Apply dual CFG
+            next_token_logits = (
+                uncond_logits
+                + cfg_scale_ref * (ref_logits - uncond_logits)
+                + cfg_scale_ins * (ins_logits - uncond_logits)
+            )
+            self._mask_reserved_codec_logits(next_token_logits)
+
+            # Apply temperature
+            if temperature is not None and temperature != 1.0:
+                next_token_logits = next_token_logits / temperature
+
+            # Sample or argmax
+            if do_sample:
+                probs = nn.functional.softmax(next_token_logits, dim=-1)
+
+                if top_k is not None and top_k > 0:
+                    top_k_probs, top_k_indices = torch.topk(
+                        probs, min(top_k, probs.size(-1))
+                    )
+                    probs = torch.zeros_like(probs).scatter_(
+                        -1, top_k_indices, top_k_probs
+                    )
+                    probs = probs / probs.sum(dim=-1, keepdim=True)
+
+                if top_p is not None and top_p < 1.0:
+                    sorted_probs, sorted_indices = torch.sort(probs, descending=True)
+                    cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
+                    sorted_indices_to_remove = cumulative_probs > top_p
+                    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[
+                        ..., :-1
+                    ].clone()
+                    sorted_indices_to_remove[..., 0] = 0
+                    indices_to_remove = sorted_indices_to_remove.scatter(
+                        -1, sorted_indices, sorted_indices_to_remove
+                    )
+                    probs = probs.masked_fill(indices_to_remove, 0.0)
+                    probs = probs / probs.sum(dim=-1, keepdim=True)
+
+                next_tokens = torch.multinomial(probs, num_samples=1)
+            else:
+                next_tokens = torch.argmax(next_token_logits, dim=-1, keepdim=True)
+
+            sequences = torch.cat([sequences, next_tokens], dim=-1)
+
+        return sequences
+
+    def _get_stopping_criteria(
+        self,
+        *args,
+        **kwargs,
+    ) -> StoppingCriteriaList:
+        criteria = super()._get_stopping_criteria(*args, **kwargs)
+
+        kept_criteria = StoppingCriteriaList()
+        for criterion in criteria:
+            if not isinstance(criterion, MaxLengthCriteria):
+                # logger.warning(
+                #     f"Breeze does not support {criterion.__class__.__name__} stopping criteria, it will be ignored."
+                # )
+                pass
+            else:
+                kept_criteria.append(criterion)
+        return kept_criteria
+
+    def _prepare_generation_config(
+        self,
+        generation_config: GenerationConfig | None,
+        use_model_defaults: bool | None = None,
+        **kwargs: Any,
+    ) -> tuple[GenerationConfig, dict]:
+        """
+        This method overrides [~generation.utils.GenerationMixin._prepare_generation_config].
+        It ensures that the depth decoder generation config is initialized and that passed args as depth_decoder_* are properly handled.
+        """
+        # Extract CFG kwargs before passing to super() to preserve them
+        cfg_kwargs = {
+            k: kwargs.pop(k) for k in list(kwargs.keys()) if k in self._cfg_kwargs
+        }
+
+        # extract depth decoder kwargs and remove them from the main kwargs
+        depth_decoder_kwargs = {
+            k[len("depth_decoder_") :]: v
+            for k, v in kwargs.items()
+            if k.startswith("depth_decoder_")
+        }
+
+        # remove the depth decoder keys from the original kwargs
+        kwargs = {k: v for k, v in kwargs.items() if not k.startswith("depth_decoder_")}
+
+        # initialize the generation config
+        generation_config, model_kwargs = super()._prepare_generation_config(
+            generation_config, use_model_defaults, **kwargs
+        )
+
+        # Restore CFG kwargs to model_kwargs
+        model_kwargs.update(cfg_kwargs)
+
+        self.depth_decoder.generation_config.update(**depth_decoder_kwargs)
+
+        # ensure the depth decoder generation config is valid
+        depth_decoder_min_new_tokens = self.depth_decoder.generation_config.min_new_tokens or (self.config.num_codebooks - 1)
+        depth_decoder_max_new_tokens = self.depth_decoder.generation_config.max_new_tokens or (self.config.num_codebooks - 1)
+
+        if {depth_decoder_min_new_tokens, depth_decoder_max_new_tokens} != {
+            self.config.num_codebooks - 1
+        }:
+            raise ValueError(
+                f"depth_decoder_generation_config's min_new_tokens ({depth_decoder_min_new_tokens}) and max_new_tokens ({depth_decoder_max_new_tokens}) must be equal to self.config.num_codebooks - 1 ({self.config.num_codebooks - 1})"
+            )
+        elif self.depth_decoder.generation_config.return_dict_in_generate:
+            logger.warning(
+                "depth_decoder_generation_config.return_dict_in_generate is set to True, but this will be ignored as the depth decoder model does not return a dictionary in generate"
+            )
+            self.depth_decoder.generation_config.return_dict_in_generate = False
+
+        self.depth_decoder.generation_config.min_new_tokens = (
+            depth_decoder_min_new_tokens
+        )
+        self.depth_decoder.generation_config.max_new_tokens = (
+            depth_decoder_max_new_tokens
+        )
+
+        # Monkey patch the get_generation_mode method to support Breeze model
+        original_get_generation_mode = generation_config.get_generation_mode
+
+        def patched_get_generation_mode(assistant_model=None):
+            generation_mode = original_get_generation_mode(assistant_model)
+            if generation_mode not in [
+                GenerationMode.GREEDY_SEARCH,
+                GenerationMode.SAMPLE,
+            ]:
+                raise ValueError(
+                    f"Generation mode {generation_mode} is not supported for Breeze model. Please set generation parameters to use greedy or sampling generation."
+                )
+
+            return generation_mode
+
+        generation_config.get_generation_mode = patched_get_generation_mode
+
+        return generation_config, model_kwargs
+
+    def _sample(
+        self,
+        input_ids: torch.LongTensor,
+        logits_processor: LogitsProcessorList,
+        stopping_criteria: StoppingCriteriaList,
+        generation_config: GenerationConfig,
+        synced_gpus: bool = False,
+        streamer: Optional["BaseStreamer"] = None,
+        **model_kwargs,
+    ) -> GenerateNonBeamOutput | torch.LongTensor:
+        """
+        This method overrides [~generation.utils.GenerationMixin._sample].
+        To ease maintenance, modifications are marked with the comment "Breeze specific".
+
+        Indeed, Breeze model requires a custom generation sampling step:
+        1. Infer the backbone model to sample the first codebook token
+        2. Call generate on the depth decoder with the first codebook token as input_ids to sample the next codebook tokens
+        3. Use these generated codebook tokens as input_ids to sample the next first codebook token using the backbone model
+        4. Repeat until stopping criteria is met
+
+        Breeze supports two stopping criteria:
+        - stop when the generated sequence is at max_length
+        - stop when all the generated codebook tokens are the codebook_eos_token_id
+        """
+        # init values
+        # *************** Breeze specific ***************
+        pad_token_id = self.config.codebook_pad_token_id
+        has_eos_stopping_criteria = generation_config._eos_token_tensor is not None
+        # ============================================
+        output_attentions = generation_config.output_attentions
+        output_hidden_states = generation_config.output_hidden_states
+        output_scores = generation_config.output_scores
+        output_logits = generation_config.output_logits
+        return_dict_in_generate = generation_config.return_dict_in_generate
+        do_sample = generation_config.do_sample
+
+        # *************** CFG specific ***************
+        # Breeze requires custom CFG because transformers' UnbatchedClassifierFreeGuidanceLogitsProcessor
+        # doesn't pass text_ids_mask/text_ids_len which Breeze model requires
+        # We use 'cfg_scale' instead of 'guidance_scale' to avoid triggering transformers' CFG
+        # cfg_scale=1.0: pure cond (normal generation with description)
+        # cfg_scale=0.0: pure uncond (generation without description)
+        # cfg_scale>1.0: enhance description influence
+        cfg_scale = model_kwargs.pop("cfg_scale", 1.0)
+        negative_prompt_ids = model_kwargs.pop("cfg_negative_prompt_ids", None)
+        negative_prompt_attention_mask = model_kwargs.pop(
+            "cfg_negative_prompt_attention_mask", None
+        )
+        negative_text_ids_mask = model_kwargs.pop("cfg_negative_text_ids_mask", None)
+        negative_text_ids_len = model_kwargs.pop("cfg_negative_text_ids_len", None)
+        negative_input_values = model_kwargs.pop("cfg_negative_input_values", None)
+        use_cfg = cfg_scale != 1.0 and negative_prompt_ids is not None
+
+        # cfg_scale=0 means pure unconditional: swap in the negative prompt as the main input
+        # so we only run a single forward pass instead of a wasteful dual pass
+        if cfg_scale == 0.0 and negative_prompt_ids is not None:
+            input_ids = negative_prompt_ids
+            model_kwargs["attention_mask"] = negative_prompt_attention_mask
+            if negative_text_ids_mask is not None:
+                model_kwargs["text_ids_mask"] = negative_text_ids_mask
+            else:
+                model_kwargs["text_ids_mask"] = negative_prompt_attention_mask.bool()
+            if negative_text_ids_len is not None:
+                model_kwargs["text_ids_len"] = negative_text_ids_len
+            else:
+                model_kwargs["text_ids_len"] = negative_prompt_attention_mask.sum(
+                    dim=1
+                ).long()
+            use_cfg = False
+        # ============================================
+
+        # *************** Dual CFG specific ***************
+        # Dual CFG: separate scales for reference audio and instruction/performance
+        # Formula: logits = uncond + cfg_ref * (ref - uncond) + cfg_ins * (ins - uncond)
+        cfg_scale_ref = model_kwargs.pop("cfg_scale_ref", None)
+        cfg_scale_ins = model_kwargs.pop("cfg_scale_ins", None)
+        uncond_prompt_ids = model_kwargs.pop("cfg_uncond_prompt_ids", None)
+        uncond_prompt_attention_mask = model_kwargs.pop(
+            "cfg_uncond_prompt_attention_mask", None
+        )
+        uncond_text_ids_mask = model_kwargs.pop("cfg_uncond_text_ids_mask", None)
+        uncond_text_ids_len = model_kwargs.pop("cfg_uncond_text_ids_len", None)
+        ref_prompt_ids = model_kwargs.pop("cfg_ref_prompt_ids", None)
+        ref_prompt_attention_mask = model_kwargs.pop(
+            "cfg_ref_prompt_attention_mask", None
+        )
+        ref_text_ids_mask = model_kwargs.pop("cfg_ref_text_ids_mask", None)
+        ref_text_ids_len = model_kwargs.pop("cfg_ref_text_ids_len", None)
+        ins_prompt_ids = model_kwargs.pop("cfg_ins_prompt_ids", None)
+        ins_prompt_attention_mask = model_kwargs.pop(
+            "cfg_ins_prompt_attention_mask", None
+        )
+        ins_text_ids_mask = model_kwargs.pop("cfg_ins_text_ids_mask", None)
+        ins_text_ids_len = model_kwargs.pop("cfg_ins_text_ids_len", None)
+        use_dual_cfg = (
+            cfg_scale_ref is not None
+            and cfg_scale_ins is not None
+            and uncond_prompt_ids is not None
+            and ref_prompt_ids is not None
+            and ins_prompt_ids is not None
+        )
+        if use_dual_cfg:
+            use_cfg = False  # disable single CFG when dual CFG is active
+        # ============================================
+
+        # init attention / hidden states / scores tuples
+        scores = () if (return_dict_in_generate and output_scores) else None
+        raw_logits = () if (return_dict_in_generate and output_logits) else None
+        decoder_attentions = (
+            () if (return_dict_in_generate and output_attentions) else None
+        )
+        decoder_hidden_states = (
+            () if (return_dict_in_generate and output_hidden_states) else None
+        )
+
+        # keep track of which sequences are already finished
+        batch_size, cur_len = input_ids.shape[:2]
+        this_peer_finished = False
+        unfinished_sequences = torch.ones(
+            batch_size, dtype=torch.long, device=input_ids.device
+        )
+        model_kwargs = self._get_initial_cache_position(
+            cur_len, input_ids.device, model_kwargs
+        )
+
+        # *************** CFG specific ***************
+        if use_cfg:
+            # Start with a copy of model_kwargs to get use_cache and other necessary params
+            # but exclude attention_mask, text_ids_mask, text_ids_len, cache_position which need to be recomputed
+            negative_model_kwargs = {
+                k: v
+                for k, v in model_kwargs.items()
+                if k
+                not in [
+                    "attention_mask",
+                    "text_ids_mask",
+                    "text_ids_len",
+                    "cache_position",
+                    "past_key_values",
+                ]
+            }
+            negative_model_kwargs["attention_mask"] = negative_prompt_attention_mask
+            # Override input_values for negative prompt if provided (e.g. prompt audio for ref audio mode)
+            if negative_input_values is not None:
+                negative_model_kwargs["input_values"] = negative_input_values
+            # Use pre-computed text_ids_mask and text_ids_len if provided, otherwise fall back to computing from attention_mask
+            # Pre-computed values ensure correct position_ids generation in text encoder (one segment per sample)
+            if negative_text_ids_mask is not None:
+                negative_model_kwargs["text_ids_mask"] = negative_text_ids_mask
+            else:
+                negative_model_kwargs["text_ids_mask"] = (
+                    negative_prompt_attention_mask.bool()
+                )
+            if negative_text_ids_len is not None:
+                negative_model_kwargs["text_ids_len"] = negative_text_ids_len
+            else:
+                negative_model_kwargs["text_ids_len"] = (
+                    negative_prompt_attention_mask.sum(dim=1).long()
+                )
+            negative_model_kwargs = self._get_initial_cache_position(
+                negative_prompt_ids.shape[1],
+                negative_prompt_ids.device,
+                negative_model_kwargs,
+            )
+        # ============================================
+
+        # *************** Dual CFG init ***************
+        if use_dual_cfg:
+            _exclude_keys = {
+                "attention_mask",
+                "text_ids_mask",
+                "text_ids_len",
+                "cache_position",
+                "past_key_values",
+            }
+
+            def _make_branch_kwargs(
+                prompt_ids, attn_mask, text_mask, text_len, include_input_values
+            ):
+                kw = {k: v for k, v in model_kwargs.items() if k not in _exclude_keys}
+                kw["attention_mask"] = attn_mask
+                kw["text_ids_mask"] = text_mask
+                kw["text_ids_len"] = text_len
+                if not include_input_values:
+                    kw.pop("input_values", None)
+                return self._get_initial_cache_position(
+                    prompt_ids.shape[1], prompt_ids.device, kw
+                )
+
+            # uncond: text only (no ref audio, no instruction)
+            uncond_model_kwargs = _make_branch_kwargs(
+                uncond_prompt_ids,
+                uncond_prompt_attention_mask,
+                uncond_text_ids_mask,
+                uncond_text_ids_len,
+                include_input_values=False,
+            )
+            # ref_cond: ref audio + text (no instruction) — inherits input_values from model_kwargs
+            ref_model_kwargs = _make_branch_kwargs(
+                ref_prompt_ids,
+                ref_prompt_attention_mask,
+                ref_text_ids_mask,
+                ref_text_ids_len,
+                include_input_values=True,
+            )
+            # ins_cond: instruction + text (no ref audio)
+            ins_model_kwargs = _make_branch_kwargs(
+                ins_prompt_ids,
+                ins_prompt_attention_mask,
+                ins_text_ids_mask,
+                ins_text_ids_len,
+                include_input_values=False,
+            )
+        # ============================================
+
+        # *************** Breeze specific ***************
+        if input_ids.ndim == 2 and model_kwargs.get("inputs_embeds") is None:
+            # in the case where the passed input_ids correspond to text tokens, i.e. don't have a third dimension for codebook ids,
+            # we need to remove the input length to the MaxLengthCriteria stopping criteria has such input are not returned
+            for criterion in stopping_criteria:
+                if isinstance(criterion, MaxLengthCriteria):
+                    criterion.max_length -= cur_len
+        # ============================================
+
+        model_forward = self.__call__
+        compile_forward = self._valid_auto_compile_criteria(
+            model_kwargs, generation_config
+        )
+        if compile_forward:
+            os.environ["TOKENIZERS_PARALLELISM"] = "0"
+            model_forward = self.get_compiled_call(generation_config.compile_config)
+
+        is_prefill = True
+        while self._has_unfinished_sequences(
+            this_peer_finished,
+            synced_gpus,
+            device=input_ids.device,
+        ):
+            # *************** Dual CFG specific ***************
+            # In dual CFG mode, skip the main forward pass entirely — only run the 3 branches.
+            # We reuse ref_outputs as `outputs` for model_kwargs update and hidden_states
+            # (backbone_last_hidden_state from `outputs` is unused since dual CFG computes its own).
+            if use_dual_cfg:
+                if is_prefill:
+
+                    def _run_branch(prompt_ids, branch_kwargs):
+                        mi = self.prepare_inputs_for_generation(
+                            prompt_ids, **branch_kwargs
+                        )
+                        mi.update({"output_hidden_states": True})
+                        out = self(**mi, return_dict=True)
+                        return out, self._update_model_kwargs_for_generation(
+                            out, branch_kwargs
+                        )
+
+                    uncond_outputs, uncond_model_kwargs = _run_branch(
+                        uncond_prompt_ids, uncond_model_kwargs
+                    )
+                    ref_outputs, ref_model_kwargs = _run_branch(
+                        ref_prompt_ids, ref_model_kwargs
+                    )
+                    ins_outputs, ins_model_kwargs = _run_branch(
+                        ins_prompt_ids, ins_model_kwargs
+                    )
+                    is_prefill = False
+                else:
+
+                    def _run_branch_decode(prompt_ids, branch_kwargs):
+                        mi = self.prepare_inputs_for_generation(
+                            prompt_ids, **branch_kwargs
+                        )
+                        mi.update({"output_hidden_states": True})
+                        out = model_forward(**mi, return_dict=True)
+                        return out, self._update_model_kwargs_for_generation(
+                            out, branch_kwargs
+                        )
+
+                    uncond_outputs, uncond_model_kwargs = _run_branch_decode(
+                        uncond_prompt_ids, uncond_model_kwargs
+                    )
+                    ref_outputs, ref_model_kwargs = _run_branch_decode(
+                        ref_prompt_ids, ref_model_kwargs
+                    )
+                    ins_outputs, ins_model_kwargs = _run_branch_decode(
+                        ins_prompt_ids, ins_model_kwargs
+                    )
+
+                # Use ref_outputs as the dummy `outputs` for the rest of the loop
+                outputs = ref_outputs
+            else:
+                # ============================================
+                # prepare model inputs
+                model_inputs = self.prepare_inputs_for_generation(
+                    input_ids, **model_kwargs
+                )
+
+                # prepare variable output controls (note: some models won't accept all output controls)
+                model_inputs.update(
+                    {"output_attentions": output_attentions}
+                    if output_attentions
+                    else {}
+                )
+                # *************** Breeze specific ***************
+                model_inputs.update({"output_hidden_states": True})
+                # ============================================
+
+                if is_prefill:
+                    outputs = self(**model_inputs, return_dict=True)
+                    # *************** CFG specific ***************
+                    if use_cfg:
+                        negative_model_inputs = self.prepare_inputs_for_generation(
+                            negative_prompt_ids, **negative_model_kwargs
+                        )
+                        negative_model_inputs.update({"output_hidden_states": True})
+                        negative_outputs = self(
+                            **negative_model_inputs, return_dict=True
+                        )
+                        negative_model_kwargs = (
+                            self._update_model_kwargs_for_generation(
+                                negative_outputs, negative_model_kwargs
+                            )
+                        )
+                    # ============================================
+                    is_prefill = False
+                else:
+                    outputs = model_forward(**model_inputs, return_dict=True)
+                    # *************** CFG specific ***************
+                    if use_cfg:
+                        negative_model_inputs = self.prepare_inputs_for_generation(
+                            negative_prompt_ids, **negative_model_kwargs
+                        )
+                        negative_model_inputs.update({"output_hidden_states": True})
+                        negative_outputs = model_forward(
+                            **negative_model_inputs, return_dict=True
+                        )
+                        negative_model_kwargs = (
+                            self._update_model_kwargs_for_generation(
+                                negative_outputs, negative_model_kwargs
+                            )
+                        )
+                    # ============================================
+
+            # synced_gpus: don't waste resources running the code we don't need; kwargs must be updated before skipping
+            model_kwargs = self._update_model_kwargs_for_generation(
+                outputs,
+                model_kwargs,
+            )
+            if synced_gpus and this_peer_finished:
+                continue
+
+            # Clone is needed to avoid keeping a hanging ref to outputs.logits which may be very large for first iteration
+            # (the clone itself is always small)
+            next_token_logits = outputs.logits[:, -1, :].clone().float()
+            next_token_logits = next_token_logits.to(input_ids.device)
+
+            # *************** CFG specific ***************
+            if use_cfg:
+                negative_logits = negative_outputs.logits[:, -1, :].clone().float()
+                negative_logits = negative_logits.to(input_ids.device)
+                # CFG formula: logits = uncond_logits + cfg_scale * (cond_logits - uncond_logits)
+                next_token_logits = negative_logits + cfg_scale * (
+                    next_token_logits - negative_logits
+                )
+            # ============================================
+            # *************** Dual CFG specific ***************
+            if use_dual_cfg:
+                uncond_logits = (
+                    uncond_outputs.logits[:, -1, :].clone().float().to(input_ids.device)
+                )
+                ref_logits = (
+                    ref_outputs.logits[:, -1, :].clone().float().to(input_ids.device)
+                )
+                ins_logits = (
+                    ins_outputs.logits[:, -1, :].clone().float().to(input_ids.device)
+                )
+                # Dual CFG: logits = uncond + cfg_ref * (ref - uncond) + cfg_ins * (ins - uncond)
+                # Note: next_token_logits from backbone is not used — we don't run a full-cond forward
+                next_token_logits = (
+                    uncond_logits
+                    + cfg_scale_ref * (ref_logits - uncond_logits)
+                    + cfg_scale_ins * (ins_logits - uncond_logits)
+                )
+            # ============================================
+
+            # pre-process distribution
+            next_token_scores = logits_processor(input_ids, next_token_logits)
+            self._mask_reserved_codec_logits(next_token_scores)
+
+            # Store scores, attentions and hidden_states when required
+            if return_dict_in_generate:
+                if output_scores:
+                    scores += (next_token_scores,)
+                if output_logits:
+                    raw_logits += (next_token_logits,)
+                if output_attentions:
+                    decoder_attentions += (outputs.attentions,)
+
+                if output_hidden_states:
+                    decoder_hidden_states += (outputs.hidden_states,)
+
+            # token selection
+            if do_sample:
+                probs = nn.functional.softmax(next_token_scores, dim=-1)
+                # TODO (joao): this OP throws "skipping cudagraphs due to ['incompatible ops']", find solution
+                # next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)
+
+                def multinomial(probs, num_samples=1):
+                    ori_shape = probs.shape[:-1]
+                    probs = probs.reshape(-1, probs.shape[-1])
+                    samples = torch.multinomial(probs, num_samples=num_samples)
+                    samples = samples.reshape([*ori_shape, -1])
+                    return samples
+
+                next_tokens = multinomial(probs, num_samples=1).squeeze(-1)
+            else:
+                next_tokens = torch.argmax(next_token_scores, dim=-1)
+
+            # *************** Breeze specific ***************
+            # Detect backbone EOS: backbone predicts token at index vocab_size
+            backbone_eos_token_id = self.config.vocab_size
+            backbone_eos_mask = next_tokens == backbone_eos_token_id  # (batch_size,)
+
+            # infer the depth decoder only for non-EOS rows
+            backbone_last_hidden_state = outputs.hidden_states[-1][:, -1, :]
+
+            # *************** CFG specific for depth decoder ***************
+            if use_cfg:
+                uncond_backbone_last_hidden_state = negative_outputs.hidden_states[-1][
+                    :, -1, :
+                ]
+            # ============================================
+            # *************** Dual CFG specific for depth decoder ***************
+            if use_dual_cfg:
+                uncond_backbone_last_hidden_state = uncond_outputs.hidden_states[-1][
+                    :, -1, :
+                ]
+                ref_backbone_last_hidden_state = ref_outputs.hidden_states[-1][:, -1, :]
+                ins_backbone_last_hidden_state = ins_outputs.hidden_states[-1][:, -1, :]
+            # ============================================
+
+            if backbone_last_hidden_state.ndim == 3:
+                # patched generate
+                (b, patch_size, hidden_size) = backbone_last_hidden_state.shape
+                backbone_last_hidden_state = backbone_last_hidden_state.reshape(
+                    b * patch_size, hidden_size
+                )
+                if use_cfg:
+                    uncond_backbone_last_hidden_state = (
+                        uncond_backbone_last_hidden_state.reshape(
+                            b * patch_size, hidden_size
+                        )
+                    )
+                if use_dual_cfg:
+                    uncond_backbone_last_hidden_state = (
+                        uncond_backbone_last_hidden_state.reshape(
+                            b * patch_size, hidden_size
+                        )
+                    )
+                    ref_backbone_last_hidden_state = (
+                        ref_backbone_last_hidden_state.reshape(
+                            b * patch_size, hidden_size
+                        )
+                    )
+                    ins_backbone_last_hidden_state = (
+                        ins_backbone_last_hidden_state.reshape(
+                            b * patch_size, hidden_size
+                        )
+                    )
+            else:
+                patch_size = 1
+
+            # Preallocate full codebook output with pad tokens.
+            # Newer checkpoints return the complete audio frame from the depth decoder,
+            # while older ones return only the continuation after the backbone's first codebook.
+            num_codebooks = self.config.num_codebooks
+            total_rows = batch_size * patch_size if patch_size > 1 else batch_size
+            codebook_ids = torch.full(
+                (total_rows, num_codebooks),
+                pad_token_id,
+                device=input_ids.device,
+                dtype=input_ids.dtype,
+            )
+
+            # Determine which rows need depth decoder (non-EOS and unfinished)
+            if patch_size > 1:
+                non_eos_mask = ~backbone_eos_mask.reshape(total_rows)
+            else:
+                non_eos_mask = ~backbone_eos_mask
+
+            # Clamp EOS rows to 0 so they are valid codebook ids for depth decoder input
+            first_codebook_ids_clamped = next_tokens.clone()
+            first_codebook_ids_clamped[backbone_eos_mask] = 0
+
+            if non_eos_mask.any():
+                if patch_size > 1:
+                    first_codebook_ids_flat = first_codebook_ids_clamped.reshape(
+                        total_rows, 1
+                    )
+                else:
+                    first_codebook_ids_flat = first_codebook_ids_clamped[..., None]
+
+                depth_decoder_input_ids = nn.functional.pad(
+                    first_codebook_ids_flat[non_eos_mask], (1, 0), value=0
+                )
+                bhs = backbone_last_hidden_state[non_eos_mask]
+
+                # *************** Dual CFG specific for depth decoder ***************
+                if use_dual_cfg:
+                    active_codebook_ids = self._depth_decoder_generate_with_dual_cfg(
+                        depth_decoder_input_ids=depth_decoder_input_ids,
+                        uncond_backbone_hidden_state=uncond_backbone_last_hidden_state[
+                            non_eos_mask
+                        ].clone(),
+                        ref_backbone_hidden_state=ref_backbone_last_hidden_state[
+                            non_eos_mask
+                        ].clone(),
+                        ins_backbone_hidden_state=ins_backbone_last_hidden_state[
+                            non_eos_mask
+                        ].clone(),
+                        cfg_scale_ref=cfg_scale_ref,
+                        cfg_scale_ins=cfg_scale_ins,
+                    )
+                elif use_cfg:
+                    active_codebook_ids = self._depth_decoder_generate_with_cfg(
+                        depth_decoder_input_ids=depth_decoder_input_ids,
+                        cond_backbone_hidden_state=bhs.clone(),
+                        uncond_backbone_hidden_state=uncond_backbone_last_hidden_state[
+                            non_eos_mask
+                        ].clone(),
+                        cfg_scale=cfg_scale,
+                    )
+                else:
+                    depth_decoder_outputs = self.depth_decoder.generate(
+                        input_ids=depth_decoder_input_ids,
+                        backbone_last_hidden_state=bhs.clone(),
+                        suppress_tokens=self._reserved_codec_token_ids(),
+                    )
+                    active_codebook_ids = (
+                        depth_decoder_outputs
+                        if isinstance(depth_decoder_outputs, torch.Tensor)
+                        else depth_decoder_outputs.sequences
+                    )
+                # ============================================
+                # Remove the placeholder in position 0. Depending on checkpoint vintage,
+                # the depth decoder may now return either the full frame or only the
+                # continuation after the backbone's first codebook.
+                active_codebook_ids = active_codebook_ids[:, 1:]
+                if active_codebook_ids.shape[-1] == num_codebooks:
+                    normalized_codebook_ids = active_codebook_ids
+                elif active_codebook_ids.shape[-1] == num_codebooks - 1:
+                    normalized_codebook_ids = torch.cat(
+                        [first_codebook_ids_flat[non_eos_mask], active_codebook_ids],
+                        dim=-1,
+                    )
+                else:
+                    raise ValueError(
+                        "Unexpected depth decoder output width "
+                        f"{active_codebook_ids.shape[-1]} after placeholder removal; "
+                        f"expected {num_codebooks} or {num_codebooks - 1}."
+                    )
+                codebook_ids[non_eos_mask] = normalized_codebook_ids
+
+            next_tokens = codebook_ids  # (total_rows, num_codebooks)
+
+            # add sequence dimension
+            next_tokens_3d = next_tokens.reshape(
+                batch_size, 1, patch_size * next_tokens.shape[-1]
+            )
+
+            # finished sentences should have their next token be a padding token
+            next_tokens_3d = next_tokens_3d * unfinished_sequences.view(
+                batch_size, 1, 1
+            ) + pad_token_id * (1 - unfinished_sequences.view(batch_size, 1, 1))
+
+            # update generated ids, model inputs, and length for next step
+            if input_ids.ndim == 2:
+                input_ids = next_tokens_3d
+            else:
+                input_ids = torch.cat([input_ids, next_tokens_3d], dim=1)
+            # ============================================
+
+            if streamer is not None:
+                streamer.put(next_tokens.cpu())
+
+            # *************** Breeze specific ***************
+            # Stop sequences where backbone predicted EOS
+            unfinished_sequences = unfinished_sequences & ~backbone_eos_mask
+            # ============================================
+            unfinished_sequences = unfinished_sequences & ~stopping_criteria(
+                input_ids, scores
+            )
+            this_peer_finished = unfinished_sequences.max() == 0
+            cur_len += 1
+
+            # This is needed to properly delete outputs.logits which may be very large for first iteration
+            # Otherwise a reference to outputs is kept which keeps the logits alive in the next iteration
+            del outputs
+
+            # *************** Breeze specific ***************
+            if not use_cfg and not use_dual_cfg:
+                depth_decoder_outputs = None
+            # ============================================
+
+            # *************** CFG specific ***************
+            if use_cfg:
+                # negative_prompt_ids starts as 2D (batch, seq_len), convert to 3D after first iteration
+                if negative_prompt_ids.ndim == 2:
+                    negative_prompt_ids = next_tokens_3d
+                else:
+                    negative_prompt_ids = torch.cat(
+                        [negative_prompt_ids, next_tokens_3d], dim=1
+                    )
+                del negative_outputs
+            # ============================================
+
+            # *************** Dual CFG specific ***************
+            if use_dual_cfg:
+
+                def _update_branch_ids(prompt_ids):
+                    if prompt_ids.ndim == 2:
+                        return next_tokens_3d
+                    return torch.cat([prompt_ids, next_tokens_3d], dim=1)
+
+                uncond_prompt_ids = _update_branch_ids(uncond_prompt_ids)
+                ref_prompt_ids = _update_branch_ids(ref_prompt_ids)
+                ins_prompt_ids = _update_branch_ids(ins_prompt_ids)
+                del uncond_outputs, ref_outputs, ins_outputs
+            # ============================================
+
+        if streamer is not None:
+            streamer.end()
+
+        if return_dict_in_generate:
+            return GenerateDecoderOnlyOutput(
+                sequences=input_ids,
+                scores=scores,
+                logits=raw_logits,
+                attentions=decoder_attentions,
+                hidden_states=decoder_hidden_states,
+                past_key_values=model_kwargs.get("past_key_values"),
+            )
+        else:
+            return input_ids
+
+    def generate(
+        self,
+        input_ids: torch.Tensor | None = None,
+        input_values: torch.Tensor | None = None,
+        input_values_cutoffs: torch.Tensor | None = None,
+        generation_config: GenerationConfig | None = None,
+        logits_processor: LogitsProcessorList | None = None,
+        stopping_criteria: StoppingCriteriaList | None = None,
+        synced_gpus: bool | None = None,
+        streamer: Optional["BaseStreamer"] = None,
+        output_audio: bool | None = False,
+        audio_tokenizer=None,
+        **kwargs,
+    ) -> GenerateNonBeamOutput | torch.LongTensor:
+        r"""
+        This method overrides [`~generation.utils.GenerationMixin.generate`] to match the specifics of the Breeze model.
+        Indeed, Breeze model requires a custom generation sampling step:
+        1. Infer the backbone model to sample the first codebook token
+        2. Call generate on the depth decoder with the first codebook token as `input_ids` to sample the next codebook tokens
+        3. Use these generated codebook tokens as `input_ids` to sample the next first codebook token using the backbone model
+        4. Repeat until stopping criteria is met
+
+        <Tip warning={true}>
+
+        Most generation-controlling parameters are set in `generation_config` which, if not passed, will be set to the
+        model's default generation configuration. You can override any `generation_config` by passing the corresponding
+        parameters to generate(), e.g. `.generate(inputs, do_sample=True)`.
+        </Tip>
+
+        Parameters:
+            inputs_ids (`torch.Tensor` of shape (batch_size, seq_length), *optional*):
+                The sequence used as a prompt for the backbone model.
+            input_values (`torch.Tensor` of shape (batch_size, channels, max_concatenated_audio_length), *optional*):
+                The batched audio input values, where each batch entry contains the concatenation of all audio segments for that entry.
+                These values will be encoded into codebook tokens using the codec model and merged with the text input ids provided in `input_ids`.
+            input_values_cutoffs (`torch.Tensor` of shape (batch_size, max_num_audio), *optional*):
+                Specify the end positions of audio segments within each batch entry, relative to the concatenated audio input.
+                If a batch entry has fewer segments than the maximum, it is padded with -1. For example, in a batch of 2 sequences
+                where the first contains 2 audio segments of length l1, and the second contains 1 audio segment of length l2,
+                the input_values_cutoffs would be: [[l1, 2 * l1], [l2, -1]].
+            generation_config ([`~generation.GenerationConfig`], *optional*):
+                The generation configuration to be used as base parametrization for the generation call. `**kwargs`
+                passed to generate matching the attributes of `generation_config` will override them. If
+                `generation_config` is not provided, the default will be used, which has the following loading
+                priority: 1) from the `generation_config.json` model file, if it exists; 2) from the model
+                configuration. Please note that unspecified parameters will inherit [`~generation.GenerationConfig`]'s
+                default values, whose documentation should be checked to parameterize generation.
+            logits_processor (`LogitsProcessorList`, *optional*):
+                Custom logits processors that complement the default logits processors built from arguments and
+                generation config. If a logit processor is passed that is already created with the arguments or a
+                generation config an error is thrown. This feature is intended for advanced users.
+            stopping_criteria (`StoppingCriteriaList`, *optional*):
+                Custom stopping criteria that complements the default stopping criteria built from arguments and a
+                generation config. If a stopping criteria is passed that is already created with the arguments or a
+                generation config an error is thrown. If your stopping criteria depends on the `scores` input, make
+                sure you pass `return_dict_in_generate=True, output_scores=True` to `generate`. This feature is
+                intended for advanced users.
+            synced_gpus (`bool`, *optional*):
+                Whether to continue running the while loop until max_length. Unless overridden, this flag will be set
+                to `True` if using `FullyShardedDataParallel` or DeepSpeed ZeRO Stage 3 with multiple GPUs to avoid
+                deadlocking if one GPU finishes generating before other GPUs. Otherwise, defaults to `False`.
+            streamer (`BaseStreamer`, *optional*):
+                Streamer object that will be used to stream the generated sequences. Generated tokens are passed
+                through `streamer.put(token_ids)` and the streamer is responsible for any further processing.
+            output_audio (`bool`, *optional*):
+                Whether to return the generated audio.
+            kwargs (`dict[str, Any]`, *optional*):
+                Ad hoc parametrization of `generation_config` and/or additional model-specific kwargs that will be
+                forwarded to the `forward` function of the model. Depth decoder specific kwargs should be prefixed with *depth_decoder_*.
+
+        Return:
+            [`BreezeGenerateOutput`] or `torch.LongTensor` or `list[torch.FloatTensor]`: A [`BreezeGenerateOutput`]
+            (if `return_dict_in_generate=True` or when `config.return_dict_in_generate=True`) or a `torch.LongTensor` when `output_audio=False`
+            or a `list[torch.FloatTensor]` otherwise.
+
+        Example:
+
+        ```python
+        >>> from models.breeze import BreezeForConditionalGeneration
+        from transformers import AutoTokenizer
+        >>> from datasets import load_dataset, Audio
+
+        >>> model_id = "/path/to/breeze-model"
+        >>> torch_device = "cuda" if torch.cuda.is_available() else "cpu"
+
+        >>> processor = AutoProcessor.from_pretrained(model_id)
+
+        >>> ds = load_dataset("hf-internal-testing/dailytalk-dummy", split="train")
+        >>> # ensure the audio is 24kHz
+        >>> ds = ds.cast_column("audio", Audio(sampling_rate=24000))
+
+        >>> conversation = []
+        >>> # prepare a conversation with text and corresponding audio
+        >>> for text, audio, speaker_id in zip(ds[:4]["text"], ds[:4]["audio"], ds[:4]["speaker_id"]):
+        ...     conversation.append(
+        ...         {
+        ...             "role": f"{speaker_id}",
+        ...             "content": [{"type": "text", "text": text}, {"type": "audio", "path": audio["array"]}],
+        ...         }
+        ...     )
+
+        >>> # text prompt
+        >>> conversation.append({"role": f"{ds[4]['speaker_id']}", "content": [{"type": "text", "text": ds[4]["text"]}]})
+
+        >>> inputs = processor.apply_chat_template(
+        ...     conversation,
+        ...     tokenize=True,
+        ...     return_dict=True,
+        ... ).to(torch_device)
+
+        >>> model = BreezeForConditionalGeneration.from_pretrained(model_id, device_map=torch_device)
+        >>> audio = model.generate(**inputs, output_audio=True)
+        >>> processor.save_audio(audio, "output.wav")
+        ```
+        """
+        generate_output = super().generate(
+            input_ids=input_ids,
+            input_values=input_values,
+            input_values_cutoffs=input_values_cutoffs,
+            generation_config=generation_config,
+            logits_processor=logits_processor,
+            stopping_criteria=stopping_criteria,
+            synced_gpus=synced_gpus,
+            streamer=streamer,
+            **kwargs,
+        )
+
+        generate_returned_dict = not isinstance(generate_output, torch.Tensor)
+        audio = None
+        if output_audio:
+            generated_audio_codes = (
+                generate_output.sequences if generate_returned_dict else generate_output
+            )
+
+            # reshape for patched audio tokens, if patch size == 1, this does nothing
+            patch_size = generated_audio_codes.shape[-1] // self.config.num_codebooks
+            generated_audio_codes = generated_audio_codes.reshape(
+                generated_audio_codes.shape[0],
+                generated_audio_codes.shape[1] * patch_size,
+                generated_audio_codes.shape[2] // patch_size,
+            )
+
+            # infer the codec model
+            audio = []
+            with torch.no_grad():
+                # =======================================
+                # TODO: @eustlb, this should be batched !!!
+                # but requires making sure batched inference of the codec model works as intended
+                for sample_index, audio_codes_batch in enumerate(generated_audio_codes):
+                    # Truncate at first pad frame (EOS rows are stored as all-pad)
+                    is_pad_frame = (
+                        audio_codes_batch == self.config.codebook_pad_token_id
+                    ).all(dim=-1)
+                    eos_idxs = is_pad_frame.nonzero()
+                    if eos_idxs.numel() != 0:
+                        cutoff_idx = eos_idxs.min()
+                    else:
+                        cutoff_idx = audio_codes_batch.shape[0]
+
+                    audio_codes_batch = audio_codes_batch[:cutoff_idx]
+
+                    if cutoff_idx == 0:
+                        # write log to file
+                        with open(
+                            "/tmp/breeze_generation_warnings.log", "a"
+                        ) as log_file:
+                            log_file.write(
+                                "[no-code] No codebook tokens were generated, generating silent audio.\n"
+                            )
+
+                        # generate silent audio if no codebook tokens were generated
+                        num_codebooks = getattr(
+                            self.codec_model,
+                            "num_codebooks",
+                            getattr(
+                                self.config,
+                                "num_codebooks",
+                                audio_codes_batch.shape[-1],
+                            ),
+                        )
+                        dummy_codes = torch.ones(
+                            1,
+                            num_codebooks,
+                            device=audio_codes_batch.device,
+                            dtype=audio_codes_batch.dtype,
+                        )
+                        logger.warning(
+                            "No codebook tokens were generated, generating silent audio."
+                        )
+                        if audio_tokenizer is not None:
+                            codec_decode_output = audio_tokenizer.decode(
+                                {"audio_codes": dummy_codes}
+                            )
+                            decode_audio = _extract_decoded_audio_tensor(
+                                codec_decode_output
+                            )
+                        else:
+                            codec_decode_output = self.codec_model.decode(
+                                dummy_codes.transpose(0, 1).unsqueeze(0)
+                            )
+                            decode_audio = _extract_decoded_audio_tensor(
+                                codec_decode_output
+                            )
+                    else:
+                        codebook_size = int(self.config.codec_config.codebook_size)
+                        invalid_codes = (audio_codes_batch < 0) | (
+                            audio_codes_batch >= codebook_size
+                        )
+                        if invalid_codes.any():
+                            frame_index, codebook_index = invalid_codes.nonzero()[
+                                0
+                            ].tolist()
+                            token_id = int(
+                                audio_codes_batch[frame_index, codebook_index]
+                            )
+                            raise ValueError(
+                                "Generated codec token is outside the decoder codebook range: "
+                                f"sample={sample_index}, frame={frame_index}, "
+                                f"codebook={codebook_index}, token={token_id}, "
+                                f"valid=[0, {codebook_size - 1}]"
+                            )
+
+                        if (
+                            audio_tokenizer is None
+                        ):  # 兼容qwen3, 因为qwen3 tokenizer没有cardinality属性
+                            if (
+                                audio_codes_batch.max().item()
+                                >= self.codec_model.quantizer.cardinality
+                            ):
+                                audio_codes_batch = torch.clamp(
+                                    audio_codes_batch,
+                                    0,
+                                    self.codec_model.quantizer.cardinality - 1,
+                                )
+
+                        if audio_tokenizer is not None:
+                            codec_decode_output = audio_tokenizer.decode(
+                                {"audio_codes": [audio_codes_batch]}
+                            )
+                            decode_audio = _extract_decoded_audio_tensor(
+                                codec_decode_output
+                            )
+                        else:
+                            codec_decode_output = self.codec_model.decode(
+                                audio_codes_batch.transpose(0, 1).unsqueeze(0)
+                            )
+                            decode_audio = _extract_decoded_audio_tensor(
+                                codec_decode_output
+                            )
+
+                        codes_min_max = (
+                            audio_codes_batch.min().item(),
+                            audio_codes_batch.max().item(),
+                        )
+                        # logger.info(f"Decoded audio codes with min/max values: {codes_min_max}")
+                    audio.append(decode_audio)
+                # =======================================
+
+        if generate_returned_dict:
+            return BreezeGenerateOutput(audio=audio, **generate_output)
+        elif output_audio:
+            return audio
+        else:
+            return generate_output

models/logits_process.py

@@ -0,0 +1,78 @@
+from __future__ import annotations
+
+import torch
+from transformers.generation.logits_process import LogitsProcessor
+
+
+def mask_invalid_codec_token_logits(
+    scores: torch.Tensor,
+    *,
+    codebook_size: int,
+    token_vocab_size: int,
+) -> torch.Tensor:
+    """Mask reserved model tokens that are not valid codec codebook entries."""
+    if codebook_size <= 0 or token_vocab_size < codebook_size:
+        raise ValueError(
+            f"Invalid codec token range: codebook_size={codebook_size}, "
+            f"token_vocab_size={token_vocab_size}"
+        )
+    if token_vocab_size > scores.shape[-1]:
+        raise ValueError(
+            f"token_vocab_size={token_vocab_size} exceeds logits size {scores.shape[-1]}"
+        )
+    if codebook_size < token_vocab_size:
+        scores[..., codebook_size:token_vocab_size] = float("-inf")
+    return scores
+
+
+class GeneratedTokenRepetitionPenaltyLogitsProcessor(LogitsProcessor):
+    """Apply HF-style repetition penalty only to generated, in-vocab tokens."""
+
+    def __init__(self, penalty: float) -> None:
+        if penalty <= 0:
+            raise ValueError("penalty must be > 0")
+        self.penalty = float(penalty)
+
+    def __call__(
+        self, input_ids: torch.LongTensor, scores: torch.FloatTensor
+    ) -> torch.FloatTensor:
+        # Breeze starts generation with 2D text prompt IDs, then replaces them
+        # with 3D generated codec frames after the first decoding step.
+        if input_ids.ndim == 2:
+            return scores
+        if input_ids.ndim != 3:
+            raise ValueError(
+                f"Expected 2D prompt IDs or 3D codec frames, got {input_ids.ndim}D"
+            )
+
+        generated_ids = input_ids[..., 0]
+        if self.penalty == 1.0 or generated_ids.numel() == 0:
+            return scores
+
+        vocab_size = scores.shape[-1]
+        valid = (generated_ids >= 0) & (generated_ids < vocab_size)
+        if not torch.any(valid):
+            return scores
+
+        # Invalid IDs use token 0 as a temporary gather/scatter target. Restore
+        # token 0 explicitly afterward so those placeholders cannot affect it.
+        safe_ids = generated_ids.masked_fill(~valid, 0)
+        selected_scores = torch.gather(scores, 1, safe_ids)
+        penalized_scores = torch.where(
+            selected_scores < 0,
+            selected_scores * self.penalty,
+            selected_scores / self.penalty,
+        )
+
+        processed = scores.clone()
+        processed.scatter_(1, safe_ids, penalized_scores)
+
+        original_zero = scores[:, 0]
+        penalized_zero = torch.where(
+            original_zero < 0,
+            original_zero * self.penalty,
+            original_zero / self.penalty,
+        )
+        generated_zero = (valid & (generated_ids == 0)).any(dim=1)
+        processed[:, 0] = torch.where(generated_zero, penalized_zero, original_zero)
+        return processed

models/stream_runtime/__init__.py

@@ -0,0 +1,15 @@
+from .inference import (
+    BACKEND_NAME,
+    ExecutionLane,
+    MultiRequestStreamRuntime,
+    QwenStreamRuntimeConfig,
+    load_tokenizer,
+)
+
+__all__ = [
+    "BACKEND_NAME",
+    "ExecutionLane",
+    "MultiRequestStreamRuntime",
+    "QwenStreamRuntimeConfig",
+    "load_tokenizer",
+]

models/stream_runtime/core/__init__.py

@@ -0,0 +1,21 @@
+from .compat import (
+    BACKEND_NAME,
+    Qwen3TTSTokenizer,
+    Qwen3TTSTokenizerV2CausalConvNet,
+    Qwen3TTSTokenizerV2CausalTransConvNet,
+    Qwen3TTSTokenizerV2ConvNeXtBlock,
+    Qwen3TTSTokenizerV2Decoder,
+    Qwen3TTSTokenizerV2DecoderDecoderBlock,
+    Qwen3TTSTokenizerV2DecoderDecoderResidualUnit,
+)
+
+__all__ = [
+    "BACKEND_NAME",
+    "Qwen3TTSTokenizer",
+    "Qwen3TTSTokenizerV2CausalConvNet",
+    "Qwen3TTSTokenizerV2CausalTransConvNet",
+    "Qwen3TTSTokenizerV2ConvNeXtBlock",
+    "Qwen3TTSTokenizerV2Decoder",
+    "Qwen3TTSTokenizerV2DecoderDecoderBlock",
+    "Qwen3TTSTokenizerV2DecoderDecoderResidualUnit",
+]

models/stream_runtime/core/compat.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+import os
+
+BACKEND_NAME = os.environ.get("QWEN_TTS_STREAM_BACKEND", "official").strip().lower()
+
+if BACKEND_NAME == "official":
+    from qwen_tts.core.tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import (
+        Qwen3TTSTokenizerV2CausalConvNet,
+        Qwen3TTSTokenizerV2CausalTransConvNet,
+        Qwen3TTSTokenizerV2ConvNeXtBlock,
+        Qwen3TTSTokenizerV2Decoder,
+        Qwen3TTSTokenizerV2DecoderDecoderBlock,
+        Qwen3TTSTokenizerV2DecoderDecoderResidualUnit,
+    )
+    from qwen_tts.inference.qwen3_tts_tokenizer import Qwen3TTSTokenizer
+else:
+    raise ImportError(
+        "Breeze streaming runtime only supports the official qwen_tts backend. "
+        "Set QWEN_TTS_STREAM_BACKEND=official or leave it unset."
+    )
+
+
+__all__ = [
+    "BACKEND_NAME",
+    "Qwen3TTSTokenizer",
+    "Qwen3TTSTokenizerV2CausalConvNet",
+    "Qwen3TTSTokenizerV2CausalTransConvNet",
+    "Qwen3TTSTokenizerV2ConvNeXtBlock",
+    "Qwen3TTSTokenizerV2Decoder",
+    "Qwen3TTSTokenizerV2DecoderDecoderBlock",
+    "Qwen3TTSTokenizerV2DecoderDecoderResidualUnit",
+]

models/stream_runtime/inference/__init__.py

@@ -0,0 +1,15 @@
+from ..core.compat import BACKEND_NAME
+from ..stream.lane import ExecutionLane
+from ..stream.runtime import (
+    MultiRequestStreamRuntime,
+    QwenStreamRuntimeConfig,
+    load_tokenizer,
+)
+
+__all__ = [
+    "BACKEND_NAME",
+    "ExecutionLane",
+    "MultiRequestStreamRuntime",
+    "QwenStreamRuntimeConfig",
+    "load_tokenizer",
+]

models/stream_runtime/stream/__init__.py

@@ -0,0 +1,13 @@
+from .lane import ExecutionLane
+from .runtime import MultiRequestStreamRuntime, QwenStreamRuntimeConfig, load_tokenizer
+from .state import RequestStatePool
+from .workspace import WorkspacePool
+
+__all__ = [
+    "ExecutionLane",
+    "MultiRequestStreamRuntime",
+    "QwenStreamRuntimeConfig",
+    "RequestStatePool",
+    "WorkspacePool",
+    "load_tokenizer",
+]

models/stream_runtime/stream/kv_cache.py

@@ -0,0 +1,104 @@
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+from torch import nn
+
+
+class _LayerKV(nn.Module):
+    def __init__(self, k: torch.Tensor, v: torch.Tensor):
+        super().__init__()
+        self.register_buffer("k", k, persistent=False)
+        self.register_buffer("v", v, persistent=False)
+
+
+class StaticShiftKVCache(nn.Module):
+    def __init__(
+        self,
+        num_layers: int,
+        window: int,
+        batch: int,
+        h_kv: int,
+        head_dim: int,
+        device,
+        dtype,
+    ):
+        super().__init__()
+        self.window = int(window)
+        self.layers = nn.ModuleList(
+            [
+                _LayerKV(
+                    k=torch.zeros(
+                        (batch, h_kv, self.window, head_dim), device=device, dtype=dtype
+                    ),
+                    v=torch.zeros(
+                        (batch, h_kv, self.window, head_dim), device=device, dtype=dtype
+                    ),
+                )
+                for _ in range(num_layers)
+            ]
+        )
+        self.is_sliding = [True] * num_layers
+
+    def get_seq_length(self) -> int:
+        return int(self.window)
+
+    def get_mask_sizes(self, cache_position: torch.Tensor, layer_idx: int):
+        q_len = int(cache_position.numel()) if cache_position is not None else 1
+        if cache_position is None or cache_position.numel() == 0:
+            return self.window, 0
+        cur_pos = int(cache_position.reshape(-1)[-1].item())
+        kv_length = min(cur_pos + 1, self.window)
+        kv_offset = max(cur_pos + 1 - kv_length, 0)
+        return kv_length, kv_offset
+
+    def update(
+        self,
+        key_states: torch.Tensor,
+        value_states: torch.Tensor,
+        layer_idx: int,
+        cache_kwargs: Any,
+    ):
+        layer = self.layers[layer_idx]
+        t_new = key_states.shape[-2]
+        if not torch._dynamo.is_compiling() and t_new > self.window:
+            raise ValueError(
+                f"StaticShiftKVCache assumes T_new <= window={self.window}, got {t_new}."
+            )
+        shift = min(int(t_new), self.window)
+        if shift < self.window:
+            layer.k[..., :-shift, :].copy_(layer.k[..., shift:, :])
+            layer.v[..., :-shift, :].copy_(layer.v[..., shift:, :])
+            layer.k[..., -shift:, :].copy_(key_states[..., -shift:, :])
+            layer.v[..., -shift:, :].copy_(value_states[..., -shift:, :])
+        else:
+            layer.k.copy_(key_states[..., -self.window :, :])
+            layer.v.copy_(value_states[..., -self.window :, :])
+        return layer.k, layer.v
+
+
+def make_fixed_attention_mask(
+    batch: int,
+    window: int,
+    cache_position: torch.Tensor,
+    device: torch.device,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    positions = cache_position.reshape(-1)
+    end_pos = positions[-1]
+    # Column j corresponds to the absolute key position:
+    #   abs_key(j) = end_pos - window + 1 + j
+    # We keep the column index static so torch.compile sees a fixed-size buffer.
+    cols = torch.arange(window, device=device, dtype=positions.dtype).view(1, window)
+    min_valid_col = torch.clamp((window - 1) - end_pos, min=0)
+    max_valid_col = positions.view(-1, 1) - end_pos + (window - 1)
+    valid = (cols >= min_valid_col) & (cols <= max_valid_col)
+    zeros = torch.zeros((positions.numel(), window), device=device, dtype=dtype)
+    neg_inf = torch.full(
+        (positions.numel(), window), torch.finfo(dtype).min, device=device, dtype=dtype
+    )
+    mask = torch.where(valid, zeros, neg_inf)
+    return mask.view(1, 1, positions.numel(), window).expand(
+        batch, 1, positions.numel(), window
+    )

models/stream_runtime/stream/lane.py

@@ -0,0 +1,435 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+from torch import nn
+
+from ..core.compat import (
+    Qwen3TTSTokenizer,
+    Qwen3TTSTokenizerV2CausalConvNet,
+    Qwen3TTSTokenizerV2CausalTransConvNet,
+    Qwen3TTSTokenizerV2ConvNeXtBlock,
+    Qwen3TTSTokenizerV2Decoder,
+    Qwen3TTSTokenizerV2DecoderDecoderBlock,
+    Qwen3TTSTokenizerV2DecoderDecoderResidualUnit,
+)
+from .kv_cache import StaticShiftKVCache, make_fixed_attention_mask
+from .state import ConvStateBlock, RequestStateSlot, TransConvStateBlock
+from .workspace import ConvWorkspaceBlock, TransConvWorkspaceBlock, WorkspaceSlot
+
+
+def _causal_conv_left_cache_len(conv: Qwen3TTSTokenizerV2CausalConvNet) -> int:
+    return int(conv.padding)
+
+
+def _tconv_left_cache_len(conv: Qwen3TTSTokenizerV2CausalTransConvNet) -> int:
+    kernel = int(conv.conv.kernel_size[0])
+    stride = int(conv.conv.stride[0])
+    return int((kernel - 1) // stride)
+
+
+def _cg_mark() -> None:
+    if hasattr(torch, "compiler") and hasattr(
+        torch.compiler, "cudagraph_mark_step_begin"
+    ):
+        torch.compiler.cudagraph_mark_step_begin()
+
+
+def _ensure_canonical_ncl(x: torch.Tensor) -> torch.Tensor:
+    """Materialize canonical contiguous NCL strides before compiled convolutions."""
+    x = x.contiguous()
+    if x.ndim != 3 or (x.stride(-1) == 1 and x.stride(-2) == x.shape[-1]):
+        return x
+    y = torch.empty_like(x, memory_format=torch.contiguous_format)
+    y.copy_(x)
+    return y
+
+
+@dataclass
+class LaneBinding:
+    request_slot: RequestStateSlot
+    workspace_slot: WorkspaceSlot
+
+
+class StaticCachedQwenCausalConv1dV2(nn.Module):
+    def __init__(
+        self, name: str, conv: Qwen3TTSTokenizerV2CausalConvNet, left_cache_len: int
+    ):
+        super().__init__()
+        self.name = name
+        self.conv = conv
+        self.left_cache_len = int(left_cache_len)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        state: dict[str, ConvStateBlock],
+        workspace: dict[str, ConvWorkspaceBlock],
+    ) -> torch.Tensor:
+        x = _ensure_canonical_ncl(x)
+        if self.left_cache_len <= 0:
+            return self.conv(x)
+        s = state[self.name]
+        w = workspace[self.name]
+        if not torch._dynamo.is_compiling():
+            if s.cache_buf.shape[-1] != self.left_cache_len:
+                raise ValueError(f"{self.name} cache length mismatch.")
+            if w.x_buf.shape[-1] != self.left_cache_len + x.shape[-1]:
+                raise ValueError(f"{self.name} x_buf length mismatch.")
+        w.x_buf[..., : self.left_cache_len].copy_(s.cache_buf)
+        w.x_buf[..., self.left_cache_len :].copy_(x)
+        conv_input = _ensure_canonical_ncl(w.x_buf)
+        y = self.conv(conv_input)
+        s.cache_buf.copy_(conv_input[..., -self.left_cache_len :])
+        return y[..., -x.shape[-1] :]
+
+
+class StaticCachedQwenTransposedConv1dV2(nn.Module):
+    def __init__(
+        self,
+        name: str,
+        conv: Qwen3TTSTokenizerV2CausalTransConvNet,
+        left_cache_len: int,
+    ):
+        super().__init__()
+        self.name = name
+        self.conv = conv
+        self.left_cache_len = int(left_cache_len)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        state: dict[str, TransConvStateBlock],
+        workspace: dict[str, TransConvWorkspaceBlock],
+    ) -> torch.Tensor:
+        x = _ensure_canonical_ncl(x)
+        if self.left_cache_len <= 0:
+            return self.conv(x)
+        s = state[self.name]
+        w = workspace[self.name]
+        if not torch._dynamo.is_compiling():
+            if s.cache_buf.shape[-1] != self.left_cache_len:
+                raise ValueError(f"{self.name} cache length mismatch.")
+            if w.x_buf.shape[-1] != self.left_cache_len + x.shape[-1]:
+                raise ValueError(f"{self.name} x_buf length mismatch.")
+        w.x_buf[..., : self.left_cache_len].copy_(s.cache_buf)
+        w.x_buf[..., self.left_cache_len :].copy_(x)
+        conv_input = _ensure_canonical_ncl(w.x_buf)
+        y = self.conv(conv_input)
+        s.cache_buf.copy_(conv_input[..., -self.left_cache_len :])
+        stride = int(self.conv.conv.stride[0])
+        prefix = self.left_cache_len * stride
+        new_len = x.shape[-1] * stride
+        return y[..., prefix : prefix + new_len]
+
+
+class StaticCachedConvNeXtBlockV2(nn.Module):
+    def __init__(self, name: str, block: Qwen3TTSTokenizerV2ConvNeXtBlock):
+        super().__init__()
+        self.block = block
+        self.dwconv = StaticCachedQwenCausalConv1dV2(
+            name=name,
+            conv=block.dwconv,
+            left_cache_len=_causal_conv_left_cache_len(block.dwconv),
+        )
+
+    def forward(
+        self, hidden_states: torch.Tensor, binding: LaneBinding
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.dwconv(
+            hidden_states, binding.request_slot.conv1d, binding.workspace_slot.conv1d
+        )
+        hidden_states = hidden_states.permute(0, 2, 1).contiguous()
+        hidden_states = self.block.norm(hidden_states)
+        hidden_states = self.block.pwconv1(hidden_states)
+        hidden_states = self.block.act(hidden_states)
+        hidden_states = self.block.pwconv2(hidden_states)
+        hidden_states = self.block.gamma * hidden_states
+        hidden_states = hidden_states.permute(0, 2, 1).contiguous()
+        return residual + hidden_states
+
+
+class StaticCachedDecoderResidualUnitV2(nn.Module):
+    def __init__(self, name: str, unit: Qwen3TTSTokenizerV2DecoderDecoderResidualUnit):
+        super().__init__()
+        self.unit = unit
+        self.conv1 = StaticCachedQwenCausalConv1dV2(
+            name=name,
+            conv=unit.conv1,
+            left_cache_len=_causal_conv_left_cache_len(unit.conv1),
+        )
+
+    def forward(self, hidden_state: torch.Tensor, binding: LaneBinding) -> torch.Tensor:
+        residual = hidden_state.contiguous()
+        hidden_state = self.unit.act1(hidden_state).contiguous()
+        hidden_state = self.conv1(
+            hidden_state, binding.request_slot.conv1d, binding.workspace_slot.conv1d
+        )
+        hidden_state = self.unit.act2(hidden_state).contiguous()
+        hidden_state = self.unit.conv2(_ensure_canonical_ncl(hidden_state))
+        return hidden_state + residual
+
+
+class StaticCachedDecoderBlockV2(nn.Module):
+    def __init__(self, name: str, block: Qwen3TTSTokenizerV2DecoderDecoderBlock):
+        super().__init__()
+        self.block = block
+        self.tconv = StaticCachedQwenTransposedConv1dV2(
+            name=f"{name}_tconv",
+            conv=block.block[1],
+            left_cache_len=_tconv_left_cache_len(block.block[1]),
+        )
+        self.residual_units = nn.ModuleList(
+            [
+                StaticCachedDecoderResidualUnitV2(f"{name}_residual_{idx}_conv1", unit)
+                for idx, unit in enumerate(block.block[2:])
+            ]
+        )
+
+    def forward(self, hidden: torch.Tensor, binding: LaneBinding) -> torch.Tensor:
+        hidden = self.block.block[0](hidden).contiguous()
+        hidden = self.tconv(
+            hidden, binding.request_slot.tconv1d, binding.workspace_slot.tconv1d
+        )
+        for unit in self.residual_units:
+            hidden = unit(hidden, binding)
+        return hidden
+
+
+def _copy_conv_state(
+    dst: dict[str, ConvStateBlock], src: dict[str, ConvStateBlock]
+) -> None:
+    for name, block in dst.items():
+        block.cache_buf.copy_(src[name].cache_buf)
+
+
+def _copy_tconv_state(
+    dst: dict[str, TransConvStateBlock], src: dict[str, TransConvStateBlock]
+) -> None:
+    for name, block in dst.items():
+        block.cache_buf.copy_(src[name].cache_buf)
+
+
+def _copy_kv_state(dst: StaticShiftKVCache, src: StaticShiftKVCache) -> None:
+    for dst_layer, src_layer in zip(dst.layers, src.layers):
+        dst_layer.k.copy_(src_layer.k)
+        dst_layer.v.copy_(src_layer.v)
+
+
+def _reset_request_state_tensors(slot: RequestStateSlot) -> None:
+    for block in slot.conv1d.values():
+        block.cache_buf.zero_()
+    for block in slot.tconv1d.values():
+        block.cache_buf.zero_()
+    for layer in slot.kv.layers:
+        layer.k.zero_()
+        layer.v.zero_()
+
+
+class _LaneCore(nn.Module):
+    def __init__(
+        self,
+        decoder: Qwen3TTSTokenizerV2Decoder,
+        chunk_frames: int,
+        binding: LaneBinding,
+    ):
+        super().__init__()
+        self.__dict__["decoder_ref"] = decoder
+        self.config = decoder.config
+        self.chunk_frames = int(chunk_frames)
+        self.binding = binding
+        self.pre_conv = StaticCachedQwenCausalConv1dV2(
+            "pre_conv",
+            decoder.pre_conv,
+            left_cache_len=_causal_conv_left_cache_len(decoder.pre_conv),
+        )
+        self.upsample = nn.ModuleList()
+        for idx, blocks in enumerate(decoder.upsample):
+            self.upsample.append(
+                nn.ModuleList(
+                    [
+                        StaticCachedQwenTransposedConv1dV2(
+                            f"upsample_{idx}_tconv",
+                            blocks[0],
+                            left_cache_len=_tconv_left_cache_len(blocks[0]),
+                        ),
+                        StaticCachedConvNeXtBlockV2(
+                            f"upsample_{idx}_dwconv", blocks[1]
+                        ),
+                    ]
+                )
+            )
+        self.decoder_pre_conv = StaticCachedQwenCausalConv1dV2(
+            "decoder_pre_conv",
+            decoder.decoder[0],
+            left_cache_len=_causal_conv_left_cache_len(decoder.decoder[0]),
+        )
+        self.decoder_blocks = nn.ModuleList(
+            [
+                StaticCachedDecoderBlockV2(f"decoder_block_{idx}", block)
+                for idx, block in enumerate(decoder.decoder[1:-2])
+            ]
+        )
+        self.decoder_snake = decoder.decoder[-2]
+        self.final_conv = StaticCachedQwenCausalConv1dV2(
+            "final_conv",
+            decoder.decoder[-1],
+            left_cache_len=_causal_conv_left_cache_len(decoder.decoder[-1]),
+        )
+        self.pre_transformer = decoder.pre_transformer
+        self.quantizer = decoder.quantizer
+
+    def forward(
+        self, codes: torch.Tensor, cache_position: torch.Tensor
+    ) -> torch.Tensor:
+        binding = self.binding
+        hidden = self.quantizer.decode(codes)
+        hidden = (
+            self.pre_conv(
+                hidden, binding.request_slot.conv1d, binding.workspace_slot.conv1d
+            )
+            .transpose(1, 2)
+            .contiguous()
+        )
+        attn_mask = make_fixed_attention_mask(
+            batch=hidden.shape[0],
+            window=self.config.sliding_window,
+            cache_position=cache_position,
+            device=hidden.device,
+            dtype=hidden.dtype,
+        )
+        hidden = self.pre_transformer(
+            inputs_embeds=hidden,
+            use_cache=True,
+            past_key_values=binding.request_slot.kv,
+            cache_position=cache_position,
+            attention_mask=attn_mask,
+        ).last_hidden_state
+        hidden = hidden.permute(0, 2, 1).contiguous()
+        for stage in self.upsample:
+            hidden = stage[0](
+                hidden, binding.request_slot.tconv1d, binding.workspace_slot.tconv1d
+            )
+            hidden = stage[1](hidden, binding)
+        wav = self.decoder_pre_conv(
+            hidden, binding.request_slot.conv1d, binding.workspace_slot.conv1d
+        )
+        for block in self.decoder_blocks:
+            wav = block(wav, binding)
+        wav = self.decoder_snake(wav).contiguous()
+        wav = self.final_conv(
+            wav, binding.request_slot.conv1d, binding.workspace_slot.conv1d
+        )
+        return wav.clamp(min=-1, max=1)
+
+
+class ExecutionLane:
+    def __init__(
+        self,
+        lane_id: int,
+        tokenizer: Qwen3TTSTokenizer,
+        chunk_frames: int,
+        hot_state: RequestStateSlot,
+        workspace: WorkspaceSlot,
+        fast: bool = False,
+    ):
+        if tokenizer.model is None:
+            raise ValueError("Tokenizer model is not loaded.")
+        self.chunk_frames = int(chunk_frames)
+        self.lane_id = int(lane_id)
+        self.binding = LaneBinding(request_slot=hot_state, workspace_slot=workspace)
+        self.core = _LaneCore(
+            tokenizer.model.decoder, chunk_frames=chunk_frames, binding=self.binding
+        )
+        num_quantizers = int(tokenizer.model.decoder.config.num_quantizers)
+        state_device = next(iter(hot_state.conv1d.values())).cache_buf.device
+        self.codes_in_buf = torch.zeros(
+            (1, num_quantizers, self.chunk_frames),
+            device=state_device,
+            dtype=torch.long,
+        )
+        self.cache_position_buf = torch.zeros(
+            (self.chunk_frames,), device=state_device, dtype=torch.long
+        )
+        self._base_cache_position = torch.arange(
+            self.chunk_frames, device=state_device, dtype=torch.long
+        )
+        self.step_model = self.core
+        self.fast = bool(fast)
+        self.cuda_graph: torch.cuda.CUDAGraph | None = None
+        self.cuda_graph_out: torch.Tensor | None = None
+
+    def load_request_state(self, state: RequestStateSlot) -> None:
+        self.binding.request_slot.next_step = int(state.next_step)
+        self.binding.request_slot.req_id = state.req_id
+        self.binding.request_slot.tail_flushed = bool(state.tail_flushed)
+        _copy_conv_state(self.binding.request_slot.conv1d, state.conv1d)
+        _copy_tconv_state(self.binding.request_slot.tconv1d, state.tconv1d)
+        _copy_kv_state(self.binding.request_slot.kv, state.kv)
+
+    def store_request_state(self, state: RequestStateSlot) -> None:
+        state.req_id = self.binding.request_slot.req_id
+        state.next_step = int(self.binding.request_slot.next_step)
+        state.tail_flushed = bool(self.binding.request_slot.tail_flushed)
+        _copy_conv_state(state.conv1d, self.binding.request_slot.conv1d)
+        _copy_tconv_state(state.tconv1d, self.binding.request_slot.tconv1d)
+        _copy_kv_state(state.kv, self.binding.request_slot.kv)
+
+    def reset_hot_state(self) -> None:
+        self.binding.request_slot.req_id = None
+        self.binding.request_slot.next_step = 0
+        self.binding.request_slot.tail_flushed = False
+        _reset_request_state_tensors(self.binding.request_slot)
+        self.codes_in_buf.zero_()
+        self.cache_position_buf.zero_()
+        for block in self.binding.workspace_slot.conv1d.values():
+            block.x_buf.zero_()
+        for block in self.binding.workspace_slot.tconv1d.values():
+            block.x_buf.zero_()
+
+    def warmup_cuda_graph(self, warmup_rounds: int = 2) -> None:
+        """Capture the stateful fixed-shape codec step after Snake compilation."""
+        if not self.fast or self.cuda_graph is not None:
+            return
+        device = self.codes_in_buf.device
+        if device.type != "cuda":
+            raise RuntimeError("codec CUDA Graph fast path requires CUDA")
+        capture_stream = torch.cuda.Stream(device=device)
+        capture_stream.wait_stream(torch.cuda.current_stream(device))
+        with torch.cuda.stream(capture_stream):
+            for _ in range(max(1, warmup_rounds)):
+                self.cuda_graph_out = self.step_model(
+                    self.codes_in_buf, self.cache_position_buf
+                )
+            capture_stream.synchronize()
+            self.cuda_graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(
+                self.cuda_graph,
+                stream=capture_stream,
+                capture_error_mode="thread_local",
+            ):
+                self.cuda_graph_out = self.step_model(
+                    self.codes_in_buf, self.cache_position_buf
+                )
+        torch.cuda.current_stream(device).wait_stream(capture_stream)
+
+    def run_step(self, codes_chunk: torch.Tensor, step_idx: int) -> torch.Tensor:
+        t_new = int(codes_chunk.shape[-1])
+        if t_new != self.chunk_frames:
+            raise ValueError(
+                f"ExecutionLane expects fixed chunk_frames={self.chunk_frames}, got step len={t_new}."
+            )
+        self.codes_in_buf.copy_(codes_chunk)
+        self.cache_position_buf.copy_(self._base_cache_position)
+        self.cache_position_buf.add_(int(step_idx))
+        _cg_mark()
+        if self.cuda_graph is None:
+            out = self.step_model(self.codes_in_buf, self.cache_position_buf)
+        else:
+            self.cuda_graph.replay()
+            assert self.cuda_graph_out is not None
+            out = self.cuda_graph_out
+        self.binding.request_slot.next_step = int(step_idx) + t_new
+        return out

models/stream_runtime/stream/runtime.py

@@ -0,0 +1,746 @@
+from __future__ import annotations
+
+import logging
+import time
+from collections import OrderedDict
+from dataclasses import dataclass
+from pathlib import Path
+
+import torch
+
+from ..core.compat import Qwen3TTSTokenizer
+from .kv_cache import StaticShiftKVCache
+from .lane import ExecutionLane, _causal_conv_left_cache_len, _tconv_left_cache_len
+from .state import (
+    ConvStateBlock,
+    RequestStatePool,
+    RequestStateSlot,
+    TransConvStateBlock,
+    reset_request_state,
+)
+from .workspace import (
+    ConvWorkspaceBlock,
+    TransConvWorkspaceBlock,
+    WorkspacePool,
+    WorkspaceSlot,
+    reset_workspace,
+)
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class QwenStreamRuntimeConfig:
+    chunk_frames: int = 1
+    # Strategy for a residual tail when len(codes) % chunk_frames != 0 and chunk_frames > 1.
+    # - "eager": decode the tail one frame at a time on a dedicated eager lane.
+    #            This avoids pad contamination at the tail boundary and is the default.
+    # - "pad":   keep the existing fixed-shape path by right-padding the tail to chunk_frames,
+    #            decoding it on the main lane, then trimming the wav back to the expected length.
+    non_integer_chunk_strategy: str = "eager"
+    num_lanes: int = 1
+    max_active_reqs: int | None = None
+    fast: bool = False
+    lifecycle_assert_mode: str = "warn"
+    tombstone_capacity: int = 1024
+    device: torch.device | None = None
+    dtype: torch.dtype = torch.float32
+
+
+@dataclass
+class TombstoneEntry:
+    reason: str
+
+
+def _tensor_bytes(tensor: torch.Tensor) -> int:
+    return int(tensor.numel() * tensor.element_size())
+
+
+def _format_bytes(num_bytes: int) -> str:
+    units = ["B", "KiB", "MiB", "GiB", "TiB"]
+    value = float(num_bytes)
+    for unit in units:
+        if value < 1024.0 or unit == units[-1]:
+            return f"{value:.2f} {unit}"
+        value /= 1024.0
+    return f"{num_bytes} B"
+
+
+def _request_state_bytes(slot: RequestStateSlot) -> int:
+    total = 0
+    for block in slot.conv1d.values():
+        total += _tensor_bytes(block.cache_buf)
+    for block in slot.tconv1d.values():
+        total += _tensor_bytes(block.cache_buf)
+    for layer in slot.kv.layers:
+        total += _tensor_bytes(layer.k)
+        total += _tensor_bytes(layer.v)
+    return total
+
+
+def _workspace_bytes(slot: WorkspaceSlot) -> int:
+    total = 0
+    for block in slot.conv1d.values():
+        total += _tensor_bytes(block.x_buf)
+    for block in slot.tconv1d.values():
+        total += _tensor_bytes(block.x_buf)
+    return total
+
+
+def _step_lengths(decoder, chunk_frames: int) -> dict[str, int]:
+    lengths: dict[str, int] = {}
+    cur = int(chunk_frames)
+    lengths["pre_conv"] = cur
+    for idx, blocks in enumerate(decoder.upsample):
+        lengths[f"upsample_{idx}_tconv"] = cur
+        stride = int(blocks[0].conv.stride[0])
+        cur = cur * stride
+        lengths[f"upsample_{idx}_dwconv"] = cur
+    lengths["decoder_pre_conv"] = cur
+    for idx, block in enumerate(decoder.decoder[1:-2]):
+        lengths[f"decoder_block_{idx}_tconv"] = cur
+        stride = int(block.block[1].conv.stride[0])
+        cur = cur * stride
+        for ridx, _unit in enumerate(block.block[2:]):
+            lengths[f"decoder_block_{idx}_residual_{ridx}_conv1"] = cur
+    lengths["final_conv"] = cur
+    return lengths
+
+
+def build_request_state_slot(
+    tokenizer: Qwen3TTSTokenizer,
+    chunk_frames: int,
+    device: torch.device,
+    dtype: torch.dtype,
+) -> RequestStateSlot:
+    decoder = tokenizer.model.decoder  # type: ignore[union-attr]
+    config = decoder.config
+    conv1d: dict[str, ConvStateBlock] = {}
+    tconv1d: dict[str, TransConvStateBlock] = {}
+
+    def add_conv(name: str, channels: int, left: int):
+        conv1d[name] = ConvStateBlock(
+            cache_buf=torch.zeros((1, channels, left), device=device, dtype=dtype),
+            left_cache_len=left,
+        )
+
+    def add_tconv(name: str, channels: int, left: int):
+        tconv1d[name] = TransConvStateBlock(
+            cache_buf=torch.zeros((1, channels, left), device=device, dtype=dtype),
+            left_cache_len=left,
+        )
+
+    add_conv(
+        "pre_conv", config.codebook_dim, _causal_conv_left_cache_len(decoder.pre_conv)
+    )
+    for idx, blocks in enumerate(decoder.upsample):
+        add_tconv(
+            f"upsample_{idx}_tconv",
+            blocks[0].conv.in_channels,
+            _tconv_left_cache_len(blocks[0]),
+        )
+        add_conv(
+            f"upsample_{idx}_dwconv",
+            blocks[1].dwconv.conv.in_channels,
+            _causal_conv_left_cache_len(blocks[1].dwconv),
+        )
+    add_conv(
+        "decoder_pre_conv",
+        config.latent_dim,
+        _causal_conv_left_cache_len(decoder.decoder[0]),
+    )
+    for idx, block in enumerate(decoder.decoder[1:-2]):
+        add_tconv(
+            f"decoder_block_{idx}_tconv",
+            block.block[1].conv.in_channels,
+            _tconv_left_cache_len(block.block[1]),
+        )
+        for ridx, unit in enumerate(block.block[2:]):
+            add_conv(
+                f"decoder_block_{idx}_residual_{ridx}_conv1",
+                unit.conv1.conv.in_channels,
+                _causal_conv_left_cache_len(unit.conv1),
+            )
+    add_conv(
+        "final_conv",
+        decoder.decoder[-1].conv.in_channels,
+        _causal_conv_left_cache_len(decoder.decoder[-1]),
+    )
+    kv = StaticShiftKVCache(
+        num_layers=config.num_hidden_layers,
+        window=config.sliding_window,
+        batch=1,
+        h_kv=config.num_key_value_heads,
+        head_dim=getattr(
+            config, "head_dim", config.hidden_size // config.num_attention_heads
+        ),
+        device=device,
+        dtype=dtype,
+    )
+    return RequestStateSlot(
+        req_id=None,
+        next_step=0,
+        tail_flushed=False,
+        conv1d=conv1d,
+        tconv1d=tconv1d,
+        kv=kv,
+    )
+
+
+def build_workspace_slot(
+    tokenizer: Qwen3TTSTokenizer,
+    chunk_frames: int,
+    device: torch.device,
+    dtype: torch.dtype,
+) -> WorkspaceSlot:
+    decoder = tokenizer.model.decoder  # type: ignore[union-attr]
+    step_lengths = _step_lengths(decoder, chunk_frames)
+    conv1d: dict[str, ConvWorkspaceBlock] = {}
+    tconv1d: dict[str, TransConvWorkspaceBlock] = {}
+
+    def add_conv(name: str, channels: int, left: int):
+        conv1d[name] = ConvWorkspaceBlock(
+            x_buf=torch.zeros(
+                (1, channels, left + step_lengths[name]), device=device, dtype=dtype
+            )
+        )
+
+    def add_tconv(name: str, channels: int, left: int):
+        tconv1d[name] = TransConvWorkspaceBlock(
+            x_buf=torch.zeros(
+                (1, channels, left + step_lengths[name]), device=device, dtype=dtype
+            )
+        )
+
+    add_conv(
+        "pre_conv",
+        decoder.config.codebook_dim,
+        _causal_conv_left_cache_len(decoder.pre_conv),
+    )
+    for idx, blocks in enumerate(decoder.upsample):
+        add_tconv(
+            f"upsample_{idx}_tconv",
+            blocks[0].conv.in_channels,
+            _tconv_left_cache_len(blocks[0]),
+        )
+        add_conv(
+            f"upsample_{idx}_dwconv",
+            blocks[1].dwconv.conv.in_channels,
+            _causal_conv_left_cache_len(blocks[1].dwconv),
+        )
+    add_conv(
+        "decoder_pre_conv",
+        decoder.config.latent_dim,
+        _causal_conv_left_cache_len(decoder.decoder[0]),
+    )
+    for idx, block in enumerate(decoder.decoder[1:-2]):
+        add_tconv(
+            f"decoder_block_{idx}_tconv",
+            block.block[1].conv.in_channels,
+            _tconv_left_cache_len(block.block[1]),
+        )
+        for ridx, unit in enumerate(block.block[2:]):
+            add_conv(
+                f"decoder_block_{idx}_residual_{ridx}_conv1",
+                unit.conv1.conv.in_channels,
+                _causal_conv_left_cache_len(unit.conv1),
+            )
+    add_conv(
+        "final_conv",
+        decoder.decoder[-1].conv.in_channels,
+        _causal_conv_left_cache_len(decoder.decoder[-1]),
+    )
+    return WorkspaceSlot(conv1d=conv1d, tconv1d=tconv1d)
+
+
+class MultiRequestStreamRuntime:
+    def __init__(self, tokenizer: Qwen3TTSTokenizer, config: QwenStreamRuntimeConfig):
+        if tokenizer.model is None:
+            raise ValueError("Tokenizer model is not loaded.")
+        if config.device is None:
+            try:
+                config.device = next(tokenizer.model.parameters()).device
+            except StopIteration:
+                config.device = torch.device("cpu")
+        self.tokenizer = tokenizer
+        self.config = config
+        if self.config.lifecycle_assert_mode not in {"off", "warn", "raise"}:
+            raise ValueError(
+                f"lifecycle_assert_mode must be one of off/warn/raise, got {self.config.lifecycle_assert_mode}"
+            )
+        if self.config.tombstone_capacity < 0:
+            raise ValueError(
+                f"tombstone_capacity must be >= 0, got {self.config.tombstone_capacity}"
+            )
+        if self.config.non_integer_chunk_strategy not in {"eager", "pad"}:
+            raise ValueError(
+                "non_integer_chunk_strategy must be one of {'eager', 'pad'}, "
+                f"got {self.config.non_integer_chunk_strategy!r}"
+            )
+        if self.config.fast and (
+            self.config.num_lanes != 1 or self.config.max_active_reqs not in {None, 1}
+        ):
+            raise ValueError("fast codec requires num_lanes=1 and max_active_reqs<=1")
+        if (
+            self.config.fast
+            and self.config.device is not None
+            and self.config.device.type == "cuda"
+        ):
+            torch.backends.cudnn.benchmark = True
+        self._compile_snakes()
+        self._req_tombstones: OrderedDict[str, TombstoneEntry] = OrderedDict()
+        self._enforce_compile_attention_backend()
+        self.request_pool = RequestStatePool(
+            create_slot=lambda: build_request_state_slot(
+                self.tokenizer,
+                self.config.chunk_frames,
+                self.config.device,
+                self.config.dtype,
+            ),
+            reset_slot=reset_request_state,
+            max_active_reqs=self.config.max_active_reqs,
+            on_evict=self._on_request_evicted,
+        )
+        self.workspace_pool = WorkspacePool(
+            create_slot=lambda: build_workspace_slot(
+                self.tokenizer,
+                self.config.chunk_frames,
+                self.config.device,
+                self.config.dtype,
+            ),
+            reset_slot=reset_workspace,
+            size=self.config.num_lanes,
+        )
+        self.lanes: list[ExecutionLane] = []
+        self.tail_eager_lanes: list[ExecutionLane] = []
+        self.req_to_lane: dict[str, int] = {}
+        self._single_lane_req_id: str | None = None
+        for lane_id in range(self.config.num_lanes):
+            workspace_idx, workspace_slot = self.workspace_pool.acquire()
+            if workspace_idx != lane_id:
+                raise RuntimeError(
+                    "Workspace pool returned unexpected lane slot order."
+                )
+            hot_state = build_request_state_slot(
+                self.tokenizer,
+                self.config.chunk_frames,
+                self.config.device,
+                self.config.dtype,
+            )
+            self.lanes.append(
+                ExecutionLane(
+                    lane_id=lane_id,
+                    tokenizer=self.tokenizer,
+                    chunk_frames=self.config.chunk_frames,
+                    hot_state=hot_state,
+                    workspace=workspace_slot,
+                    fast=self.config.fast,
+                )
+            )
+            if (
+                self.config.chunk_frames > 1
+                and self.config.non_integer_chunk_strategy == "eager"
+            ):
+                eager_tail_hot_state = build_request_state_slot(
+                    self.tokenizer,
+                    chunk_frames=1,
+                    device=self.config.device,
+                    dtype=self.config.dtype,
+                )
+                eager_tail_workspace = build_workspace_slot(
+                    self.tokenizer,
+                    chunk_frames=1,
+                    device=self.config.device,
+                    dtype=self.config.dtype,
+                )
+                # Residual tails are rare and only happen at request shutdown. Keep this fallback
+                # eager so the default path avoids padding a future code into the tail boundary.
+                self.tail_eager_lanes.append(
+                    ExecutionLane(
+                        lane_id=lane_id,
+                        tokenizer=self.tokenizer,
+                        chunk_frames=1,
+                        hot_state=eager_tail_hot_state,
+                        workspace=eager_tail_workspace,
+                        fast=False,
+                    )
+                )
+        self.samples_per_code = self._validate_and_get_samples_per_code()
+        self._log_runtime_config()
+        self._maybe_warmup_fast_codec()
+
+    def _validate_and_get_samples_per_code(self) -> int:
+        theoretical_chunk_out = _step_lengths(
+            self.tokenizer.model.decoder, self.config.chunk_frames
+        )["final_conv"]  # type: ignore[union-attr]
+        if theoretical_chunk_out % self.config.chunk_frames != 0:
+            logger.error(
+                "breeze_codec invalid chunk->wav ratio: theoretical_chunk_out=%d chunk_frames=%d is not divisible",
+                theoretical_chunk_out,
+                self.config.chunk_frames,
+            )
+        theoretical_spc = theoretical_chunk_out // self.config.chunk_frames
+        num_quantizers = int(self.tokenizer.model.decoder.config.num_quantizers)  # type: ignore[union-attr]
+        dummy_codes = torch.zeros(
+            (1, num_quantizers, self.config.chunk_frames),
+            device=self.config.device,
+            dtype=torch.long,
+        )
+        with torch.inference_mode():
+            observed = self.tokenizer.model.decoder(dummy_codes)  # type: ignore[union-attr]
+        observed_chunk_out = int(observed.shape[-1])
+        if observed_chunk_out % self.config.chunk_frames != 0:
+            logger.error(
+                "breeze_codec observed non-integer chunk->wav ratio: observed_chunk_out=%d chunk_frames=%d",
+                observed_chunk_out,
+                self.config.chunk_frames,
+            )
+        observed_spc = observed_chunk_out // self.config.chunk_frames
+        if observed_spc != theoretical_spc:
+            logger.error(
+                "breeze_codec chunk->wav ratio mismatch: theoretical=%d observed=%d",
+                theoretical_spc,
+                observed_spc,
+            )
+        else:
+            logger.info(
+                "breeze_codec chunk->wav ratio verified: chunk_frames=%d chunk_output=%d samples_per_code=%d",
+                self.config.chunk_frames,
+                observed_chunk_out,
+                observed_spc,
+            )
+        return int(observed_spc)
+
+    def _compile_snakes(self) -> None:
+        if not self.config.fast:
+            return
+        torch._dynamo.config.recompile_limit = max(
+            int(torch._dynamo.config.recompile_limit), 64
+        )
+        torch._dynamo.config.accumulated_recompile_limit = max(
+            int(torch._dynamo.config.accumulated_recompile_limit), 256
+        )
+        compiled = 0
+        for module in self.tokenizer.model.decoder.modules():
+            if module.__class__.__name__ != "SnakeBeta":
+                continue
+            module.forward = torch.compile(
+                module.forward,
+                mode="max-autotune-no-cudagraphs",
+                fullgraph=True,
+                dynamic=True,
+            )
+            compiled += 1
+        logger.info("breeze_codec compiled %d SnakeBeta modules", compiled)
+
+    def _enforce_compile_attention_backend(self) -> None:
+        if not self.config.fast:
+            return
+        decoder = self.tokenizer.model.decoder  # type: ignore[union-attr]
+        pre_transformer = getattr(decoder, "pre_transformer", None)
+        if pre_transformer is None or not hasattr(pre_transformer, "config"):
+            return
+        current_impl = getattr(pre_transformer.config, "_attn_implementation", None)
+        if current_impl != "eager":
+            logger.warning(
+                "breeze_codec forcing pre_transformer attention backend to eager for fast codec. previous=%s",
+                current_impl,
+            )
+            pre_transformer.config._attn_implementation = "eager"
+        else:
+            logger.info(
+                "breeze_codec pre_transformer attention backend already eager for fast codec."
+            )
+
+    def _log_runtime_config(self) -> None:
+        sample_request_state = build_request_state_slot(
+            self.tokenizer,
+            self.config.chunk_frames,
+            self.config.device,
+            self.config.dtype,
+        )
+        sample_workspace = build_workspace_slot(
+            self.tokenizer,
+            self.config.chunk_frames,
+            self.config.device,
+            self.config.dtype,
+        )
+        current_impl = getattr(
+            self.tokenizer.model.decoder.pre_transformer.config,
+            "_attn_implementation",
+            "unknown",
+        )  # type: ignore[union-attr]
+        logger.info(
+            "breeze_codec runtime init: chunk_frames=%d num_lanes=%d device=%s dtype=%s fast=%s attn_impl=%s non_integer_chunk_strategy=%s lifecycle_assert_mode=%s tombstone_capacity=%d max_active_reqs=%s",
+            self.config.chunk_frames,
+            self.config.num_lanes,
+            self.config.device,
+            self.config.dtype,
+            self.config.fast,
+            current_impl,
+            self.config.non_integer_chunk_strategy,
+            self.config.lifecycle_assert_mode,
+            self.config.tombstone_capacity,
+            self.config.max_active_reqs,
+        )
+        logger.info(
+            "breeze_codec memory: per_request_state=%s per_workspace=%s total_request_capacity=%s total_workspace_capacity=%s samples_per_code=%s",
+            _format_bytes(_request_state_bytes(sample_request_state)),
+            _format_bytes(_workspace_bytes(sample_workspace)),
+            _format_bytes(
+                _request_state_bytes(sample_request_state)
+                * (self.config.max_active_reqs or 1)
+            ),
+            _format_bytes(_workspace_bytes(sample_workspace) * self.config.num_lanes),
+            self.samples_per_code,
+        )
+
+    def _sync_device(self) -> None:
+        if self.config.device is not None and self.config.device.type == "cuda":
+            torch.cuda.synchronize(self.config.device)
+
+    def _handle_lifecycle_assert(self, message: str) -> None:
+        mode = self.config.lifecycle_assert_mode
+        if mode == "off":
+            return
+        if mode == "warn":
+            logger.warning(message)
+            return
+        raise RuntimeError(message)
+
+    def _record_tombstone(self, req_id: str, reason: str) -> None:
+        if self.config.tombstone_capacity == 0:
+            return
+        self._req_tombstones.pop(req_id, None)
+        self._req_tombstones[req_id] = TombstoneEntry(reason=reason)
+        while len(self._req_tombstones) > self.config.tombstone_capacity:
+            self._req_tombstones.popitem(last=False)
+
+    def _clear_tombstone(self, req_id: str) -> None:
+        self._req_tombstones.pop(req_id, None)
+
+    def _on_request_evicted(self, req_id: str) -> None:
+        self.req_to_lane.pop(req_id, None)
+        if self._single_lane_req_id == req_id:
+            self._single_lane_req_id = None
+        self._record_tombstone(req_id, "evicted_lru")
+
+    def _maybe_warmup_fast_codec(self) -> None:
+        if not self.config.fast:
+            return
+        self._sync_device()
+        t0 = time.perf_counter()
+        with torch.inference_mode():
+            for lane in self.lanes:
+                lane.reset_hot_state()
+                lane.warmup_cuda_graph(warmup_rounds=2)
+                lane.reset_hot_state()
+        self._sync_device()
+        logger.info(
+            "breeze_codec fast codec warmup done: elapsed_ms=%.3f",
+            (time.perf_counter() - t0) * 1000.0,
+        )
+
+    def create_request(self, req_id: str, reset: bool = True) -> RequestStateSlot:
+        slot = self.request_pool.get(req_id)
+        if reset:
+            reset_request_state(slot)
+            slot.req_id = req_id
+        elif slot.req_id is None:
+            slot.req_id = req_id
+        return slot
+
+    def reset_request(self, req_id: str) -> RequestStateSlot:
+        return self.create_request(req_id, reset=True)
+
+    def release_request(self, req_id: str) -> None:
+        was_active = req_id in self.request_pool.active_req_ids()
+        self.req_to_lane.pop(req_id, None)
+        if self._single_lane_req_id == req_id:
+            self._single_lane_req_id = None
+        self.request_pool.pop(req_id)
+        if was_active:
+            self._record_tombstone(req_id, "closed")
+
+    # Public request-centric API
+    def open_request(
+        self, req_id: str, reset: bool = True, is_first_decode: bool | None = None
+    ) -> RequestStateSlot:
+        is_active = req_id in self.request_pool.active_req_ids()
+        tombstone = self._req_tombstones.get(req_id)
+        if is_first_decode is True and is_active:
+            self._handle_lifecycle_assert(
+                f"open_request received is_first_decode=True for an already active request: req_id={req_id}"
+            )
+        if is_first_decode is False and not is_active:
+            reason = tombstone.reason if tombstone is not None else "missing"
+            self._handle_lifecycle_assert(
+                f"open_request received is_first_decode=False for a non-active request: req_id={req_id} tombstone_reason={reason}"
+            )
+        self._clear_tombstone(req_id)
+        return self.create_request(req_id=req_id, reset=reset)
+
+    def close_request(self, req_id: str) -> None:
+        self.release_request(req_id)
+
+    def _select_lane(self, req_id: str) -> ExecutionLane:
+        lane_idx = self.req_to_lane.get(req_id)
+        if lane_idx is None:
+            lane_idx = len(self.req_to_lane) % self.config.num_lanes
+            self.req_to_lane[req_id] = lane_idx
+        return self.lanes[lane_idx]
+
+    def _select_tail_eager_lane(self, req_id: str) -> ExecutionLane:
+        lane_idx = self.req_to_lane.get(req_id)
+        if lane_idx is None:
+            raise RuntimeError(
+                f"tail eager lane requested before main lane selection: req_id={req_id}"
+            )
+        return self.tail_eager_lanes[lane_idx]
+
+    def _normalize_codes(self, codes_chunk: torch.Tensor) -> torch.Tensor:
+        if codes_chunk.ndim != 3:
+            raise ValueError(
+                f"Expected codes chunk shape [B, Q, T], got {tuple(codes_chunk.shape)}"
+            )
+        if codes_chunk.shape[0] != 1:
+            raise ValueError(
+                f"Only batch_size=1 is supported, got {codes_chunk.shape[0]}"
+            )
+        if codes_chunk.device != self.config.device:
+            codes_chunk = codes_chunk.to(self.config.device)
+        if codes_chunk.dtype != torch.long:
+            codes_chunk = codes_chunk.to(torch.long)
+        return codes_chunk.contiguous()
+
+    def decode_request_chunk(
+        self, req_id: str, codes_chunk: torch.Tensor, reset: bool = False
+    ) -> torch.Tensor:
+        state = self.create_request(req_id, reset=reset)
+        codes_chunk = self._normalize_codes(codes_chunk)
+        if codes_chunk.shape[-1] == 0:
+            return torch.zeros(
+                (1, 1, 0), device=self.config.device, dtype=self.config.dtype
+            )
+        if state.tail_flushed:
+            logger.error(
+                "decode() called after residual tail flush; request must be reset or closed before reuse. req_id=%s",
+                req_id,
+            )
+            return torch.zeros(
+                (1, 1, 0), device=self.config.device, dtype=self.config.dtype
+            )
+        lane = self._select_lane(req_id)
+        self.request_pool.pin(req_id)
+        try:
+            sticky_lane = self.config.fast
+            if reset or not sticky_lane or self._single_lane_req_id != req_id:
+                lane.load_request_state(state)
+                if sticky_lane:
+                    self._single_lane_req_id = req_id
+            outs = []
+            state_stored = False
+            with torch.inference_mode():
+                full_len = (
+                    codes_chunk.shape[-1] // self.config.chunk_frames
+                ) * self.config.chunk_frames
+                for offset in range(0, full_len, self.config.chunk_frames):
+                    step_codes = codes_chunk[
+                        ..., offset : offset + self.config.chunk_frames
+                    ]
+                    outs.append(lane.run_step(step_codes, state.next_step).clone())
+                    state.next_step += int(step_codes.shape[-1])
+                    lane.binding.request_slot.next_step = state.next_step
+                tail = codes_chunk.shape[-1] - full_len
+                if tail > 0:
+                    tail_codes = codes_chunk[..., full_len:]
+                    if self.config.non_integer_chunk_strategy == "eager":
+                        # Persist the fixed-chunk lane state first, then finish the residual tail
+                        # on a dedicated eager chunk=1 lane. This keeps the default strategy free
+                        # of future-code padding at the tail boundary.
+                        lane.store_request_state(state)
+                        state_stored = True
+                        tail_lane = self._select_tail_eager_lane(req_id)
+                        tail_lane.load_request_state(state)
+                        for offset in range(tail):
+                            step_codes = tail_codes[..., offset : offset + 1]
+                            outs.append(
+                                tail_lane.run_step(step_codes, state.next_step).clone()
+                            )
+                            state.next_step += 1
+                            tail_lane.binding.request_slot.next_step = state.next_step
+                        state.tail_flushed = True
+                        tail_lane.binding.request_slot.tail_flushed = True
+                        tail_lane.store_request_state(state)
+                    else:
+                        pad = torch.zeros(
+                            (
+                                tail_codes.shape[0],
+                                tail_codes.shape[1],
+                                self.config.chunk_frames - tail,
+                            ),
+                            device=tail_codes.device,
+                            dtype=tail_codes.dtype,
+                        )
+                        padded = torch.cat([tail_codes, pad], dim=-1)
+                        padded_wav = lane.run_step(padded, state.next_step).clone()
+                        trim_len = int(tail * self.samples_per_code)
+                        outs.append(padded_wav[..., :trim_len])
+                        state.next_step += tail
+                        state.tail_flushed = True
+                        lane.binding.request_slot.next_step = state.next_step
+                        lane.binding.request_slot.tail_flushed = True
+                    logger.warning(
+                        "Residual tail decode detected; request must be reset or closed after this call. req_id=%s tail_frames=%d strategy=%s",
+                        req_id,
+                        tail,
+                        self.config.non_integer_chunk_strategy,
+                    )
+            if not state_stored and not sticky_lane:
+                lane.store_request_state(state)
+            if not outs:
+                return torch.zeros(
+                    (1, 1, 0), device=self.config.device, dtype=self.config.dtype
+                )
+            return torch.cat(outs, dim=-1)
+        finally:
+            self.request_pool.unpin(req_id)
+
+    def decode_request_sequence(
+        self, req_id: str, codes: torch.Tensor, reset: bool = True
+    ) -> torch.Tensor:
+        return self.decode_request_chunk(req_id=req_id, codes_chunk=codes, reset=reset)
+
+    def get_request_step(self, req_id: str) -> int:
+        slot = self.request_pool.get(req_id)
+        return int(slot.next_step)
+
+    def decode(
+        self, req_id: str, codes_chunk: torch.Tensor, reset: bool = False
+    ) -> torch.Tensor:
+        return self.decode_request_chunk(
+            req_id=req_id, codes_chunk=codes_chunk, reset=reset
+        )
+
+    def decode_sequence(
+        self, req_id: str, codes: torch.Tensor, reset: bool = True
+    ) -> torch.Tensor:
+        return self.decode_request_sequence(req_id=req_id, codes=codes, reset=reset)
+
+    def get_step(self, req_id: str) -> int:
+        return self.get_request_step(req_id)
+
+
+def load_tokenizer(model_path: str | Path) -> Qwen3TTSTokenizer:
+    model_path = Path(model_path)
+    if not model_path.exists():
+        raise FileNotFoundError(f"model path not found: {model_path}")
+    tok = Qwen3TTSTokenizer.from_pretrained(
+        str(model_path),
+        torch_dtype=torch.float32,
+        load_feature_extractor=False,
+    )
+    if tok.model is None:
+        raise RuntimeError("Tokenizer model failed to load.")
+    tok.model.eval()
+    return tok

models/stream_runtime/stream/state.py

@@ -0,0 +1,165 @@
+from __future__ import annotations
+
+import logging
+from collections import OrderedDict
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Generic, TypeVar
+
+import torch
+
+from .kv_cache import StaticShiftKVCache
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T")
+
+
+@dataclass
+class ConvStateBlock:
+    cache_buf: torch.Tensor
+    left_cache_len: int
+
+
+@dataclass
+class TransConvStateBlock:
+    cache_buf: torch.Tensor
+    left_cache_len: int
+
+
+@dataclass
+class RequestStateSlot:
+    req_id: str | None
+    next_step: int
+    tail_flushed: bool
+    conv1d: dict[str, ConvStateBlock]
+    tconv1d: dict[str, TransConvStateBlock]
+    kv: StaticShiftKVCache
+
+
+def reset_request_state(slot: RequestStateSlot) -> None:
+    slot.req_id = None
+    slot.next_step = 0
+    slot.tail_flushed = False
+    for block in slot.conv1d.values():
+        block.cache_buf.zero_()
+    for block in slot.tconv1d.values():
+        block.cache_buf.zero_()
+    for layer in slot.kv.layers:
+        layer.k.zero_()
+        layer.v.zero_()
+
+
+class RequestStatePool(Generic[T]):
+    def __init__(
+        self,
+        create_slot: Callable[[], T],
+        reset_slot: Callable[[T], None],
+        max_active_reqs: int | None = None,
+        on_evict: Callable[[str], None] | None = None,
+    ):
+        if max_active_reqs is not None and max_active_reqs <= 0:
+            raise ValueError(
+                f"max_active_reqs must be > 0 or None, got {max_active_reqs}"
+            )
+        self._create_slot = create_slot
+        self._reset_slot = reset_slot
+        self.max_active_reqs = max_active_reqs
+        self._on_evict = on_evict
+        self._req_to_slot: OrderedDict[str, int] = OrderedDict()
+        self._slot_to_req: list[str | None] = []
+        self._slots: list[T] = []
+        self._pinned: set[str] = set()
+        self._dynamic: dict[str, T] = {}
+        if self.max_active_reqs is not None:
+            self._slots = [self._create_slot() for _ in range(self.max_active_reqs)]
+            self._slot_to_req = [None for _ in range(self.max_active_reqs)]
+
+    def __len__(self) -> int:
+        return (
+            len(self._req_to_slot)
+            if self.max_active_reqs is not None
+            else len(self._dynamic)
+        )
+
+    def active_req_ids(self) -> list[str]:
+        if self.max_active_reqs is None:
+            return list(self._dynamic.keys())
+        return list(self._req_to_slot.keys())
+
+    def num_free_slots(self) -> int | None:
+        if self.max_active_reqs is None:
+            return None
+        return sum(1 for req_id in self._slot_to_req if req_id is None)
+
+    def pin(self, req_id: str) -> None:
+        self._pinned.add(req_id)
+
+    def unpin(self, req_id: str) -> None:
+        self._pinned.discard(req_id)
+
+    def _find_free_slot(self) -> int | None:
+        for idx, req_id in enumerate(self._slot_to_req):
+            if req_id is None:
+                return idx
+        return None
+
+    def _evict_lru(self) -> int:
+        victim_req_id: str | None = None
+        for req_id in self._req_to_slot.keys():
+            if req_id not in self._pinned:
+                victim_req_id = req_id
+                break
+        if victim_req_id is None:
+            logger.error(
+                "RequestStatePool found no evictable request. max_active_reqs=%s active_req_ids=%s pinned=%s",
+                self.max_active_reqs,
+                list(self._req_to_slot.keys()),
+                sorted(self._pinned),
+            )
+            raise RuntimeError("No evictable request state slot available.")
+        slot_idx = self._req_to_slot.pop(victim_req_id)
+        self._slot_to_req[slot_idx] = None
+        self._reset_slot(self._slots[slot_idx])
+        if self._on_evict is not None:
+            self._on_evict(victim_req_id)
+        logger.info(
+            "Evicted request state due to LRU capacity limit: req_id=%s slot_idx=%d",
+            victim_req_id,
+            slot_idx,
+        )
+        return slot_idx
+
+    def get(self, req_id: str) -> T:
+        if self.max_active_reqs is None:
+            value = self._dynamic.get(req_id)
+            if value is None:
+                value = self._create_slot()
+                self._dynamic[req_id] = value
+            return value
+        slot_idx = self._req_to_slot.get(req_id)
+        if slot_idx is None:
+            slot_idx = self._find_free_slot()
+            if slot_idx is None:
+                slot_idx = self._evict_lru()
+            self._slot_to_req[slot_idx] = req_id
+            self._req_to_slot[req_id] = slot_idx
+            self._reset_slot(self._slots[slot_idx])
+        self._req_to_slot.move_to_end(req_id)
+        return self._slots[slot_idx]
+
+    def pop(self, req_id: str) -> None:
+        self._pinned.discard(req_id)
+        if self.max_active_reqs is None:
+            self._dynamic.pop(req_id, None)
+            return
+        slot_idx = self._req_to_slot.pop(req_id, None)
+        if slot_idx is None:
+            return
+        self._reset_slot(self._slots[slot_idx])
+        self._slot_to_req[slot_idx] = None
+
+    def get_slot_index(self, req_id: str) -> int | None:
+        if self.max_active_reqs is None:
+            return None
+        return self._req_to_slot.get(req_id)

models/stream_runtime/stream/workspace.py

@@ -0,0 +1,67 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Generic, TypeVar
+
+import torch
+
+W = TypeVar("W")
+
+
+@dataclass
+class ConvWorkspaceBlock:
+    x_buf: torch.Tensor
+
+
+@dataclass
+class TransConvWorkspaceBlock:
+    x_buf: torch.Tensor
+
+
+@dataclass
+class WorkspaceSlot:
+    conv1d: dict[str, ConvWorkspaceBlock]
+    tconv1d: dict[str, TransConvWorkspaceBlock]
+
+
+def reset_workspace(slot: WorkspaceSlot) -> None:
+    for block in slot.conv1d.values():
+        block.x_buf.zero_()
+    for block in slot.tconv1d.values():
+        block.x_buf.zero_()
+
+
+class WorkspacePool(Generic[W]):
+    def __init__(
+        self, create_slot: Callable[[], W], reset_slot: Callable[[W], None], size: int
+    ):
+        if size <= 0:
+            raise ValueError(f"workspace pool size must be > 0, got {size}")
+        self._reset_slot = reset_slot
+        self._slots: list[W] = [create_slot() for _ in range(size)]
+        self._free: list[int] = list(range(size))
+
+    def __len__(self) -> int:
+        return len(self._slots)
+
+    def num_free_slots(self) -> int:
+        return len(self._free)
+
+    def acquire(self) -> tuple[int, W]:
+        if not self._free:
+            raise RuntimeError("No free workspace slot available.")
+        slot_idx = self._free.pop(0)
+        slot = self._slots[slot_idx]
+        self._reset_slot(slot)
+        return slot_idx, slot
+
+    def release(self, slot_idx: int) -> None:
+        if slot_idx in self._free:
+            return
+        self._reset_slot(self._slots[slot_idx])
+        self._free.append(slot_idx)
+        self._free.sort()
+
+    def get_slot(self, slot_idx: int) -> W:
+        return self._slots[slot_idx]

models/t5gemma2_compat.py

@@ -0,0 +1,813 @@
+# Local compatibility shim for loading T5Gemma 2 text encoders under
+# transformers 4.57.x, which does not ship the t5gemma2 architecture.
+
+
+import torch
+from torch import nn
+from transformers import PretrainedConfig
+from transformers.activations import ACT2FN
+from transformers.modeling_layers import GradientCheckpointingLayer
+from transformers.modeling_outputs import BaseModelOutput
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+
+
+class T5Gemma2TextConfig(PretrainedConfig):
+    model_type = "t5gemma2_text"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        vocab_size: int = 262208,
+        hidden_size: int = 2304,
+        intermediate_size: int = 9216,
+        num_hidden_layers: int = 26,
+        num_attention_heads: int = 8,
+        num_key_value_heads: int = 4,
+        head_dim: int = 256,
+        hidden_activation: str = "gelu_pytorch_tanh",
+        max_position_embeddings: int = 131072,
+        initializer_range: float = 0.02,
+        rms_norm_eps: float = 1e-6,
+        use_cache: bool = True,
+        rope_parameters: dict | None = None,
+        attention_bias: bool = False,
+        attention_dropout: float = 0.0,
+        dropout_rate: float = 0.0,
+        query_pre_attn_scalar: int = 256,
+        sliding_window: int | None = 4096,
+        layer_types: list[str] | None = None,
+        final_logit_softcapping: float | None = None,
+        attn_logit_softcapping: float | None = None,
+        eoi_token_index: int = 256000,
+        pad_token_id: int | None = 0,
+        eos_token_id: int | None = 1,
+        bos_token_id: int | None = 2,
+        tie_word_embeddings: bool = True,
+        **kwargs,
+    ):
+        sliding_window_pattern = kwargs.pop(
+            "_sliding_window_pattern", kwargs.pop("sliding_window_pattern", 6)
+        )
+        rope_scaling = kwargs.pop("rope_scaling", None)
+        super().__init__(
+            pad_token_id=pad_token_id,
+            eos_token_id=eos_token_id,
+            bos_token_id=bos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_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.hidden_activation = hidden_activation
+        self.max_position_embeddings = max_position_embeddings
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        self.attention_bias = attention_bias
+        self.attention_dropout = attention_dropout
+        self.dropout_rate = dropout_rate
+        self.query_pre_attn_scalar = query_pre_attn_scalar
+        self.sliding_window = sliding_window
+        self.final_logit_softcapping = final_logit_softcapping
+        self.attn_logit_softcapping = attn_logit_softcapping
+        self.eoi_token_index = eoi_token_index
+
+        if layer_types is None:
+            layer_types = [
+                "sliding_attention"
+                if (i + 1) % sliding_window_pattern
+                else "full_attention"
+                for i in range(num_hidden_layers)
+            ]
+        self.layer_types = layer_types
+        self.rope_parameters = self._normalize_rope_parameters(
+            rope_parameters, rope_scaling
+        )
+
+    @staticmethod
+    def _normalize_rope_parameters(
+        rope_parameters: dict | None, rope_scaling: dict | None
+    ) -> dict:
+        params = {
+            "sliding_attention": {"rope_type": "default", "rope_theta": 10000.0},
+            "full_attention": {"rope_type": "default", "rope_theta": 1000000.0},
+        }
+        if rope_parameters:
+            for key, value in rope_parameters.items():
+                params[key] = (
+                    dict(value) if value is not None else {"rope_type": "default"}
+                )
+        if rope_scaling:
+            params["full_attention"].update(rope_scaling)
+        params["sliding_attention"].setdefault("rope_type", "default")
+        params["sliding_attention"].setdefault("rope_theta", 10000.0)
+        params["full_attention"].setdefault("rope_type", "default")
+        params["full_attention"].setdefault("rope_theta", 1000000.0)
+        return params
+
+
+class T5Gemma2RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.zeros(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        output = x.float() * torch.rsqrt(
+            x.float().pow(2).mean(-1, keepdim=True) + self.eps
+        )
+        output = output * (1.0 + self.weight.float())
+        return output.type_as(x)
+
+
+class T5Gemma2MLP(nn.Module):
+    def __init__(self, config: T5Gemma2TextConfig):
+        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
+        )
+        self.act_fn = ACT2FN[config.hidden_activation]
+        self.dropout = nn.Dropout(config.dropout_rate)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.act_fn(self.gate_proj(x)) * self.up_proj(x)
+        hidden_states = self.dropout(hidden_states)
+        return self.down_proj(hidden_states)
+
+
+class T5Gemma2RotaryEmbedding(nn.Module):
+    def __init__(self, config: T5Gemma2TextConfig, device=None):
+        super().__init__()
+        self.config = config
+        self.layer_types = sorted(set(config.layer_types))
+        for layer_type in self.layer_types:
+            inv_freq, attention_scaling = self._compute_inv_freq(
+                config, layer_type, device=device
+            )
+            self.register_buffer(f"{layer_type}_inv_freq", inv_freq, persistent=False)
+            setattr(self, f"{layer_type}_attention_scaling", attention_scaling)
+
+    @staticmethod
+    def _compute_inv_freq(
+        config: T5Gemma2TextConfig, layer_type: str, device=None
+    ) -> tuple[torch.Tensor, float]:
+        rope_params = config.rope_parameters[layer_type]
+        rope_type = rope_params.get("rope_type", "default")
+        base = rope_params["rope_theta"]
+        dim = (
+            getattr(config, "head_dim", None)
+            or config.hidden_size // config.num_attention_heads
+        )
+        inv_freq = 1.0 / (
+            base ** (torch.arange(0, dim, 2, device=device, dtype=torch.float32) / dim)
+        )
+        if rope_type == "linear":
+            inv_freq = inv_freq / float(rope_params["factor"])
+        elif rope_type != "default":
+            raise ValueError(
+                f"Unsupported T5Gemma2 rope_type in local compat shim: {rope_type}"
+            )
+        return inv_freq, 1.0
+
+    def forward(self, x: torch.Tensor, position_ids: torch.LongTensor, layer_type: str):
+        inv_freq = getattr(self, f"{layer_type}_inv_freq")
+        attention_scaling = getattr(self, f"{layer_type}_attention_scaling")
+        inv_freq_expanded = (
+            inv_freq[None, :, None]
+            .float()
+            .expand(position_ids.shape[0], -1, 1)
+            .to(x.device)
+        )
+        position_ids_expanded = position_ids[:, None, :].float()
+
+        device_type = (
+            x.device.type
+            if isinstance(x.device.type, str) and x.device.type != "mps"
+            else "cpu"
+        )
+        with torch.autocast(device_type=device_type, enabled=False):
+            freqs = (
+                inv_freq_expanded.float() @ position_ids_expanded.float()
+            ).transpose(1, 2)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            cos = emb.cos() * attention_scaling
+            sin = emb.sin() * attention_scaling
+
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+def rotate_half(x: torch.Tensor) -> torch.Tensor:
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    unsqueeze_dim=1,
+):
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(
+        batch, num_key_value_heads, n_rep, slen, head_dim
+    )
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+def eager_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: torch.Tensor | None,
+    dropout: float = 0.0,
+    scaling: float | None = None,
+    softcap: float | None = None,
+    **kwargs,
+):
+    if scaling is None:
+        scaling = module.head_dim**-0.5
+    key_states = repeat_kv(key, module.num_key_value_groups)
+    value_states = repeat_kv(value, module.num_key_value_groups)
+    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+    if softcap is not None:
+        attn_weights = torch.tanh(attn_weights / softcap) * softcap
+    if attention_mask is not None:
+        attn_weights = attn_weights + attention_mask[:, :, :, : key_states.shape[-2]]
+    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(
+        query.dtype
+    )
+    attn_weights = nn.functional.dropout(
+        attn_weights, p=dropout, training=module.training
+    )
+    attn_output = torch.matmul(attn_weights, value_states)
+    attn_output = attn_output.transpose(1, 2).contiguous()
+    return attn_output, attn_weights
+
+
+def _padding_mask_to_additive_mask(
+    attention_mask: torch.Tensor | None,
+    query: torch.Tensor,
+    layer_type: str,
+    sliding_window: int | None,
+) -> torch.Tensor | None:
+    batch_size, _, seq_len, _ = query.shape
+    device = query.device
+    dtype = query.dtype
+
+    has_padding = attention_mask is not None and not bool(
+        attention_mask.to(torch.bool).all()
+    )
+    if layer_type == "full_attention" and not has_padding:
+        return None
+
+    if attention_mask is None:
+        valid_kv = torch.ones(batch_size, seq_len, device=device, dtype=torch.bool)
+    else:
+        valid_kv = attention_mask.to(device=device, dtype=torch.bool)
+    allowed = valid_kv[:, None, None, :].expand(batch_size, 1, seq_len, seq_len)
+
+    if layer_type == "sliding_attention":
+        if sliding_window is None:
+            raise ValueError(
+                "T5Gemma2 sliding_attention requires config.sliding_window"
+            )
+        q_idx = torch.arange(seq_len, device=device)[:, None]
+        kv_idx = torch.arange(seq_len, device=device)[None, :]
+        left_window_size = (sliding_window + 1) // 2
+        right_window_size = sliding_window // 2 + 1
+        dist = q_idx - kv_idx
+        local = ((dist >= 0) & (dist < left_window_size)) | (
+            (dist < 0) & (-dist < right_window_size)
+        )
+        allowed = allowed & local[None, None, :, :]
+
+    additive_mask = torch.zeros(
+        (batch_size, 1, seq_len, seq_len), device=device, dtype=dtype
+    )
+    additive_mask.masked_fill_(~allowed, torch.finfo(dtype).min)
+    return additive_mask
+
+
+def t5gemma2_flash_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: torch.Tensor | None,
+    dropout: float = 0.0,
+    scaling: float | None = None,
+    sliding_window: int | None = None,
+    softcap: float | None = None,
+    **kwargs,
+):
+    if query.device.type != "cuda" or query.dtype not in (
+        torch.float16,
+        torch.bfloat16,
+    ):
+        additive_mask = _padding_mask_to_additive_mask(
+            attention_mask,
+            query,
+            module.layer_type,
+            module.config.sliding_window,
+        )
+        return eager_attention_forward(
+            module,
+            query,
+            key,
+            value,
+            additive_mask,
+            dropout=dropout,
+            scaling=scaling,
+            softcap=softcap,
+            **kwargs,
+        )
+
+    try:
+        from flash_attn import flash_attn_func, flash_attn_varlen_func
+        from flash_attn.bert_padding import pad_input, unpad_input
+    except ImportError as exc:
+        raise ImportError(
+            "flash_attn is required for T5Gemma2 flash_attention_2"
+        ) from exc
+
+    query = query.transpose(1, 2).contiguous()
+    key = key.transpose(1, 2).contiguous()
+    value = value.transpose(1, 2).contiguous()
+    batch_size, seq_len = query.shape[:2]
+    dropout_p = dropout if module.training else 0.0
+    softcap_value = 0.0 if softcap is None else softcap
+
+    window_size = (-1, -1)
+    if module.layer_type == "sliding_attention":
+        if sliding_window is None:
+            raise ValueError(
+                "T5Gemma2 sliding_attention requires config.sliding_window"
+            )
+        left_window_size = (sliding_window + 1) // 2
+        right_window_size = sliding_window // 2 + 1
+        window_size = (left_window_size - 1, right_window_size - 1)
+
+    has_padding = attention_mask is not None and not bool(
+        attention_mask.to(torch.bool).all()
+    )
+    if has_padding:
+        attention_mask = attention_mask.to(device=query.device, dtype=torch.bool)
+        query_unpad, indices_q, cu_seqlens_q, max_seqlen_q, _ = unpad_input(
+            query, attention_mask
+        )
+        key_unpad, _, cu_seqlens_k, max_seqlen_k, _ = unpad_input(key, attention_mask)
+        value_unpad, _, _, _, _ = unpad_input(value, attention_mask)
+        attn_output_unpad = flash_attn_varlen_func(
+            query_unpad,
+            key_unpad,
+            value_unpad,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            max_seqlen_q,
+            max_seqlen_k,
+            dropout_p=dropout_p,
+            softmax_scale=scaling,
+            causal=False,
+            window_size=window_size,
+            softcap=softcap_value,
+        )
+        attn_output = pad_input(attn_output_unpad, indices_q, batch_size, seq_len)
+    else:
+        attn_output = flash_attn_func(
+            query,
+            key,
+            value,
+            dropout_p=dropout_p,
+            softmax_scale=scaling,
+            causal=False,
+            window_size=window_size,
+            softcap=softcap_value,
+        )
+
+    return attn_output, None
+
+
+class T5Gemma2SelfAttention(nn.Module):
+    def __init__(self, config: T5Gemma2TextConfig, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.layer_type = config.layer_types[layer_idx]
+        self.head_dim = getattr(
+            config, "head_dim", config.hidden_size // config.num_attention_heads
+        )
+        self.num_key_value_groups = (
+            config.num_attention_heads // config.num_key_value_heads
+        )
+        self.scaling = config.query_pre_attn_scalar**-0.5
+        self.attention_dropout = config.attention_dropout
+        self.is_causal = False
+        self.sliding_window = (
+            config.sliding_window if self.layer_type == "sliding_attention" else None
+        )
+
+        self.q_proj = nn.Linear(
+            config.hidden_size,
+            config.num_attention_heads * self.head_dim,
+            bias=config.attention_bias,
+        )
+        self.k_proj = nn.Linear(
+            config.hidden_size,
+            config.num_key_value_heads * self.head_dim,
+            bias=config.attention_bias,
+        )
+        self.v_proj = nn.Linear(
+            config.hidden_size,
+            config.num_key_value_heads * self.head_dim,
+            bias=config.attention_bias,
+        )
+        self.o_proj = nn.Linear(
+            config.num_attention_heads * self.head_dim,
+            config.hidden_size,
+            bias=config.attention_bias,
+        )
+        self.q_norm = T5Gemma2RMSNorm(dim=self.head_dim, eps=config.rms_norm_eps)
+        self.k_norm = T5Gemma2RMSNorm(dim=self.head_dim, eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: torch.Tensor | None = None,
+        **kwargs,
+    ):
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+        query_states = self.q_norm(query_states)
+        key_states = self.k_norm(key_states)
+
+        cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(
+            query_states, key_states, cos, sin
+        )
+
+        attn_impl = getattr(self.config, "_attn_implementation", None) or getattr(
+            self.config,
+            "preferred_attn_implementation",
+            "flash_attention_2",
+        )
+        if attn_impl == "flash_attention_2":
+            attention_interface = t5gemma2_flash_attention_forward
+        else:
+            attention_interface = (
+                eager_attention_forward
+                if attn_impl == "eager"
+                else ALL_ATTENTION_FUNCTIONS[attn_impl]
+            )
+        attn_output, attn_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            dropout=self.attention_dropout if self.training else 0.0,
+            scaling=self.scaling,
+            sliding_window=self.sliding_window,
+            softcap=self.config.attn_logit_softcapping,
+            **kwargs,
+        )
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+
+
+class T5Gemma2EncoderLayer(GradientCheckpointingLayer):
+    def __init__(self, config: T5Gemma2TextConfig, layer_idx: int):
+        super().__init__()
+        self.attention_type = config.layer_types[layer_idx]
+        self.self_attn = T5Gemma2SelfAttention(config=config, layer_idx=layer_idx)
+        self.pre_self_attn_layernorm = T5Gemma2RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+        self.post_self_attn_layernorm = T5Gemma2RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+        self.mlp = T5Gemma2MLP(config)
+        self.pre_feedforward_layernorm = T5Gemma2RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+        self.post_feedforward_layernorm = T5Gemma2RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+        self.dropout = nn.Dropout(config.dropout_rate)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: torch.Tensor | None = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.pre_self_attn_layernorm(hidden_states)
+        hidden_states, _ = self.self_attn(
+            hidden_states=hidden_states,
+            position_embeddings=position_embeddings,
+            attention_mask=attention_mask,
+            **kwargs,
+        )
+        hidden_states = self.post_self_attn_layernorm(hidden_states)
+        hidden_states = residual + self.dropout(hidden_states)
+
+        residual = hidden_states
+        hidden_states = self.pre_feedforward_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = self.post_feedforward_layernorm(hidden_states)
+        hidden_states = residual + self.dropout(hidden_states)
+        return hidden_states
+
+
+class T5Gemma2TextScaledWordEmbedding(nn.Embedding):
+    def __init__(
+        self,
+        num_embeddings: int,
+        embedding_dim: int,
+        padding_idx: int | None,
+        embed_scale: float = 1.0,
+        eoi_token_index: int = 256000,
+        device=None,
+        dtype=None,
+    ):
+        super().__init__(
+            num_embeddings, embedding_dim, padding_idx, device=device, dtype=dtype
+        )
+        self.scalar_embed_scale = embed_scale
+        self.register_buffer("embed_scale", torch.tensor(embed_scale), persistent=False)
+        self.eoi_token_index = eoi_token_index
+        self.eoi_embedding = nn.Parameter(
+            torch.zeros(self.embedding_dim, device=device, dtype=dtype)
+        )
+
+    def forward(self, input_ids: torch.Tensor):
+        input_embeddings = super().forward(input_ids) * self.embed_scale.to(
+            self.weight.dtype
+        )
+        eoi_mask = input_ids == self.eoi_token_index
+        return torch.where(
+            eoi_mask.unsqueeze(-1),
+            self.eoi_embedding.to(input_embeddings.dtype),
+            input_embeddings,
+        )
+
+
+class T5Gemma2TextEncoder(PreTrainedModel):
+    config_class = T5Gemma2TextConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["T5Gemma2EncoderLayer"]
+    _supports_flash_attn = True
+    _supports_sdpa = True
+    _supports_attention_backend = True
+
+    def __init__(self, config: T5Gemma2TextConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+        self.embed_tokens = T5Gemma2TextScaledWordEmbedding(
+            config.vocab_size,
+            config.hidden_size,
+            self.padding_idx,
+            embed_scale=config.hidden_size**0.5,
+            eoi_token_index=config.eoi_token_index,
+        )
+        self.norm = T5Gemma2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.layers = nn.ModuleList(
+            [
+                T5Gemma2EncoderLayer(config, layer_idx)
+                for layer_idx in range(config.num_hidden_layers)
+            ]
+        )
+        self.dropout = nn.Dropout(config.dropout_rate)
+        self.rotary_emb = T5Gemma2RotaryEmbedding(config)
+        self.gradient_checkpointing = False
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embed_tokens
+
+    def set_input_embeddings(self, new_embeddings):
+        self.embed_tokens = new_embeddings
+        self.vocab_size = new_embeddings.num_embeddings
+
+    def _get_resized_embeddings(
+        self,
+        old_embeddings: nn.Embedding,
+        new_num_tokens: int | None = None,
+        pad_to_multiple_of: int | None = None,
+        mean_resizing: bool = True,
+    ) -> nn.Embedding:
+        if new_num_tokens is None:
+            return old_embeddings
+        if pad_to_multiple_of is not None:
+            new_num_tokens = (
+                (new_num_tokens + pad_to_multiple_of - 1) // pad_to_multiple_of
+            ) * pad_to_multiple_of
+
+        old_num_tokens, old_embedding_dim = old_embeddings.weight.shape
+        if old_num_tokens == new_num_tokens:
+            return old_embeddings
+        padding_idx = old_embeddings.padding_idx
+        if padding_idx is not None and padding_idx >= new_num_tokens:
+            padding_idx = None
+        new_embeddings = T5Gemma2TextScaledWordEmbedding(
+            new_num_tokens,
+            old_embedding_dim,
+            padding_idx,
+            embed_scale=getattr(
+                old_embeddings, "scalar_embed_scale", self.config.hidden_size**0.5
+            ),
+            eoi_token_index=getattr(
+                old_embeddings, "eoi_token_index", self.config.eoi_token_index
+            ),
+            device=old_embeddings.weight.device,
+            dtype=old_embeddings.weight.dtype,
+        )
+        self._init_weights(new_embeddings)
+        num_tokens_to_copy = min(old_num_tokens, new_num_tokens)
+        new_embeddings.weight.data[:num_tokens_to_copy] = old_embeddings.weight.data[
+            :num_tokens_to_copy
+        ]
+        if hasattr(old_embeddings, "eoi_embedding"):
+            new_embeddings.eoi_embedding.data.copy_(old_embeddings.eoi_embedding.data)
+        return new_embeddings
+
+    def _init_weights(self, module):
+        std = self.config.initializer_range
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, T5Gemma2TextScaledWordEmbedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+            module.eoi_embedding.data.zero_()
+            module.embed_scale.fill_(module.scalar_embed_scale)
+        elif isinstance(module, T5Gemma2RMSNorm):
+            module.weight.data.zero_()
+
+    @staticmethod
+    def _build_additive_attention_mask(
+        attention_mask: torch.Tensor | None,
+        inputs_embeds: torch.Tensor,
+        layer_type: str,
+        sliding_window: int | None,
+    ) -> torch.Tensor | None:
+        batch_size, seq_len = inputs_embeds.shape[:2]
+        device = inputs_embeds.device
+
+        has_padding = attention_mask is not None
+        if layer_type == "full_attention" and not has_padding:
+            return None
+
+        if attention_mask is None:
+            valid_kv = torch.ones(batch_size, seq_len, device=device, dtype=torch.bool)
+        else:
+            valid_kv = attention_mask.to(device=device, dtype=torch.bool)
+        allowed = valid_kv[:, None, None, :].expand(batch_size, 1, seq_len, seq_len)
+
+        if layer_type == "sliding_attention":
+            if sliding_window is None:
+                raise ValueError(
+                    "T5Gemma2 sliding_attention requires config.sliding_window"
+                )
+            q_idx = torch.arange(seq_len, device=device)[:, None]
+            kv_idx = torch.arange(seq_len, device=device)[None, :]
+            left_window_size = (sliding_window + 1) // 2
+            right_window_size = sliding_window // 2 + 1
+            dist = q_idx - kv_idx
+            local = ((dist >= 0) & (dist < left_window_size)) | (
+                (dist < 0) & (-dist < right_window_size)
+            )
+            allowed = allowed & local[None, None, :, :]
+
+        min_value = torch.finfo(inputs_embeds.dtype).min
+        additive_mask = torch.zeros(
+            (batch_size, 1, seq_len, seq_len), device=device, dtype=inputs_embeds.dtype
+        )
+        additive_mask.masked_fill_(~allowed, min_value)
+        return additive_mask
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        output_hidden_states: bool | None = None,
+        return_dict: bool | None = None,
+        token_type_ids: torch.Tensor | None = None,
+        **kwargs,
+    ) -> BaseModelOutput:
+        if (input_ids is None) == (inputs_embeds is None):
+            raise ValueError(
+                "You must specify exactly one of input_ids or inputs_embeds"
+            )
+        kwargs.pop("past_key_values", None)
+        kwargs.pop("use_cache", None)
+
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+        if position_ids is None:
+            position_ids = torch.arange(
+                0, inputs_embeds.shape[1], device=inputs_embeds.device
+            ).unsqueeze(0)
+
+        hidden_states = self.dropout(inputs_embeds)
+        all_hidden_states = () if output_hidden_states else None
+
+        position_embeddings = {
+            layer_type: self.rotary_emb(hidden_states, position_ids, layer_type)
+            for layer_type in set(self.config.layer_types)
+        }
+        attn_impl = getattr(self.config, "_attn_implementation", None) or getattr(
+            self.config,
+            "preferred_attn_implementation",
+            "flash_attention_2",
+        )
+        if attn_impl == "flash_attention_2":
+            padding_mask = None
+            if getattr(self, "_force_static_attention_mask", False) or attention_mask is not None and not bool(
+                attention_mask.to(torch.bool).all()
+            ):
+                padding_mask = attention_mask
+            attn_masks = {
+                layer_type: padding_mask for layer_type in set(self.config.layer_types)
+            }
+        else:
+            attn_masks = {
+                layer_type: self._build_additive_attention_mask(
+                    attention_mask,
+                    inputs_embeds,
+                    layer_type,
+                    self.config.sliding_window,
+                )
+                for layer_type in set(self.config.layer_types)
+            }
+
+        for layer_module in self.layers[: self.config.num_hidden_layers]:
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+            hidden_states = layer_module(
+                hidden_states,
+                position_embeddings[layer_module.attention_type],
+                attn_masks[layer_module.attention_type],
+                output_attentions=False,
+                **kwargs,
+            )
+
+        hidden_states = self.norm(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        if not return_dict:
+            return tuple(v for v in [hidden_states, all_hidden_states] if v is not None)
+        return BaseModelOutput(
+            last_hidden_state=hidden_states, hidden_states=all_hidden_states
+        )

models/text_encoder_graph.py

@@ -0,0 +1,164 @@
+"""Batch-1 static CUDA Graph cache for the HF text encoder."""
+
+from __future__ import annotations
+
+import threading
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+
+
+@dataclass
+class _GraphRecord:
+    graph: torch.cuda.CUDAGraph
+    stream: torch.cuda.Stream
+    input_ids: torch.Tensor
+    attention_mask: torch.Tensor
+    position_ids: torch.Tensor
+    output: torch.Tensor
+
+
+class TextEncoderGraphCache:
+    """Capture fixed padded text-encoder buckets and replay them without CPU dispatch."""
+
+    def __init__(
+        self, text_encoder: torch.nn.Module, *, token_granularity: int = 32
+    ) -> None:
+        if token_granularity <= 0:
+            raise ValueError("token_granularity must be > 0")
+        self.text_encoder = text_encoder
+        self.text_encoder._force_static_attention_mask = True
+        self.text_encoder.config._attn_implementation = "sdpa"
+        self.token_granularity = int(token_granularity)
+        self._records: dict[tuple[int, int], _GraphRecord] = {}
+        self._lock = threading.RLock()
+        self._graph_pool = torch.cuda.graph_pool_handle()
+        self.captures = 0
+        self.replays = 0
+        self._frozen = False
+
+    def _bucket(self, length: int) -> int:
+        return (
+            (int(length) + self.token_granularity - 1) // self.token_granularity
+        ) * self.token_granularity
+
+    @staticmethod
+    def _copy_inputs(
+        segments: Sequence[torch.Tensor],
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        position_ids: torch.Tensor,
+    ) -> None:
+        input_ids.zero_()
+        attention_mask.zero_()
+        position_ids.zero_()
+        for row, segment in enumerate(segments):
+            length = int(segment.shape[0])
+            input_ids[row, :length].copy_(segment)
+            attention_mask[row, :length].fill_(1)
+            position_ids[row, :length].copy_(
+                torch.arange(length, device=segment.device, dtype=torch.long)
+            )
+
+    @torch.inference_mode()
+    def __call__(
+        self, segments: Sequence[torch.Tensor]
+    ) -> tuple[list[torch.Tensor], list[Any]]:
+        if not segments:
+            return [], []
+        lengths = [int(segment.shape[0]) for segment in segments]
+        if any(length <= 0 for length in lengths):
+            raise ValueError("text encoder graph does not support empty segments")
+        batch_size = len(segments)
+        max_length = self._bucket(max(lengths))
+        key = (batch_size, max_length)
+
+        with self._lock:
+            record = self._records.get(key)
+            if record is None:
+                if self._frozen:
+                    raise RuntimeError(
+                        f"text encoder CUDA graph {key} was not declared in the warmup profile"
+                    )
+                device = segments[0].device
+                static_ids = torch.zeros(
+                    (batch_size, max_length), dtype=segments[0].dtype, device=device
+                )
+                static_mask = torch.zeros(
+                    (batch_size, max_length), dtype=torch.long, device=device
+                )
+                static_positions = torch.zeros_like(static_mask)
+                self._copy_inputs(segments, static_ids, static_mask, static_positions)
+
+                capture_stream = torch.cuda.Stream(device=device)
+                capture_stream.wait_stream(torch.cuda.current_stream(device))
+                with torch.cuda.stream(capture_stream):
+                    for _ in range(3):
+                        static_output = self.text_encoder(
+                            input_ids=static_ids,
+                            attention_mask=static_mask,
+                            position_ids=static_positions,
+                            output_hidden_states=False,
+                        ).last_hidden_state
+                capture_stream.synchronize()
+
+                graph = torch.cuda.CUDAGraph()
+                with torch.cuda.graph(
+                    graph,
+                    stream=capture_stream,
+                    pool=self._graph_pool,
+                    capture_error_mode="thread_local",
+                ):
+                    static_output = self.text_encoder(
+                        input_ids=static_ids,
+                        attention_mask=static_mask,
+                        position_ids=static_positions,
+                        output_hidden_states=False,
+                    ).last_hidden_state
+                record = _GraphRecord(
+                    graph=graph,
+                    stream=capture_stream,
+                    input_ids=static_ids,
+                    attention_mask=static_mask,
+                    position_ids=static_positions,
+                    output=static_output,
+                )
+                self._records[key] = record
+                self.captures += 1
+
+            self._copy_inputs(
+                segments,
+                record.input_ids,
+                record.attention_mask,
+                record.position_ids,
+            )
+            current_stream = torch.cuda.current_stream(segments[0].device)
+            record.stream.wait_stream(current_stream)
+            record.graph.replay()
+            current_stream.wait_stream(record.stream)
+            self.replays += 1
+            hidden_states = [
+                record.output[row, :length] for row, length in enumerate(lengths)
+            ]
+            return hidden_states, []
+
+    @property
+    def graph_keys(self) -> tuple[tuple[int, int], ...]:
+        return tuple(sorted(self._records))
+
+    @torch.inference_mode()
+    def warmup_graph(
+        self, *, batch_size: int, token_length: int, device: torch.device
+    ) -> None:
+        if self._frozen:
+            raise RuntimeError("text encoder CUDA graph cache is already frozen")
+        segments = [
+            torch.zeros(token_length, dtype=torch.long, device=device)
+            for _ in range(batch_size)
+        ]
+        self(segments)
+
+    def freeze(self) -> None:
+        self._frozen = True

models/warmup_profile.py

@@ -0,0 +1,278 @@
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+
+@dataclass(frozen=True)
+class TextEncoderWarmupGraph:
+    batch_size: int
+    token_length: int
+
+
+@dataclass(frozen=True)
+class BackbonePrefillWarmupGraph:
+    branch_batch_size: int
+    sequence_length: int
+
+
+@dataclass(frozen=True)
+class SyntheticWarmupRequest:
+    template: str
+    text: str
+    instruction: str
+    speaker: str
+    seed: int
+
+
+@dataclass(frozen=True)
+class FastStreamingWarmupProfile:
+    schema_version: int
+    name: str
+    cfg_scales: tuple[float, ...]
+    text_encoder_graphs: tuple[TextEncoderWarmupGraph, ...]
+    backbone_prefill_graphs: tuple[BackbonePrefillWarmupGraph, ...]
+    backbone_decode_branch_batch_sizes: tuple[int, ...]
+    depth_decoder_batch_sizes: tuple[int, ...]
+    codec_chunk_frames: int
+    codec_num_lanes: int
+    warmup_request: SyntheticWarmupRequest
+    freeze_after_warmup: bool
+    source: str | None = None
+
+    @property
+    def cfg_modes(self) -> tuple[str, ...]:
+        return tuple(
+            "single_cfg" if cfg_scale != 1.0 else "no_cfg"
+            for cfg_scale in self.cfg_scales
+        )
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "schema_version": self.schema_version,
+            "name": self.name,
+            "service": {
+                "concurrency": 1,
+                "cfg_scales": list(self.cfg_scales),
+                "cfg_modes": list(self.cfg_modes),
+                "freeze_after_warmup": self.freeze_after_warmup,
+            },
+            "stages": {
+                "text_encoder": {
+                    "graphs": [
+                        {
+                            "batch_size": graph.batch_size,
+                            "token_length": graph.token_length,
+                        }
+                        for graph in self.text_encoder_graphs
+                    ]
+                },
+                "backbone_prefill": {
+                    "graphs": [
+                        {
+                            "branch_batch_size": graph.branch_batch_size,
+                            "sequence_length": graph.sequence_length,
+                        }
+                        for graph in self.backbone_prefill_graphs
+                    ]
+                },
+                "backbone_decode": {
+                    "graphs": [
+                        {"branch_batch_size": batch_size}
+                        for batch_size in self.backbone_decode_branch_batch_sizes
+                    ]
+                },
+                "depth_decoder": {
+                    "graphs": [
+                        {"batch_size": batch_size}
+                        for batch_size in self.depth_decoder_batch_sizes
+                    ]
+                },
+                "codec": {
+                    "graphs": [
+                        {
+                            "num_lanes": self.codec_num_lanes,
+                            "chunk_frames": self.codec_chunk_frames,
+                        }
+                    ]
+                },
+            },
+            "warmup_request": {
+                "template": self.warmup_request.template,
+                "text": self.warmup_request.text,
+                "instruction": self.warmup_request.instruction,
+                "speaker": self.warmup_request.speaker,
+                "seed": self.warmup_request.seed,
+            },
+        }
+
+
+def _positive_int(value: Any, path: str) -> int:
+    value = int(value)
+    if value <= 0:
+        raise ValueError(f"{path} must be > 0")
+    return value
+
+
+def _graphs(stages: dict[str, Any], stage: str) -> list[dict[str, Any]]:
+    value = stages.get(stage)
+    if not isinstance(value, dict):
+        raise ValueError(f"stages.{stage} must be an object")
+    graphs = value.get("graphs")
+    if not isinstance(graphs, list) or not graphs:
+        raise ValueError(f"stages.{stage}.graphs must be a non-empty list")
+    if not all(isinstance(graph, dict) for graph in graphs):
+        raise ValueError(f"stages.{stage}.graphs entries must be objects")
+    return graphs
+
+
+def parse_warmup_profile(
+    payload: dict[str, Any], *, source: str | None = None
+) -> FastStreamingWarmupProfile:
+    if int(payload.get("schema_version", 0)) != 1:
+        raise ValueError("warmup profile schema_version must be 1")
+    name = str(payload.get("name", "")).strip()
+    if not name:
+        raise ValueError("warmup profile name must not be empty")
+
+    service = payload.get("service")
+    if not isinstance(service, dict):
+        raise ValueError("service must be an object")
+    concurrency = _positive_int(service.get("concurrency", 0), "service.concurrency")
+    if concurrency != 1:
+        raise ValueError("fast streaming warmup supports only service.concurrency=1")
+    raw_cfg_scales = service.get("cfg_scales")
+    if not isinstance(raw_cfg_scales, list) or not raw_cfg_scales:
+        raise ValueError("service.cfg_scales must be a non-empty list")
+    cfg_scales = tuple(sorted({float(value) for value in raw_cfg_scales}))
+    if any(cfg_scale < 1.0 for cfg_scale in cfg_scales):
+        raise ValueError("warmup profiles require service.cfg_scales values >= 1")
+    expected_branch_batches = tuple(
+        sorted({1 if cfg_scale == 1.0 else 2 for cfg_scale in cfg_scales})
+    )
+    freeze_after_warmup = bool(service.get("freeze_after_warmup", True))
+
+    stages = payload.get("stages")
+    if not isinstance(stages, dict):
+        raise ValueError("stages must be an object")
+
+    text_graphs = tuple(
+        TextEncoderWarmupGraph(
+            batch_size=_positive_int(
+                graph.get("batch_size"), "text_encoder.batch_size"
+            ),
+            token_length=_positive_int(
+                graph.get("token_length"), "text_encoder.token_length"
+            ),
+        )
+        for graph in _graphs(stages, "text_encoder")
+    )
+    if any(graph.token_length % 32 for graph in text_graphs):
+        raise ValueError("text_encoder token_length values must be multiples of 32")
+    prefill_graphs = tuple(
+        BackbonePrefillWarmupGraph(
+            branch_batch_size=_positive_int(
+                graph.get("branch_batch_size"), "backbone_prefill.branch_batch_size"
+            ),
+            sequence_length=_positive_int(
+                graph.get("sequence_length"), "backbone_prefill.sequence_length"
+            ),
+        )
+        for graph in _graphs(stages, "backbone_prefill")
+    )
+    if any(graph.sequence_length % 32 for graph in prefill_graphs):
+        raise ValueError(
+            "backbone_prefill sequence_length values must be multiples of 32"
+        )
+
+    decode_batches = tuple(
+        sorted(
+            {
+                _positive_int(
+                    graph.get("branch_batch_size"),
+                    "backbone_decode.branch_batch_size",
+                )
+                for graph in _graphs(stages, "backbone_decode")
+            }
+        )
+    )
+    if decode_batches != expected_branch_batches:
+        raise ValueError(
+            "backbone_decode branch batches must match cfg_scales; "
+            f"expected {expected_branch_batches}, got {decode_batches}"
+        )
+    if {graph.branch_batch_size for graph in prefill_graphs} != set(decode_batches):
+        raise ValueError(
+            "backbone_prefill must declare graphs for every backbone_decode branch batch"
+        )
+
+    depth_batches = tuple(
+        _positive_int(graph.get("batch_size"), "depth_decoder.batch_size")
+        for graph in _graphs(stages, "depth_decoder")
+    )
+    if tuple(sorted(set(depth_batches))) != expected_branch_batches:
+        raise ValueError(
+            "depth_decoder batch sizes must match cfg_scales; "
+            f"expected {expected_branch_batches}"
+        )
+
+    codec_graphs = _graphs(stages, "codec")
+    if len(codec_graphs) != 1:
+        raise ValueError("codec must declare exactly one graph")
+    codec_num_lanes = _positive_int(codec_graphs[0].get("num_lanes"), "codec.num_lanes")
+    codec_chunk_frames = _positive_int(
+        codec_graphs[0].get("chunk_frames"), "codec.chunk_frames"
+    )
+    if codec_num_lanes != 1:
+        raise ValueError("single-concurrency service requires codec.num_lanes=1")
+
+    warmup_request_payload = payload.get("warmup_request")
+    if not isinstance(warmup_request_payload, dict):
+        raise ValueError("warmup_request must be an object")
+    warmup_request = SyntheticWarmupRequest(
+        template=str(warmup_request_payload.get("template", "")).strip(),
+        text=str(warmup_request_payload.get("text", "")).strip(),
+        instruction=str(warmup_request_payload.get("instruction", "")).strip(),
+        speaker=str(warmup_request_payload.get("speaker", "S0")).strip(),
+        seed=int(warmup_request_payload.get("seed", 42)),
+    )
+    if not warmup_request.template or not warmup_request.text:
+        raise ValueError(
+            "warmup_request.template and warmup_request.text must not be empty"
+        )
+
+    return FastStreamingWarmupProfile(
+        schema_version=1,
+        name=name,
+        cfg_scales=cfg_scales,
+        text_encoder_graphs=tuple(
+            sorted(
+                set(text_graphs),
+                key=lambda graph: (graph.batch_size, graph.token_length),
+            )
+        ),
+        backbone_prefill_graphs=tuple(
+            sorted(
+                set(prefill_graphs),
+                key=lambda graph: (graph.branch_batch_size, graph.sequence_length),
+            )
+        ),
+        backbone_decode_branch_batch_sizes=decode_batches,
+        depth_decoder_batch_sizes=depth_batches,
+        codec_chunk_frames=codec_chunk_frames,
+        codec_num_lanes=codec_num_lanes,
+        warmup_request=warmup_request,
+        freeze_after_warmup=freeze_after_warmup,
+        source=source,
+    )
+
+
+def load_warmup_profile(path: str | Path) -> FastStreamingWarmupProfile:
+    path = Path(path).resolve()
+    with path.open("r", encoding="utf-8") as handle:
+        payload = json.load(handle)
+    if not isinstance(payload, dict):
+        raise ValueError("warmup profile root must be an object")
+    return parse_warmup_profile(payload, source=str(path))

packages.txt

@@ -0,0 +1,2 @@
+ffmpeg
+libsndfile1

requirements.txt

@@ -0,0 +1,6 @@
+torch==2.9.1
+torchaudio==2.9.1
+qwen-tts==0.1.1
+transformers==4.57.3
+numpy>=2.0,<3
+soundfile>=0.13