Base files

.github/workflows/ci.yaml

@@ -0,0 +1,22 @@
+name: Continuous Integration
+
+on:
+  pull_request:
+    branches:
+      - main
+
+jobs:
+  lint_and_format:
+    runs-on: ubuntu-latest
+    name: Lint and Format
+    steps:
+      - uses: actions/checkout@v4
+      - uses: astral-sh/ruff-action@v3
+        with:
+          version: latest
+
+      - name: Check Lint using Ruff
+        run: ruff check
+
+      - name: Check Format using Ruff
+        run: ruff format --check --diff

.gitignore

@@ -0,0 +1,22 @@
+# Python-generated files
+__pycache__/
+*.py[oc]
+build/
+dist/
+wheels/
+*.egg-info
+
+# Virtual environments
+.venv
+.idea/
+.gradio
+
+**/*.pth
+**/*.safetensors
+**/*.mp3
+!example_prompt.mp3
+**/*.txt
+
+.ruff_cache
+.ipynb_checkpoints
+config.json

.python-version

@@ -0,0 +1 @@
+3.10

LICENSE

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

app.py

@@ -0,0 +1,444 @@
+import argparse
+import contextlib
+import io
+import random
+import tempfile
+import time
+from pathlib import Path
+from typing import Optional, Tuple
+
+import gradio as gr
+import numpy as np
+import soundfile as sf
+import torch
+
+from dia.model import Dia
+
+
+# --- Global Setup ---
+parser = argparse.ArgumentParser(description="Gradio interface for Nari TTS")
+parser.add_argument("--device", type=str, default=None, help="Force device (e.g., 'cuda', 'mps', 'cpu')")
+parser.add_argument("--share", action="store_true", help="Enable Gradio sharing")
+
+args = parser.parse_args()
+
+
+# Determine device
+if args.device:
+    device = torch.device(args.device)
+elif torch.cuda.is_available():
+    device = torch.device("cuda")
+# Simplified MPS check for broader compatibility
+elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+    # Basic check is usually sufficient, detailed check can be problematic
+    device = torch.device("mps")
+else:
+    device = torch.device("cpu")
+
+print(f"Using device: {device}")
+
+# Load Nari model and config
+print("Loading Nari model...")
+try:
+    dtype_map = {
+        "cpu": "float32",
+        "mps": "float32",  # Apple M series – better with float32
+        "cuda": "float16",  # NVIDIA – better with float16
+    }
+
+    dtype = dtype_map.get(device.type, "float16")
+    print(f"Using device: {device}, attempting to load model with {dtype}")
+    model = Dia.from_pretrained("nari-labs/Dia-1.6B-0626", compute_dtype=dtype, device=device)
+except Exception as e:
+    print(f"Error loading Nari model: {e}")
+    raise
+
+
+def set_seed(seed: int):
+    """Sets the random seed for reproducibility."""
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+
+def run_inference(
+    text_input: str,
+    audio_prompt_text_input: str,
+    audio_prompt_input: Optional[Tuple[int, np.ndarray]],
+    max_new_tokens: int,
+    cfg_scale: float,
+    temperature: float,
+    top_p: float,
+    cfg_filter_top_k: int,
+    speed_factor: float,
+    seed: Optional[int] = None,
+):
+    """
+    Runs Nari inference using the globally loaded model and provided inputs.
+    Uses temporary files for text and audio prompt compatibility with inference.generate.
+    """
+    global model, device  # Access global model, config, device
+    console_output_buffer = io.StringIO()
+
+    with contextlib.redirect_stdout(console_output_buffer):
+        # Prepend transcript text if audio_prompt provided
+        if audio_prompt_input and audio_prompt_text_input and not audio_prompt_text_input.isspace():
+            text_input = audio_prompt_text_input + "\n" + text_input
+            text_input = text_input.strip()
+
+        if audio_prompt_input and (not audio_prompt_text_input or audio_prompt_text_input.isspace()):
+            raise gr.Error("Audio Prompt Text input cannot be empty.")
+
+        if not text_input or text_input.isspace():
+            raise gr.Error("Text input cannot be empty.")
+
+        # Preprocess Audio
+        temp_txt_file_path = None
+        temp_audio_prompt_path = None
+        output_audio = (44100, np.zeros(1, dtype=np.float32))
+
+        try:
+            prompt_path_for_generate = None
+            if audio_prompt_input is not None:
+                sr, audio_data = audio_prompt_input
+                # Check if audio_data is valid
+                if audio_data is None or audio_data.size == 0 or audio_data.max() == 0:  # Check for silence/empty
+                    gr.Warning("Audio prompt seems empty or silent, ignoring prompt.")
+                else:
+                    # Save prompt audio to a temporary WAV file
+                    with tempfile.NamedTemporaryFile(mode="wb", suffix=".wav", delete=False) as f_audio:
+                        temp_audio_prompt_path = f_audio.name  # Store path for cleanup
+
+                        # Basic audio preprocessing for consistency
+                        # Convert to float32 in [-1, 1] range if integer type
+                        if np.issubdtype(audio_data.dtype, np.integer):
+                            max_val = np.iinfo(audio_data.dtype).max
+                            audio_data = audio_data.astype(np.float32) / max_val
+                        elif not np.issubdtype(audio_data.dtype, np.floating):
+                            gr.Warning(f"Unsupported audio prompt dtype {audio_data.dtype}, attempting conversion.")
+                            # Attempt conversion, might fail for complex types
+                            try:
+                                audio_data = audio_data.astype(np.float32)
+                            except Exception as conv_e:
+                                raise gr.Error(f"Failed to convert audio prompt to float32: {conv_e}")
+
+                        # Ensure mono (average channels if stereo)
+                        if audio_data.ndim > 1:
+                            if audio_data.shape[0] == 2:  # Assume (2, N)
+                                audio_data = np.mean(audio_data, axis=0)
+                            elif audio_data.shape[1] == 2:  # Assume (N, 2)
+                                audio_data = np.mean(audio_data, axis=1)
+                            else:
+                                gr.Warning(
+                                    f"Audio prompt has unexpected shape {audio_data.shape}, taking first channel/axis."
+                                )
+                                audio_data = (
+                                    audio_data[0] if audio_data.shape[0] < audio_data.shape[1] else audio_data[:, 0]
+                                )
+                            audio_data = np.ascontiguousarray(audio_data)  # Ensure contiguous after slicing/mean
+
+                        # Write using soundfile
+                        try:
+                            sf.write(
+                                temp_audio_prompt_path, audio_data, sr, subtype="FLOAT"
+                            )  # Explicitly use FLOAT subtype
+                            prompt_path_for_generate = temp_audio_prompt_path
+                            print(f"Created temporary audio prompt file: {temp_audio_prompt_path} (orig sr: {sr})")
+                        except Exception as write_e:
+                            print(f"Error writing temporary audio file: {write_e}")
+                            raise gr.Error(f"Failed to save audio prompt: {write_e}")
+
+            # Set and Display Generation Seed
+            if seed is None or seed < 0:
+                seed = random.randint(0, 2**32 - 1)
+                print(f"\nNo seed provided, generated random seed: {seed}\n")
+            else:
+                print(f"\nUsing user-selected seed: {seed}\n")
+            set_seed(seed)
+
+            # Run Generation
+            print(f'Generating speech: \n"{text_input}"\n')
+
+            start_time = time.time()
+
+            # Use torch.inference_mode() context manager for the generation call
+            with torch.inference_mode():
+                output_audio_np = model.generate(
+                    text_input,
+                    max_tokens=max_new_tokens,
+                    cfg_scale=cfg_scale,
+                    temperature=temperature,
+                    top_p=top_p,
+                    cfg_filter_top_k=cfg_filter_top_k,  # Pass the value here
+                    use_torch_compile=False,  # Keep False for Gradio stability
+                    audio_prompt=prompt_path_for_generate,
+                    verbose=True,
+                )
+
+            end_time = time.time()
+            print(f"Generation finished in {end_time - start_time:.2f} seconds.\n")
+
+            # 4. Convert Codes to Audio
+            if output_audio_np is not None:
+                # Get sample rate from the loaded DAC model
+                output_sr = 44100
+
+                # --- Slow down audio ---
+                original_len = len(output_audio_np)
+                # Ensure speed_factor is positive and not excessively small/large to avoid issues
+                speed_factor = max(0.1, min(speed_factor, 5.0))
+                target_len = int(original_len / speed_factor)  # Target length based on speed_factor
+                if target_len != original_len and target_len > 0:  # Only interpolate if length changes and is valid
+                    x_original = np.arange(original_len)
+                    x_resampled = np.linspace(0, original_len - 1, target_len)
+                    resampled_audio_np = np.interp(x_resampled, x_original, output_audio_np)
+                    output_audio = (
+                        output_sr,
+                        resampled_audio_np.astype(np.float32),
+                    )  # Use resampled audio
+                    print(
+                        f"Resampled audio from {original_len} to {target_len} samples for {speed_factor:.2f}x speed."
+                    )
+                else:
+                    output_audio = (
+                        output_sr,
+                        output_audio_np,
+                    )  # Keep original if calculation fails or no change
+                    print(f"Skipping audio speed adjustment (factor: {speed_factor:.2f}).")
+                # --- End slowdown ---
+
+                print(f"Audio conversion successful. Final shape: {output_audio[1].shape}, Sample Rate: {output_sr}")
+
+                # Explicitly convert to int16 to prevent Gradio warning
+                if output_audio[1].dtype == np.float32 or output_audio[1].dtype == np.float64:
+                    audio_for_gradio = np.clip(output_audio[1], -1.0, 1.0)
+                    audio_for_gradio = (audio_for_gradio * 32767).astype(np.int16)
+                    output_audio = (output_sr, audio_for_gradio)
+                    print("Converted audio to int16 for Gradio output.")
+
+            else:
+                print("\nGeneration finished, but no valid tokens were produced.")
+                # Return default silence
+                gr.Warning("Generation produced no output.")
+
+        except Exception as e:
+            print(f"Error during inference: {e}")
+            import traceback
+
+            traceback.print_exc()
+            # Re-raise as Gradio error to display nicely in the UI
+            raise gr.Error(f"Inference failed: {e}")
+
+        finally:
+            # Cleanup Temporary Files defensively
+            if temp_txt_file_path and Path(temp_txt_file_path).exists():
+                try:
+                    Path(temp_txt_file_path).unlink()
+                    print(f"Deleted temporary text file: {temp_txt_file_path}")
+                except OSError as e:
+                    print(f"Warning: Error deleting temporary text file {temp_txt_file_path}: {e}")
+            if temp_audio_prompt_path and Path(temp_audio_prompt_path).exists():
+                try:
+                    Path(temp_audio_prompt_path).unlink()
+                    print(f"Deleted temporary audio prompt file: {temp_audio_prompt_path}")
+                except OSError as e:
+                    print(f"Warning: Error deleting temporary audio prompt file {temp_audio_prompt_path}: {e}")
+
+        # After generation, capture the printed output
+        console_output = console_output_buffer.getvalue()
+
+    return output_audio, seed, console_output
+
+
+# --- Create Gradio Interface ---
+css = """
+#col-container {max-width: 90%; margin-left: auto; margin-right: auto;}
+"""
+# Attempt to load default text from example.txt
+default_text = "[S1] Dia is an open weights text to dialogue model. \n[S2] You get full control over scripts and voices. \n[S1] Wow. Amazing. (laughs) \n[S2] Try it now on Git hub or Hugging Face."
+example_txt_path = Path("./example.txt")
+if example_txt_path.exists():
+    try:
+        default_text = example_txt_path.read_text(encoding="utf-8").strip()
+        if not default_text:  # Handle empty example file
+            default_text = "Example text file was empty."
+    except Exception as e:
+        print(f"Warning: Could not read example.txt: {e}")
+
+
+# Build Gradio UI
+with gr.Blocks(css=css, theme="gradio/dark") as demo:
+    gr.Markdown("# Nari Text-to-Speech Synthesis")
+
+    with gr.Row(equal_height=False):
+        with gr.Column(scale=1):
+            with gr.Accordion("Audio Reference Prompt (Optional)", open=False):
+                audio_prompt_input = gr.Audio(
+                    label="Audio Prompt (Optional)",
+                    show_label=True,
+                    sources=["upload", "microphone"],
+                    type="numpy",
+                )
+                audio_prompt_text_input = gr.Textbox(
+                    label="Transcript of Audio Prompt (Required if using Audio Prompt)",
+                    placeholder="Enter text here...",
+                    value="",
+                    lines=5,  # Increased lines
+                )
+            text_input = gr.Textbox(
+                label="Text To Generate",
+                placeholder="Enter text here...",
+                value=default_text,
+                lines=5,  # Increased lines
+            )
+            with gr.Accordion("Generation Parameters", open=False):
+                max_new_tokens = gr.Slider(
+                    label="Max New Tokens (Audio Length)",
+                    minimum=860,
+                    maximum=3072,
+                    value=model.config.decoder_config.max_position_embeddings,  # Use config default if available, else fallback
+                    step=50,
+                    info="Controls the maximum length of the generated audio (more tokens = longer audio).",
+                )
+                cfg_scale = gr.Slider(
+                    label="CFG Scale (Guidance Strength)",
+                    minimum=1.0,
+                    maximum=5.0,
+                    value=3.0,  # Default from inference.py
+                    step=0.1,
+                    info="Higher values increase adherence to the text prompt.",
+                )
+                temperature = gr.Slider(
+                    label="Temperature (Randomness)",
+                    minimum=1.0,
+                    maximum=2.5,
+                    value=1.8,  # Default from inference.py
+                    step=0.05,
+                    info="Lower values make the output more deterministic, higher values increase randomness.",
+                )
+                top_p = gr.Slider(
+                    label="Top P (Nucleus Sampling)",
+                    minimum=0.70,
+                    maximum=1.0,
+                    value=0.95,  # Default from inference.py
+                    step=0.01,
+                    info="Filters vocabulary to the most likely tokens cumulatively reaching probability P.",
+                )
+                cfg_filter_top_k = gr.Slider(
+                    label="CFG Filter Top K",
+                    minimum=15,
+                    maximum=100,
+                    value=45,
+                    step=1,
+                    info="Top k filter for CFG guidance.",
+                )
+                speed_factor_slider = gr.Slider(
+                    label="Speed Factor",
+                    minimum=0.8,
+                    maximum=1.0,
+                    value=1.0,
+                    step=0.02,
+                    info="Adjusts the speed of the generated audio (1.0 = original speed).",
+                )
+                seed_input = gr.Number(
+                    label="Generation Seed (Optional)",
+                    value=-1,
+                    precision=0,  # No decimal points
+                    step=1,
+                    interactive=True,
+                    info="Set a generation seed for reproducible outputs. Leave empty or -1 for random seed.",
+                )
+
+            run_button = gr.Button("Generate Audio", variant="primary")
+
+        with gr.Column(scale=1):
+            audio_output = gr.Audio(
+                label="Generated Audio",
+                type="numpy",
+                autoplay=False,
+            )
+            seed_output = gr.Textbox(label="Generation Seed", interactive=False)
+            console_output = gr.Textbox(label="Console Output Log", lines=10, interactive=False)
+
+    # Link button click to function
+    run_button.click(
+        fn=run_inference,
+        inputs=[
+            text_input,
+            audio_prompt_text_input,
+            audio_prompt_input,
+            max_new_tokens,
+            cfg_scale,
+            temperature,
+            top_p,
+            cfg_filter_top_k,
+            speed_factor_slider,
+            seed_input,
+        ],
+        outputs=[
+            audio_output,
+            seed_output,
+            console_output,
+        ],  # Add status_output here if using it
+        api_name="generate_audio",
+    )
+
+    # Add examples (ensure the prompt path is correct or remove it if example file doesn't exist)
+    example_prompt_path = "./example_prompt.mp3"  # Adjust if needed
+    examples_list = [
+        [
+            "[S1] Oh fire! Oh my goodness! What's the procedure? What to we do people? The smoke could be coming through an air duct! \n[S2] Oh my god! Okay.. it's happening. Everybody stay calm! \n[S1] What's the procedure... \n[S2] Everybody stay fucking calm!!!... Everybody fucking calm down!!!!! \n[S1] No! No! If you touch the handle, if its hot there might be a fire down the hallway! ",
+            None,
+            3072,
+            3.0,
+            1.8,
+            0.95,
+            45,
+            1.0,
+        ],
+        [
+            "[S1] Open weights text to dialogue model. \n[S2] You get full control over scripts and voices. \n[S1] I'm biased, but I think we clearly won. \n[S2] Hard to disagree. (laughs) \n[S1] Thanks for listening to this demo. \n[S2] Try it now on Git hub and Hugging Face. \n[S1] If you liked our model, please give us a star and share to your friends. \n[S2] This was Nari Labs.",
+            example_prompt_path if Path(example_prompt_path).exists() else None,
+            3072,
+            3.0,
+            1.8,
+            0.95,
+            45,
+            1.0,
+        ],
+    ]
+
+    if examples_list:
+        gr.Examples(
+            examples=examples_list,
+            inputs=[
+                text_input,
+                audio_prompt_input,
+                max_new_tokens,
+                cfg_scale,
+                temperature,
+                top_p,
+                cfg_filter_top_k,
+                speed_factor_slider,
+                seed_input,
+            ],
+            outputs=[audio_output],
+            fn=run_inference,
+            cache_examples=False,
+            label="Examples (Click to Run)",
+        )
+    else:
+        gr.Markdown("_(No examples configured or example prompt file missing)_")
+
+# --- Launch the App ---
+if __name__ == "__main__":
+    print("Launching Gradio interface...")
+
+    # set `GRADIO_SERVER_NAME`, `GRADIO_SERVER_PORT` env vars to override default values
+    # use `GRADIO_SERVER_NAME=0.0.0.0` for Docker
+    demo.launch(share=args.share)

