dia

app.py

@@ -8,6 +8,8 @@ import tempfile
 import os
 import re
 import spaces
+from csm1b_dv import get_csm1b_tab
+from dia_1_6B_dv import get_dia_1_6B_tab
 
 # Global variables to persist across requests
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -465,214 +467,8 @@ with gr.Blocks(
     }
     """
 ) as app:
-    with gr.Tab("🎙️ CSM-1B"):
-        gr.Markdown("# 🎙️  CSM-1B Text-to-Speech Synthesis")
-        gr.Markdown("**CSM (Conversational Speech Model)** is a speech generation model from [Sesame](sesame.com) that generates **RVQ audio codes** from text and audio inputs. The model architecture employs a [Llama](https://www.llama.com/) backbone and a smaller audio decoder that produces [Mimi](https://huggingface.co/kyutai/mimi) audio codes. This demo uses a **fine-tuned version** of the model for **Dhivehi speech synthesis**.")
-        # Model selection
-        with gr.Row():
-            model_dropdown = gr.Dropdown(
-                choices=list(MODELS.keys()),
-                value=list(MODELS.keys())[0],
-                label="🤖 Select Model"
-            )
-            model_info = gr.Textbox(
-                value="No model loaded - select a model to load",
-                label="Model Status",
-                interactive=False
-            )
-        
-        with gr.Tabs():
-            # Simple Generation
-            with gr.TabItem("🎯 Simple Generation"):
-                gr.Markdown("### Generate speech from text without context")
-                
-                with gr.Row():
-                    with gr.Column():
-                        simple_text = gr.Textbox(
-                            label="Text to Generate (Dhivehi)",
-                            placeholder="މަންމަ ކިހާއިރެއް ދަރިފުޅުގެ އިންތިޒާރުގަ އިންނަތާ",
-                            value="މަންމަ ކިހާއިރެއް ދަރިފުޅުގެ އިންތިޒާރުގަ އިންނަތާ",
-                            lines=3,
-                            elem_classes=["dhivehi-text"]
-                        )
-                        # Generate speaker choices from current model's speaker prompts
-                        current_prompts = get_current_speaker_prompts()
-                        speaker_choices = [(f"{prompt_data['speaker_id']}: {prompt_data['name']}", prompt_data['speaker_id'])
-                                         for prompt_key, prompt_data in current_prompts.items()]
-                        simple_speaker = gr.Radio(
-                            choices=[choice[0] for choice in speaker_choices],
-                            label="Speaker",
-                            value=speaker_choices[0][0] if speaker_choices else "0: Speaker"
-                        )
-                        simple_btn = gr.Button("🎵 Generate", variant="primary")
-                    
-                    with gr.Column():
-                        simple_audio = gr.Audio(label="Generated Audio")
-                        simple_status = gr.Textbox(label="Status", interactive=False)
-                
-                def simple_generate_with_mapping(text, speaker_display, selected_model):
-                    # Extract speaker ID from display text (e.g., "0: Female Speaker 01" -> "0")
-                    speaker_id = speaker_display.split(":")[0]
-                    return generate_simple_audio(text, speaker_id, selected_model)
-                
-                simple_btn.click(
-                    simple_generate_with_mapping,
-                    inputs=[simple_text, simple_speaker, model_dropdown],
-                    outputs=[simple_audio, simple_status]
-                )
-            
-            # Context Generation
-            with gr.TabItem("🎭 Context Generation"):
-                gr.Markdown("### Generate speech with voice prompt")
-                
-                with gr.Row():
-                    with gr.Column():
-                        context_text = gr.Textbox(
-                            label="Speaker prompt",
-                            placeholder="މަންމަ ކިހާއިރެއް ދަރިފުޅުގެ އިންތިޒާރުގަ އިންނަތާ",
-                            value="",
-                            lines=2,
-                            elem_classes=["dhivehi-text"]
-                        )
-                        context_audio = gr.Audio(
-                            label="Speaker Prompt",
-                            type="filepath"
-                        )
-                        target_text = gr.Textbox(
-                            label="Text to Generate",
-                            placeholder="މަންމަ ކިހާއިރެއް ދަރިފުޅުގެ އިންތިޒާރުގަ އިންނަތާ",
-                            value="މަންމަ ކިހާއިރެއް ދަރިފުޅުގެ އިންތިޒާރުގަ އިންނަތާ",
-                            lines=3,
-                            elem_classes=["dhivehi-text"]
-                        )
-                        # Generate speaker choices for context generation
-                        context_speaker = gr.Radio(
-                            choices=[choice[0] for choice in speaker_choices],
-                            label="Speaker",
-                            value=speaker_choices[0][0] if speaker_choices else "0: Speaker"
-                        )
-                        context_btn = gr.Button("🎵 Generate with Context", variant="primary")
-                    
-                    with gr.Column():
-                        context_audio_out = gr.Audio(label="Generated Audio")
-                        context_status = gr.Textbox(label="Status", interactive=False)
-                
-                def context_generate_with_mapping(text, speaker_display, context_text_val, context_audio_val, selected_model):
-                    # Extract speaker ID from display text
-                    speaker_id = speaker_display.split(":")[0]
-                    return generate_context_audio(text, speaker_id, context_text_val, context_audio_val, selected_model)
-                
-                context_btn.click(
-                    context_generate_with_mapping,
-                    inputs=[target_text, context_speaker, context_text, context_audio, model_dropdown],
-                    outputs=[context_audio_out, context_status]
-                )
-            
-            # Conversation Generation
-            with gr.TabItem("💬 Conversation"):
-                gr.Markdown("### Generate dual-speaker conversations")
-                
-                with gr.Row():
-                    speaker_a = gr.Dropdown(
-                        choices=get_speaker_choices(),
-                        label="Speaker A",
-                        value=get_speaker_choices()[0]
-                    )
-                    speaker_b = gr.Dropdown(
-                        choices=get_speaker_choices(),
-                        label="Speaker B",
-                        value=get_speaker_choices()[1] if len(get_speaker_choices()) > 1 else get_speaker_choices()[0]
-                    )
-                
-                with gr.Accordion("🎵 Audio Style References", open=False):
-                    with gr.Row():
-                        with gr.Column():
-                            gr.Markdown("**Speaker A Prompt**")
-                            speaker_a_audio = gr.Audio(
-                                type="filepath",
-                                label="Speaker A Audio Style"
-                            )
-                            speaker_a_text = gr.Textbox(
-                                label="Speaker A Prompt Text",
-                                lines=2,
-                                placeholder="މަންމަ ކިހާއިރެއް ދަރިފުޅުގެ އިންތިޒާރުގަ އިންނަތާ",
-                                elem_classes=["dhivehi-text"]
-                            )
-                        
-                        with gr.Column():
-                            gr.Markdown("**Speaker B Prompt**")
-                            speaker_b_audio = gr.Audio(
-                                type="filepath",
-                                label="Speaker B Speaker Prompt"
-                            )
-                            speaker_b_text = gr.Textbox(
-                                label="Speaker B Speaker Prompt",
-                                lines=2,
-                                placeholder="މަންމަ ކިހާއިރެއް ދަރިފުޅުގެ އިންތިޒާރުގަ އިންނަތާ",
-                                elem_classes=["dhivehi-text"]
-                            )
-                
-                with gr.Accordion("⚙️ Options", open=False):
-                    use_style = gr.Checkbox(
-                        label="Use audio style references",
-                        value=False
-                    )
-                    split_sentences_checkbox = gr.Checkbox(
-                        label="Split sentences",
-                        value=True
-                    )
-                
-                dialogue_text = gr.Textbox(
-                    lines=6,
-                    placeholder="ދަރިފުޅު މިއަދު ހާދަ ލަސްތިވީ.. މަންމަ ކިހާއިރެއް ދަރިފުޅުގެ އިންތިޒާރުގަ އިންނަތާ...... ކޮބާ ތިޔަ ރިޕޯޓް ފޮތް؟\nދަބަހުގަ އެބައޮތް!.\nވަދެބަލަ އެތެރެއަށް!... ދީބަލަ ރިޕޯޓްފޮތް މަންމަ ބަލައިލަން!",
-                    value="ދަރިފުޅު މިއަދު ހާދަ ލަސްތިވީ.. މަންމަ ކިހާއިރެއް ދަރިފުޅުގެ އިންތިޒާރުގަ އިންނަތާ...... ކޮބާ ތިޔަ ރިޕޯޓް ފޮތް؟\nދަބަހުގަ އެބައޮތް!.\nވަދެބަލަ އެތެރެއަށް!... ދީބަލަ ރިޕޯޓްފޮތް މަންމަ ބަލައިލަން!",
-                    label="Dialogue Lines (one per line)",
-                    elem_classes=["dhivehi-text"]
-                )
-                
-                conv_btn = gr.Button("🎵 Generate Conversation", variant="primary")
-                conv_audio = gr.Audio(label="Generated Conversation")
-                conv_status = gr.Textbox(label="Status", interactive=False)
-                
-                def conversation_generate_with_model(speaker_a_val, speaker_b_val, speaker_a_audio_val, speaker_a_text_val,
-                                                   speaker_b_audio_val, speaker_b_text_val, dialogue_text_val,
-                                                   split_sentences_flag, use_style_flag, selected_model):
-                    return generate_conversation(speaker_a_val, speaker_b_val, speaker_a_audio_val, speaker_a_text_val,
-                                               speaker_b_audio_val, speaker_b_text_val, dialogue_text_val,
-                                               split_sentences_flag, use_style_flag, selected_model)
-                
-                conv_btn.click(
-                    conversation_generate_with_model,
-                    inputs=[speaker_a, speaker_b, speaker_a_audio, speaker_a_text,
-                           speaker_b_audio, speaker_b_text, dialogue_text,
-                           split_sentences_checkbox, use_style, model_dropdown],
-                    outputs=[conv_audio, conv_status]
-                )
-        
-        # Wire up model change
-        model_dropdown.change(
-            change_model_and_update_ui,
-            inputs=[model_dropdown],
-            outputs=[model_info, speaker_a, speaker_b, simple_speaker, context_speaker]
-        )
-        
-        gr.Markdown("""
-        ---
-        **Tips:**
-        - Simple: Basic text-to-speech
-        - Context: Use reference audio for voice consistency
-        - Conversation: Multi-speaker dialogues with style control
-        
-        **Issues:**
-        - Context: Context breaks sometimes. Adding multiple context audio seems to make it work, or adding previous generation to the context helps.
-        - Audio: Sometimes the generated audio is not in sync with the text.
-        - Long sentences: Generated long sentences seems sped up.
-        - Repeating words: Generated text sometimes repeats words.
-        """)
-    with gr.Tab("🎙️ Dia-1.6B"):
-        gr.Markdown("# 🎙️  Dia-1.6B Text-to-Speech Synthesis")
-        gr.Markdown("Dia is a 1.6B parameter text to speech model created by [Nari Labs](https://huggingface.co/nari-labs/Dia-1.6B). This demo uses a fine-tuned version of the model with **Dhivehi (Thaana) voices**. Also supports mixed languages and can generate non-verbal, dialogue and voice-clone.")
-        gr.Markdown("Check back later...")
+    get_csm1b_tab()
+    get_dia_1_6B_tab()
 
 if __name__ == "__main__":
-    app.launch(share=False)
+    app.launch(share=False,server_name="172.20.0.246",server_port=7860)

csm1b_dv.py

@@ -0,0 +1,508 @@
+import gradio as gr
+import torch
+from transformers import AutoProcessor, CsmForConditionalGeneration
+import librosa
+import numpy as np
+import soundfile as sf
+import tempfile
+import os
+import re
+import spaces
+
+# Global variables to persist across requests
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model = None
+processor = None
+current_model_id = None
+
+# Model configuration
+MODELS = {
+    "CSM 1B Dhivehi 2-Speakers": "alakxender/csm-1b-dhivehi-2-speakers",
+    "CSM 1B Dhivehi 5-Speaker": "alakxender/csm-1b-dhivehi-5-spk-pv",
+}
+
+# Model-specific speaker prompts
+MODEL_SPEAKER_PROMPTS = {
+    "alakxender/csm-1b-dhivehi-2-speakers": {
+        "female_01": {
+            "text": None,
+            "audio": None,
+            "speaker_id": "0",
+            "name": "Female Speaker 01"
+        },
+        "male_01": {
+            "text": None,
+            "audio": None,
+            "speaker_id": "1",
+            "name": "Male Speaker 01"
+        }
+    },
+    "alakxender/csm-1b-dhivehi-5-spk-pv": {
+        "female_01": {
+            "text": None,
+            "audio": None,
+            "speaker_id": "0",
+            "name": "Female Speaker 01"
+        },
+        "male_01": {
+            "text": None,
+            "audio": None,
+            "speaker_id": "1",
+            "name": "Male Speaker 01"
+        },
+        "female_02": {
+            "text": None,
+            "audio": None,
+            "speaker_id": "2",
+            "name": "Female Speaker 02"
+        },
+        "male_02": {
+            "text": None,
+            "audio": None,
+            "speaker_id": "4",
+            "name": "Male Speaker 02"
+        },
+        "female_03": {
+            "text": None,
+            "audio": None,
+            "speaker_id": "3",
+            "name": "Female Speaker 03"
+        }
+    }
+}
+
+@spaces.GPU
+def load_model(model_name):
+    global model, processor, current_model_id
+    if model_name not in MODELS:
+        return False
+    model_id = MODELS[model_name]
+    if current_model_id == model_id:
+        return True
+    try:
+        if model is not None:
+            del model
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+        processor = AutoProcessor.from_pretrained(model_id)
+        if hasattr(processor.tokenizer, "init_kwargs"):
+            processor.tokenizer.init_kwargs.pop("pad_to_multiple_of", None)
+        model = CsmForConditionalGeneration.from_pretrained(
+            model_id,
+            device_map=device,
+            torch_dtype=torch.float32
+        )
+        current_model_id = model_id
+        return True
+    except Exception:
+        return False
+
+def get_current_speaker_prompts():
+    if current_model_id in MODEL_SPEAKER_PROMPTS:
+        return MODEL_SPEAKER_PROMPTS[current_model_id]
+    return {
+        "female_01": {"text": "", "audio": None, "speaker_id": "0", "name": "Female Speaker 01"},
+        "male_01": {"text": "", "audio": None, "speaker_id": "1", "name": "Male Speaker 01"}
+    }
+
+def get_model_info():
+    if current_model_id:
+        model_name = next((name for name, id in MODELS.items() if id == current_model_id), "Unknown")
+        return f"Current Model: {model_name}"
+    return "No model loaded"
+
+def get_speaker_choices():
+    prompts = get_current_speaker_prompts()
+    return list(prompts.keys())
+
+def load_audio_file(filepath, target_sr=24000):
+    if filepath is None or not os.path.exists(filepath):
+        return None
+    try:
+        audio, sr = librosa.load(filepath, sr=target_sr)
+        return audio.astype(np.float32)
+    except Exception:
+        return None
+
+@spaces.GPU
+def generate_simple_audio(text, speaker_id, model_name=None):
+    global model, processor
+    if not text.strip():
+        return None, "Please enter some text."
+    if model_name and model_name in MODELS:
+        load_model(model_name)
+    elif model is None or processor is None:
+        load_model(list(MODELS.keys())[0])
+    try:
+        formatted_text = f"[{speaker_id}]{text}"
+        inputs = processor(formatted_text, add_special_tokens=True).to(device)
+        with torch.no_grad():
+            audio = model.generate(**inputs, output_audio=True)
+        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
+        processor.save_audio(audio, temp_file.name)
+        return temp_file.name, "Audio generated successfully!"
+    except Exception as e:
+        return None, f"Error: {str(e)}"
+
+@spaces.GPU
+def generate_context_audio(text, speaker_id, context_text, context_audio_file, model_name=None):
+    global model, processor
+    if not text.strip():
+        return None, "Please enter text to generate."
+    if not context_text.strip() or context_audio_file is None:
+        return None, "Please provide both context text and audio."
+    if model_name and model_name in MODELS:
+        load_model(model_name)
+    elif model is None or processor is None:
+        load_model(list(MODELS.keys())[0])
+    try:
+        context_audio = load_audio_file(context_audio_file)
+        if context_audio is None:
+            return None, "Failed to load context audio."
+        conversation = [
+            {
+                "role": str(speaker_id),
+                "content": [
+                    {"type": "text", "text": context_text},
+                    {"type": "audio", "path": context_audio}
+                ]
+            },
+            {
+                "role": str(speaker_id),
+                "content": [{"type": "text", "text": text}]
+            }
+        ]
+        inputs = processor.apply_chat_template(
+            conversation, tokenize=True, return_dict=True
+        ).to(device)
+        with torch.no_grad():
+            audio = model.generate(**inputs, output_audio=True)
+        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
+        processor.save_audio(audio, temp_file.name)
+        return temp_file.name, "Audio generated with context!"
+    except Exception as e:
+        return None, f"Error: {str(e)}"
+
+def trim_silence(audio_array, top_db=20):
+    try:
+        trimmed, _ = librosa.effects.trim(audio_array, top_db=top_db)
+        return trimmed
+    except:
+        return audio_array
+
+def split_sentences(text):
+    sentences = re.split(r'[.؟!،]', text)
+    return [s.strip() for s in sentences if s.strip()]
+
+def extract_audio_output(audio_output):
+    global processor
+    try:
+        if hasattr(audio_output, 'audio_values'):
+            return audio_output.audio_values.cpu().squeeze().numpy()
+        elif isinstance(audio_output, torch.Tensor):
+            return audio_output.cpu().squeeze().numpy()
+        else:
+            temp_file = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
+            processor.save_audio(audio_output, temp_file.name)
+            audio, _ = librosa.load(temp_file.name, sr=24000)
+            os.unlink(temp_file.name)
+            return audio
+    except Exception:
+        return np.array([])
+
+@spaces.GPU
+def generate_conversation(speaker_a, speaker_b, speaker_a_audio, speaker_a_text,
+                        speaker_b_audio, speaker_b_text, dialogue_text,
+                        split_sentences_flag, use_style, model_name=None, progress=gr.Progress()):
+    global model, processor
+    if model_name and model_name in MODELS:
+        load_model(model_name)
+    elif model is None or processor is None:
+        load_model(list(MODELS.keys())[0])
+    try:
+        lines = [line.strip() for line in dialogue_text.strip().split("\n") if line.strip()]
+        if not lines:
+            return None, "Please enter dialogue text."
+        current_prompts = get_current_speaker_prompts()
+        speaker_a_id = current_prompts.get(speaker_a, {}).get("speaker_id", "0")
+        speaker_b_id = current_prompts.get(speaker_b, {}).get("speaker_id", "1")
+        speakers = [speaker_a_id, speaker_b_id]
+        progress(0, desc="Preparing...")
+        speaker_contexts = {}
+        if use_style:
+            if speaker_a in current_prompts and speaker_a_audio and speaker_a_text:
+                audio_data = load_audio_file(speaker_a_audio)
+                if audio_data is not None:
+                    speaker_contexts[speaker_a_id] = {
+                        "audio": trim_silence(audio_data),
+                        "text": speaker_a_text.strip()
+                    }
+            if speaker_b in current_prompts and speaker_b_audio and speaker_b_text:
+                audio_data = load_audio_file(speaker_b_audio)
+                if audio_data is not None:
+                    speaker_contexts[speaker_b_id] = {
+                        "audio": trim_silence(audio_data),
+                        "text": speaker_b_text.strip()
+                    }
+        text_units = []
+        if split_sentences_flag:
+            for i, line in enumerate(lines):
+                speaker_role = speakers[i % 2]
+                sentences = split_sentences(line)
+                for sentence in sentences:
+                    text_units.append((sentence, speaker_role))
+        else:
+            for i, line in enumerate(lines):
+                speaker_role = speakers[i % 2]
+                text_units.append((line, speaker_role))
+        if not text_units:
+            return None, "No text to generate."
+        audio_segments = []
+        progress(0.1, desc="Generating audio...")
+        for i, (text, role) in enumerate(text_units):
+            progress(0.1 + (i / len(text_units)) * 0.8,
+                    desc=f"Generating {i+1}/{len(text_units)}: Speaker {role}")
+            conversation = []
+            if use_style and role in speaker_contexts:
+                is_first = True
+                for prev_i in range(i):
+                    if text_units[prev_i][1] == role:
+                        is_first = False
+                        break
+                if is_first:
+                    context = speaker_contexts[role]
+                    conversation.append({
+                        "role": role,
+                        "content": [
+                            {"type": "text", "text": context["text"]},
+                            {"type": "audio", "path": context["audio"]}
+                        ]
+                    })
+            conversation.append({
+                "role": role,
+                "content": [{"type": "text", "text": text}]
+            })
+            inputs = processor.apply_chat_template(
+                conversation, tokenize=True, return_dict=True
+            ).to(device)
+            with torch.no_grad():
+                audio_output = model.generate(**inputs, output_audio=True)
+            audio_segment = extract_audio_output(audio_output)
+            if len(audio_segment) > 0:
+                audio_segment = trim_silence(audio_segment)
+                audio_segments.append(audio_segment)
+        if not audio_segments:
+            return None, "No audio generated."
+        progress(0.9, desc="Finalizing...")
+        if len(audio_segments) == 1:
+            final_audio = audio_segments[0]
+        else:
+            final_audio = np.concatenate(audio_segments)
+        if np.max(np.abs(final_audio)) > 0:
+            final_audio = final_audio / np.max(np.abs(final_audio)) * 0.95
+        output_file = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
+        sf.write(output_file.name, final_audio, 24000)
+        duration = len(final_audio) / 24000
+        progress(1.0, desc="Complete!")
+        return output_file.name, f"Generated {len(text_units)} segments, {duration:.1f}s total"
+    except Exception as e:
+        return None, f"Error: {str(e)}"
+
+def change_model_and_update_ui(model_name):
+    success = load_model(model_name)
+    if not success:
+        return (
+            gr.update(value="Error loading model"),
+            gr.update(),
+            gr.update(),
+            gr.update(),
+            gr.update()
+        )
+    choices = get_speaker_choices()
+    current_prompts = get_current_speaker_prompts()
+    new_speaker_choices = [(f"{prompt_data['speaker_id']}: {prompt_data['name']}", prompt_data['speaker_id'])
+                          for prompt_key, prompt_data in current_prompts.items()]
+    display_choices = [choice[0] for choice in new_speaker_choices]
+    return (
+        gr.update(value=get_model_info()),
+        gr.update(choices=choices, value=choices[0]),
+        gr.update(choices=choices, value=choices[1] if len(choices) > 1 else choices[0]),
+        gr.update(choices=display_choices, value=display_choices[0] if display_choices else "0: Speaker"),
+        gr.update(choices=display_choices, value=display_choices[0] if display_choices else "0: Speaker")
+    )
+
+def get_csm1b_tab():
+    with gr.Tab("🎙️ CSM-1B"):
+        gr.Markdown("# 🎙️  CSM-1B Text-to-Speech Synthesis")
+        gr.Markdown("**CSM (Conversational Speech Model)** is a speech generation model from [Sesame](sesame.com) that generates **RVQ audio codes** from text and audio inputs. The model architecture employs a [Llama](https://www.llama.com/) backbone and a smaller audio decoder that produces [Mimi](https://huggingface.co/kyutai/mimi) audio codes. This demo uses a **fine-tuned version** of the model for **Dhivehi speech synthesis**.")
+        with gr.Row():
+            model_dropdown = gr.Dropdown(
+                choices=list(MODELS.keys()),
+                value=list(MODELS.keys())[0],
+                label="🤖 Select Model"
+            )
+            model_info = gr.Textbox(
+                value="No model loaded - select a model to load",
+                label="Model Status",
+                interactive=False
+            )
+        with gr.Tabs():
+            with gr.TabItem("🎯 Simple Generation"):
+                gr.Markdown("### Generate speech from text without context")
+                with gr.Row():
+                    with gr.Column():
+                        simple_text = gr.Textbox(
+                            label="Text to Generate (Dhivehi)",
+                            placeholder="މަންމަ ކިހާއިރެއް ދަރިފުޅުގެ އިންތިޒާރުގަ އިންނަތާ",
+                            value="މަންމަ ކިހާއިރެއް ދަރިފުޅުގެ އިންތިޒާރުގަ އިންނަތާ",
+                            lines=3,
+                            elem_classes=["dhivehi-text"]
+                        )
+                        current_prompts = get_current_speaker_prompts()
+                        speaker_choices = [(f"{prompt_data['speaker_id']}: {prompt_data['name']}", prompt_data['speaker_id'])
+                                         for prompt_key, prompt_data in current_prompts.items()]
+                        simple_speaker = gr.Radio(
+                            choices=[choice[0] for choice in speaker_choices],
+                            label="Speaker",
+                            value=speaker_choices[0][0] if speaker_choices else "0: Speaker"
+                        )
+                        simple_btn = gr.Button("🎵 Generate", variant="primary")
+                    with gr.Column():
+                        simple_audio = gr.Audio(label="Generated Audio")
+                        simple_status = gr.Textbox(label="Status", interactive=False)
+                def simple_generate_with_mapping(text, speaker_display, selected_model):
+                    speaker_id = speaker_display.split(":")[0]
+                    return generate_simple_audio(text, speaker_id, selected_model)
+                simple_btn.click(
+                    simple_generate_with_mapping,
+                    inputs=[simple_text, simple_speaker, model_dropdown],
+                    outputs=[simple_audio, simple_status]
+                )
+            with gr.TabItem("🎭 Context Generation"):
+                gr.Markdown("### Generate speech with voice prompt")
+                with gr.Row():
+                    with gr.Column():
+                        context_text = gr.Textbox(
+                            label="Speaker prompt",
+                            placeholder="މަންމަ ކިހާއިރެއް ދަރިފުޅުގެ އިންތިޒާރުގަ އިންނަތާ",
+                            value="",
+                            lines=2,
+                            elem_classes=["dhivehi-text"]
+                        )
+                        context_audio = gr.Audio(
+                            label="Speaker Prompt",
+                            type="filepath"
+                        )
+                        target_text = gr.Textbox(
+                            label="Text to Generate",
+                            placeholder="މަންމަ ކިހާއިރެއް ދަރިފުޅުގެ އިންތިޒާރުގަ އިންނަތާ",
+                            value="މަންމަ ކިހާއިރެއް ދަރިފުޅުގެ އިންތިޒާރުގަ އިންނަތާ",
+                            lines=3,
+                            elem_classes=["dhivehi-text"]
+                        )
+                        context_speaker = gr.Radio(
+                            choices=[choice[0] for choice in speaker_choices],
+                            label="Speaker",
+                            value=speaker_choices[0][0] if speaker_choices else "0: Speaker"
+                        )
+                        context_btn = gr.Button("🎵 Generate with Context", variant="primary")
+                    with gr.Column():
+                        context_audio_out = gr.Audio(label="Generated Audio")
+                        context_status = gr.Textbox(label="Status", interactive=False)
+                def context_generate_with_mapping(text, speaker_display, context_text_val, context_audio_val, selected_model):
+                    speaker_id = speaker_display.split(":")[0]
+                    return generate_context_audio(text, speaker_id, context_text_val, context_audio_val, selected_model)
+                context_btn.click(
+                    context_generate_with_mapping,
+                    inputs=[target_text, context_speaker, context_text, context_audio, model_dropdown],
+                    outputs=[context_audio_out, context_status]
+                )
+            with gr.TabItem("💬 Conversation"):
+                gr.Markdown("### Generate dual-speaker conversations")
+                with gr.Row():
+                    speaker_a = gr.Dropdown(
+                        choices=get_speaker_choices(),
+                        label="Speaker A",
+                        value=get_speaker_choices()[0]
+                    )
+                    speaker_b = gr.Dropdown(
+                        choices=get_speaker_choices(),
+                        label="Speaker B",
+                        value=get_speaker_choices()[1] if len(get_speaker_choices()) > 1 else get_speaker_choices()[0]
+                    )
+                with gr.Accordion("🎵 Audio Style References", open=False):
+                    with gr.Row():
+                        with gr.Column():
+                            gr.Markdown("**Speaker A Prompt**")
+                            speaker_a_audio = gr.Audio(
+                                type="filepath",
+                                label="Speaker A Audio Style"
+                            )
+                            speaker_a_text = gr.Textbox(
+                                label="Speaker A Prompt Text",
+                                lines=2,
+                                placeholder="މަންމަ ކިހާއިރެއް ދަރިފުޅުގެ އިންތިޒާރުގަ އިންނަތާ",
+                                elem_classes=["dhivehi-text"]
+                            )
+                        with gr.Column():
+                            gr.Markdown("**Speaker B Prompt**")
+                            speaker_b_audio = gr.Audio(
+                                type="filepath",
+                                label="Speaker B Speaker Prompt"
+                            )
+                            speaker_b_text = gr.Textbox(
+                                label="Speaker B Speaker Prompt",
+                                lines=2,
+                                placeholder="މަންމަ ކިހާއިރެއް ދަރިފުޅުގެ އިންތިޒާރުގަ އިންނަތާ",
+                                elem_classes=["dhivehi-text"]
+                            )
+                with gr.Accordion("⚙️ Options", open=False):
+                    use_style = gr.Checkbox(
+                        label="Use audio style references",
+                        value=False
+                    )
+                    split_sentences_checkbox = gr.Checkbox(
+                        label="Split sentences",
+                        value=True
+                    )
+                dialogue_text = gr.Textbox(
+                    lines=6,
+                    placeholder="ދަރިފުޅު މިއަދު ހާދަ ލަސްތިވީ.. މަންމަ ކިހާއިރެއް ދަރިފުޅުގެ އިންތިޒާރުގަ އިންނަތާ...... ކޮބާ ތިޔަ ރިޕޯޓް ފޮތް؟\nދަބަހުގަ އެބައޮތް!.",
+                    value="ދަރިފުޅު މިއަދު ހާދަ ލަސްތިވީ.. މަންމަ ކިހާއިރެއް ދަރިފުޅުގެ އިންތިޒާރުގަ އިންނަތާ...... ކޮބާ ތިޔަ ރިޕޯޓް ފޮތް؟\nދަބަހުގަ އެބައޮތް!.",
+                    label="Dialogue Lines (one per line)",
+                    elem_classes=["dhivehi-text"]
+                )
+                conv_btn = gr.Button("🎵 Generate Conversation", variant="primary")
+                conv_audio = gr.Audio(label="Generated Conversation")
+                conv_status = gr.Textbox(label="Status", interactive=False)
+                def conversation_generate_with_model(speaker_a_val, speaker_b_val, speaker_a_audio_val, speaker_a_text_val,
+                                                   speaker_b_audio_val, speaker_b_text_val, dialogue_text_val,
+                                                   split_sentences_flag, use_style_flag, selected_model):
+                    return generate_conversation(speaker_a_val, speaker_b_val, speaker_a_audio_val, speaker_a_text_val,
+                                               speaker_b_audio_val, speaker_b_text_val, dialogue_text_val,
+                                               split_sentences_flag, use_style_flag, selected_model)
+                conv_btn.click(
+                    conversation_generate_with_model,
+                    inputs=[speaker_a, speaker_b, speaker_a_audio, speaker_a_text,
+                           speaker_b_audio, speaker_b_text, dialogue_text,
+                           split_sentences_checkbox, use_style, model_dropdown],
+                    outputs=[conv_audio, conv_status]
+                )
+        model_dropdown.change(
+            change_model_and_update_ui,
+            inputs=[model_dropdown],
+            outputs=[model_info, speaker_a, speaker_b, simple_speaker, context_speaker]
+        )
+        gr.Markdown("""
+        ---
+        **Tips:**
+        - Simple: Basic text-to-speech
+        - Context: Use reference audio for voice consistency
+        - Conversation: Multi-speaker dialogues with style control
+        **Issues:**
+        - Context: Context breaks sometimes. Adding multiple context audio seems to make it work, or adding previous generation to the context helps.
+        - Audio: Sometimes the generated audio is not in sync with the text.
+        - Long sentences: Generated long sentences seems sped up.
+        - Repeating words: Generated text sometimes repeats words.
+        """)
+    # No explicit return needed for context manager pattern 

dia/__init__.py

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

dia/audio.py

@@ -0,0 +1,280 @@
+import typing as tp
+
+import torch
+
+from .config import DataConfig
+
+
+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
+
+
+@torch.no_grad()
+@torch.inference_mode()
+def audio_to_codebook(
+    model,
+    input_values,
+    data_config: DataConfig,
+    padding_mask=None,
+    sample_rate=44100,
+):
+    """
+    Encodes the input audio waveform into discrete codes.
+
+    Args:
+        model: The model to use for encoding.
+        input_values (`torch.Tensor` of shape `(batch_size, channels, sequence_length)`):
+            Float values of the input audio waveform.
+        padding_mask (`torch.Tensor` of shape `(batch_size, channels, sequence_length)`):
+            Padding mask used to pad the `input_values`.
+        sample_rate (`int`, *optional*) :
+            Signal sampling_rate
+
+    Returns:
+        A list of frames containing the discrete encoded codes for the input audio waveform, along with rescaling
+        factors for each chunk when `normalize` is True. Each frames is a tuple `(codebook, scale)`, with
+        `codebook` of shape `[batch_size, num_codebooks, frames]`.
+        Scale is not used here.
+
+    """
+    audio_data = model.preprocess(input_values, sample_rate)
+
+    if padding_mask is None:
+        padding_mask = torch.ones_like(input_values).bool()
+
+    _, encoded_frame, _, _, _ = model.encode(audio_data, n_quantizers=None)  # 1, C, T
+    seq_length = encoded_frame.shape[2]
+
+    t_idx_BxTxC, indices_BTCx3 = build_delay_indices(
+        B=1,
+        T=seq_length,
+        C=data_config.channels,
+        delay_pattern=data_config.delay_pattern,
+    )
+
+    encoded_frame = apply_audio_delay(
+        audio_BxTxC=encoded_frame.transpose(1, 2),  # 1, T, C
+        pad_value=data_config.audio_pad_value,
+        bos_value=data_config.audio_bos_value,
+        precomp=(t_idx_BxTxC, indices_BTCx3),
+    )
+
+    return encoded_frame
+
+
+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
+
+
+@torch.no_grad()
+@torch.inference_mode()
+def decode(
+    model,
+    audio_codes,
+):
+    """
+    Decodes the given frames into an output audio waveform
+    """
+    if len(audio_codes) != 1:
+        raise ValueError(f"Expected one frame, got {len(audio_codes)}")
+
+    try:
+        audio_values = model.quantizer.from_codes(audio_codes)
+        audio_values = model.decode(audio_values[0])
+
+        return audio_values
+    except Exception as e:
+        print(f"Error in decode method: {str(e)}")
+        raise
+
+
+def codebook_to_audio(generated_codes: torch.Tensor, model, delay_pattern, B=1, T=2600, C=9):
+    """Process a single codebook file to generate audio"""
+    # Remove BOS token
+    generated_codes = generated_codes[:, 1:]
+
+    if generated_codes.shape[1] > T:
+        generated_codes = generated_codes[:, :T]
+
+    seq_length = generated_codes.shape[1]
+
+    # Build revert indices
+    t_idx_BxTxC, indices_BTCx3 = build_revert_indices(B=B, T=seq_length, C=C, delay_pattern=delay_pattern)
+
+    # Transpose and add batch dimension
+    audio_BxTxC = generated_codes.transpose(1, 0).unsqueeze(0)
+    reverted_codebook = revert_audio_delay(
+        audio_BxTxC=audio_BxTxC,
+        pad_value=0,
+        precomp=(t_idx_BxTxC, indices_BTCx3),
+        T=seq_length,
+    )
+    reverted_codebook = reverted_codebook[:, :-30, :]
+
+    codebook = reverted_codebook.transpose(1, 2)
+
+    min_valid_index = 0
+    max_valid_index = 1023
+    invalid_mask = (codebook < min_valid_index) | (codebook > max_valid_index)
+
+    num_invalid = torch.sum(invalid_mask).item()
+    if num_invalid > 0:
+        print(f"Warning: Clamping {num_invalid} indices outside range [{min_valid_index}, {max_valid_index}] to 0.")
+
+    # Set invalid values to 0 (modify the tensor in-place)
+    codebook[invalid_mask] = 0
+    audio_array = decode(model, codebook)
+
+    return audio_array

dia/config.py

@@ -0,0 +1,206 @@
+"""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 typing import Annotated
+
+from pydantic import BaseModel, BeforeValidator, Field
+
+
+class DataConfig(BaseModel, frozen=True):
+    """Configuration for data loading and preprocessing.
+
+    Attributes:
+        text_length: Maximum length of text sequences (must be multiple of 128).
+        audio_length: Maximum length of audio sequences (must be multiple of 128).
+        channels: Number of audio channels.
+        text_pad_value: Value used for padding text sequences.
+        audio_eos_value: Value representing the end of audio sequences.
+        audio_bos_value: Value representing the beginning of audio sequences.
+        audio_pad_value: Value used for padding audio sequences.
+        delay_pattern: List of delay values for each audio channel.
+    """
+
+    text_length: Annotated[int, BeforeValidator(lambda x: (x + 127) // 128 * 128)] = Field(gt=0, multiple_of=128)
+    audio_length: Annotated[int, BeforeValidator(lambda x: (x + 127) // 128 * 128)] = Field(gt=0, multiple_of=128)
+    channels: int = Field(default=9, gt=0, multiple_of=1)
+    text_pad_value: int = Field(default=0)
+    audio_eos_value: int = Field(default=1024)
+    audio_pad_value: int = Field(default=1025)
+    audio_bos_value: int = Field(default=1026)
+    delay_pattern: list[Annotated[int, Field(ge=0)]] = Field(default_factory=lambda: [0, 8, 9, 10, 11, 12, 13, 14, 15])
+
+    def __hash__(self) -> int:
+        """Generate a hash based on all fields of the config."""
+        return hash(
+            (
+                self.text_length,
+                self.audio_length,
+                self.channels,
+                self.text_pad_value,
+                self.audio_pad_value,
+                self.audio_bos_value,
+                self.audio_eos_value,
+                tuple(self.delay_pattern),
+            )
+        )
+
+
+class EncoderConfig(BaseModel, frozen=True):
+    """Configuration for the encoder component of the Dia model.
+
+    Attributes:
+        n_layer: Number of transformer layers.
+        n_embd: Embedding dimension.
+        n_hidden: Hidden dimension size in the MLP layers.
+        n_head: Number of attention heads.
+        head_dim: Dimension per attention head.
+        mlp_activations: List of activation functions for the MLP layers.
+        use_pre_norm: Whether to use pre-normalization (LayerNorm before attention/MLP).
+    """
+
+    n_layer: int = Field(gt=0)
+    n_embd: int = Field(gt=0)
+    n_hidden: int = Field(gt=0)
+    n_head: int = Field(gt=0)
+    head_dim: int = Field(gt=0)
+    mlp_activations: list[str] = Field(default=["silu", "linear"])
+    use_pre_norm: bool = Field(default=False)
+
+
+class DecoderConfig(BaseModel, frozen=True):
+    """Configuration for the decoder component of the Dia model.
+
+    Attributes:
+        n_layer: Number of transformer layers.
+        n_embd: Embedding dimension.
+        n_hidden: Hidden dimension size in the MLP layers.
+        gqa_query_heads: Number of query heads for grouped-query self-attention.
+        kv_heads: Number of key/value heads for grouped-query self-attention.
+        gqa_head_dim: Dimension per query head for grouped-query self-attention.
+        cross_query_heads: Number of query heads for cross-attention.
+        cross_head_dim: Dimension per cross-attention head.
+        mlp_activations: List of activation functions for the MLP layers.
+        use_pre_norm: Whether to use pre-normalization.
+    """
+
+    n_layer: int = Field(gt=0)
+    n_embd: int = Field(gt=0)
+    n_hidden: int = Field(gt=0)
+    gqa_query_heads: int = Field(gt=0)
+    kv_heads: int = Field(gt=0)
+    gqa_head_dim: int = Field(gt=0)
+    cross_query_heads: int = Field(gt=0)
+    cross_head_dim: int = Field(gt=0)
+    mlp_activations: list[str] = Field(default=["silu", "linear"])
+    use_pre_norm: bool = Field(default=False)
+
+
+class ModelConfig(BaseModel, frozen=True):
+    """Main configuration container for the Dia model architecture.
+
+    Attributes:
+        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.
+        dropout: Dropout probability applied within the model.
+        normalization_layer_epsilon: Epsilon value for normalization layers (e.g., LayerNorm).
+        weight_dtype: Data type for model weights (e.g., "float32", "bfloat16").
+        rope_min_timescale: Minimum timescale for Rotary Positional Embeddings (RoPE).
+        rope_max_timescale: Maximum timescale for Rotary Positional Embeddings (RoPE).
+    """
+
+    encoder: EncoderConfig
+    decoder: DecoderConfig
+    src_vocab_size: int = Field(default=128, gt=0)
+    tgt_vocab_size: int = Field(default=1028, gt=0)
+    dropout: float = Field(default=0.0, ge=0.0, lt=1.0)
+    normalization_layer_epsilon: float = Field(default=1.0e-5, ge=0.0)
+    weight_dtype: str = Field(default="float32", description="Weight precision")
+    rope_min_timescale: int = Field(default=1, description="Timescale For global Attention")
+    rope_max_timescale: int = Field(default=10_000, description="Timescale For global Attention")
+
+
+class TrainingConfig(BaseModel, frozen=True):
+    """Training process configuration and hyperparameters.
+
+    Note: This configuration currently only includes precision settings.
+    Other training parameters (like batch size, learning rate, optimizer settings)
+    are assumed to be handled externally.
+
+    Attributes:
+        dtype: Data type for activations during training (e.g., "bfloat16", "float32").
+        logits_dot_in_fp32: Whether to compute the final logits dot product in fp32 for stability.
+    """
+
+    dtype: str = Field(default="bfloat16", description="Activation precision")
+    logits_dot_in_fp32: bool = Field(default=False)
+
+
+class DiaConfig(BaseModel, frozen=True):
+    """Master configuration for the Dia model.
+
+    Combines all sub-configurations into a single validated object.
+
+    Attributes:
+        version: Configuration version string.
+        model: Model architecture configuration.
+        training: Training process configuration (precision settings).
+        data: Data loading and processing configuration.
+    """
+
+    version: str = Field(default="1.0")
+    model: ModelConfig
+    training: TrainingConfig
+    data: DataConfig
+
+    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,909 @@
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+from torch.nn import RMSNorm
+
+from .config import DiaConfig
+
+
+def _normalize_axes(axes: tuple[int, ...], ndim: int) -> tuple[int, ...]:
+    return tuple(ax if ax >= 0 else ndim + ax for ax in axes)
+
+
+def _str_to_dtype(dtype_str: str) -> torch.dtype | None:
+    # Allow None for default behavior
+    if dtype_str is None or dtype_str.lower() == "none":
+        return None
+    if dtype_str == "float32":
+        return torch.float32
+    elif dtype_str == "float16":
+        return torch.float16
+    elif dtype_str == "bfloat16":
+        return torch.bfloat16
+    else:
+        raise ValueError(f"Unsupported dtype string: {dtype_str}")
+
+
+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,),
+        dtype: torch.dtype | None = None,
+        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.dtype = dtype
+        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))
+        self.register_parameter("bias", None)
+
+    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.float(),
+            self.weight.float(),
+            dims=(norm_axis, kernel_contract_axes),
+        ).to(inputs.dtype)
+        return output
+
+
+def get_activation_fn(activation_string: str) -> nn.Module:  # Return Module instance
+    """Maps activation string to PyTorch activation function module."""
+    if activation_string == "gelu":
+        return nn.GELU()
+    elif activation_string == "relu":
+        return nn.ReLU()
+    elif activation_string == "silu" or activation_string == "swish":
+        return nn.SiLU()
+    elif activation_string == "linear":
+        return nn.Identity()
+    else:
+        raise ValueError(f"Unsupported activation function: {activation_string}")
+
+
+class MlpBlock(nn.Module):
+    """MLP block using DenseGeneral."""
+
+    def __init__(
+        self,
+        config: DiaConfig,
+        embed_dim: int,
+        intermediate_dim: int,
+        dropout_rate: float,
+        activations: list[str] = ["silu", "linear"],
+        use_pre_norm: bool = False,
+    ):
+        super().__init__()
+        self.use_pre_norm = use_pre_norm
+        num_activations = len(activations)
+        compute_dtype = _str_to_dtype(config.training.dtype)
+        weight_dtype = _str_to_dtype(config.model.weight_dtype)
+        self.dtype = compute_dtype
+        # Assume default device for now, could be passed in config
+
+        if use_pre_norm:
+            self.pre_norm = RMSNorm(
+                embed_dim,
+                eps=config.model.normalization_layer_epsilon,
+                dtype=torch.float32,
+            )
+
+        self.wi_fused = DenseGeneral(
+            in_shapes=(embed_dim,),
+            out_features=(
+                num_activations,
+                intermediate_dim,
+            ),
+            axis=(-1,),
+            dtype=compute_dtype,
+            weight_dtype=weight_dtype,
+        )
+
+        self.activation_fn_0 = get_activation_fn(activations[0])  # silu
+        self.activation_fn_1 = get_activation_fn(activations[1])  # linear
+
+        self.dropout = nn.Dropout(dropout_rate)
+
+        # Output layer using DenseGeneral
+        self.wo = DenseGeneral(
+            in_shapes=(intermediate_dim,),
+            out_features=(embed_dim,),
+            axis=(-1,),
+            dtype=compute_dtype,
+            weight_dtype=weight_dtype,
+        )
+
+    def forward(self, x: torch.Tensor, deterministic: bool) -> torch.Tensor:
+        """Forward pass."""
+        if self.use_pre_norm and hasattr(self, "pre_norm"):
+            x = self.pre_norm(x)
+
+        fused_x = self.wi_fused(x)
+
+        gate_input = fused_x[..., 0, :]
+        up_input = fused_x[..., 1, :]
+
+        gate = self.activation_fn_0(gate_input)
+        up = self.activation_fn_1(up_input)
+        hidden = torch.mul(gate, up).to(self.dtype)
+
+        if not deterministic:
+            hidden = self.dropout(hidden)
+
+        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: int = 1,
+        max_timescale: int = 10000,
+        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.dtype = dtype
+
+        half_embedding_dim = embedding_dims // 2
+        fraction = (2.0 * torch.arange(0, half_embedding_dim)) / embedding_dims
+        self.register_buffer(
+            "timescale",
+            self.min_timescale * (self.max_timescale / self.min_timescale) ** fraction,
+            persistent=False,
+        )
+
+    def extra_repr(self) -> str:
+        s = f"{self.timescale.shape}"
+        return s
+
+    def forward(self, inputs: torch.Tensor, position: torch.Tensor):
+        """Applies RoPE."""
+        position = position.unsqueeze(-1).unsqueeze(-1)
+        timescale = self.timescale.to(inputs.device)
+        sinusoid_inp = position / timescale
+        sin = torch.sin(sinusoid_inp).to(inputs.dtype)
+        cos = torch.cos(sinusoid_inp).to(inputs.dtype)
+        first_half, second_half = torch.chunk(inputs, 2, dim=-1)
+        first_part = first_half * cos - second_half * sin
+        second_part = second_half * cos + first_half * sin
+        return torch.cat((first_part, second_part), dim=-1)
+
+
+class KVCache:
+    def __init__(self, num_heads, max_len, head_dim, device, k=None, v=None):
+        self.k = torch.zeros((2, num_heads, max_len, head_dim), device=device) if k is None else k
+        self.v = torch.zeros((2, num_heads, max_len, head_dim), device=device) if v is None else v
+        self.current_idx = 0
+        self.max_len = max_len
+
+    def get_kv_for_attention(self, current_k, current_v):
+        if self.current_idx == 0:
+            return current_k, current_v
+        else:
+            past_k = self.k[:, :, : self.current_idx, :]
+            past_v = self.v[:, :, : self.current_idx, :]
+            attn_k = torch.cat((past_k, current_k), dim=2)
+            attn_v = torch.cat((past_v, current_v), dim=2)
+            return attn_k, attn_v
+
+    def update_cache(self, k, v):
+        assert self.current_idx < self.max_len
+        self.k[:, :, self.current_idx : self.current_idx + 1, :] = k
+        self.v[:, :, self.current_idx : self.current_idx + 1, :] = v
+        self.current_idx += 1
+
+    def prefill_kv(self, k, v):
+        prefill_len = k.shape[2]
+        assert prefill_len <= self.max_len
+        self.k[:, :, :prefill_len, :] = k
+        self.v[:, :, :prefill_len, :] = v
+        self.current_idx = prefill_len
+
+
+class Attention(nn.Module):
+    """Attention using DenseGeneral."""
+
+    def __init__(
+        self,
+        config: DiaConfig,
+        q_embed_dim: int,
+        kv_embed_dim: int,
+        num_query_heads: int,
+        num_kv_heads: int,
+        head_dim: int,
+        dropout_rate: float,
+        is_cross_attn: bool = False,
+        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.is_cross_attn = is_cross_attn
+        self.dropout_rate = dropout_rate
+        compute_dtype = _str_to_dtype(config.training.dtype)
+        weight_dtype = _str_to_dtype(config.model.weight_dtype)
+        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,),
+            dtype=compute_dtype,
+            weight_dtype=weight_dtype,
+        )
+        self.k_proj = DenseGeneral(
+            in_shapes=(kv_embed_dim,),
+            out_features=(num_kv_heads, head_dim),
+            axis=(-1,),
+            dtype=compute_dtype,
+            weight_dtype=weight_dtype,
+        )
+        self.v_proj = DenseGeneral(
+            in_shapes=(kv_embed_dim,),
+            out_features=(num_kv_heads, head_dim),
+            axis=(-1,),
+            dtype=compute_dtype,
+            weight_dtype=weight_dtype,
+        )
+        self.o_proj = DenseGeneral(
+            in_shapes=(num_query_heads, head_dim),
+            out_features=(self.output_dim,),
+            axis=(-2, -1),
+            dtype=compute_dtype,
+            weight_dtype=weight_dtype,
+        )
+
+        # --- Rotary Embedding ---
+        self.rotary_emb = RotaryEmbedding(
+            embedding_dims=self.head_dim,
+            min_timescale=config.model.rope_min_timescale,
+            max_timescale=config.model.rope_max_timescale,
+            dtype=compute_dtype,
+        )
+
+    def forward(
+        self,
+        Xq: torch.Tensor,  # (B, T, D) T = 1 in AR generation
+        Xkv: torch.Tensor,  # (B, S, E) S = 1 in AR generation
+        q_positions: torch.Tensor,  # (B, T)
+        kv_positions: torch.Tensor | None = None,  # (B, S)
+        deterministic: bool = True,
+        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,  # True only when prefilling KV Cache
+    ) -> 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.
+            deterministic: If True, disable dropout.
+            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 = Xq.dtype
+
+        Xq_BxTxNxH = self.q_proj(Xq)
+        Xq_BxTxNxH = self.rotary_emb(Xq_BxTxNxH, position=q_positions)
+        Xq_BxNxTxH = Xq_BxTxNxH.transpose(1, 2)
+
+        # Input values into attention calculation
+        attn_k: torch.Tensor | None = None
+        attn_v: torch.Tensor | None = None
+        new_kv_cache: tuple[torch.Tensor, torch.Tensor] | None = None
+
+        # Decoder Cross Attention
+        if self.is_cross_attn:
+            # Directly use cache (no need to check index)
+            attn_k, attn_v = cache.k, cache.v
+            if attn_k.shape[1] != self.num_query_heads or attn_v.shape[1] != self.num_query_heads:
+                raise ValueError(
+                    f"Cross-attention cache head dimension ({attn_k.shape[1]}) "
+                    f"does not match num_query_heads ({self.num_query_heads}). "
+                    "Cache should be pre-repeated for GQA."
+                )
+        # Self Attention
+        else:
+            Xk_BxSxKxH = self.k_proj(Xkv)  # (B, S, K, H)
+            Xv_BxSxKxH = self.v_proj(Xkv)  # (B, S, K, H)
+            Xk_BxSxKxH = self.rotary_emb(Xk_BxSxKxH, position=kv_positions)  # (B, S, K, H)
+
+            Xk_BxKxSxH = Xk_BxSxKxH.transpose(1, 2)  # (B, K, S, H)
+            Xv_BxKxSxH = Xv_BxSxKxH.transpose(1, 2)  # (B, K, S, H)
+            # S=1 for Decode Step
+
+            if self.num_gqa_groups > 1:
+                Xk_BxNxSxH = Xk_BxKxSxH.repeat_interleave(self.num_gqa_groups, dim=1)
+                Xv_BxNxSxH = Xv_BxKxSxH.repeat_interleave(self.num_gqa_groups, dim=1)
+            else:
+                Xk_BxNxSxH = Xk_BxKxSxH
+                Xv_BxNxSxH = Xv_BxKxSxH
+
+            # Encoder Self Attention
+            if cache is None:
+                attn_k = Xk_BxNxSxH
+                attn_v = Xv_BxNxSxH
+            # Decoder Self Attention
+            else:
+                # In prefill mode, we fill in cache until prefill length
+                if prefill:
+                    attn_k, attn_v = Xk_BxNxSxH, Xv_BxNxSxH
+                    cache.prefill_kv(attn_k, attn_v)
+                # In decode step, we add current K/V to cache step by step
+                else:
+                    new_kv_cache = Xk_BxNxSxH, Xv_BxNxSxH
+                    attn_k, attn_v = cache.get_kv_for_attention(Xk_BxNxSxH, Xv_BxNxSxH)
+
+        attn_output = F.scaled_dot_product_attention(
+            Xq_BxNxTxH,
+            attn_k,
+            attn_v,
+            attn_mask=attn_mask,
+            dropout_p=self.dropout_rate if not deterministic else 0.0,
+            scale=1.0,
+        )
+
+        attn_output = attn_output.transpose(1, 2).contiguous()  # (B, T, N, H)
+        output = self.o_proj(attn_output)
+
+        return output.to(original_dtype), new_kv_cache
+
+
+class EncoderLayer(nn.Module):
+    """Transformer Encoder Layer using DenseGeneral."""
+
+    def __init__(self, config: DiaConfig):
+        super().__init__()
+        self.config = config
+        model_config = config.model
+        enc_config = config.model.encoder
+        embed_dim = enc_config.n_embd
+
+        self.pre_sa_norm = RMSNorm(
+            embed_dim,
+            eps=model_config.normalization_layer_epsilon,
+            dtype=torch.float32,
+        )
+        self.self_attention = Attention(
+            config=config,
+            q_embed_dim=embed_dim,
+            kv_embed_dim=embed_dim,
+            num_query_heads=enc_config.n_head,
+            num_kv_heads=enc_config.n_head,
+            head_dim=enc_config.head_dim,
+            dropout_rate=model_config.dropout,
+            is_cross_attn=False,
+            out_embed_dim=embed_dim,
+        )
+        self.post_sa_norm = RMSNorm(
+            embed_dim,
+            eps=model_config.normalization_layer_epsilon,
+            dtype=torch.float32,
+        )
+        self.mlp = MlpBlock(
+            config=config,
+            embed_dim=embed_dim,
+            intermediate_dim=enc_config.n_hidden,
+            activations=enc_config.mlp_activations,
+            dropout_rate=model_config.dropout,
+            use_pre_norm=enc_config.use_pre_norm,
+        )
+        self.dropout = nn.Dropout(model_config.dropout)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        src_positions: torch.Tensor | None = None,
+        deterministic: bool = True,
+        attn_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        residual = x
+        x_norm = self.pre_sa_norm(x)
+
+        sa_out, _ = self.self_attention(
+            Xq=x_norm,
+            Xkv=x_norm,
+            q_positions=src_positions,
+            kv_positions=src_positions,
+            deterministic=deterministic,
+            attn_mask=attn_mask,
+        )
+        x = residual + sa_out
+
+        residual = x
+        x_norm = self.post_sa_norm(x)
+        mlp_out = self.mlp(x_norm, deterministic=deterministic)
+        x = residual + mlp_out
+
+        if not deterministic:
+            x = self.dropout(x)
+        return x
+
+
+class Encoder(nn.Module):
+    """Transformer Encoder Stack using DenseGeneral."""
+
+    def __init__(self, config: DiaConfig):
+        super().__init__()
+        self.config = config
+        model_config = config.model
+        enc_config = config.model.encoder
+        compute_dtype = _str_to_dtype(config.training.dtype)
+
+        self.embedding = nn.Embedding(
+            model_config.src_vocab_size,
+            enc_config.n_embd,
+            dtype=compute_dtype,
+        )
+        self.dropout = nn.Dropout(model_config.dropout)
+        self.layers = nn.ModuleList([EncoderLayer(config=config) for _ in range(enc_config.n_layer)])
+        self.norm = RMSNorm(
+            enc_config.n_embd,
+            eps=model_config.normalization_layer_epsilon,
+            dtype=torch.float32,
+        )
+
+    def forward(
+        self,
+        x_ids: torch.Tensor,
+        src_positions: torch.Tensor | None = None,
+        deterministic: bool = True,
+        attn_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        x = self.embedding(x_ids)
+
+        if not deterministic:
+            x = self.dropout(x)
+
+        for layer in self.layers:
+            x = layer(
+                x,
+                src_positions=src_positions,
+                deterministic=deterministic,
+                attn_mask=attn_mask,
+            )
+        x = self.norm(x)
+        if not deterministic:
+            x = self.dropout(x)
+        return x
+
+
+class DecoderLayer(nn.Module):
+    """Transformer Decoder Layer using DenseGeneral."""
+
+    def __init__(self, config: DiaConfig):
+        super().__init__()
+        self.config = config
+        model_config = config.model
+        dec_config = config.model.decoder
+        enc_config = config.model.encoder
+        dec_embed_dim = dec_config.n_embd
+        enc_embed_dim = enc_config.n_embd
+
+        # Norms
+        self.pre_sa_norm = RMSNorm(
+            dec_embed_dim,
+            eps=model_config.normalization_layer_epsilon,
+            dtype=torch.float32,
+        )
+        self.pre_ca_norm = RMSNorm(
+            dec_embed_dim,
+            eps=model_config.normalization_layer_epsilon,
+            dtype=torch.float32,
+        )
+        self.pre_mlp_norm = RMSNorm(
+            dec_embed_dim,
+            eps=model_config.normalization_layer_epsilon,
+            dtype=torch.float32,
+        )
+
+        # Self-Attention (GQA) with Causal Masking
+        self.self_attention = Attention(
+            config=config,
+            q_embed_dim=dec_embed_dim,
+            kv_embed_dim=dec_embed_dim,
+            num_query_heads=dec_config.gqa_query_heads,
+            num_kv_heads=dec_config.kv_heads,
+            head_dim=dec_config.gqa_head_dim,
+            dropout_rate=model_config.dropout,
+            is_cross_attn=False,
+            out_embed_dim=dec_embed_dim,
+        )
+        # Cross-Attention (MHA)
+        self.cross_attention = Attention(
+            config=config,
+            q_embed_dim=dec_embed_dim,
+            kv_embed_dim=enc_embed_dim,  # Note kv_embed_dim
+            num_query_heads=dec_config.cross_query_heads,
+            num_kv_heads=dec_config.cross_query_heads,
+            head_dim=dec_config.cross_head_dim,
+            dropout_rate=model_config.dropout,
+            is_cross_attn=True,
+            out_embed_dim=dec_embed_dim,
+        )
+        # MLP
+        self.mlp = MlpBlock(
+            config=config,
+            embed_dim=dec_embed_dim,
+            intermediate_dim=dec_config.n_hidden,
+            activations=dec_config.mlp_activations,
+            dropout_rate=model_config.dropout,
+            use_pre_norm=dec_config.use_pre_norm,
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        encoder_out: torch.Tensor,
+        tgt_positions: torch.Tensor,
+        src_positions: torch.Tensor | None,
+        deterministic: bool,
+        self_attn_mask: torch.Tensor,
+        cross_attn_mask: torch.Tensor,
+        self_attn_cache: KVCache,
+        cross_attn_cache: KVCache,
+        prefill: bool = False,
+    ) -> torch.Tensor:
+        residual = x
+        x_norm = self.pre_sa_norm(x)
+
+        sa_out, new_kv_cache = self.self_attention(
+            Xq=x_norm,  # (2, 1, D)
+            Xkv=x_norm,  # (2, 1, D)
+            q_positions=tgt_positions,  # (2, 1)
+            kv_positions=tgt_positions,  # (2, 1)
+            deterministic=deterministic,
+            attn_mask=self_attn_mask,  # (2, 1, 1, S_max)
+            cache=self_attn_cache,
+            prefill=prefill,
+        )
+
+        x = residual + sa_out
+
+        # 2. Cross-Attention
+        residual = x
+        x_norm = self.pre_ca_norm(x)
+        ca_out, _ = self.cross_attention(
+            Xq=x_norm,
+            Xkv=encoder_out,
+            q_positions=tgt_positions,
+            kv_positions=src_positions,
+            deterministic=deterministic,
+            attn_mask=cross_attn_mask,
+            cache=cross_attn_cache,
+        )
+        x = residual + ca_out
+
+        # 3. MLP
+        residual = x
+        x_norm = self.pre_mlp_norm(x)
+        mlp_out = self.mlp(x_norm, deterministic=deterministic)
+        x = residual + mlp_out
+
+        return x, new_kv_cache
+
+
+class Decoder(nn.Module):
+    """Transformer Decoder Stack using DenseGeneral."""
+
+    def __init__(self, config: DiaConfig):
+        super().__init__()
+        self.config = config
+        model_config = config.model
+        dec_config = config.model.decoder
+        train_config = config.training
+        data_config = config.data
+        compute_dtype = _str_to_dtype(config.training.dtype)
+        weight_dtype = _str_to_dtype(config.model.weight_dtype)
+        self.num_channels = data_config.channels
+        self.num_layers = dec_config.n_layer
+
+        self.embeddings = nn.ModuleList(
+            [
+                nn.Embedding(model_config.tgt_vocab_size, dec_config.n_embd, dtype=compute_dtype)
+                for _ in range(self.num_channels)
+            ]
+        )
+        self.dropout = nn.Dropout(model_config.dropout)
+        self.layers = nn.ModuleList([DecoderLayer(config=config) for _ in range(self.num_layers)])
+        self.norm = RMSNorm(
+            dec_config.n_embd,
+            eps=model_config.normalization_layer_epsilon,
+            dtype=torch.float32,
+        )
+
+        # Final Logits Projection using DenseGeneral
+        self.logits_dense = DenseGeneral(
+            in_shapes=(dec_config.n_embd,),
+            out_features=(self.num_channels, model_config.tgt_vocab_size),
+            axis=(-1,),
+            dtype=(torch.float32 if train_config.logits_dot_in_fp32 else compute_dtype),
+            weight_dtype=weight_dtype,
+        )
+        self.logits_in_fp32 = train_config.logits_dot_in_fp32
+
+    def precompute_cross_attention_kv(
+        self,
+        max_len: int,
+        encoder_out: torch.Tensor,  # (B, S, E)
+        src_positions: torch.Tensor | None,  # (B, S)
+    ) -> 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(encoder_out)
+            v_proj = cross_attn_module.v_proj(encoder_out)
+
+            k_proj = cross_attn_module.rotary_emb(k_proj, position=src_positions)
+            k = k_proj.transpose(1, 2)
+            v = v_proj.transpose(1, 2)
+
+            per_layer_kv_cache.append(
+                KVCache(
+                    cross_attn_module.num_kv_heads,
+                    max_len,
+                    cross_attn_module.head_dim,
+                    k.device,
+                    k=k,
+                    v=v,
+                )
+            )
+
+        return per_layer_kv_cache
+
+    def decode_step(
+        self,
+        tgt_ids_Bx1xC: torch.Tensor,  # [B, 1, C]
+        tgt_pos_Bx1: torch.Tensor,  # [B, 1]
+        encoder_out: torch.Tensor,  # [B, S, E]
+        self_attn_mask: Any,  # None
+        cross_attn_mask: torch.Tensor,  # [B, 1, 1, S]
+        self_attention_cache: list[KVCache],
+        cross_attention_cache: list[KVCache],
+    ) -> 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.
+        """
+        assert self_attn_mask is None, "Self-attention mask should be None, kept for pattern"
+
+        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
+
+        new_cache = []
+
+        for i, layer in enumerate(self.layers):
+            self_cache = self_attention_cache[i]
+            cross_cache = cross_attention_cache[i]
+            x, new_kv_cache = layer(
+                x,  # (2, 1, D)
+                encoder_out,  # (2, S, E)
+                src_positions=None,  # CA KV is already computed
+                tgt_positions=tgt_pos_Bx1,  # (2, 1)
+                deterministic=True,
+                self_attn_mask=None,
+                cross_attn_mask=cross_attn_mask,
+                self_attn_cache=self_cache,
+                cross_attn_cache=cross_cache,
+            )
+            new_cache.append(new_kv_cache)
+
+        x = self.norm(x)
+        logits_Bx1xCxV = self.logits_dense(x)
+
+        return logits_Bx1xCxV.to(torch.float32), new_cache
+
+    def forward(
+        self,
+        tgt_ids_BxTxC: torch.Tensor,
+        encoder_out: torch.Tensor,
+        tgt_positions: torch.Tensor,
+        src_positions: torch.Tensor,
+        deterministic: bool,
+        self_attn_mask: torch.Tensor,
+        cross_attn_mask: torch.Tensor,
+        self_attention_cache: list[KVCache],
+        cross_attention_cache: list[KVCache],
+    ) -> 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).
+            deterministic: Disable dropout if True.
+            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
+
+        if not deterministic:
+            x = self.dropout(x)
+
+        for i, layer in enumerate(self.layers):
+            x, _ = layer(
+                x,
+                encoder_out,
+                tgt_positions=tgt_positions,
+                src_positions=src_positions,
+                deterministic=deterministic,
+                self_attn_mask=self_attn_mask,
+                cross_attn_mask=cross_attn_mask,
+                self_attn_cache=self_attention_cache[i],
+                cross_attn_cache=cross_attention_cache[i],
+                prefill=True,
+            )
+
+        # Final Norm
+        x = self.norm(x)
+        logits_BxTxCxV = self.logits_dense(x)
+
+        return logits_BxTxCxV.to(torch.float32)
+
+
+class DiaModel(nn.Module):
+    """PyTorch Dia Model using DenseGeneral."""
+
+    def __init__(self, config: DiaConfig):
+        super().__init__()
+        self.config = config
+        self.encoder = Encoder(config)
+        self.decoder = Decoder(config)
+        #self._init_weights()
+
+    
+    def _init_weights(self):
+        for module in self.modules():
+            if isinstance(module, (torch.nn.Linear, torch.nn.Conv1d)):
+                torch.nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    torch.nn.init.zeros_(module.bias)
+            elif isinstance(module, torch.nn.Embedding):
+                torch.nn.init.xavier_uniform_(module.weight)
+            elif isinstance(module, torch.nn.LayerNorm) or isinstance(module, torch.nn.modules.normalization.RMSNorm):
+                if hasattr(module, 'weight') and module.weight is not None:
+                    torch.nn.init.ones_(module.weight)
+                if hasattr(module, 'bias') and module.bias is not None:
+                    torch.nn.init.zeros_(module.bias)
+
+    def forward(
+        self,
+        src_BxS: torch.Tensor,
+        tgt_BxTxC: torch.Tensor,
+        src_positions: torch.Tensor | None = None,
+        tgt_positions: torch.Tensor | None = None,
+        enc_self_attn_mask: torch.Tensor | None = None,
+        dec_self_attn_mask: torch.Tensor | None = None,
+        dec_cross_attn_mask: torch.Tensor | None = None,
+        enable_dropout: bool = True,
+    ):
+        deterministic = not enable_dropout
+
+        # --- Encoder Pass ---
+        encoder_out = self.encoder(
+            x_ids=src_BxS,
+            src_positions=src_positions,
+            deterministic=deterministic,
+            attn_mask=enc_self_attn_mask,
+        )
+
+        B, T, C = tgt_BxTxC.shape  # Batch size, target sequence length, channels
+        device = tgt_BxTxC.device
+
+        self_attention_cache = [
+            KVCache(
+                num_heads=self.decoder.layers[i].self_attention.num_query_heads,  # ✅ FIXED: use query heads!
+                max_len=T,
+                head_dim=self.decoder.layers[i].self_attention.head_dim,
+                device=device,
+            )
+            for i in range(self.decoder.num_layers)
+        ]
+
+        cross_attention_cache = self.decoder.precompute_cross_attention_kv(
+            max_len=encoder_out.shape[1],
+            encoder_out=encoder_out,
+            src_positions=src_positions,
+        )
+
+        # --- Decoder Pass ---
+        logits = self.decoder(
+            tgt_ids_BxTxC=tgt_BxTxC,
+            encoder_out=encoder_out,
+            tgt_positions=tgt_positions,
+            src_positions=src_positions,
+            deterministic=deterministic,
+            self_attn_mask=dec_self_attn_mask,
+            cross_attn_mask=dec_cross_attn_mask,
+            self_attention_cache=self_attention_cache,
+            cross_attention_cache=cross_attention_cache
+        )
+
+        return logits

dia/model.py

@@ -0,0 +1,460 @@
+import dac
+import numpy as np
+import torch
+import torchaudio
+from huggingface_hub import hf_hub_download
+
+from .audio import audio_to_codebook, codebook_to_audio
+from .config import DiaConfig
+from .layers import DiaModel, KVCache
+
+
+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,
+    use_cfg_filter: bool,
+    cfg_filter_top_k: int | None = None,
+) -> torch.Tensor:
+    if temperature == 0.0:
+        return torch.argmax(logits_BCxV, dim=-1)
+
+    logits_BCxV = logits_BCxV / temperature
+    if use_cfg_filter and cfg_filter_top_k is not None:
+        _, top_k_indices_BCxV = torch.topk(logits_BCxV, k=cfg_filter_top_k, dim=-1)
+        mask = torch.ones_like(logits_BCxV, dtype=torch.bool)
+        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)
+
+        # Calculate indices to remove based on top_p
+        sorted_indices_to_remove_BCxV = cumulative_probs_BCxV > top_p
+        # Shift the mask to the right to keep the first token above the threshold
+        sorted_indices_to_remove_BCxV[..., 1:] = sorted_indices_to_remove_BCxV[..., :-1].clone()
+        sorted_indices_to_remove_BCxV[..., 0] = 0  # Always keep the most probable token
+
+        indices_to_remove_BCxV = torch.zeros_like(sorted_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 Dia:
+    def __init__(self, config: DiaConfig, device: torch.device | None = None):
+        """Initializes the Dia model.
+
+        Args:
+            config: The configuration object for the model.
+            device: The device to load the model onto. If None, will automatically select the best available device.
+
+        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()
+        self.model = DiaModel(config)
+        self.dac_model = None
+
+    @classmethod
+    def from_local(cls, config_path: str, checkpoint_path: str, device: torch.device | None = None) -> "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.
+            device: The device to load the model onto. If None, will automatically select the best available device.
+
+        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, device)
+
+        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()
+        dia._load_dac_model()
+        return dia
+
+    @classmethod
+    def from_pretrained(
+        cls, model_name: str = "nari-labs/Dia-1.6B", device: torch.device | None = None
+    ) -> "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., "NariLabs/Dia-1.6B").
+            device: The device to load the model onto. If None, will automatically select the best available device.
+
+        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.
+        """
+        config_path = hf_hub_download(repo_id=model_name, filename="config.json")
+        checkpoint_path = hf_hub_download(repo_id=model_name, filename="dia-v0_1.pth")
+        return cls.from_local(config_path, checkpoint_path, device)
+
+    def _load_dac_model(self):
+        try:
+            dac_model_path = dac.utils.download()
+            dac_model = dac.DAC.load(dac_model_path).to(self.device)
+        except Exception as e:
+            raise RuntimeError("Failed to load DAC model") from e
+        self.dac_model = dac_model
+
+    def _create_attn_mask(
+        self,
+        q_padding_mask_1d: torch.Tensor,
+        k_padding_mask_1d: torch.Tensor,
+        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
+        assert B1 == B2, "Query and key batch dimensions must match"
+
+        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)
+        # This implementation follows Jax TPU splash attention kernel
+        mask = non_pad_attends_non_pad | pad_attends_pad  # Shape [B, Tq, Tk]
+
+        if is_causal:
+            # Ensure causality for self-attention (Tq == Tk)
+            assert Tq == Tk, "Causal mask requires query and key sequence lengths to be equal"
+            # Standard lower-triangular causal mask (True means allow)
+            causal_mask_2d = torch.tril(torch.ones((Tq, Tk), dtype=torch.bool, device=self.device))  # Shape [Tq, Tk]
+            causal_mask = mask & causal_mask_2d  # Shape [B, Tq, Tk]
+            return causal_mask.unsqueeze(1)  # Shape [B, 1, Tq, Tk] for broadcasting across heads
+        else:
+            # For cross-attention or non-causal self-attention
+            return mask.unsqueeze(1)  # Shape [B, 1, Tq, Tk] for broadcasting across heads
+
+    def _prepare_text_input(self, text: str) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Encodes text prompt, pads, and creates attention mask and positions."""
+        text_pad_value = self.config.data.text_pad_value
+        max_len = self.config.data.text_length
+
+        byte_text = text.encode("utf-8")
+        
+        
+        replaced_bytes = byte_text
+
+        LANG2BYTE = {
+            "en": 3,
+            "de": 4,
+            "fr": 5,
+            "es": 6,
+            "it": 7,
+            "nl": 14,
+            "pl": 15,
+            "pt": 16,
+            "tr": 17,
+            "hu": 18,
+        }
+
+        for lang, byte_val in LANG2BYTE.items():
+            tag = f"[{lang}]".encode("ascii")        # e.g. b"[de]"
+            code = bytes([byte_val])                 # e.g. b"\x04"
+            replaced_bytes = replaced_bytes.replace(tag, code)
+        text_tokens = list(replaced_bytes)
+
+        current_len = len(text_tokens)
+        padding_needed = max_len - current_len
+        if padding_needed <= 0:
+            text_tokens = text_tokens[:max_len]
+            padded_text_np = np.array(text_tokens, dtype=np.uint8)
+        else:
+            padded_text_np = np.pad(
+                text_tokens,
+                (0, padding_needed),
+                mode="constant",
+                constant_values=text_pad_value,
+            ).astype(np.uint8)
+
+        src_tokens = torch.from_numpy(padded_text_np).to(torch.long).to(self.device).unsqueeze(0)  # [1, S]
+        src_positions = torch.arange(max_len, device=self.device).to(torch.long).unsqueeze(0)  # [1, S]
+
+        src_padding_mask = (src_tokens != text_pad_value).to(self.device)  # [1, S]
+
+        enc_self_attn_mask = self._create_attn_mask(src_padding_mask, src_padding_mask, is_causal=False)  # [1, S, S]
+
+        return src_tokens, src_positions, src_padding_mask, enc_self_attn_mask
+
+    @torch.inference_mode()
+    def generate(
+        self,
+        text: str,
+        max_tokens: int | None = None,
+        cfg_scale: float = 3.0,
+        temperature: float = 1.3,
+        top_p: float = 0.95,
+        use_cfg_filter: bool = True,
+        use_torch_compile: bool = False,
+        cfg_filter_top_k: int = 35,
+        audio_prompt_path: str | None = None,
+    ) -> np.ndarray:
+        """
+        Generates audio from a text prompt (and optional audio prompt) using the Nari model.
+
+        Returns:
+            A tensor of generated audio codes (shape: [max_tokens, num_channels]).
+        """
+        num_channels = self.config.data.channels
+        audio_bos_value = self.config.data.audio_bos_value
+        audio_eos_value = self.config.data.audio_eos_value
+        audio_pad_value = self.config.data.audio_pad_value
+        delay_pattern = self.config.data.delay_pattern
+        max_tokens = self.config.data.audio_length if max_tokens is None else max_tokens
+        delay_tensor = torch.tensor(delay_pattern, dtype=torch.long, device=self.device)
+        max_delay_pattern = max(delay_pattern)
+        self.model.eval()
+
+        (
+            cond_src_BxS,
+            cond_src_positions_BxS,
+            cond_src_padding_mask_BxS,
+            cond_enc_self_attn_mask_Bx1xSxS,
+        ) = self._prepare_text_input(text)
+
+        unc_src_BxS = torch.zeros_like(cond_src_BxS)
+        src_BxS = torch.cat([unc_src_BxS, cond_src_BxS], dim=0)
+        src_positions_BxS = cond_src_positions_BxS.expand(2, -1)
+        src_padding_mask_BxS = cond_src_padding_mask_BxS.expand(2, -1)
+        enc_self_attn_mask_Bx1xSxS = cond_enc_self_attn_mask_Bx1xSxS.expand(2, -1, -1, -1)
+
+        # 2. Encoder Pass
+        # with torch.autocast(device_type="cuda", dtype=forward_dtype):
+        encoder_out = self.model.encoder(
+            x_ids=src_BxS,
+            src_positions=src_positions_BxS,
+            deterministic=True,
+            attn_mask=enc_self_attn_mask_Bx1xSxS,
+        )  # Shape: (B, S, E)
+
+        # 3. Prepare Decoder Inputs
+        # 3-1. Allocate KV Cache (Static)
+        decoder_cross_attention_cache: list[KVCache] = self.model.decoder.precompute_cross_attention_kv(
+            max_tokens, encoder_out, src_positions_BxS
+        )
+
+        decoder_self_attention_cache: list[KVCache] = []
+        for _ in range(self.model.decoder.num_layers):
+            decoder_self_attention_cache.append(
+                KVCache(
+                    self.config.model.decoder.gqa_query_heads,
+                    max_tokens,
+                    self.config.model.decoder.gqa_head_dim,
+                    self.device,
+                )
+            )
+
+        # 3-2. Initialize Decoder Inputs
+        generated_BxTxC = torch.full(
+            (2, 1, num_channels),
+            fill_value=audio_bos_value,
+            dtype=torch.long,
+            device=self.device,
+        )
+
+        current_step = 0
+        prompt_len_inc_bos = 1  # Start with BOS length
+
+        # 3-3. Load Audio Prompt (if provided)
+        if audio_prompt_path is not None:
+            audio_prompt, sr = torchaudio.load(audio_prompt_path, channels_first=True)  # C, T
+            if sr != 44100:  # Resample to 44.1kHz
+                audio_prompt = torchaudio.functional.resample(audio_prompt, sr, 44100)
+            audio_prompt = audio_prompt.to(self.device).unsqueeze(0)  # 1, C, T
+            audio_prompt = audio_to_codebook(self.dac_model, audio_prompt, data_config=self.config.data)
+            generated_BxTxC = torch.cat([generated_BxTxC, audio_prompt.expand(2, -1, -1)], dim=1)
+
+            prefill_len = generated_BxTxC.shape[1]
+            prompt_len_inc_bos = prefill_len
+            prefill_tgt_pos = torch.arange(prefill_len, device=self.device).unsqueeze(0).expand(2, -1)
+            prefill_tgt_padding_mask = (generated_BxTxC != audio_pad_value).any(dim=2)
+
+            prefill_self_attn_mask = self._create_attn_mask(
+                prefill_tgt_padding_mask,
+                prefill_tgt_padding_mask,
+                is_causal=True,
+            )
+            prefill_cross_attn_mask = self._create_attn_mask(
+                prefill_tgt_padding_mask,
+                src_padding_mask_BxS,
+                is_causal=False,
+            )
+
+            _ = self.model.decoder.forward(
+                tgt_ids_BxTxC=generated_BxTxC,
+                encoder_out=encoder_out,
+                tgt_positions=prefill_tgt_pos,
+                src_positions=src_positions_BxS,
+                deterministic=True,
+                self_attn_mask=prefill_self_attn_mask,
+                cross_attn_mask=prefill_cross_attn_mask,
+                self_attention_cache=decoder_self_attention_cache,
+                cross_attention_cache=decoder_cross_attention_cache,
+            )
+
+            current_step = prefill_len - 1
+
+        # 4. Autoregressive Generation Loop
+        eos_detected_channel_0 = False
+        eos_countdown = -1
+        extra_steps_after_eos = 30
+        # Make generated_BxTxC a fixed size tensor
+        # Length is either 1 + max tokens or 1 + prompt len + max tokens
+        generated_BxTxC = torch.cat(
+            [
+                generated_BxTxC,
+                torch.full(
+                    (2, max_tokens, num_channels),
+                    fill_value=-1,
+                    dtype=torch.long,
+                    device=self.device,
+                ),
+            ],
+            dim=1,
+        )
+
+        decode_step = self.model.decoder.decode_step
+        if use_torch_compile:
+            decode_step = torch.compile(
+                self.model.decoder.decode_step,
+                mode="default",
+            )
+
+        tgt_padding_mask = (
+            (generated_BxTxC[:, -1, :].unsqueeze(1) != audio_pad_value).any(dim=2).to(self.device)
+        )  # [B, 1]
+        # Generated tokens are never PAD, so we use fixed mask
+        decoder_cross_attn_mask = self._create_attn_mask(
+            tgt_padding_mask,  # Query mask [B, 1]
+            src_padding_mask_BxS,  # Key mask [B, S]
+            is_causal=False,
+        )  # [B, 1, 1, S]
+
+        for step in range(current_step, current_step + max_tokens):
+            tgt_ids_Bx1xC = generated_BxTxC[:, step, :].unsqueeze(1)
+            tgt_pos_Bx1 = torch.full(
+                (2, 1),
+                fill_value=step,
+                dtype=torch.long,
+                device=self.device,
+            )
+
+            logits_Bx1xCxV, new_cache = decode_step(
+                tgt_ids_Bx1xC=tgt_ids_Bx1xC,
+                tgt_pos_Bx1=tgt_pos_Bx1,
+                encoder_out=encoder_out,
+                self_attn_mask=None,
+                cross_attn_mask=decoder_cross_attn_mask,
+                self_attention_cache=decoder_self_attention_cache,
+                cross_attention_cache=decoder_cross_attention_cache,
+            )
+
+            for i, layer_cache in enumerate(decoder_self_attention_cache):
+                layer_cache.update_cache(new_cache[i][0], new_cache[i][1])
+
+            V = self.config.model.tgt_vocab_size
+            logits_last_BxCxV = logits_Bx1xCxV[:, -1, :, :]  # B, C, V
+            uncond_logits_CxV = logits_last_BxCxV[0, :, :]
+            cond_logits_CxV = logits_last_BxCxV[1, :, :]
+
+            cfg_logits_CxV = cond_logits_CxV + cfg_scale * (cond_logits_CxV - uncond_logits_CxV)
+
+            logits_CxV = cfg_logits_CxV.reshape((-1, V))  # C, V
+            logits_CxV[:, 1025:] = -torch.inf
+
+            # Sample next token
+            pred_C = _sample_next_token(
+                logits_CxV.float(),
+                temperature=temperature,
+                top_p=top_p,
+                use_cfg_filter=use_cfg_filter,
+                cfg_filter_top_k=cfg_filter_top_k,
+            )
+
+            generation_step_index = step - current_step
+            if audio_prompt_path is None:
+                pred_C = torch.where(
+                    generation_step_index >= delay_tensor,
+                    pred_C,
+                    audio_bos_value,
+                )
+
+            generated_BxTxC[:, step + 1, :] = pred_C.unsqueeze(0).expand(2, -1)
+
+            if not eos_detected_channel_0 and pred_C[0] == audio_eos_value:
+                eos_detected_channel_0 = True
+                eos_countdown = extra_steps_after_eos
+
+            if eos_countdown > 0:
+                step_after_eos = max_delay_pattern - eos_countdown
+                for i, d in enumerate(delay_pattern):
+                    if step_after_eos == d:
+                        generated_BxTxC[:, step + 1, i] = audio_eos_value
+                    elif step_after_eos > d:
+                        generated_BxTxC[:, step + 1, i] = audio_pad_value
+                eos_countdown -= 1
+                if eos_countdown == 0:
+                    break
+
+            generation_step_index = step - current_step + 1
+
+        output_codes = generated_BxTxC[:, prompt_len_inc_bos : step + 1, :]
+
+        generated_codes = output_codes[0]
+
+        audio = codebook_to_audio(
+            generated_codes.transpose(1, 0), self.dac_model, delay_pattern, B=1, T=max_tokens, C=num_channels
+        )
+        return audio.squeeze().cpu().numpy()