cli.py

@@ -0,0 +1,144 @@
+import argparse
+import os
+import random
+
+import numpy as np
+import soundfile as sf
+import torch
+
+from dia.model import Dia
+
+
+def set_seed(seed: int):
+    """Sets the random seed for reproducibility."""
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)
+    # Ensure deterministic behavior for cuDNN (if used)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Generate audio using the Dia model.")
+
+    parser.add_argument("text", type=str, help="Input text for speech generation.")
+    parser.add_argument(
+        "--output", type=str, required=True, help="Path to save the generated audio file (e.g., output.wav)."
+    )
+
+    parser.add_argument(
+        "--repo-id",
+        type=str,
+        default="nari-labs/Dia-1.6B-0626",
+        help="Hugging Face repository ID (e.g., nari-labs/Dia-1.6B-0626).",
+    )
+    parser.add_argument(
+        "--local-paths", action="store_true", help="Load model from local config and checkpoint files."
+    )
+
+    parser.add_argument(
+        "--config", type=str, help="Path to local config.json file (required if --local-paths is set)."
+    )
+    parser.add_argument(
+        "--checkpoint", type=str, help="Path to local model checkpoint .pth file (required if --local-paths is set)."
+    )
+    parser.add_argument(
+        "--audio-prompt", type=str, default=None, help="Path to an optional audio prompt WAV file for voice cloning."
+    )
+
+    gen_group = parser.add_argument_group("Generation Parameters")
+    gen_group.add_argument(
+        "--max-tokens",
+        type=int,
+        default=None,
+        help="Maximum number of audio tokens to generate (defaults to config value).",
+    )
+    gen_group.add_argument(
+        "--cfg-scale", type=float, default=3.0, help="Classifier-Free Guidance scale (default: 3.0)."
+    )
+    gen_group.add_argument(
+        "--temperature", type=float, default=1.3, help="Sampling temperature (higher is more random, default: 0.7)."
+    )
+    gen_group.add_argument("--top-p", type=float, default=0.95, help="Nucleus sampling probability (default: 0.95).")
+
+    infra_group = parser.add_argument_group("Infrastructure")
+    infra_group.add_argument("--seed", type=int, default=None, help="Random seed for reproducibility.")
+    infra_group.add_argument(
+        "--device",
+        type=str,
+        default="cuda" if torch.cuda.is_available() else "cpu",
+        help="Device to run inference on (e.g., 'cuda', 'cpu', default: auto).",
+    )
+
+    args = parser.parse_args()
+
+    # Validation for local paths
+    if args.local_paths:
+        if not args.config:
+            parser.error("--config is required when --local-paths is set.")
+        if not args.checkpoint:
+            parser.error("--checkpoint is required when --local-paths is set.")
+        if not os.path.exists(args.config):
+            parser.error(f"Config file not found: {args.config}")
+        if not os.path.exists(args.checkpoint):
+            parser.error(f"Checkpoint file not found: {args.checkpoint}")
+
+    # Set seed if provided
+    if args.seed is not None:
+        set_seed(args.seed)
+        print(f"Using user-selected seed: {args.seed}")
+
+    # Determine device
+    device = torch.device(args.device)
+    print(f"Using device: {device}")
+
+    # Load model
+    print("Loading model...")
+    if args.local_paths:
+        print(f"Loading from local paths: config='{args.config}', checkpoint='{args.checkpoint}'")
+        try:
+            model = Dia.from_local(args.config, args.checkpoint, device=device)
+        except Exception as e:
+            print(f"Error loading local model: {e}")
+            exit(1)
+    else:
+        print(f"Loading from Hugging Face Hub: repo_id='{args.repo_id}'")
+        try:
+            model = Dia.from_pretrained(args.repo_id, device=device)
+        except Exception as e:
+            print(f"Error loading model from Hub: {e}")
+            exit(1)
+    print("Model loaded.")
+
+    # Generate audio
+    print("Generating audio...")
+    try:
+        sample_rate = 44100  # Default assumption
+
+        output_audio = model.generate(
+            text=args.text,
+            audio_prompt=args.audio_prompt,
+            max_tokens=args.max_tokens,
+            cfg_scale=args.cfg_scale,
+            temperature=args.temperature,
+            top_p=args.top_p,
+        )
+        print("Audio generation complete.")
+
+        print(f"Saving audio to {args.output}...")
+        os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True)
+
+        sf.write(args.output, output_audio, sample_rate)
+        print(f"Audio successfully saved to {args.output}")
+
+    except Exception as e:
+        print(f"Error during audio generation or saving: {e}")
+        exit(1)
+
+
+if __name__ == "__main__":
+    main()

dia/__init__.py

@@ -0,0 +1,6 @@
+from .model import Dia
+
+
+__all__ = [
+    "Dia",
+]

dia/audio.py

@@ -0,0 +1,163 @@
+import typing as tp
+
+import torch
+
+
+def build_delay_indices(B: int, T: int, C: int, delay_pattern: tp.List[int]) -> tp.Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Precompute (t_idx_BxTxC, indices_BTCx3) so that out[t, c] = in[t - delay[c], c].
+    Negative t_idx => BOS; t_idx >= T => PAD.
+    """
+    delay_arr = torch.tensor(delay_pattern, dtype=torch.int32)
+
+    t_idx_BxT = torch.broadcast_to(
+        torch.arange(T, dtype=torch.int32)[None, :],
+        [B, T],
+    )
+    t_idx_BxTx1 = t_idx_BxT[..., None]
+    t_idx_BxTxC = t_idx_BxTx1 - delay_arr.view(1, 1, C)
+
+    b_idx_BxTxC = torch.broadcast_to(
+        torch.arange(B, dtype=torch.int32).view(B, 1, 1),
+        [B, T, C],
+    )
+    c_idx_BxTxC = torch.broadcast_to(
+        torch.arange(C, dtype=torch.int32).view(1, 1, C),
+        [B, T, C],
+    )
+
+    # We must clamp time indices to [0..T-1] so gather_nd equivalent won't fail
+    t_clamped_BxTxC = torch.clamp(t_idx_BxTxC, 0, T - 1)
+
+    indices_BTCx3 = torch.stack(
+        [
+            b_idx_BxTxC.reshape(-1),
+            t_clamped_BxTxC.reshape(-1),
+            c_idx_BxTxC.reshape(-1),
+        ],
+        dim=1,
+    ).long()  # Ensure indices are long type for indexing
+
+    return t_idx_BxTxC, indices_BTCx3
+
+
+def apply_audio_delay(
+    audio_BxTxC: torch.Tensor,
+    pad_value: int,
+    bos_value: int,
+    precomp: tp.Tuple[torch.Tensor, torch.Tensor],
+) -> torch.Tensor:
+    """
+    Applies the delay pattern to batched audio tokens using precomputed indices,
+    inserting BOS where t_idx < 0 and PAD where t_idx >= T.
+
+    Args:
+        audio_BxTxC: [B, T, C] int16 audio tokens (or int32/float)
+        pad_value: the padding token
+        bos_value: the BOS token
+        precomp:  (t_idx_BxTxC, indices_BTCx3) from build_delay_indices
+
+    Returns:
+        result_BxTxC: [B, T, C] delayed audio tokens
+    """
+    device = audio_BxTxC.device  # Get device from input tensor
+    t_idx_BxTxC, indices_BTCx3 = precomp
+    t_idx_BxTxC = t_idx_BxTxC.to(device)  # Move precomputed indices to device
+    indices_BTCx3 = indices_BTCx3.to(device)
+
+    # Equivalent of tf.gather_nd using advanced indexing
+    # Ensure indices are long type if not already (build_delay_indices should handle this)
+    gathered_flat = audio_BxTxC[indices_BTCx3[:, 0], indices_BTCx3[:, 1], indices_BTCx3[:, 2]]
+    gathered_BxTxC = gathered_flat.view(audio_BxTxC.shape)
+
+    # Create masks on the correct device
+    mask_bos = t_idx_BxTxC < 0  # => place bos_value
+    mask_pad = t_idx_BxTxC >= audio_BxTxC.shape[1]  # => place pad_value
+
+    # Create scalar tensors on the correct device
+    bos_tensor = torch.tensor(bos_value, dtype=audio_BxTxC.dtype, device=device)
+    pad_tensor = torch.tensor(pad_value, dtype=audio_BxTxC.dtype, device=device)
+
+    # If mask_bos, BOS; else if mask_pad, PAD; else original gather
+    # All tensors should now be on the same device
+    result_BxTxC = torch.where(mask_bos, bos_tensor, torch.where(mask_pad, pad_tensor, gathered_BxTxC))
+
+    return result_BxTxC
+
+
+def build_revert_indices(B: int, T: int, C: int, delay_pattern: tp.List[int]) -> tp.Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Precompute indices for the revert operation using PyTorch.
+
+    Returns:
+        A tuple (t_idx_BxTxC, indices_BTCx3) where:
+            - t_idx_BxTxC is a tensor of shape [B, T, C] computed as time indices plus the delay.
+            - indices_BTCx3 is a tensor of shape [B*T*C, 3] used for gathering, computed from:
+                batch indices, clamped time indices, and channel indices.
+    """
+    # Use default device unless specified otherwise; assumes inputs might define device later
+    device = None  # Or determine dynamically if needed, e.g., from a model parameter
+
+    delay_arr = torch.tensor(delay_pattern, dtype=torch.int32, device=device)
+
+    t_idx_BT1 = torch.broadcast_to(torch.arange(T, device=device).unsqueeze(0), [B, T])
+    t_idx_BT1 = t_idx_BT1.unsqueeze(-1)
+
+    t_idx_BxTxC = torch.minimum(
+        t_idx_BT1 + delay_arr.view(1, 1, C),
+        torch.tensor(T - 1, device=device),
+    )
+    b_idx_BxTxC = torch.broadcast_to(torch.arange(B, device=device).view(B, 1, 1), [B, T, C])
+    c_idx_BxTxC = torch.broadcast_to(torch.arange(C, device=device).view(1, 1, C), [B, T, C])
+
+    indices_BTCx3 = torch.stack(
+        [
+            b_idx_BxTxC.reshape(-1),
+            t_idx_BxTxC.reshape(-1),
+            c_idx_BxTxC.reshape(-1),
+        ],
+        axis=1,
+    ).long()  # Ensure indices are long type
+
+    return t_idx_BxTxC, indices_BTCx3
+
+
+def revert_audio_delay(
+    audio_BxTxC: torch.Tensor,
+    pad_value: int,
+    precomp: tp.Tuple[torch.Tensor, torch.Tensor],
+    T: int,
+) -> torch.Tensor:
+    """
+    Reverts a delay pattern from batched audio tokens using precomputed indices (PyTorch version).
+
+    Args:
+        audio_BxTxC: Input delayed audio tensor
+        pad_value: Padding value for out-of-bounds indices
+        precomp: Precomputed revert indices tuple containing:
+            - t_idx_BxTxC: Time offset indices tensor
+            - indices_BTCx3: Gather indices tensor for original audio
+        T: Original sequence length before padding
+
+    Returns:
+        Reverted audio tensor with same shape as input
+    """
+    t_idx_BxTxC, indices_BTCx3 = precomp
+    device = audio_BxTxC.device  # Get device from input tensor
+
+    # Move precomputed indices to the same device as audio_BxTxC if they aren't already
+    t_idx_BxTxC = t_idx_BxTxC.to(device)
+    indices_BTCx3 = indices_BTCx3.to(device)
+
+    # Using PyTorch advanced indexing (equivalent to tf.gather_nd or np equivalent)
+    gathered_flat = audio_BxTxC[indices_BTCx3[:, 0], indices_BTCx3[:, 1], indices_BTCx3[:, 2]]
+    gathered_BxTxC = gathered_flat.view(audio_BxTxC.size())  # Use .size() for robust reshaping
+
+    # Create pad_tensor on the correct device
+    pad_tensor = torch.tensor(pad_value, dtype=audio_BxTxC.dtype, device=device)
+    # Create T tensor on the correct device for comparison
+    T_tensor = torch.tensor(T, device=device)
+
+    result_BxTxC = torch.where(t_idx_BxTxC >= T_tensor, pad_tensor, gathered_BxTxC)  # Changed np.where to torch.where
+
+    return result_BxTxC

dia/config.py

@@ -0,0 +1,176 @@
+"""Configuration management module for the Dia model.
+
+This module provides comprehensive configuration management for the Dia model,
+utilizing Pydantic for validation. It defines configurations for data processing,
+model architecture (encoder and decoder), and training settings.
+
+Key components:
+- DataConfig: Parameters for data loading and preprocessing.
+- EncoderConfig: Architecture details for the encoder module.
+- DecoderConfig: Architecture details for the decoder module.
+- ModelConfig: Combined model architecture settings.
+- TrainingConfig: Training hyperparameters and settings.
+- DiaConfig: Master configuration combining all components.
+"""
+
+import os
+
+from pydantic import BaseModel, Field
+
+
+class EncoderConfig(BaseModel, frozen=True):
+    """Configuration for the encoder component of the Dia model.
+
+    Attributes:
+        model_type: Type of the model, defaults to "dia_encoder".
+        hidden_size: Size of the encoder layers, defaults to 1024.
+        intermediate_size: Size of the "intermediate" (i.e., feed-forward) layer in the encoder, defaults to 4096.
+        num_hidden_layers: Number of hidden layers in the encoder, defaults to 12.
+        num_attention_heads: Number of attention heads in the encoder, defaults to 16.
+        num_key_value_heads: Number of key-value heads in the encoder, defaults to 16.
+        head_dim: Dimension of each attention head, defaults to 128.
+        hidden_act: Activation function in the encoder, defaults to "silu".
+        max_position_embeddings: Maximum number of position embeddings, defaults to 1024.
+        initializer_range: Range for initializing weights, defaults to 0.02.
+        norm_eps: Epsilon value for normalization layers, defaults to 1e-5.
+        rope_theta: Theta value for RoPE, defaults to 10000.0.
+        rope_scaling: Optional scaling factor for RoPE.
+        vocab_size: Vocabulary size, defaults to 256.
+    """
+
+    head_dim: int = Field(default=128, gt=0)
+    hidden_act: str = Field(default="silu")
+    hidden_size: int = Field(default=1024, gt=0)
+    initializer_range: float = Field(default=0.02)
+    intermediate_size: int = Field(default=4096, gt=0)
+    max_position_embeddings: int = Field(default=1024, gt=0)
+    model_type: str = Field(default="dia_encoder")
+    norm_eps: float = Field(default=1e-5)
+    num_attention_heads: int = Field(default=16, gt=0)
+    num_hidden_layers: int = Field(default=12, gt=0)
+    num_key_value_heads: int = Field(default=16, gt=0)
+    rope_scaling: float | None = Field(default=None)
+    rope_theta: float = Field(default=10000.0)
+    vocab_size: int = Field(default=256, gt=0)
+
+
+class DecoderConfig(BaseModel, frozen=True):
+    """Configuration for the decoder component of the Dia model.
+
+    Attributes:
+        model_type: Type of the model, defaults to "dia_decoder".
+        hidden_size: Size of the decoder layers, defaults to 2048.
+        intermediate_size: Size of the "intermediate" (i.e., feed-forward) layer in the decoder, defaults to 8192.
+        num_hidden_layers: Number of hidden layers in the decoder, defaults to 18.
+        num_attention_heads: Number of attention heads in the decoder, defaults to 16.
+        num_key_value_heads: Number of key-value heads in the decoder, defaults to 4.
+        head_dim: Dimension of each attention head, defaults to 128.
+        cross_hidden_size: Size of the cross-attention layers, defaults to 1024.
+        cross_num_attention_heads: Number of attention heads in the cross-attention mechanism, defaults to 16.
+        cross_num_key_value_heads: Number of key-value heads in the cross-attention mechanism, defaults to 16.
+        cross_head_dim: Dimension of each cross-attention head, defaults to 128.
+        hidden_act: Activation function in the decoder, defaults to "silu".
+        max_position_embeddings: Maximum number of position embeddings in the decoder, defaults to 3072.
+        initializer_range: Range for initializing weights in the decoder, defaults to 0.02.
+        norm_eps: Epsilon value for normalization layers in the decoder, defaults to 1e-5.
+        rope_theta: Theta value for RoPE in the decoder, defaults to 10000.0.
+        rope_scaling: Optional scaling factor for RoPE in the decoder.
+        vocab_size: Vocabulary size for the decoder, defaults to 1028.
+        num_channels: Number of channels in the decoder, defaults to 9.
+    """
+
+    cross_head_dim: int = Field(default=128, gt=0)
+    cross_hidden_size: int = Field(default=1024, gt=0)
+    cross_num_attention_heads: int = Field(default=16, gt=0)
+    cross_num_key_value_heads: int = Field(default=16, gt=0)
+    head_dim: int = Field(default=128, gt=0)
+    hidden_act: str = Field(default="silu")
+    hidden_size: int = Field(default=2048, gt=0)
+    initializer_range: float = Field(default=0.02)
+    intermediate_size: int = Field(default=8192, gt=0)
+    max_position_embeddings: int = Field(default=3072, gt=0)
+    model_type: str = Field(default="dia_decoder")
+    norm_eps: float = Field(default=1e-5)
+    num_attention_heads: int = Field(default=16, gt=0)
+    num_channels: int = Field(default=9, gt=0)
+    num_hidden_layers: int = Field(default=18, gt=0)
+    num_key_value_heads: int = Field(default=4, gt=0)
+    rope_scaling: float | None = Field(default=None)
+    rope_theta: float = Field(default=10000.0)
+    vocab_size: int = Field(default=1028, gt=0)
+
+
+class DiaConfig(BaseModel, frozen=True):
+    """Main configuration container for the Dia model architecture.
+
+    Attributes:
+        model_type: Type of the model, defaults to "dia".
+        is_encoder_decoder: Flag indicating if the model is an encoder-decoder type, defaults to True.
+        encoder: Configuration for the encoder component.
+        decoder: Configuration for the decoder component.
+        src_vocab_size: Size of the source (text) vocabulary.
+        tgt_vocab_size: Size of the target (audio code) vocabulary.
+        initializer_range: Range for initializing weights, defaults to 0.02.
+        norm_eps: Epsilon value for normalization layers, defaults to 1e-5.
+        torch_dtype: Data type for model weights in PyTorch, defaults to "float32".
+        bos_token_id: Beginning-of-sequence token ID, defaults to 1026.
+        eos_token_id: End-of-sequence token ID, defaults to 1024.
+        pad_token_id: Padding token ID, defaults to 1025.
+        rope_theta: Theta value for RoPE, defaults to 10000.0.
+        rope_scaling: Optional scaling factor for RoPE.
+        transformers_version: Version of the transformers library, defaults to "4.53.0.dev0".
+        architectures: List of model architectures, defaults to ["DiaForConditionalGeneration"].
+        delay_pattern: List of delay values for each audio channel, defaults to [0,8,9,10,11,12,13,14,15].
+    """
+
+    architectures: list[str] = Field(default_factory=lambda: ["DiaForConditionalGeneration"])
+    bos_token_id: int = Field(default=1026)
+    decoder_config: DecoderConfig
+    delay_pattern: list[int] = Field(default_factory=lambda: [0, 8, 9, 10, 11, 12, 13, 14, 15])
+    encoder_config: EncoderConfig
+    eos_token_id: int = Field(default=1024)
+    initializer_range: float = Field(default=0.02)
+    is_encoder_decoder: bool = Field(default=True)
+    model_type: str = Field(default="dia")
+    norm_eps: float = Field(default=1e-5)
+    pad_token_id: int = Field(default=1025)
+    torch_dtype: str = Field(default="float32")
+    transformers_version: str = Field(default="4.53.0.dev0")
+
+    def save(self, path: str) -> None:
+        """Save the current configuration instance to a JSON file.
+
+        Ensures the parent directory exists and the file has a .json extension.
+
+        Args:
+            path: The target file path to save the configuration.
+
+        Raises:
+            ValueError: If the path is not a file with a .json extension.
+        """
+        os.makedirs(os.path.dirname(path), exist_ok=True)
+        config_json = self.model_dump_json(indent=2)
+        with open(path, "w") as f:
+            f.write(config_json)
+
+    @classmethod
+    def load(cls, path: str) -> "DiaConfig | None":
+        """Load and validate a Dia configuration from a JSON file.
+
+        Args:
+            path: The path to the configuration file.
+
+        Returns:
+            A validated DiaConfig instance if the file exists and is valid,
+            otherwise None if the file is not found.
+
+        Raises:
+            ValueError: If the path does not point to an existing .json file.
+            pydantic.ValidationError: If the JSON content fails validation against the DiaConfig schema.
+        """
+        try:
+            with open(path, "r") as f:
+                content = f.read()
+            return cls.model_validate_json(content)
+        except FileNotFoundError:
+            return None

dia/layers.py