dia_1_6B_dv.py

@@ -0,0 +1,390 @@
+import tempfile
+import time
+from pathlib import Path
+from typing import Optional, Tuple
+import spaces
+
+import gradio as gr
+import numpy as np
+import soundfile as sf
+import torch
+
+from dia.model import Dia
+
+# Model selection
+DIA_MODELS = {
+    "Dhivehi Dia-1.6B": "alakxender/Dia-1.6B-dhivehi-ep1",
+    #"Dhivehi 18k": "alakxender/Dia-1.6B-dhivehi-18k"
+}
+
+dia_models = {}
+
+def load_dia_model(model_id):
+    if model_id not in dia_models:
+        print(f"Loading model {model_id}")
+        dia_models[model_id] = Dia.from_pretrained(model_id)
+    print(f"Loaded model {model_id}")
+    return dia_models[model_id]
+
+@spaces.GPU
+def run_inference(
+    text_input: str,
+    audio_prompt_input: Optional[Tuple[int, np.ndarray]],
+    transcription_input: Optional[str],
+    max_new_tokens: int,
+    cfg_scale: float,
+    temperature: float,
+    top_p: float,
+    cfg_filter_top_k: int,
+    speed_factor: float,
+    model_name: str,
+):
+    model_id = DIA_MODELS[model_name]
+    model = load_dia_model(model_id)
+    if not text_input or text_input.isspace():
+        raise gr.Error("Text input cannot be empty.")
+    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
+            duration_sec = len(audio_data) / float(sr) if sr else 0
+            if duration_sec > 10.0:
+                raise gr.Error("Audio prompt must be 10 seconds or shorter.")
+            if (
+                audio_data is None or audio_data.size == 0 or audio_data.max() == 0
+            ):
+                gr.Warning("Audio prompt seems empty or silent, ignoring prompt.")
+            else:
+                with tempfile.NamedTemporaryFile(
+                    mode="wb", suffix=".wav", delete=False
+                ) as f_audio:
+                    temp_audio_prompt_path = f_audio.name
+                    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."
+                        )
+                        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}"
+                            )
+                    if audio_data.ndim > 1:
+                        if audio_data.shape[0] == 2:
+                            audio_data = np.mean(audio_data, axis=0)
+                        elif audio_data.shape[1] == 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)
+                    try:
+                        sf.write(
+                            temp_audio_prompt_path, audio_data, sr, subtype="FLOAT"
+                        )
+                        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}")
+        start_time = time.time()
+        with torch.inference_mode():
+            combined_text = (
+                text_input.strip() + "\n" + transcription_input.strip()
+                if transcription_input and not transcription_input.isspace()
+                else text_input
+            )
+            output_audio_np = model.generate(
+                combined_text,
+                max_tokens=max_new_tokens,
+                cfg_scale=cfg_scale,
+                temperature=temperature,
+                top_p=top_p,
+                cfg_filter_top_k=cfg_filter_top_k,
+                use_torch_compile=False,
+                audio_prompt_path=prompt_path_for_generate,
+            )
+        end_time = time.time()
+        print(f"Generation finished in {end_time - start_time:.2f} seconds.")
+        if output_audio_np is not None:
+            output_sr = 44100
+            original_len = len(output_audio_np)
+            speed_factor = max(0.1, min(speed_factor, 5.0))
+            target_len = int(original_len / speed_factor)
+            if target_len != original_len and target_len > 0:
+                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),
+                )
+                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,
+                )
+                print(f"Skipping audio speed adjustment (factor: {speed_factor:.2f}).")
+            print(
+                f"Audio conversion successful. Final shape: {output_audio[1].shape}, Sample Rate: {output_sr}"
+            )
+            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.")
+            gr.Warning("Generation produced no output.")
+    except Exception as e:
+        print(f"Error during inference: {e}")
+        import traceback
+        traceback.print_exc()
+        raise gr.Error(f"Inference failed: {e}")
+    finally:
+        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}"
+                )
+    return output_audio
+
+def get_dia_1_6B_tab():
+    css = """
+    #col-container {max-width: 90%; margin-left: auto; margin-right: auto;}
+    .dhivehi-text-nofont textarea {
+        font-size: 18px !important;
+        line-height: 1.8 !important;
+        direction: rtl !important;
+        text-align: right !important;
+    }
+    .dhivehi-text-nofont input {
+        font-size: 18px !important;
+        direction: rtl !important;
+        text-align: right !important;
+    }
+    """
+    default_text = ""
+    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:
+                default_text = "Example text file was empty."
+        except Exception as e:
+            print(f"Warning: Could not read example.txt: {e}")
+    with gr.Tab("🎙️ Dia-1.6B"):
+        gr.Markdown("# Dia Text-to-Speech Synthesis (Dia-1.6B)")
+        with gr.Row(equal_height=False):
+            with gr.Column(scale=1):
+                model_dropdown = gr.Dropdown(
+                    choices=list(DIA_MODELS.keys()),
+                    value=list(DIA_MODELS.keys())[0],
+                    label="Select Dia Model"
+                )
+                text_input = gr.Textbox(
+                    label="Input Text",
+                    placeholder="ލިޔެލަން",
+                    value=default_text,
+                    lines=5,
+                    elem_classes=["dhivehi-text-nofont"]
+                )
+                audio_prompt_input = gr.Audio(
+                    label="Audio Prompt (≤ 10 s, Optional)",
+                    show_label=True,
+                    sources=["upload", "microphone"],
+                    type="numpy",
+                )
+                transcription_input = gr.Textbox(
+                    label="Audio Prompt Transcription (Optional)",
+                    placeholder="ޓްރާންސްކްރިޕްޓް ލިޔެލަން",
+                    lines=3,
+                    elem_classes=["dhivehi-text-nofont"]
+                )
+                with gr.Accordion("Generation Parameters", open=False):
+                    default_model = load_dia_model(DIA_MODELS[list(DIA_MODELS.keys())[0]])
+                    max_new_tokens = gr.Slider(
+                        label="Max New Tokens (Audio Length)",
+                        minimum=860,
+                        maximum=3072,
+                        value=getattr(getattr(default_model.config, 'data', None), 'audio_length', 1536),
+                        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,
+                        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,
+                        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,
+                        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).",
+                    )
+                generate_btn = gr.Button("Generate Audio", variant="primary")
+            with gr.Column(scale=1):
+                audio_output = gr.Audio(
+                    label="Generated Audio",
+                    type="numpy",
+                    autoplay=False,
+                )
+        generate_btn.click(
+            run_inference,
+            inputs=[
+                text_input,
+                audio_prompt_input,
+                transcription_input,
+                max_new_tokens,
+                cfg_scale,
+                temperature,
+                top_p,
+                cfg_filter_top_k,
+                speed_factor_slider,
+                model_dropdown,
+            ],
+            outputs=[audio_output],
+        )
+        # Examples (optional, can be extended)
+        examples_list = [
+            [
+                """[S1] އައްސަލާމު އަލައިކުމް. (clears throat) Good morning!
+[S2] How are you today?
+[S1] I'm fine, thanks. ކިހިނެއް ހާލު؟
+[S2] (coughs) ކުޑަކޮށް ބަލިކޮށް މިއުޅެނީ
+[S1] Oh okay. Get well soon...
+[S2] Thanks! See you later... ފަހުން ދިމާވެލާނީ
+[S1]""",
+                None,
+                "",
+                1536,
+                5.0,
+                2.5,
+                0.95,
+                45,
+                1.0,
+                list(DIA_MODELS.keys())[0],
+            ],
+            ["""[FEMALE-01] [S1] ގައުމަށް އައި މިނިވަން ނޫރާނީ... [S2] ދައުރުން މި ހަނދާންތައް އާކުރަނީ... [S1] އައުދާނަ އިތުރު އަބުޠާލުންނަށް... [S2] ޒިކުރާގެ މަލުން މި ވެދުން ކުރަނީ.""",
+                None,
+                "",
+                1536,
+                3.0,
+                1.8,
+                0.95,
+                45,
+                0.96,
+                list(DIA_MODELS.keys())[0]
+            ],
+            ["""[MALE-01] [S1] މާޒީގެ އުޖާލާ މަންޒަރުތައް... [S2] މާރީތި އުފާވެރި ކުރެހުންތައް... [S2] ދާތީ އަދު ހާމަ ވަމުން ކުލަތައް... [S2] ތާރީޚު އަލުން މި އިޢާދަ ވަނީ.""",
+                None,
+                "",
+                1536,
+                3.0,
+                1.8,
+                0.95,
+                45,
+                0.96,
+                list(DIA_MODELS.keys())[0]
+            ],
+        ]
+        if examples_list:
+            gr.Examples(
+                examples=examples_list,
+                inputs=[
+                    text_input,
+                    audio_prompt_input,
+                    transcription_input,
+                    max_new_tokens,
+                    cfg_scale,
+                    temperature,
+                    top_p,
+                    cfg_filter_top_k,
+                    speed_factor_slider,
+                    model_dropdown,
+                ],
+                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)_")
+    gr.Markdown(
+        "---\n"
+        "**General Guidelines:**\n"
+    "- Keep input text length moderate\n"
+    "  - Short input (corresponding to under 5s of audio) will sound unnatural\n"
+    "  - Very long input (corresponding to over 20s of audio) will make the speech unnaturally fast\n\n"
+    "- Use non-verbal tags sparingly, from the list in the README. Overusing or using unlisted non-verbals may cause weird artifacts\n\n"
+    "- Always begin input text with [S1], and always alternate between [S1] and [S2] (i.e. [S1]... [S1]... is not good)\n\n"
+    "**When using audio prompts (voice cloning):**\n"
+    "- Provide the transcript of the to-be cloned audio before the generation text\n"
+    "- Transcript must use [S1], [S2] speaker tags correctly:\n"
+    "  - Single speaker: [S1]...\n"
+    "  - Two speakers: [S1]... [S2]...\n"
+    "- Duration of the to-be cloned audio should be 5~10 seconds for the best results\n"
+    "  - (Keep in mind: 1 second ≈ 86 tokens)\n"
+    "- Put [S1] or [S2] (the second-to-last speaker's tag) at the end of the audio to improve audio quality at the end"
+    )
+    # No explicit return needed for context manager pattern 

example.txt

@@ -0,0 +1,4 @@
+[S1] [dv] Hmm... (coughs) sorry.... ކެއްސާވަރުން!...
+[S2] Good Morning!...
+[S1] ބާއްޖަވެރި ހެދުނެއް ދުރާގާތުން އަޑުއަހާ އެންމެންނަށް (laughs)
+[S2]