@@ -0,0 +1,888 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from huggingface_hub import PyTorchModelHubMixin
+from torch import Tensor
+from torch.nn import RMSNorm
+
+from .config import DecoderConfig, DiaConfig, EncoderConfig
+from .state import DecoderInferenceState, EncoderInferenceState, KVCache
+
+
+def _normalize_axes(axes: tuple[int, ...], ndim: int) -> tuple[int, ...]:
+    return tuple(ax if ax >= 0 else ndim + ax for ax in axes)
+
+
+class DenseGeneral(nn.Module):
+    """
+    PyTorch equivalent of flax.linen.DenseGeneral with shapes defined at init.
+    Stores weights (`kernel`) in the same layout as Jax and uses torch.tensordot
+    for the generalized matrix multiplication. Weight/bias shapes are calculated
+    and parameters created during initialization based on config.
+    `load_weights` validates shapes and copies data.
+    Attributes:
+        axis (Tuple[int, ...]): Input axis or axes to contract.
+        in_shapes (Tuple[int, ...]): Sizes of the input dimensions specified by `axis`.
+        out_features (Tuple[int, ...]): Shape of the output features (non-contracted dims).
+        use_bias (bool): Whether to add a bias term.
+        weight (nn.Parameter): The kernel parameter.
+        bias (Optional[nn.Parameter]): The bias parameter (if use_bias=True).
+    """
+
+    def __init__(
+        self,
+        in_shapes: tuple[int, ...],
+        out_features: tuple[int, ...],
+        axis: tuple[int, ...] = (-1,),
+        weight_dtype: torch.dtype | None = None,
+        device: torch.device | None = None,
+    ):
+        super().__init__()
+        self.in_shapes = in_shapes
+        self.out_features = out_features
+        self.axis = axis
+        self.kernel_shape = self.in_shapes + self.out_features
+
+        factory_kwargs = {"device": device, "dtype": weight_dtype}
+        self.weight = nn.Parameter(torch.empty(self.kernel_shape, **factory_kwargs))
+
+    def forward(self, inputs: Tensor) -> Tensor:
+        norm_axis = _normalize_axes(self.axis, inputs.ndim)
+        kernel_contract_axes = tuple(range(len(norm_axis)))
+
+        output = torch.tensordot(
+            inputs.to(self.weight.dtype),
+            self.weight,
+            dims=(norm_axis, kernel_contract_axes),
+        ).to(inputs.dtype)
+        return output
+
+
+class MlpBlock(nn.Module):
+    """MLP block using DenseGeneral."""
+
+    def __init__(self, embed_dim: int, intermediate_dim: int, compute_dtype: torch.dtype):
+        super().__init__()
+        self.dtype = compute_dtype
+
+        self.wi_fused = DenseGeneral(
+            in_shapes=(embed_dim,),
+            out_features=(2, intermediate_dim),
+            axis=(-1,),
+            weight_dtype=compute_dtype,
+        )
+
+        self.wo = DenseGeneral(
+            in_shapes=(intermediate_dim,),
+            out_features=(embed_dim,),
+            axis=(-1,),
+            weight_dtype=compute_dtype,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass."""
+        fused_x = self.wi_fused(x)
+
+        gate = fused_x[..., 0, :]
+        up = fused_x[..., 1, :]
+
+        hidden = torch.mul(F.silu(gate), up).to(self.dtype)
+
+        output = self.wo(hidden)
+        return output
+
+
+class RotaryEmbedding(nn.Module):
+    """Rotary Position Embedding (RoPE) implementation in PyTorch."""
+
+    def __init__(
+        self,
+        embedding_dims: int,
+        min_timescale: float = 1.0,
+        max_timescale: float = 10000.0,
+        dtype: torch.dtype = torch.float32,
+    ):
+        super().__init__()
+        if embedding_dims % 2 != 0:
+            raise ValueError("Embedding dim must be even for RoPE.")
+        self.embedding_dims = embedding_dims
+        self.min_timescale = min_timescale
+        self.max_timescale = max_timescale
+        self.compute_dtype = dtype
+
+        half_embedding_dim = embedding_dims // 2
+        fraction = (2.0 * torch.arange(0, half_embedding_dim)) / embedding_dims
+        timescale = (self.min_timescale * (self.max_timescale / self.min_timescale) ** fraction).to(torch.float32)
+        self.register_buffer("timescale", timescale, persistent=False)
+
+    def forward(self, inputs: torch.Tensor, position: torch.Tensor):
+        """Applies RoPE."""
+        position = position.unsqueeze(-1).unsqueeze(-1)
+        sinusoid_inp = position / self.timescale
+        sin = torch.sin(sinusoid_inp)
+        cos = torch.cos(sinusoid_inp)
+        first_half, second_half = torch.chunk(inputs.to(torch.float32), 2, dim=-1)
+        first_part = first_half * cos - second_half * sin
+        second_part = second_half * cos + first_half * sin
+        return torch.cat(
+            (first_part.to(self.compute_dtype), second_part.to(self.compute_dtype)),
+            dim=-1,
+        )
+
+    def apply_rope(self, inputs: torch.Tensor, sin: torch.Tensor, cos: torch.Tensor):
+        first_half, second_half = torch.chunk(inputs.to(torch.float32), 2, dim=-1)
+        first_part = first_half * cos - second_half * sin
+        second_part = second_half * cos + first_half * sin
+        return torch.cat((first_part.to(self.compute_dtype), second_part.to(self.compute_dtype)), dim=-1)
+
+
+def custom_scaled_dot_product_attention(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attn_mask: torch.Tensor | None = None,
+    scale: float = 1.0,
+    is_causal: bool = False,
+    num_gqa_groups: int = 1,
+) -> torch.Tensor:
+    """
+    Custom scaled dot-product attention with GQA support for MPS compatibility.
+
+    Args:
+        query: (B, N_q, T, H) - Query tensor, N_q = num_query_heads
+        key: (B, N_kv, S, H) - Key tensor, N_kv = num_kv_heads
+        value: (B, N_kv, S, H) - Value tensor
+        attn_mask: (B, 1, T, S) - Attention mask, optional
+        scale: Scaling factor for attention scores
+        is_causal: If True, apply causal masking
+        num_gqa_groups: Number of query groups per KV head (N_q / N_kv)
+
+    Returns:
+        output: (B, N_q, T, H) - Attention output
+    """
+    B, N_q, T, H = query.shape
+    _, N_kv, S, _ = key.shape
+
+    # For GQA, repeat key and value tensors to match query heads
+    if num_gqa_groups > 1:
+        key = key.repeat_interleave(num_gqa_groups, dim=1)  # (B, N_q, S, H)
+        value = value.repeat_interleave(num_gqa_groups, dim=1)  # (B, N_q, S, H)
+
+    # Compute attention scores: (B, N_q, T, H) @ (B, N_q, H, S) -> (B, N_q, T, S)
+    scores = torch.matmul(query, key.transpose(-1, -2)) * scale
+
+    # Apply causal mask if needed
+    if is_causal:
+        causal_mask = torch.tril(torch.ones(T, S, dtype=torch.bool, device=query.device))
+        scores = scores.masked_fill(~causal_mask, float("-inf"))
+
+    # Apply attention mask if provided
+    if attn_mask is not None:
+        scores = scores.masked_fill(~attn_mask, float("-inf"))
+
+    # Softmax over the last dimension (S)
+    attn_weights = F.softmax(scores, dim=-1)
+
+    # Compute output: (B, N_q, T, S) @ (B, N_q, S, H) -> (B, N_q, T, H)
+    output = torch.matmul(attn_weights, value)
+
+    return output
+
+
+class CrossAttention(nn.Module):
+    """Cross-Attention using DenseGeneral."""
+
+    def __init__(
+        self,
+        config: EncoderConfig | DecoderConfig,
+        q_embed_dim: int,
+        kv_embed_dim: int,
+        num_query_heads: int,
+        num_kv_heads: int,
+        head_dim: int,
+        compute_dtype: torch.dtype,
+        out_embed_dim: int | None = None,
+    ):
+        super().__init__()
+        self.num_query_heads = num_query_heads
+        self.num_kv_heads = num_kv_heads
+        self.head_dim = head_dim
+        self.output_dim = out_embed_dim if out_embed_dim is not None else q_embed_dim
+        self.projected_query_dim = num_query_heads * head_dim
+        if num_query_heads % num_kv_heads != 0:
+            raise ValueError(f"num_query_heads ({num_query_heads}) must be divisible by num_kv_heads ({num_kv_heads})")
+        self.num_gqa_groups = num_query_heads // num_kv_heads
+
+        # --- Projection Layers using DenseGeneral ---
+        self.q_proj = DenseGeneral(
+            in_shapes=(q_embed_dim,),
+            out_features=(num_query_heads, head_dim),
+            axis=(-1,),
+            weight_dtype=compute_dtype,
+        )
+        self.k_proj = DenseGeneral(
+            in_shapes=(kv_embed_dim,),
+            out_features=(num_kv_heads, head_dim),
+            axis=(-1,),
+            weight_dtype=compute_dtype,
+        )
+        self.v_proj = DenseGeneral(
+            in_shapes=(kv_embed_dim,),
+            out_features=(num_kv_heads, head_dim),
+            axis=(-1,),
+            weight_dtype=compute_dtype,
+        )
+        self.o_proj = DenseGeneral(
+            in_shapes=(num_query_heads, head_dim),
+            out_features=(self.output_dim,),
+            axis=(-2, -1),
+            weight_dtype=compute_dtype,
+        )
+
+        # --- Rotary Embedding ---
+        self.rotary_emb = RotaryEmbedding(
+            embedding_dims=self.head_dim,
+            max_timescale=config.rope_theta,
+            dtype=compute_dtype,
+        )
+
+    def forward(
+        self,
+        Xq: torch.Tensor,  # (B, T, D) T = 1 in AR generation
+        q_positions: torch.Tensor,  # (B, T)
+        kv_positions: torch.Tensor | None = None,  # (B, S)
+        attn_mask: torch.Tensor | None = None,  # None in Decoder Self Attention, Valid mask in Others
+        cache: KVCache | None = None,  # None in Encoder, KVCache in Decoder
+        is_causal: bool = False,
+    ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor] | None]:
+        """
+        Performs attention calculation with optional KV caching.
+
+        Args:
+            Xq: Query tensor (B, T, D). T=1 during single-step decoding.
+            Xkv: Key/Value source tensor (B, S, E). S=1 during single-step decoding for self-attn.
+            q_positions: Positions for queries (B, T).
+            kv_positions: Positions for keys/values (B, S). If None, uses q_positions.
+            attn_mask: Attention mask.
+            cache: KVCache.
+
+        Returns:
+            A tuple containing:
+            - output: The attention output tensor (B, T, output_dim).
+            - present_kv: The K/V state to be cached for the next step ((B, N, S_new, H), (B, N, S_new, H)). For self-attn, S_new = S_past + S. For cross-attn, S_new = S_kv.
+        """
+        if kv_positions is None:
+            kv_positions = q_positions
+        original_dtype = Xq.dtype
+
+        Xq_BxTxNxH = self.q_proj(Xq)
+        Xq_BxNxTxH = Xq_BxTxNxH.transpose(1, 2)
+
+        attn_k: torch.Tensor | None = cache.k if cache is not None else None
+        attn_v: torch.Tensor | None = cache.v if cache is not None else None
+
+        # Use custom attention for MPS backend, otherwise use optimized PyTorch function
+        is_mps = Xq.device.type == "mps" and torch.backends.mps.is_available()
+        if is_mps:
+            attn_output = custom_scaled_dot_product_attention(
+                query=Xq_BxNxTxH,
+                key=attn_k,
+                value=attn_v,
+                attn_mask=attn_mask if not is_causal else None,
+                scale=1.0,
+                is_causal=is_causal,
+                num_gqa_groups=self.num_gqa_groups,
+            )
+        else:
+            attn_output = F.scaled_dot_product_attention(
+                Xq_BxNxTxH,
+                attn_k,
+                attn_v,
+                attn_mask=attn_mask if not is_causal else None,
+                scale=1.0,
+                enable_gqa=self.num_gqa_groups > 1,
+                is_causal=is_causal,
+            )
+
+        attn_output = attn_output.transpose(1, 2).contiguous()  # (B, T, N, H)
+        output = self.o_proj(attn_output)
+
+        return output.to(original_dtype)
+
+
+class FusedQKV(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        bias: bool = False,
+        num_q_heads: int = 1,
+        q_head_dim: int = 1,
+        num_kv_heads: int = 1,
+        kv_head_dim: int = 1,
+    ):
+        super().__init__()
+        self.num_q_heads = num_q_heads
+        self.q_head_dim = q_head_dim
+        self.num_kv_heads = num_kv_heads
+        self.kv_head_dim = kv_head_dim
+        self.q_output_dim = num_q_heads * q_head_dim
+        self.kv_output_dim = num_kv_heads * kv_head_dim
+        self.linear = nn.Linear(in_features, out_features, bias=bias)
+
+    def forward(self, inputs: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        x = self.linear(inputs)
+
+        q, k, v = x.split([self.q_output_dim, self.kv_output_dim, self.kv_output_dim], dim=-1)
+
+        q = q.reshape(q.shape[:-1] + (self.num_q_heads, self.q_head_dim))
+        k = k.reshape(k.shape[:-1] + (self.num_kv_heads, self.kv_head_dim))
+        v = v.reshape(v.shape[:-1] + (self.num_kv_heads, self.kv_head_dim))
+
+        return q, k, v
+
+
+class SelfAttention(nn.Module):
+    """Attention using DenseGeneral."""
+
+    def __init__(
+        self,
+        config: EncoderConfig | DecoderConfig,
+        q_embed_dim: int,
+        kv_embed_dim: int,
+        num_query_heads: int,
+        num_kv_heads: int,
+        head_dim: int,
+        compute_dtype: torch.dtype,
+        out_embed_dim: int | None = None,
+    ):
+        super().__init__()
+        self.num_query_heads = num_query_heads
+        self.num_kv_heads = num_kv_heads
+        self.head_dim = head_dim
+        self.output_dim = out_embed_dim if out_embed_dim is not None else q_embed_dim
+        self.projected_query_dim = num_query_heads * head_dim
+        if num_query_heads % num_kv_heads != 0:
+            raise ValueError(f"num_query_heads ({num_query_heads}) must be divisible by num_kv_heads ({num_kv_heads})")
+        self.num_gqa_groups = num_query_heads // num_kv_heads
+        self.kv_embed_dim = kv_embed_dim
+        self.q_embed_dim = q_embed_dim
+
+        # --- Projection Layers using DenseGeneral ---
+        self.q_proj = DenseGeneral(
+            in_shapes=(q_embed_dim,),
+            out_features=(num_query_heads, head_dim),
+            axis=(-1,),
+            weight_dtype=compute_dtype,
+        )
+        self.k_proj = DenseGeneral(
+            in_shapes=(kv_embed_dim,),
+            out_features=(num_kv_heads, head_dim),
+            axis=(-1,),
+            weight_dtype=compute_dtype,
+        )
+        self.v_proj = DenseGeneral(
+            in_shapes=(kv_embed_dim,),
+            out_features=(num_kv_heads, head_dim),
+            axis=(-1,),
+            weight_dtype=compute_dtype,
+        )
+        self.o_proj = DenseGeneral(
+            in_shapes=(num_query_heads, head_dim),
+            out_features=(self.output_dim,),
+            axis=(-2, -1),
+            weight_dtype=compute_dtype,
+        )
+
+        # --- Rotary Embedding ---
+        self.rotary_emb = RotaryEmbedding(
+            embedding_dims=self.head_dim,
+            max_timescale=config.rope_theta,
+            dtype=compute_dtype,
+        )
+
+        self.is_fused_qkv = False
+
+    def get_linear_weight(self, dense: DenseGeneral):
+        W_dg = dense.weight.data
+
+        out_features = 1
+        input_features = 1
+        for dim in dense.out_features:
+            out_features *= dim
+        for dim in dense.in_shapes:
+            input_features *= dim
+
+        W_dg_reshaped_for_linear_T = W_dg.reshape(input_features, out_features)
+        linear_weight = W_dg_reshaped_for_linear_T.transpose(0, 1).contiguous()
+        return linear_weight
+
+    def patch_fused_qkv(self):
+        q_proj_weight = self.get_linear_weight(self.q_proj)
+        k_proj_weight = self.get_linear_weight(self.k_proj)
+        v_proj_weight = self.get_linear_weight(self.v_proj)
+
+        self.qkv = FusedQKV(
+            self.kv_embed_dim,
+            (self.num_query_heads * self.head_dim + 2 * (self.num_kv_heads * self.head_dim)),
+            bias=False,
+            num_q_heads=self.num_query_heads,
+            q_head_dim=self.head_dim,
+            num_kv_heads=self.num_kv_heads,
+            kv_head_dim=self.head_dim,
+        )
+        self.qkv.linear.weight.data = torch.cat([q_proj_weight, k_proj_weight, v_proj_weight], dim=0)
+
+        # print(f"qkv.weight.shape: {self.qkv.linear.weight.shape}")
+        self.is_fused_qkv = True
+
+    def forward(
+        self,
+        X: torch.Tensor,  # (B, T, D) T = 1 in AR generation
+        q_positions: torch.Tensor,  # (B, T)
+        kv_positions: torch.Tensor | None = None,  # (B, S)
+        attn_mask: torch.Tensor | None = None,  # None in Decoder Self Attention, Valid mask in Others
+        cache: KVCache | None = None,  # None in Encoder, KVCache in Decoder
+        prefill: bool = False,
+        is_causal: bool = False,
+        current_idx: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor] | None]:
+        """
+        Performs attention calculation with optional KV caching.
+        Args:
+            Xq: Query tensor (B, T, D). T=1 during single-step decoding.
+            Xkv: Key/Value source tensor (B, S, E). S=1 during single-step decoding for self-attn.
+            q_positions: Positions for queries (B, T).
+            kv_positions: Positions for keys/values (B, S). If None, uses q_positions.
+            attn_mask: Attention mask.
+            cache: KVCache.
+            prefill: If True, use prefill mode.
+        Returns:
+            A tuple containing:
+            - output: The attention output tensor (B, T, output_dim).
+            - present_kv: The K/V state to be cached for the next step ((B, N, S_new, H), (B, N, S_new, H)). For self-attn, S_new = S_past + S. For cross-attn, S_new = S_kv.
+        """
+        if kv_positions is None:
+            kv_positions = q_positions
+
+        original_dtype = X.dtype
+
+        if self.is_fused_qkv:
+            Xq_BxTxNxH, Xk_BxSxKxH, Xv_BxSxKxH = self.qkv(X)
+        else:
+            Xq_BxTxNxH = self.q_proj(X)
+            Xk_BxSxKxH = self.k_proj(X)
+            Xv_BxSxKxH = self.v_proj(X)
+
+        position = q_positions.unsqueeze(-1).unsqueeze(-1)
+        sinusoid_inp = position / self.rotary_emb.timescale
+        sin = torch.sin(sinusoid_inp)
+        cos = torch.cos(sinusoid_inp)
+
+        Xq_BxTxNxH = self.rotary_emb.apply_rope(Xq_BxTxNxH, sin, cos)
+        Xk_BxSxKxH = self.rotary_emb.apply_rope(Xk_BxSxKxH, sin, cos)
+
+        Xq_BxNxTxH = Xq_BxTxNxH.transpose(1, 2)
+
+        attn_k: torch.Tensor | None = cache.k if cache is not None else None
+        attn_v: torch.Tensor | None = cache.v if cache is not None else None
+
+        Xk_BxKxSxH = Xk_BxSxKxH.transpose(1, 2)  # (B, K, S, H)
+        Xv_BxKxSxH = Xv_BxSxKxH.transpose(1, 2)  # (B, K, S, H)
+
+        if cache is None:
+            attn_k = Xk_BxKxSxH
+            attn_v = Xv_BxKxSxH
+        elif prefill:
+            attn_k, attn_v = Xk_BxKxSxH, Xv_BxKxSxH
+            cache.prefill(attn_k, attn_v)
+        else:
+            attn_k, attn_v = cache.update(Xk_BxKxSxH, Xv_BxKxSxH, current_idx)
+
+        # Use custom attention for MPS backend, otherwise use optimized PyTorch function
+        is_mps = Xv_BxSxKxH.device.type == "mps" and torch.backends.mps.is_available()
+        if is_mps:
+            attn_output = custom_scaled_dot_product_attention(
+                query=Xq_BxNxTxH,
+                key=attn_k,
+                value=attn_v,
+                attn_mask=attn_mask if not is_causal else None,
+                scale=1.0,
+                is_causal=is_causal,
+                num_gqa_groups=self.num_gqa_groups,
+            )
+        else:
+            attn_output = F.scaled_dot_product_attention(
+                Xq_BxNxTxH,
+                attn_k,
+                attn_v,
+                attn_mask=attn_mask if not is_causal else None,
+                scale=1.0,
+                enable_gqa=self.num_gqa_groups > 1,
+                is_causal=is_causal,
+            )
+
+        attn_output = attn_output.transpose(1, 2).contiguous()  # (B, T, N, H)
+        output = self.o_proj(attn_output)
+
+        return output.to(original_dtype)
+
+
+class EncoderLayer(nn.Module):
+    """Transformer Encoder Layer using DenseGeneral."""
+
+    def __init__(self, config: DiaConfig, compute_dtype: torch.dtype):
+        super().__init__()
+        self.config = config
+        enc_config = config.encoder_config
+        embed_dim = enc_config.hidden_size
+        self.compute_dtype = compute_dtype
+
+        self.pre_sa_norm = RMSNorm(
+            embed_dim,
+            eps=enc_config.norm_eps,
+            dtype=torch.float32,
+        )
+        self.self_attention = SelfAttention(
+            enc_config,
+            q_embed_dim=embed_dim,
+            kv_embed_dim=embed_dim,
+            num_query_heads=enc_config.num_attention_heads,
+            num_kv_heads=enc_config.num_key_value_heads,
+            head_dim=enc_config.head_dim,
+            compute_dtype=compute_dtype,
+            out_embed_dim=embed_dim,
+        )
+        self.post_sa_norm = RMSNorm(
+            embed_dim,
+            eps=enc_config.norm_eps,
+            dtype=torch.float32,
+        )
+        self.mlp = MlpBlock(
+            embed_dim=embed_dim,
+            intermediate_dim=enc_config.intermediate_size,
+            compute_dtype=compute_dtype,
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        state: EncoderInferenceState,
+    ) -> torch.Tensor:
+        residual = x
+        x_norm = self.pre_sa_norm(x).to(self.compute_dtype)
+
+        sa_out = self.self_attention(
+            X=x_norm,
+            q_positions=state.positions,
+            kv_positions=state.positions,
+            attn_mask=state.attn_mask,
+        )
+        x = residual + sa_out
+
+        residual = x
+        x_norm = self.post_sa_norm(x).to(self.compute_dtype)
+        mlp_out = self.mlp(x_norm)
+        x = residual + mlp_out
+
+        return x
+
+
+class Encoder(nn.Module):
+    """Transformer Encoder Stack using DenseGeneral."""
+
+    def __init__(self, config: DiaConfig, compute_dtype: torch.dtype):
+        super().__init__()
+        self.config = config
+        enc_config = config.encoder_config
+        self.compute_dtype = compute_dtype
+
+        self.embedding = nn.Embedding(
+            enc_config.vocab_size,
+            enc_config.hidden_size,
+            dtype=compute_dtype,
+        )
+        self.layers = nn.ModuleList([EncoderLayer(config, compute_dtype) for _ in range(enc_config.num_hidden_layers)])
+        self.norm = RMSNorm(
+            enc_config.hidden_size,
+            eps=enc_config.norm_eps,
+            dtype=torch.float32,
+        )
+
+    def forward(
+        self,
+        x_ids: torch.Tensor,
+        state: EncoderInferenceState,
+    ) -> torch.Tensor:
+        x = self.embedding(x_ids)
+
+        for layer in self.layers:
+            x = layer(x, state)
+
+        x = self.norm(x).to(self.compute_dtype)
+        return x
+
+
+class DecoderLayer(nn.Module):
+    """Transformer Decoder Layer using DenseGeneral."""
+
+    def __init__(self, config: DiaConfig, compute_dtype: torch.dtype):
+        super().__init__()
+        self.config = config
+        dec_config = config.decoder_config
+        enc_config = config.encoder_config
+        dec_embed_dim = dec_config.hidden_size
+        enc_embed_dim = enc_config.hidden_size
+        self.compute_dtype = compute_dtype
+
+        # Norms
+        self.pre_sa_norm = RMSNorm(
+            dec_embed_dim,
+            eps=dec_config.norm_eps,
+            dtype=torch.float32,
+        )
+        self.pre_ca_norm = RMSNorm(
+            dec_embed_dim,
+            eps=dec_config.norm_eps,
+            dtype=torch.float32,
+        )
+        self.pre_mlp_norm = RMSNorm(
+            dec_embed_dim,
+            eps=dec_config.norm_eps,
+            dtype=torch.float32,
+        )
+
+        # Self-Attention (GQA) with Causal Masking
+        self.self_attention = SelfAttention(
+            dec_config,
+            q_embed_dim=dec_embed_dim,
+            kv_embed_dim=dec_embed_dim,
+            num_query_heads=dec_config.num_attention_heads,
+            num_kv_heads=dec_config.num_key_value_heads,
+            head_dim=dec_config.head_dim,
+            compute_dtype=compute_dtype,
+            out_embed_dim=dec_embed_dim,
+        )
+        # Cross-Attention (MHA)
+        self.cross_attention = CrossAttention(
+            dec_config,
+            q_embed_dim=dec_embed_dim,
+            kv_embed_dim=enc_embed_dim,  # Note kv_embed_dim
+            num_query_heads=dec_config.cross_num_attention_heads,
+            num_kv_heads=dec_config.cross_num_key_value_heads,
+            head_dim=dec_config.cross_head_dim,
+            compute_dtype=compute_dtype,
+            out_embed_dim=dec_embed_dim,
+        )
+        # MLP
+        self.mlp = MlpBlock(
+            embed_dim=dec_embed_dim,
+            intermediate_dim=dec_config.intermediate_size,
+            compute_dtype=compute_dtype,
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        state: DecoderInferenceState,
+        self_attn_cache: KVCache | None = None,
+        cross_attn_cache: KVCache | None = None,
+        prefill: bool = False,
+        current_idx: int = 0,
+    ) -> torch.Tensor:
+        residual = x
+        x_norm = self.pre_sa_norm(x).to(self.compute_dtype)
+
+        self_attn_mask = state.casual_attn_mask[None, None, current_idx]
+
+        sa_out = self.self_attention(
+            X=x_norm,  # (2, 1, D)
+            q_positions=state.dec_positions,  # (2, 1)
+            kv_positions=state.dec_positions,  # (2, 1)
+            attn_mask=self_attn_mask,
+            cache=self_attn_cache,
+            prefill=prefill,
+            is_causal=prefill,
+            current_idx=current_idx,
+        )
+
+        x = residual + sa_out
+
+        residual = x
+        x_norm = self.pre_ca_norm(x).to(self.compute_dtype)
+        ca_out = self.cross_attention(
+            Xq=x_norm,
+            q_positions=state.dec_positions,
+            kv_positions=state.enc_positions,
+            attn_mask=state.cross_attn_mask,
+            cache=cross_attn_cache,
+        )
+        x = residual + ca_out
+
+        residual = x
+        x_norm = self.pre_mlp_norm(x).to(self.compute_dtype)
+        mlp_out = self.mlp(x_norm)
+        x = residual + mlp_out
+
+        return x
+
+
+class Decoder(nn.Module):
+    """Transformer Decoder Stack using DenseGeneral."""
+
+    def __init__(self, config: DiaConfig, compute_dtype: torch.dtype):
+        super().__init__()
+        self.config = config
+        dec_config = config.decoder_config
+        self.num_channels = dec_config.num_channels
+        self.num_layers = dec_config.num_hidden_layers
+
+        self.embeddings = nn.ModuleList(
+            [
+                nn.Embedding(dec_config.vocab_size, dec_config.hidden_size, dtype=compute_dtype)
+                for _ in range(self.num_channels)
+            ]
+        )
+        self.layers = nn.ModuleList(
+            [DecoderLayer(config=config, compute_dtype=compute_dtype) for _ in range(self.num_layers)]
+        )
+
+        self.norm = RMSNorm(
+            dec_config.hidden_size,
+            eps=dec_config.norm_eps,
+            dtype=torch.float32,
+        )
+
+        self.logits_dense = DenseGeneral(
+            in_shapes=(dec_config.hidden_size,),
+            out_features=(self.num_channels, dec_config.vocab_size),
+            axis=(-1,),
+            weight_dtype=compute_dtype,
+        )
+
+    def precompute_cross_attn_cache(
+        self,
+        enc_out: torch.Tensor,  # (B, S, E)
+    ) -> list[KVCache]:
+        """
+        Computes the Key and Value tensors for cross-attention for each layer from the encoder output.
+        """
+        per_layer_kv_cache: list[KVCache] = []
+
+        for layer in self.layers:
+            cross_attn_module = layer.cross_attention
+            k_proj = cross_attn_module.k_proj(enc_out)
+            v_proj = cross_attn_module.v_proj(enc_out)
+
+            k = k_proj.transpose(1, 2)
+            v = v_proj.transpose(1, 2)
+
+            per_layer_kv_cache.append(KVCache.from_kv(k, v))
+
+        return per_layer_kv_cache
+
+    def decode_step(
+        self,
+        tgt_ids_Bx1xC: torch.Tensor,  # [B, 1, C]
+        state: DecoderInferenceState,
+        current_idx: int,
+    ) -> torch.Tensor:
+        """
+        Performs a single decoding step, managing KV caches layer by layer.
+        Returns:
+            A tuple containing:
+            - logits_Bx1xCV: The final output logits for the current step (B, 1, C*V), cast to float32.
+        """
+
+        x = None
+        for i in range(self.num_channels):
+            channel_tokens = tgt_ids_Bx1xC[..., i]
+            channel_embed = self.embeddings[i](channel_tokens)
+            x = channel_embed if x is None else x + channel_embed
+
+        for i, layer in enumerate(self.layers):
+            self_cache = state.self_attn_cache[i]
+            cross_cache = state.cross_attn_cache[i]
+            x = layer(
+                x,  # (2, 1, D)
+                state,
+                self_attn_cache=self_cache,
+                cross_attn_cache=cross_cache,
+                current_idx=current_idx,
+            )
+
+        x = self.norm(x)
+        logits_Bx1xCxV = self.logits_dense(x)
+
+        return logits_Bx1xCxV.to(torch.float32)
+
+    def forward(self, tgt_ids_BxTxC: torch.Tensor, state: DecoderInferenceState) -> torch.Tensor:
+        """
+        Forward pass for the Decoder stack, managing KV caches.
+        Args:
+            tgt_ids_BxTxC: Target token IDs (B, T, C).
+            encoder_out: Output from the encoder (B, S, E).
+            tgt_positions: Positions for target sequence (B, T).
+            src_positions: Positions for source sequence (B, S).
+            self_attn_mask: Mask for self-attention.
+            cross_attn_mask: Mask for cross-attention.
+            past_key_values: List containing the self-attention KV cache for each layer
+                             from the previous decoding step. `len(past_key_values)` should
+                             equal `num_layers`.
+            precomputed_cross_attn_kv: A single tuple containing the pre-computed K/V cache
+                                      derived from `encoder_out`. This is passed identically
+                                      to all layers.
+        Returns:
+            A tuple containing:
+            - logits: The final output logits (B, T, C * V), cast to float32.
+            - present_key_values: A list containing the updated self-attention KV cache
+                                 for each layer for the *current* decoding step.
+        """
+        _, _, num_channels_in = tgt_ids_BxTxC.shape
+        assert num_channels_in == self.num_channels, "Input channels mismatch"
+
+        # Embeddings
+        x = None
+        for i in range(self.num_channels):
+            channel_tokens = tgt_ids_BxTxC[..., i]
+            channel_embed = self.embeddings[i](channel_tokens)
+            x = channel_embed if x is None else x + channel_embed
+
+        for i, layer in enumerate(self.layers):
+            self_cache = state.self_attn_cache[i]
+            cross_cache = state.cross_attn_cache[i]
+            x = layer(
+                x,
+                state,
+                self_attn_cache=self_cache,
+                cross_attn_cache=cross_cache,
+                prefill=True,
+            )
+
+        # Final Norm
+        x = self.norm(x)
+        logits_BxTxCxV = self.logits_dense(x)
+
+        return logits_BxTxCxV.to(torch.float32)
+
+
+class DiaModel(
+    nn.Module,
+    PyTorchModelHubMixin,
+    repo_url="https://github.com/nari-labs/dia",
+    pipeline_tag="text-to-speech",
+    license="apache-2.0",
+    coders={
+        DiaConfig: (
+            lambda x: x.model_dump(),
+            lambda data: DiaConfig.model_validate(data),
+        ),
+    },
+):
+    """PyTorch Dia Model using DenseGeneral."""
+
+    def __init__(self, config: DiaConfig, compute_dtype: torch.dtype):
+        super().__init__()
+        self.config = config
+        self.encoder = Encoder(config, compute_dtype)
+        self.decoder = Decoder(config, compute_dtype)

dia/model.py

@@ -0,0 +1,802 @@
+import time
+from enum import Enum
+from typing import Callable
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+import torchaudio
+
+from .audio import apply_audio_delay, build_delay_indices, build_revert_indices, revert_audio_delay
+from .config import DiaConfig
+from .layers import DiaModel
+from .state import DecoderInferenceState, DecoderOutput, EncoderInferenceState
+
+
+DEFAULT_SAMPLE_RATE = 44100
+SAMPLE_RATE_RATIO = 512
+
+
+def _get_default_device():
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+        return torch.device("mps")
+    return torch.device("cpu")
+
+
+def _sample_next_token(
+    logits_BCxV: torch.Tensor,
+    temperature: float,
+    top_p: float,
+    top_k: int | None,
+    audio_eos_value: int,
+) -> torch.Tensor:
+    if temperature == 0.0:
+        return torch.argmax(logits_BCxV, dim=-1)
+
+    logits_BCxV = logits_BCxV / temperature
+
+    if audio_eos_value is not None and audio_eos_value >= 0:
+        top_logit_indices_BC = torch.argmax(logits_BCxV, dim=-1)
+        eos_not_highest_mask_BC = top_logit_indices_BC != audio_eos_value
+        mask_eos_unless_highest_BCxV = torch.zeros_like(logits_BCxV, dtype=torch.bool)
+        mask_eos_unless_highest_BCxV[eos_not_highest_mask_BC, audio_eos_value] = True
+        logits_BCxV = logits_BCxV.masked_fill(mask_eos_unless_highest_BCxV, -torch.inf)
+        eos_highest_mask_BC = top_logit_indices_BC == audio_eos_value
+        mask_eos_highest_BCxV = torch.zeros_like(logits_BCxV, dtype=torch.bool)
+        mask_eos_highest_BCxV[eos_highest_mask_BC, :audio_eos_value] = True
+        logits_BCxV = logits_BCxV.masked_fill(mask_eos_highest_BCxV, -torch.inf)
+
+    if top_k is not None:
+        _, top_k_indices_BCxV = torch.topk(logits_BCxV, k=top_k, dim=-1)
+        mask = torch.ones_like(logits_BCxV, dtype=torch.bool)
+        mask = mask.scatter(dim=-1, index=top_k_indices_BCxV, value=False)
+        logits_BCxV = logits_BCxV.masked_fill(mask, -torch.inf)
+
+    if top_p < 1.0:
+        probs_BCxV = torch.softmax(logits_BCxV, dim=-1)
+        sorted_probs_BCxV, sorted_indices_BCxV = torch.sort(probs_BCxV, dim=-1, descending=True)
+        cumulative_probs_BCxV = torch.cumsum(sorted_probs_BCxV, dim=-1)
+
+        sorted_indices_to_remove_BCxV = cumulative_probs_BCxV > top_p
+        sorted_indices_to_remove_BCxV = torch.roll(sorted_indices_to_remove_BCxV, shifts=1, dims=-1)
+        sorted_indices_to_remove_BCxV[..., 0] = torch.zeros_like(sorted_indices_to_remove_BCxV[..., 0])
+
+        indices_to_remove_BCxV = torch.zeros_like(sorted_indices_to_remove_BCxV)
+        indices_to_remove_BCxV = indices_to_remove_BCxV.scatter(
+            dim=-1, index=sorted_indices_BCxV, src=sorted_indices_to_remove_BCxV
+        )
+        logits_BCxV = logits_BCxV.masked_fill(indices_to_remove_BCxV, -torch.inf)
+
+    final_probs_BCxV = torch.softmax(logits_BCxV, dim=-1)
+
+    sampled_indices_BC = torch.multinomial(final_probs_BCxV, num_samples=1)
+    sampled_indices_C = sampled_indices_BC.squeeze(-1)
+    return sampled_indices_C
+
+
+class ComputeDtype(str, Enum):
+    FLOAT32 = "float32"
+    FLOAT16 = "float16"
+    BFLOAT16 = "bfloat16"
+
+    def to_dtype(self) -> torch.dtype:
+        if self == ComputeDtype.FLOAT32:
+            return torch.float32
+        elif self == ComputeDtype.FLOAT16:
+            return torch.float16
+        elif self == ComputeDtype.BFLOAT16:
+            return torch.bfloat16
+        else:
+            raise ValueError(f"Unsupported compute dtype: {self}")
+
+
+class Dia:
+    def __init__(
+        self,
+        config: DiaConfig,
+        compute_dtype: str | ComputeDtype = ComputeDtype.FLOAT32,
+        device: torch.device | None = None,
+        load_dac: bool = True,
+    ):
+        """Initializes the Dia model.
+
+        Args:
+            config: The configuration object for the model.
+            compute_dtype: The computation dtype to use.
+            device: The device to load the model onto. If None, will automatically select the best available device.
+            load_dac: Whether to load the DAC model.
+
+        Raises:
+            RuntimeError: If there is an error loading the DAC model.
+        """
+        super().__init__()
+        self.config = config
+        self.device = device if device is not None else _get_default_device()
+        if isinstance(compute_dtype, str):
+            compute_dtype = ComputeDtype(compute_dtype)
+        self.compute_dtype = compute_dtype.to_dtype()
+        self.model: DiaModel = DiaModel(config, self.compute_dtype)
+        self.dac_model = None
+        self._compiled_step = None
+        self.load_dac = load_dac
+
+        if not self.load_dac:
+            print("Warning: DAC model will not be loaded. This is not recommended.")
+
+        if torch.cuda.is_available():
+            torch.backends.cuda.matmul.allow_tf32 = True
+
+    @classmethod
+    def from_local(
+        cls,
+        config_path: str,
+        checkpoint_path: str,
+        compute_dtype: str | ComputeDtype = ComputeDtype.FLOAT32,
+        device: torch.device | None = None,
+        load_dac: bool = True,
+    ) -> "Dia":
+        """Loads the Dia model from local configuration and checkpoint files.
+
+        Args:
+            config_path: Path to the configuration JSON file.
+            checkpoint_path: Path to the model checkpoint (.pth) file.
+            compute_dtype: The computation dtype to use.
+            device: The device to load the model onto. If None, will automatically select the best available device.
+            load_dac: Whether to load the DAC model.
+
+        Returns:
+            An instance of the Dia model loaded with weights and set to eval mode.
+
+        Raises:
+            FileNotFoundError: If the config or checkpoint file is not found.
+            RuntimeError: If there is an error loading the checkpoint.
+        """
+        config = DiaConfig.load(config_path)
+        if config is None:
+            raise FileNotFoundError(f"Config file not found at {config_path}")
+
+        dia = cls(config, compute_dtype, device, load_dac)
+
+        try:
+            state_dict = torch.load(checkpoint_path, map_location=dia.device)
+            dia.model.load_state_dict(state_dict)
+        except FileNotFoundError:
+            raise FileNotFoundError(f"Checkpoint file not found at {checkpoint_path}")
+        except Exception as e:
+            raise RuntimeError(f"Error loading checkpoint from {checkpoint_path}") from e
+
+        dia.model.to(dia.device)
+        dia.model.eval()
+        if load_dac:
+            dia._load_dac_model()
+        return dia
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        model_name: str = "nari-labs/Dia-1.6B-0626",
+        compute_dtype: str | ComputeDtype = ComputeDtype.FLOAT32,
+        device: torch.device | None = None,
+        load_dac: bool = True,
+    ) -> "Dia":
+        """Loads the Dia model from a Hugging Face Hub repository.
+
+        Downloads the configuration and checkpoint files from the specified
+        repository ID and then loads the model.
+
+        Args:
+            model_name: The Hugging Face Hub repository ID (e.g., "nari-labs/Dia-1.6B-0626").
+            compute_dtype: The computation dtype to use.
+            device: The device to load the model onto. If None, will automatically select the best available device.
+            load_dac: Whether to load the DAC model.
+
+        Returns:
+            An instance of the Dia model loaded with weights and set to eval mode.
+
+        Raises:
+            FileNotFoundError: If config or checkpoint download/loading fails.
+            RuntimeError: If there is an error loading the checkpoint.
+        """
+        if isinstance(compute_dtype, str):
+            compute_dtype = ComputeDtype(compute_dtype)
+
+        # Load model directly using DiaModel's from_pretrained which handles HF download
+        try:
+            loaded_model = DiaModel.from_pretrained(model_name, compute_dtype=compute_dtype.to_dtype())
+        except Exception as e:
+            raise RuntimeError(f"Error loading model from Hugging Face Hub ({model_name})") from e
+
+        config = loaded_model.config  # Get config from the loaded model
+        dia = cls(config, compute_dtype, device, load_dac)
+
+        dia.model = loaded_model  # Assign the already loaded model
+        dia.model.to(dia.device)
+        dia.model.eval()
+        if load_dac:
+            dia._load_dac_model()
+        return dia
+
+    def _load_dac_model(self):
+        """Loads the Descript Audio Codec (DAC) model.
+
+        Downloads the DAC model if necessary and loads it onto the specified device.
+        Sets the DAC model to evaluation mode.
+
+        Raises:
+            RuntimeError: If downloading or loading the DAC model fails.
+        """
+        import dac
+
+        try:
+            dac_model_path = dac.utils.download()
+            dac_model = dac.DAC.load(dac_model_path).to(self.device)
+            dac_model.eval()  # Ensure DAC is in eval mode
+        except Exception as e:
+            raise RuntimeError("Failed to load DAC model") from e
+        self.dac_model = dac_model
+
+    def _encode_text(self, text: str) -> torch.Tensor:
+        """Encodes the input text string into a tensor of token IDs using byte-level encoding.
+
+        Special tokens [S1] and [S2] are replaced by their byte values. The resulting
+        sequence is truncated to the maximum configured text length.
+
+        Args:
+            text: The input text string.
+
+        Returns:
+            A tensor containing the encoded byte token IDs.
+        """
+        max_len = self.config.encoder_config.max_position_embeddings
+
+        byte_text = text.encode("utf-8")
+        # Replace special tokens with their byte values if needed by the specific tokenizer/config
+        # Assuming byte values 1 and 2 are correct placeholders based on original code
+        replaced_bytes = byte_text.replace(b"[S1]", b"\x01").replace(b"[S2]", b"\x02")
+        text_tokens = list(replaced_bytes)
+        return torch.tensor(
+            text_tokens[:max_len],
+            dtype=torch.long,
+            device=self.device,
+        )
+
+    def _pad_text_input(self, text_tokens: list[torch.Tensor]) -> torch.Tensor:
+        """Pads the text input to the maximum length."""
+        text_pad_value = 0
+        max_len = self.config.encoder_config.max_position_embeddings
+        batch_size = len(text_tokens)
+
+        src_tokens = torch.full(
+            (batch_size, 1, max_len),
+            fill_value=text_pad_value,
+            dtype=torch.long,
+            device=self.device,
+        )
+        for i in range(batch_size):
+            current_len = len(text_tokens[i])
+            src_tokens[i, 0, :current_len] = text_tokens[i]
+        return src_tokens
+
+    def _prepare_audio_prompt(self, audio_prompts: list[torch.Tensor | None]) -> tuple[torch.Tensor, list[int]]:
+        """Prepares the audio prompt tensor for the decoder.
+
+        Handles padding, adds the beginning-of-sequence (BOS) token, applies the
+        delay pattern, and determines the number of prefill steps for each item
+        in the batch.
+
+        Args:
+            audio_prompts: A list of audio prompt tensors (encoded DAC frames) or None.
+                           Each tensor should have shape [T, C].
+
+        Returns:
+            A tuple containing:
+                - delayed_batch (torch.Tensor): The prepared audio prompt tensor with
+                  delays applied, shape [B, T_max_padded, C].
+                - prefill_steps (list[int]): A list containing the number of valid
+                  tokens (including BOS) for each prompt in the batch.
+        """
+        num_channels = self.config.decoder_config.num_channels
+        audio_bos_value = self.config.bos_token_id
+        delay_pattern = self.config.delay_pattern
+        max_delay_pattern = max(delay_pattern)
+        batch_size = len(audio_prompts)
+
+        max_len = max(p.shape[0] if p is not None else 0 for p in audio_prompts) + max_delay_pattern
+        prefill_steps = []
+
+        prefill = torch.full(
+            (batch_size, max_len, num_channels),
+            fill_value=-1,
+            dtype=torch.int,
+            device=self.device,
+        )
+
+        prefill[:, 0, :] = audio_bos_value
+
+        for i in range(batch_size):
+            prompt = audio_prompts[i]
+            if prompt is not None:
+                prompt = prompt.to(device=self.device, dtype=torch.int)
+                prefill[i, 1 : prompt.shape[0] + 1, :] = prompt
+                prefill_steps.append(prompt.shape[0] + 1)
+            else:
+                prefill_steps.append(1)
+
+        delay_precomp = build_delay_indices(
+            B=batch_size,
+            T=max_len,
+            C=num_channels,
+            delay_pattern=delay_pattern,
+        )
+
+        delayed_batch = apply_audio_delay(
+            audio_BxTxC=prefill,
+            pad_value=-1,
+            bos_value=audio_bos_value,
+            precomp=delay_precomp,
+        )
+
+        return delayed_batch, prefill_steps
+
+    def _prepare_generation(
+        self,
+        text: torch.Tensor,
+        audio_prompts: list[torch.Tensor | None],
+        max_tokens: int | None = None,
+        attn_fn: Callable = F.scaled_dot_product_attention,
+    ):
+        """Initializes the model state for generation.
+
+        Encodes the text input (conditional and unconditional), prepares the
+        encoder and decoder states (including KV caches and cross-attention),
+        prepares the audio prompt, and performs the initial decoder prefill steps
+        based on the audio prompts.
+
+        Args:
+            text: The padded text input tensor, shape [B, 1, T_text].
+            audio_prompts: A list of prepared audio prompt tensors or None.
+
+        Returns:
+            A tuple containing:
+                - dec_state (DecoderInferenceState): The initialized decoder state.
+                - dec_output (DecoderOutput): The initialized decoder output manager,
+                  containing the prefilled audio tokens.
+        """
+        batch_size = text.shape[0]
+
+        enc_input_uncond = torch.zeros_like(text)
+        enc_input_cond = text
+        stacked_inputs = torch.stack([enc_input_uncond, enc_input_cond], dim=1)
+        enc_input = stacked_inputs.view(2 * batch_size, -1)
+
+        enc_state = EncoderInferenceState.new(self.config, enc_input_cond)
+        encoder_out = self.model.encoder(enc_input, enc_state)
+
+        dec_cross_attn_cache = self.model.decoder.precompute_cross_attn_cache(encoder_out)
+        dec_state = DecoderInferenceState.new(
+            self.config,
+            enc_state,
+            encoder_out,
+            dec_cross_attn_cache,
+            self.compute_dtype,
+            max_generation_length=max_tokens,
+        )
+        prefill, prefill_steps = self._prepare_audio_prompt(audio_prompts)
+
+        dec_output = DecoderOutput.new(batch_size, self.config, self.device)
+        dec_output.prefill(prefill, prefill_steps)
+
+        dec_step = min(prefill_steps) - 1
+        if dec_step > 0:
+            dec_state.prepare_step(0, dec_step)
+            tokens_BxTxC = dec_output.get_tokens_at(0, dec_step).repeat_interleave(2, dim=0)
+            self.model.decoder.forward(tokens_BxTxC, dec_state)
+
+        return dec_state, dec_output
+
+    def _decoder_step(
+        self,
+        tokens_Bx1xC: torch.Tensor,
+        dec_state: DecoderInferenceState,
+        cfg_scale: float,
+        temperature: float,
+        top_p: float,
+        top_k: int,
+        current_idx: int,
+    ) -> torch.Tensor:
+        """Performs a single step of the decoder inference.
+
+        Takes the tokens from the previous step, runs them through the decoder
+        (for both conditional and unconditional paths), applies classifier-free
+        guidance (CFG), samples the next token using temperature, top-p, and top-k
+        sampling, and applies constraints (e.g., preventing EOS in certain channels).
+
+        Args:
+            tokens_Bx1xC: The input tokens for the current step, shape [2*B, 1, C].
+                         Repeated for CFG (unconditional and conditional).
+            dec_state: The current state of the decoder (KV caches, etc.).
+            cfg_scale: The scale factor for classifier-free guidance.
+            temperature: The temperature for sampling.
+            top_p: The cumulative probability threshold for top-p sampling.
+            top_k: The number of top logits to consider for top-k sampling.
+            current_idx: The current generation step index.
+
+        Returns:
+            torch.Tensor: The sampled next tokens for each item in the batch,
+                          shape [B, C].
+        """
+        B = tokens_Bx1xC.shape[0] // 2
+
+        audio_eos_value = self.config.eos_token_id
+        logits_Bx1xCxV = self.model.decoder.decode_step(tokens_Bx1xC, dec_state, current_idx)
+
+        logits_last_2BxCxV = logits_Bx1xCxV[:, -1]
+        logits_last_Bx2xCxV = logits_last_2BxCxV.view(B, 2, *logits_last_2BxCxV.shape[1:])
+
+        uncond_logits_BxCxV = logits_last_Bx2xCxV[:, 0, :, :]  # Shape [B, C, V]
+        cond_logits_BxCxV = logits_last_Bx2xCxV[:, 1, :, :]  # Shape [B, C, V]
+        logits_BxCxV = cond_logits_BxCxV + cfg_scale * (cond_logits_BxCxV - uncond_logits_BxCxV)
+
+        _, top_k_indices_BxCxk = torch.topk(logits_BxCxV, k=top_k, dim=-1)
+        mask_BxCxV = torch.ones_like(logits_BxCxV, dtype=torch.bool)
+        mask_BxCxV = mask_BxCxV.scatter(dim=-1, index=top_k_indices_BxCxk, value=False)
+        logits_BxCxV = cond_logits_BxCxV.masked_fill(mask_BxCxV, -torch.inf)
+
+        logits_BxCxV[:, :, audio_eos_value + 1 :] = torch.full_like(
+            logits_BxCxV[:, :, audio_eos_value + 1 :],
+            fill_value=-torch.inf,
+        )
+        logits_BxCxV[:, 1:, audio_eos_value:] = torch.full_like(
+            logits_BxCxV[:, 1:, audio_eos_value:],
+            fill_value=-torch.inf,
+        )
+
+        flat_logits_BCxV = logits_BxCxV.view(B * self.config.decoder_config.num_channels, -1)
+
+        pred_BC = _sample_next_token(
+            flat_logits_BCxV.float(),
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
+            audio_eos_value=audio_eos_value,
+        )
+
+        pred_BxC = pred_BC.view(B, self.config.decoder_config.num_channels)
+        return pred_BxC
+
+    def _generate_output(self, generated_codes: torch.Tensor, lengths_Bx: torch.Tensor) -> list[np.ndarray]:
+        """Converts generated delayed codes into audio waveforms.
+
+        Reverts the delay pattern applied during generation, decodes the resulting
+        codebook using the DAC model (if loaded), and returns a list of audio
+        waveforms as NumPy arrays. If DAC is not loaded, returns the raw codebook indices.
+
+        Args:
+            generated_codes: The tensor of generated audio codes with delays,
+                             shape [B, T_gen, C].
+            lengths_Bx: A tensor containing the valid length of generated codes
+                        (excluding padding and BOS/EOS markers) for each item
+                        in the batch, shape [B].
+
+        Returns:
+            A list of NumPy arrays, where each array represents the generated audio
+            waveform for one item in the batch. If DAC is not loaded, returns the
+            raw, reverted codebook indices as NumPy arrays.
+        """
+        num_channels = self.config.decoder_config.num_channels
+        batch_size = generated_codes.shape[0]
+        seq_length = generated_codes.shape[1]
+        delay_pattern = self.config.delay_pattern
+        audio_pad_value = self.config.pad_token_id
+        max_delay_pattern = max(delay_pattern)
+
+        revert_precomp = build_revert_indices(
+            B=batch_size,
+            T=seq_length,
+            C=num_channels,
+            delay_pattern=delay_pattern,
+        )
+
+        codebook = revert_audio_delay(
+            audio_BxTxC=generated_codes,
+            pad_value=audio_pad_value,
+            precomp=revert_precomp,
+            T=seq_length,
+        )[:, :-max_delay_pattern, :]
+
+        min_valid_index = 0
+        max_valid_index = 1023
+        invalid_mask = (codebook < min_valid_index) | (codebook > max_valid_index)
+        codebook[invalid_mask] = 0
+
+        audios = []
+
+        if self.load_dac:
+            for i in range(batch_size):
+                audio = self._decode(codebook[i, : lengths_Bx[i], :])
+                audio_np = audio.cpu().numpy()
+                audios.append(audio_np)
+        else:
+            for i in range(batch_size):
+                audios.append(codebook[i, : lengths_Bx[i], :].cpu().numpy())
+        return audios
+
+    @torch.no_grad()
+    @torch.inference_mode()
+    def _encode(self, audio: torch.Tensor) -> torch.Tensor:
+        """
+        Encodes the given audio waveform into a tensor of DAC codebook indices
+        """
+        audio = audio.unsqueeze(0)
+        audio_data = self.dac_model.preprocess(audio, DEFAULT_SAMPLE_RATE)
+        _, encoded_frame, _, _, _ = self.dac_model.encode(audio_data)
+        encoded_frame: torch.Tensor
+        return encoded_frame.squeeze(0).transpose(0, 1)
+
+    @torch.no_grad()
+    @torch.inference_mode()
+    def _decode(self, audio_codes: torch.Tensor) -> torch.Tensor:
+        """
+        Decodes the given frames into an output audio waveform
+        """
+        audio_codes = audio_codes.unsqueeze(0).transpose(1, 2)
+        audio_values, _, _ = self.dac_model.quantizer.from_codes(audio_codes)
+        audio_values = self.dac_model.decode(audio_values)
+        audio_values: torch.Tensor
+        return audio_values.squeeze()
+
+    def load_audio(self, audio_path: str) -> torch.Tensor:
+        """Loads and preprocesses an audio file for use as a prompt.
+
+        Loads the audio file, resamples it to the target sample rate if necessary,
+        preprocesses it using the DAC model's preprocessing, and encodes it into
+        DAC codebook indices.
+
+        Args:
+            audio_path: Path to the audio file.
+
+        Returns:
+            torch.Tensor: The encoded audio prompt as DAC codebook indices,
+                          shape [T, C].
+
+        Raises:
+            RuntimeError: If the DAC model is not loaded (`load_dac=False` during init).
+            FileNotFoundError: If the audio file cannot be found.
+            Exception: If there's an error during loading or processing.
+        """
+        if self.dac_model is None:
+            raise RuntimeError("DAC model is required for loading audio prompts but was not loaded.")
+        audio, sr = torchaudio.load(audio_path, channels_first=True)  # C, T
+        if sr != DEFAULT_SAMPLE_RATE:
+            audio = torchaudio.functional.resample(audio, sr, DEFAULT_SAMPLE_RATE)
+        # Convert to mono if stereo
+        if audio.shape[0] > 1:
+            audio = torch.mean(audio, dim=0, keepdim=True)  # Average channels to get mono
+        return self._encode(audio.to(self.device))
+
+    def save_audio(self, path: str, audio: np.ndarray):
+        """Saves the generated audio waveform to a file.
+
+        Uses the soundfile library to write the NumPy audio array to the specified
+        path with the default sample rate.
+
+        Args:
+            path: The path where the audio file will be saved.
+            audio: The audio waveform as a NumPy array.
+        """
+        import soundfile as sf
+
+        sf.write(path, audio, DEFAULT_SAMPLE_RATE)
+
+    @torch.inference_mode()
+    def generate(
+        self,
+        text: str | list[str],
+        max_tokens: int = 3072,
+        cfg_scale: float = 3.0,
+        temperature: float = 1.2,
+        top_p: float = 0.95,
+        use_torch_compile: bool = False,
+        cfg_filter_top_k: int = 45,
+        audio_prompt: list[str | torch.Tensor | None] | str | torch.Tensor | None = None,
+        audio_prompt_path: list[str | torch.Tensor | None] | str | torch.Tensor | None = None,
+        use_cfg_filter: bool | None = None,
+        verbose: bool = False,
+    ) -> np.ndarray | list[np.ndarray]:
+        """Generates audio corresponding to the input text.
+
+        Args:
+            text: The input text prompt, or a list of text prompts for batch generation.
+            max_tokens: The maximum number of audio tokens to generate per prompt.
+                        Defaults to the model's configured audio length if None.
+            cfg_scale: The scale factor for classifier-free guidance (CFG). Higher values
+                       lead to stronger guidance towards the text prompt.
+            temperature: The temperature for sampling. Higher values increase randomness.
+            top_p: The cumulative probability threshold for nucleus (top-p) sampling.
+            use_torch_compile: Whether to compile the generation steps using torch.compile.
+                               Can significantly speed up generation after the initial
+                               compilation overhead. Defaults to False.
+            cfg_filter_top_k: The number of top logits to consider during CFG filtering.
+                              (Note: This parameter name might be slightly misleading based
+                              on the code; it's used in the `_sample_next_token` function.)
+            audio_prompt: An audio prompt or list of prompts to condition the generation.
+                          Can be a file path (str), a pre-loaded tensor (DAC codes), or None.
+                          If a list, its length must match the batch size of the text input.
+            audio_prompt_path: (Deprecated) Use `audio_prompt` instead.
+            use_cfg_filter: (Deprecated) This parameter is no longer used.
+            verbose: If True, prints progress information during generation, including
+                     speed metrics.
+
+        Returns:
+            If a single text prompt was provided, returns a NumPy array containing the
+            generated audio waveform.
+            If a list of text prompts was provided, returns a list of NumPy arrays,
+            each corresponding to a prompt in the input list. Returns None for a
+            sequence if no audio was generated for it.
+        """
+        batch_size = len(text) if isinstance(text, list) else 1
+        audio_eos_value = self.config.eos_token_id
+        audio_pad_value = self.config.pad_token_id
+        delay_pattern = self.config.delay_pattern
+        max_delay_pattern = max(delay_pattern)
+        delay_pattern_Cx = torch.tensor(delay_pattern, device=self.device, dtype=torch.long)
+        self.model.eval()
+
+        if audio_prompt_path:
+            print("Warning: audio_prompt_path is deprecated. Use audio_prompt instead.")
+            audio_prompt = audio_prompt_path
+        if use_cfg_filter is not None:
+            print("Warning: use_cfg_filter is deprecated.")
+
+        if verbose:
+            total_start_time = time.time()
+
+        if use_torch_compile and not hasattr(self, "_compiled"):
+            # Compilation can take about a minute.
+            self._prepare_generation = torch.compile(self._prepare_generation, dynamic=True, fullgraph=True)
+            self._decoder_step = torch.compile(self._decoder_step, fullgraph=True, mode="max-autotune")
+            self._compiled = True
+
+        if isinstance(audio_prompt, list):
+            audio_prompt = [self.load_audio(p) if isinstance(p, str) else p for p in audio_prompt]
+        elif isinstance(audio_prompt, str):
+            audio_prompt = [self.load_audio(audio_prompt)]
+        elif isinstance(audio_prompt, torch.Tensor):
+            audio_prompt = [audio_prompt]
+        elif audio_prompt is None:
+            audio_prompt = [None] * batch_size
+
+        assert len(audio_prompt) == batch_size, "Number of audio prompts must match batch size"
+
+        if isinstance(text, list):
+            text = [self._encode_text(t) for t in text]
+        else:
+            text = [self._encode_text(text)]
+        text = self._pad_text_input(text)
+
+        dec_state, dec_output = self._prepare_generation(text, audio_prompt, max_tokens=max_tokens)
+        dec_step = min(dec_output.prefill_steps) - 1
+        current_idx = torch.tensor([dec_step], device=self.device)
+
+        eos_detected_Bx = torch.zeros((batch_size,), dtype=torch.bool, device=self.device)
+        eos_countdown_Bx = torch.full((batch_size,), -1, dtype=torch.long, device=self.device)
+        finished_step_Bx = torch.full((batch_size,), -1, dtype=torch.long, device=self.device)
+
+        bos_over = False
+
+        if verbose:
+            print("generate: starting generation loop")
+            if use_torch_compile:
+                print("generate: using use_torch_compile=True, the first step may be slow")
+            start_time = time.time()
+
+        # --- Generation Loop ---
+        while dec_step < max_tokens:
+            if (eos_countdown_Bx == 0).all():
+                break
+
+            current_step_idx = dec_step + 1
+            torch.compiler.cudagraph_mark_step_begin()
+            dec_state.prepare_step(dec_step)
+            tokens_Bx1xC = dec_output.get_tokens_at(dec_step).repeat_interleave(2, dim=0)  # Repeat for CFG
+
+            pred_BxC = self._decoder_step(
+                tokens_Bx1xC,
+                dec_state,
+                cfg_scale,
+                temperature,
+                top_p,
+                cfg_filter_top_k,
+                current_idx,
+            )
+
+            current_idx += 1
+
+            active_mask_Bx = eos_countdown_Bx != 0
+            eos_trigger_Bx = torch.zeros_like(active_mask_Bx)
+            if active_mask_Bx.any():
+                is_eos_token = (~eos_detected_Bx[active_mask_Bx]) & (pred_BxC[active_mask_Bx, 0] == audio_eos_value)
+                is_max_len = current_step_idx >= max_tokens - max_delay_pattern
+                eos_trigger_Bx[active_mask_Bx] = is_eos_token | is_max_len
+            eos_detected_Bx |= eos_trigger_Bx
+            start_countdown_mask_Bx = eos_trigger_Bx & (eos_countdown_Bx < 0)
+            if start_countdown_mask_Bx.any():
+                eos_countdown_Bx[start_countdown_mask_Bx] = max_delay_pattern
+                finished_step_Bx[start_countdown_mask_Bx] = current_step_idx
+
+            padding_mask_Bx = eos_countdown_Bx > 0
+            if padding_mask_Bx.any():
+                pred_active_BxC = pred_BxC[padding_mask_Bx].clone()
+                countdown_active_Bx = eos_countdown_Bx[padding_mask_Bx]
+                step_after_eos_Bx = max_delay_pattern - countdown_active_Bx
+                step_after_eos_Bx_ = step_after_eos_Bx.unsqueeze(1)
+                delay_pattern_Cx_ = delay_pattern_Cx.unsqueeze(0)
+                eos_mask_NxC = step_after_eos_Bx_ == delay_pattern_Cx_
+                pad_mask_NxC = step_after_eos_Bx_ > delay_pattern_Cx_
+                pred_active_BxC[eos_mask_NxC] = audio_eos_value
+                pred_active_BxC[pad_mask_NxC] = audio_pad_value
+                pred_BxC[padding_mask_Bx] = pred_active_BxC
+                eos_countdown_Bx[padding_mask_Bx] -= 1
+
+            # --- Update BOS flag (Original) ---
+            if not bos_over:
+                bos_over = all(
+                    dec_step - prefill_step > max_delay_pattern for prefill_step in dec_output.prefill_steps
+                )
+
+            dec_output.update_one(pred_BxC, current_step_idx, not bos_over)
+
+            dec_step += 1
+
+            if verbose and dec_step % 86 == 0:
+                duration = time.time() - start_time
+                if duration > 0:
+                    print(
+                        f"generate step {dec_step}: speed={86 * batch_size / duration:.3f} tokens/s, realtime factor={batch_size / duration:.3f}x"
+                    )
+                start_time = time.time()
+
+        # --- Finalize and Extract Output ---
+        final_step = dec_step + 1
+
+        finished_step_Bx[finished_step_Bx == -1] = final_step - max_delay_pattern
+
+        prefill_steps_tensor = torch.tensor(dec_output.prefill_steps, device=self.device)
+        lengths_Bx = finished_step_Bx - prefill_steps_tensor
+        lengths_Bx = torch.clamp(lengths_Bx, min=0)
+
+        max_len = lengths_Bx.max().item() + max_delay_pattern
+        outputs = []
+
+        if max_len > 0:
+            num_channels = self.config.decoder_config.num_channels
+            audio_pad_value = self.config.pad_token_id
+            generated_codes = torch.full(
+                (batch_size, max_len, num_channels),
+                fill_value=audio_pad_value,
+                dtype=torch.long,
+                device=self.device,
+            )
+
+            for i in range(batch_size):
+                start_step = dec_output.prefill_steps[i]
+                actual_len = lengths_Bx[i].item() + max_delay_pattern
+                if actual_len > 0:
+                    tokens_to_copy = dec_output.generated_tokens[i, start_step : start_step + actual_len, :]
+                    generated_codes[i, :actual_len, :] = tokens_to_copy
+
+            if verbose:
+                avg_steps = lengths_Bx.float().mean().item()
+                total_duration = time.time() - total_start_time
+                print(f"generate: avg steps={avg_steps:.1f}, total duration={total_duration:.3f}s")
+
+            del dec_state
+
+            outputs = self._generate_output(generated_codes, lengths_Bx)
+        else:
+            print("Warning: Nothing generated for any sequence in the batch.")
+            outputs = [None] * batch_size
+
+        return outputs if batch_size > 1 else outputs[0]

dia/state.py

@@ -0,0 +1,217 @@
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+
+from .config import DiaConfig
+
+
+def create_attn_mask(
+    q_padding_mask_1d: torch.Tensor,
+    k_padding_mask_1d: torch.Tensor,
+    device: torch.device,
+    is_causal: bool = False,
+) -> torch.Tensor:
+    """
+    Creates the attention mask (self or cross) mimicking JAX segment ID logic.
+    """
+    # B1, Tq = q_padding_mask_1d.shape
+    # B2, Tk = k_padding_mask_1d.shape
+
+    p_mask_q = q_padding_mask_1d.unsqueeze(2)  # Shape [B, Tq, 1]
+    p_mask_k = k_padding_mask_1d.unsqueeze(1)  # Shape [B, 1, Tk]
+
+    # Condition A: Non-padding query attends to non-padding key
+    non_pad_attends_non_pad = p_mask_q & p_mask_k  # Shape [B, Tq, Tk]
+
+    # Condition B: Padding query attends to padding key
+    pad_attends_pad = (~p_mask_q) & (~p_mask_k)  # Shape [B, Tq, Tk]
+
+    # Combine: True if padding status is compatible (both non-pad OR both pad)
+    mask = non_pad_attends_non_pad | pad_attends_pad  # Shape [B, Tq, Tk]
+
+    if is_causal:
+        # assert Tq == Tk, "Causal mask requires query and key sequence lengths to be equal"
+        causal_mask_2d = torch.tril(torch.ones_like(mask[0], dtype=torch.bool, device=device))  # Shape [B, Tq, Tk]
+        causal_mask = mask & causal_mask_2d  # Shape [B, Tq, Tk]
+        return causal_mask.unsqueeze(1)  # Shape [B, 1, Tq, Tk]
+    else:
+        return mask.unsqueeze(1)  # Shape [B, 1, Tq, Tk]
+
+
+@dataclass
+class EncoderInferenceState:
+    """Parameters specifically for encoder inference."""
+
+    max_seq_len: int
+    device: torch.device
+    positions: torch.Tensor
+    padding_mask: torch.Tensor
+    attn_mask: torch.Tensor
+
+    @classmethod
+    def new(cls, config: DiaConfig, cond_src: torch.Tensor) -> "EncoderInferenceState":
+        """Creates EtorchrInferenceParams from DiaConfig and a device."""
+        device = cond_src.device
+
+        positions = torch.arange(
+            config.encoder_config.max_position_embeddings, dtype=torch.float32, device=device
+        ).unsqueeze(0)
+        padding_mask = (cond_src.squeeze(1) != 0).to(device).repeat_interleave(2, dim=0)
+        attn_mask = create_attn_mask(padding_mask, padding_mask, device, is_causal=False)
+
+        return cls(
+            max_seq_len=config.encoder_config.max_position_embeddings,
+            device=device,
+            positions=positions,
+            padding_mask=padding_mask,
+            attn_mask=attn_mask,
+        )
+
+
+class KVCache(torch.nn.Module):
+    k: torch.Tensor
+    v: torch.Tensor
+
+    def __init__(
+        self,
+        batch_size: int,
+        num_heads: int,
+        max_len: int,
+        head_dim: int,
+        dtype: torch.dtype,
+        device: torch.device,
+        k: torch.Tensor | None = None,
+        v: torch.Tensor | None = None,
+    ):
+        k = torch.zeros((2 * batch_size, num_heads, max_len, head_dim), dtype=dtype, device=device) if k is None else k
+        v = torch.zeros((2 * batch_size, num_heads, max_len, head_dim), dtype=dtype, device=device) if v is None else v
+        super().__init__()
+
+        self.register_buffer("k", k)
+        self.register_buffer("v", v)
+
+    @classmethod
+    def from_kv(cls, k: torch.Tensor, v: torch.Tensor) -> "KVCache":
+        return cls(
+            batch_size=k.shape[0] // 2,
+            num_heads=k.shape[1],
+            max_len=k.shape[2],
+            head_dim=k.shape[3],
+            dtype=k.dtype,
+            device=k.device,
+            k=k,
+            v=v,
+        )
+
+    def update(self, k: torch.Tensor, v: torch.Tensor, current_idx: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        k_out, v_out = self.k, self.v
+        k_out[:, :, current_idx, :] = k
+        v_out[:, :, current_idx, :] = v
+        return self.k, self.v
+
+    def prefill(self, k: torch.Tensor, v: torch.Tensor):
+        prefill_len = k.shape[2]
+        self.k[:, :, :prefill_len, :] = k
+        self.v[:, :, :prefill_len, :] = v
+
+
+@dataclass
+class DecoderInferenceState:
+    """Parameters specifically for decoder inference."""
+
+    device: torch.device
+    dtype: torch.dtype
+    enc_out: torch.Tensor
+    enc_positions: torch.Tensor
+    dec_positions: torch.Tensor
+    self_attn_cache: list[KVCache]
+    cross_attn_cache: list[KVCache]
+    casual_attn_mask: torch.Tensor
+    cross_attn_mask: torch.Tensor
+
+    @classmethod
+    def new(
+        cls,
+        config: DiaConfig,
+        enc_state: EncoderInferenceState,
+        enc_out: torch.Tensor,
+        dec_cross_attn_cache: list[KVCache],
+        compute_dtype: torch.dtype,
+        max_generation_length: Optional[int] = None,
+    ) -> "DecoderInferenceState":
+        """Creates DecoderInferenceParams from DiaConfig and a device."""
+        device = enc_out.device
+        max_audio_len = max_generation_length or config.decoder_config.max_position_embeddings
+        batch_size = enc_out.shape[0] // 2
+
+        dec_positions = torch.full((2 * batch_size, 1), fill_value=0, dtype=torch.int32, device=device)
+        causal_mask = torch.tril(torch.ones(max_audio_len, max_audio_len, dtype=torch.bool, device=device))
+        dec_mask = torch.ones((2 * batch_size, 1), dtype=torch.bool, device=device)
+        cross_attn_mask = create_attn_mask(dec_mask, enc_state.padding_mask, device, is_causal=False)
+
+        self_attn_cache = [
+            KVCache(
+                batch_size,
+                config.decoder_config.num_key_value_heads,
+                max_audio_len,
+                config.decoder_config.head_dim,
+                compute_dtype,
+                device,
+            )
+            for _ in range(config.decoder_config.num_hidden_layers)
+        ]
+
+        return cls(
+            device=device,
+            dtype=compute_dtype,
+            enc_out=enc_out,
+            enc_positions=enc_state.positions,
+            dec_positions=dec_positions,
+            self_attn_cache=self_attn_cache,
+            cross_attn_cache=dec_cross_attn_cache,
+            casual_attn_mask=causal_mask,
+            cross_attn_mask=cross_attn_mask,
+        )
+
+    def prepare_step(self, step_from: int, step_to: int | None = None) -> None:
+        if step_to is None:
+            step_to = step_from + 1
+        self.dec_positions = torch.arange(step_from, step_to, dtype=torch.int32, device=self.device).unsqueeze(0)
+
+
+@dataclass
+class DecoderOutput:
+    generated_tokens: torch.Tensor
+    prefill_steps: list[int]
+
+    @classmethod
+    def new(cls, batch_size: int, config: DiaConfig, device: torch.device) -> "DecoderOutput":
+        max_audio_len = config.decoder_config.max_position_embeddings
+        return cls(
+            generated_tokens=torch.full(
+                (batch_size, max_audio_len, config.decoder_config.num_channels),
+                fill_value=-1,
+                dtype=torch.int,
+                device=device,
+            ),
+            prefill_steps=[],
+        )
+
+    def get_tokens_at(self, step_from: int, step_to: int | None = None) -> torch.Tensor:
+        if step_to is None:
+            step_to = step_from + 1
+        return self.generated_tokens[:, step_from:step_to, :]
+
+    def update_one(self, dec_out: torch.Tensor, step: int, apply_mask: bool = False):
+        dec_out = dec_out.to(self.generated_tokens.dtype)
+        if apply_mask:
+            mask = self.generated_tokens[:, step, :] == -1
+            self.generated_tokens[:, step, :] = torch.where(mask, dec_out, self.generated_tokens[:, step, :])
+        else:
+            self.generated_tokens[:, step, :] = dec_out
+
+    def prefill(self, dec_out: torch.Tensor, prefill_steps: list[int]):
+        length = dec_out.shape[1]
+        self.generated_tokens[:, :length, :] = dec_out
+        self.prefill_steps = prefill_steps

docker/Dockerfile.cpu

@@ -0,0 +1,48 @@
+# Dockerfile.cpu - CPU-only deployment for DIA
+# --------------------------------------------------
+# Build: docker build . -f docker/Dockerfile.cpu -t dia-cpu
+# Run:   docker run --rm -p 7860:7860 dia-cpu
+
+FROM python:3.10-slim
+
+# Set non-interactive frontend
+ENV DEBIAN_FRONTEND=noninteractive
+
+# Install venv, and system dependencies
+RUN apt-get update && apt-get install -y \
+    python3-venv \
+    libsndfile1 \
+    ffmpeg \
+    curl \
+    && apt-get clean && rm -rf /var/lib/apt/lists/*
+
+# Create non-root user and set up directories
+RUN useradd -m -u 1001 appuser && \
+    mkdir -p /app/outputs /app && \
+    chown -R appuser:appuser /app
+
+USER appuser
+WORKDIR /app
+
+# Copy all code (including pyproject.toml)
+COPY --chown=appuser:appuser . .
+
+# Create and activate virtual environment
+RUN python3 -m venv /app/venv
+ENV PATH="/app/venv/bin:$PATH"
+
+# Install all project dependencies (CPU-only PyTorch)
+RUN pip install --upgrade pip && \
+    pip install torch torchaudio --index-url https://download.pytorch.org/whl/cpu && \
+    pip install --no-cache-dir -e .[dev]
+
+# Set environment variables
+ENV PYTHONUNBUFFERED=1 \
+    PYTHONPATH=/app
+
+# Expose Gradio default port
+ENV GRADIO_SERVER_NAME="0.0.0.0"
+EXPOSE 7860
+
+# Entrypoint
+CMD ["python3", "app.py"]

docker/Dockerfile.gpu

@@ -0,0 +1,49 @@
+# Dockerfile.gpu - GPU deployment for DIA
+# --------------------------------------------------
+# Build: docker build . -f docker/Dockerfile.gpu -t dia-gpu
+# Run:   docker run --rm --gpus all -p 7860:7860 dia-gpu
+# Requires NVIDIA Container Toolkit on host.
+
+FROM pytorch/pytorch:2.1.2-cuda12.1-cudnn8-runtime
+
+# Set non-interactive frontend
+ENV DEBIAN_FRONTEND=noninteractive
+
+# Install venv, and system dependencies
+RUN apt-get update && apt-get install -y \
+    python3-venv \
+    libsndfile1 \
+    ffmpeg \
+    curl \
+    && apt-get clean && rm -rf /var/lib/apt/lists/*
+
+# Create non-root user and set up directories
+RUN useradd -m -u 1001 appuser && \
+    mkdir -p /app/outputs /app && \
+    chown -R appuser:appuser /app
+
+USER appuser
+WORKDIR /app
+
+# Copy all code (including pyproject.toml)
+COPY --chown=appuser:appuser . .
+
+# Create and activate virtual environment
+RUN python3 -m venv /app/venv
+ENV PATH="/app/venv/bin:$PATH"
+
+# Install all project dependencies
+RUN pip install --upgrade pip && pip install --no-cache-dir .
+
+# Set environment variables
+ENV PYTHONUNBUFFERED=1 \
+    PYTHONPATH=/app \
+    USE_GPU=true \
+    LD_LIBRARY_PATH=/usr/local/cuda/lib64:/usr/local/cuda-12.1/lib64:${LD_LIBRARY_PATH}
+
+# Expose Gradio default port
+ENV GRADIO_SERVER_NAME="0.0.0.0"
+EXPOSE 7860
+
+# Entrypoint
+CMD ["python3", "app.py"]

example/benchmark.py

@@ -0,0 +1,39 @@
+from random import choice
+
+import torch
+
+from dia.model import Dia
+
+
+torch._inductor.config.coordinate_descent_tuning = True
+torch._inductor.config.triton.unique_kernel_names = True
+torch._inductor.config.fx_graph_cache = True
+
+# debugging
+torch._logging.set_logs(graph_breaks=True, recompiles=True)
+
+model_name = "nari-labs/Dia-1.6B-0626"
+compute_dtype = "float16"
+
+model = Dia.from_pretrained(model_name, compute_dtype=compute_dtype)
+
+
+test_cases = [
+    "[S1] Dia is an open weights text to dialogue model.",
+    "[S1] Dia is an open weights text to dialogue model. [S2] You get full control over scripts and voices. [S1] Wow. Amazing. (laughs) [S2] Try it now on Git hub or Hugging Face.",
+    "[S1] torch.compile is a new feature in PyTorch that allows you to compile your model with a single line of code.",
+    "[S1] torch.compile is a new feature in PyTorch that allows you to compile your model with a single line of code. [S2] It is a new feature in PyTorch that allows you to compile your model with a single line of code.",
+]
+
+
+# Wram up
+for _ in range(2):
+    text = choice(test_cases)
+    output = model.generate(text, audio_prompt="./example_prompt.mp3", use_torch_compile=True, verbose=True)
+    output = model.generate(text, use_torch_compile=True, verbose=True)
+
+# Benchmark
+for _ in range(10):
+    text = choice(test_cases)
+    output = model.generate(text, use_torch_compile=True, verbose=True)
+    output = model.generate(text, audio_prompt="./example_prompt.mp3", use_torch_compile=True, verbose=True)

example/simple-cpu.py

@@ -0,0 +1,19 @@
+import torch
+
+from dia.model import Dia
+
+
+# Select device: CPU
+device = torch.device("cpu")
+print(f"Using device: {device}")
+
+# Load model
+model = Dia.from_pretrained(
+    "nari-labs/Dia-1.6B-0626", compute_dtype="float32", device=device
+)  # Float32 works better than float16 on CPU - you can also test with float16
+
+text = "[S1] Dia is an open weights text to dialogue model. [S2] You get full control over scripts and voices. [S1] Wow. Amazing. (laughs) [S2] Try it now on Git hub or Hugging Face."
+
+output = model.generate(text, use_torch_compile=False, verbose=True)
+
+model.save_audio("simple.mp3", output)

example/simple-mac.py

@@ -0,0 +1,12 @@
+from dia.model import Dia
+
+
+model = Dia.from_pretrained("nari-labs/Dia-1.6B-0626", compute_dtype="float16")
+
+text = "[S1] Dia is an open weights text to dialogue model. [S2] You get full control over scripts and voices. [S1] Wow. Amazing. (laughs) [S2] Try it now on Git hub or Hugging Face."
+
+# It is important to set the `use_torch_compile` argument to `False` when using Dia on MacOS.
+# This is because the `torch.compile` function is not supported on MacOS.
+output = model.generate(text, use_torch_compile=False, verbose=True)
+
+model.save_audio("simple.mp3", output)

example/simple.py

@@ -0,0 +1,18 @@
+from dia.model import Dia
+
+
+model = Dia.from_pretrained("nari-labs/Dia-1.6B-0626", compute_dtype="float16")
+
+text = "[S1] Dia is an open weights text to dialogue model. [S2] You get full control over scripts and voices. [S1] Wow. Amazing. (laughs) [S2] Try it now on Git hub or Hugging Face."
+
+output = model.generate(
+    text,
+    use_torch_compile=False,
+    verbose=True,
+    cfg_scale=3.0,
+    temperature=1.8,
+    top_p=0.90,
+    cfg_filter_top_k=50,
+)
+
+model.save_audio("simple.mp3", output)

example/simple_batch.py

@@ -0,0 +1,12 @@
+from dia.model import Dia
+
+
+model = Dia.from_pretrained("nari-labs/Dia-1.6B-0626", compute_dtype="float16")
+
+text = "[S1] Dia is an open weights text to dialogue model. [S2] You get full control over scripts and voices. [S1] Wow. Amazing. (laughs) [S2] Try it now on Git hub or Hugging Face."
+texts = [text for _ in range(10)]
+
+output = model.generate(texts, use_torch_compile=True, verbose=True, max_tokens=1500)
+
+for i, o in enumerate(output):
+    model.save_audio(f"simple_{i}.mp3", o)

example/voice_clone.py

@@ -0,0 +1,31 @@
+from dia.model import Dia
+
+
+model = Dia.from_pretrained("nari-labs/Dia-1.6B-0626", compute_dtype="float16")
+
+# You should put the transcript of the voice you want to clone
+# We will use the audio created by running simple.py as an example.
+# Note that you will be REQUIRED TO RUN simple.py for the script to work as-is.
+clone_from_text = "[S1] Dia is an open weights text to dialogue model. [S2] You get full control over scripts and voices. [S1] Wow. Amazing. (laughs) [S2] Try it now on Git hub or Hugging Face."
+clone_from_audio = "simple.mp3"
+
+# For your custom needs, replace above with below and add your audio file to this directory:
+# clone_from_text = "[S1] ... [S2] ... [S1] ... corresponding to your_audio_name.mp3"
+# clone_from_audio = "your_audio_name.mp3"
+
+# Text to generate
+text_to_generate = "[S1] Hello, how are you? [S2] I'm good, thank you. [S1] What's your name? [S2] My name is Dia. [S1] Nice to meet you. [S2] Nice to meet you too."
+
+# It will only return the audio from the text_to_generate
+output = model.generate(
+    clone_from_text + text_to_generate,
+    audio_prompt=clone_from_audio,
+    use_torch_compile=False,
+    verbose=True,
+    cfg_scale=4.0,
+    temperature=1.8,
+    top_p=0.90,
+    cfg_filter_top_k=50,
+)
+
+model.save_audio("voice_clone.mp3", output)

example/voice_clone_batch.py

@@ -0,0 +1,26 @@
+from dia.model import Dia
+
+
+model = Dia.from_pretrained("nari-labs/Dia-1.6B-0626", compute_dtype="float16")
+
+# You should put the transcript of the voice you want to clone
+# We will use the audio created by running simple.py as an example.
+# Note that you will be REQUIRED TO RUN simple.py for the script to work as-is.
+clone_from_text = "[S1] Dia is an open weights text to dialogue model. [S2] You get full control over scripts and voices. [S1] Wow. Amazing. (laughs) [S2] Try it now on Git hub or Hugging Face."
+
+# For your custom needs, replace above with below and add your audio file to this directory:
+# clone_from_text = "[S1] ... [S2] ... [S1] ... corresponding to your_audio_name.mp3"
+# clone_from_audio = "your_audio_name.mp3"
+
+# Text to generate
+text_to_generate = "[S1] Dia is an open weights text to dialogue model. [S2] You get full control over scripts and voices. [S1] Wow. Amazing. (laughs) [S2] Try it now on Git hub or Hugging Face."
+
+clone_from_audios = [f"simple_{i}.mp3" for i in range(10)]
+
+texts = [clone_from_text + text_to_generate for _ in range(10)]
+
+# It will only return the audio from the text_to_generate
+output = model.generate(texts, audio_prompt=clone_from_audios, use_torch_compile=True, verbose=True, max_tokens=2000)
+
+for i, o in enumerate(output):
+    model.save_audio(f"voice_clone_{i}.mp3", o)

example_prompt.mp3

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

hf.py

@@ -0,0 +1,17 @@
+from transformers import AutoProcessor, DiaForConditionalGeneration
+
+
+torch_device = "cuda"
+model_checkpoint = "nari-labs/Dia-1.6B-0626"
+
+text = [
+    "[S1] Dia is an open weights text to dialogue model. [S2] You get full control over scripts and voices. [S1] Wow. Amazing. (laughs) [S2] Try it now on Git hub or Hugging Face."
+]
+processor = AutoProcessor.from_pretrained(model_checkpoint)
+inputs = processor(text=text, padding=True, return_tensors="pt").to(torch_device)
+
+model = DiaForConditionalGeneration.from_pretrained(model_checkpoint).to(torch_device)
+outputs = model.generate(**inputs, max_new_tokens=3072, guidance_scale=3.0, temperature=1.8, top_p=0.90, top_k=45)
+
+outputs = processor.batch_decode(outputs)
+processor.save_audio(outputs, "example.mp3")

pyproject.toml

@@ -0,0 +1,66 @@
+[project]
+name = "nari-tts"
+version = "0.1.0"
+description = "Dia - A text-to-speech model for dialogue generation"
+readme = "README.md"
+requires-python = ">=3.10"
+license = {file = "LICENSE"}
+authors = [
+    {name = "Nari Labs", email = "contact@narilabs.ai"}
+]
+dependencies = [
+    "descript-audio-codec>=1.0.0",
+    "gradio>=5.25.2",
+    "huggingface-hub>=0.30.2",
+    "numpy>=2.2.4",
+    "pydantic>=2.11.3",
+    "safetensors>=0.5.3",
+    "soundfile>=0.13.1",
+    "torch==2.6.0",
+    "torchaudio==2.6.0",
+    "triton==3.2.0 ; sys_platform == 'linux'",
+    "triton-windows==3.2.0.post18 ; sys_platform == 'win32'",
+]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project.urls]
+"Homepage" = "https://github.com/nari-labs/dia"
+"Bug Tracker" = "https://github.com/nari-labs/dia/issues"
+
+[tool.hatch.build.targets.wheel]
+packages = ["dia"]
+
+[tool.ruff]
+# Never enforce `E501` (line length violations).
+lint.ignore = ["C901", "E501", "E741", "W605"]
+lint.select = ["C", "E", "F", "I", "W"]
+line-length = 119
+
+# Ignore import violations in all `__init__.py` files.
+[tool.ruff.lint.per-file-ignores]
+"__init__.py" = ["E402", "F401", "F403", "F811"]
+
+[tool.ruff.lint.isort]
+lines-after-imports = 2
+
+[tool.uv.sources]
+torch = [
+  { index = "pytorch-cu126", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
+]
+torchaudio = [
+  { index = "pytorch-cu126", marker = "sys_platform == 'linux' or sys_platform == 'win32'" },
+]
+
+[[tool.uv.index]]
+name = "pytorch-cu126"
+url = "https://download.pytorch.org/whl/cu126"
+explicit = true
+
+[dependency-groups]
+dev = [
+    "ninja>=1.11.1.4",
+    "packaging>=25.0",
+]