Add LoRA fine-tuning scripts (#5)

FINETUNING.md

@@ -0,0 +1,50 @@
+# Finetuning
+
+> [!IMPORTANT]
+> Finetuning is still **very experimental** and not well tested yet!
+
+More instructions coming soon.
+
+Join the [Discord](https://discord.gg/ZDEYTTRxWG) for support.
+
+Example script:
+
+```bash
+python -m vibevoice.finetune.train_vibevoice \
+    --model_name_or_path vibevoice/VibeVoice-1.5B \
+    --dataset_name vibevoice/jenny_vibevoice_formatted \
+    --text_column_name text \
+    --audio_column_name audio \
+    --voice_prompts_column_name audio \
+    --output_dir finetune_vibevoice_zac \
+    --per_device_train_batch_size 8 \
+    --gradient_accumulation_steps 16 \
+    --learning_rate 2.5e-5 \
+    --num_train_epochs 1 \
+    --logging_steps 10 \
+    --save_steps 100 \
+    --eval_steps 100 \
+    --report_to wandb \
+    --remove_unused_columns False \
+    --bf16 True \
+    --do_train \
+    --gradient_clipping \
+    --gradient_checkpointing False \
+    --ddpm_batch_mul 4 \
+    --diffusion_loss_weight 1.4 \
+    --train_diffusion_head True \
+    --ce_loss_weight 0.04 \
+    --voice_prompt_drop_rate 0.2 \
+    --lora_target_modules q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj \
+    --lr_scheduler_type cosine \
+    --warmup_ratio 0.03 \
+    --max_grad_norm 0.8
+```
+
+## Notes
+
+- Currently, only single-speaker finetuning is supported. Podcast fine-tuning is not supported yet.
+- This is an unofficial finetuning implementation, it has not been validated by the original authors.
+- The `voice_prompts_column_name` parameter is currently set to `audio` in the example above, which means the same audio file is used for both training data and voice prompts. This is appropriate when you don't have separate voice prompt files. However, if your dataset includes dedicated voice prompt files (short audio clips that capture the target speaker's voice characteristics), you should specify a different column name that contains these separate voice prompt files. For podcast-style training (once it is supported), the model typically uses the first utterance from each speaker within the podcast episode as the voice prompt, meaning the voice prompt is extracted from the beginning of the same audio file used for training.
+- The dataset text/transcript must be in the format of "Speaker X: text", even if there is only one speaker. Example: `Speaker 1: Hello, how are you?`
+- The default dataset is the Jenny (Dioco) dataset. This is a small dataset for testing purposes and each segment is only a few seconds long. The model may struggle to generate long audio with this dataset.

README.md

@@ -8,7 +8,11 @@
 [![Technical Report](https://img.shields.io/badge/Technical-Report-red)](https://arxiv.org/pdf/2508.19205)
 [![Colab](https://img.shields.io/badge/Colab-Demo-orange?logo=googlecolab)](https://colab.research.google.com/github/vibevoice-community/VibeVoice/blob/main/demo/VibeVoice_colab.ipynb)
 
-https://discord.gg/ZDEYTTRxWG
+## Community
+
+**Join the unofficial Discord community: https://discord.gg/ZDEYTTRxWG** - share samples, ask questions, discuss fine-tuning, etc.
+
+## Overview
 
 VibeVoice is a novel framework designed for generating **expressive**, **long-form**, **multi-speaker** conversational audio, such as podcasts, from text. It addresses significant challenges in traditional Text-to-Speech (TTS) systems, particularly in scalability, speaker consistency, and natural turn-taking.
 
@@ -16,7 +20,9 @@ A core innovation of VibeVoice is its use of continuous speech tokenizers (Acous
 
 The model can synthesize speech up to **90 minutes** long with up to **4 distinct speakers**, surpassing the typical 1-2 speaker limits of many prior models.
 
-[Examples](./EXAMPLES.md)
+## [Examples](./EXAMPLES.md)
+
+## Evaluation
 
 <p align="left">
   <img src="Figures/MOS-preference.png" alt="MOS Preference Results" height="260px">
@@ -24,83 +30,72 @@ The model can synthesize speech up to **90 minutes** long with up to **4 distinc
 </p>
 
 
-### Updates
+## Updates
 
-- **[2025-09-06]** Unofficial training/fine-tuning code coming soon!
 - **[2025-09-05]** Microsoft repo restored (without code) with statement about responsible AI use.
 - **[2025-09-04]** Community backup created after Microsoft removed original repo and models.
 - **[2025-08-26]** The [VibeVoice-7B](https://huggingface.co/vibevoice/VibeVoice-7B) model weights are open-sourced!
 - **[2025-08-28]** [Colab Notebook](https://colab.research.google.com/github/microsoft-community/VibeVoice/blob/main/demo/VibeVoice_colab.ipynb) available. Only VibeVoice-1.5B is supported due to GPU memory limitations.
 
-### Roadmap
+## Roadmap
 
 - [ ] Unofficial/community training code
 - [ ] HF Transformers integration ([PR](https://github.com/huggingface/transformers/pull/40546))
 - [ ] VibePod: End-to-end solution that creates podcasts from documents, webpages, or even a simple topic.
 
+## Model Zoo
 
-## Models
 | Model | Context Length | Generation Length |  Weight |
 |-------|----------------|----------|----------|
-| VibeVoice-0.5B-Streaming | - | - | On the way |
 | VibeVoice-1.5B | 64K | ~90 min | [HF link](https://huggingface.co/vibevoice/VibeVoice-1.5B) |
 | VibeVoice-Large| 32K | ~45 min | [HF link](https://huggingface.co/vibevoice/VibeVoice-7B) |
 
 ## Installation
-We recommend to use NVIDIA Deep Learning Container to manage the CUDA environment. 
 
-1. Launch docker
-```bash
-# NVIDIA PyTorch Container 24.07 / 24.10 / 24.12 verified. 
-# Later versions are also compatible.
-sudo docker run --privileged --net=host --ipc=host --ulimit memlock=-1:-1 --ulimit stack=-1:-1 --gpus all --rm -it nvcr.io/nvidia/pytorch:24.07-py3
-
-## If flash attention is not included in your docker environment, you need to install it manually
-## Refer to https://github.com/Dao-AILab/flash-attention for installation instructions
-# pip install flash-attn --no-build-isolation
-```
-
-2. Install from github
 ```bash
 git clone https://github.com/vibevoice-community/VibeVoice.git
 cd VibeVoice/
 
-pip install -e .
+uv pip install -e .
 ```
 
-## Usages
+## Usage
 
 ### 🚨 Tips
+
 We observed users may encounter occasional instability when synthesizing Chinese speech. We recommend:
 
 - Using English punctuation even for Chinese text, preferably only commas and periods.
 - Using the Large model variant, which is considerably more stable.
 - If you found the generated voice speak too fast. Please try to chunk your text with multiple speaker turns with same speaker label.
 
-We'd like to thank [PsiPi](https://huggingface.co/PsiPi) for sharing an interesting way for emotion control. Detials can be found via [discussion12](https://huggingface.co/microsoft/VibeVoice-1.5B/discussions/12).
+We'd like to thank [PsiPi](https://huggingface.co/PsiPi) for sharing an interesting way for emotion control. Details can be found via [discussion #12](https://huggingface.co/microsoft/VibeVoice-1.5B/discussions/12).
+
+**Option 1: Launch Gradio demo**
 
-### Usage 1: Launch Gradio demo
 ```bash
-apt update && apt install ffmpeg -y # for demo
+python demo/gradio_demo.py --model_path vibevoice/VibeVoice-1.5B --share
-
+# or python demo/gradio_demo.py --model_path vibevoice/VibeVoice-Large --share
-# For 1.5B model
-python demo/gradio_demo.py --model_path microsoft/VibeVoice-1.5B --share
-
-# For Large model
-python demo/gradio_demo.py --model_path microsoft/VibeVoice-Large --share
 ```
 
-### Usage 2: Inference from files directly
+**Option 2: Inference from files directly**
+
 ```bash
 # We provide some LLM generated example scripts under demo/text_examples/ for demo
 # 1 speaker
-python demo/inference_from_file.py --model_path microsoft/VibeVoice-Large --txt_path demo/text_examples/1p_abs.txt --speaker_names Alice
+python demo/inference_from_file.py --model_path vibevoice/VibeVoice-Large --txt_path demo/text_examples/1p_abs.txt --speaker_names Alice
 
 # or more speakers
-python demo/inference_from_file.py --model_path microsoft/VibeVoice-Large --txt_path demo/text_examples/2p_music.txt --speaker_names Alice Frank
+python demo/inference_from_file.py --model_path vibevoice/VibeVoice-Large --txt_path demo/text_examples/2p_music.txt --speaker_names Alice Frank
 ```
 
+
+## [Finetuning](./FINETUNING.md)
+
+NOTE: Finetuning is still **very experimental** and not well tested yet!
+
 ## FAQ
+
 #### Q1: Is this a pretrained model?
 **A:** Yes, it's a pretrained model without any post-training or benchmark-specific optimizations. In a way, this makes VibeVoice very versatile and fun to use.
 
@@ -127,15 +122,14 @@ In fact, we intentionally decided not to denoise our training data because we th
 #### Q6: Instability of cross-lingual transfer.
 **A:** The model does exhibit strong cross-lingual transfer capabilities, including the preservation of accents, but its performance can be unstable. This is an emergent ability of the model that we have not specifically optimized. It's possible that a satisfactory result can be achieved through repeated sampling.
 
-## Risks and limitations
+## Credits
 
-While efforts have been made to optimize it through various techniques, it may still produce outputs that are unexpected, biased, or inaccurate. VibeVoice inherits any biases, errors, or omissions produced by its base model (specifically, Qwen2.5 1.5b in this release).
+- Thanks to [Microsoft](https://github.com/microsoft/VibeVoice) for the original VibeVoice implementation.
-Potential for Deepfakes and Disinformation: High-quality synthetic speech can be misused to create convincing fake audio content for impersonation, fraud, or spreading disinformation. Users must ensure transcripts are reliable, check content accuracy, and avoid using generated content in misleading ways. Users are expected to use the generated content and to deploy the models in a lawful manner, in full compliance with all applicable laws and regulations in the relevant jurisdictions. It is best practice to disclose the use of AI when sharing AI-generated content.
+- Huge shoutout to [Juan Pablo Gallego](https://github.com/jpgallegoar) from [VoicePowered AI](https://www.voicepowered.ai/) for the unofficial training/fine-tuning code.
+- Thanks to [PsiPi](https://huggingface.co/PsiPi) for sharing an interesting way for emotion control. Details can be found via [discussion #12](https://huggingface.co/microsoft/VibeVoice-1.5B/discussions/12).
 
-English and Chinese only: Transcripts in languages other than English or Chinese may result in unexpected audio outputs.
+## License
 
-Non-Speech Audio: The model focuses solely on speech synthesis and does not handle background noise, music, or other sound effects.
+The source code and models are licensed under the MIT License. See the [LICENSE](./LICENSE) file for details.
 
-Overlapping Speech: The current model does not explicitly model or generate overlapping speech segments in conversations.
+Note: Microsoft has removed the original repo and models. This fork is based off of the MIT-licensed code from Microsoft.
-
-We do not recommend using VibeVoice in commercial or real-world applications without further testing and development. This model is intended for research and development purposes only. Please use responsibly.

pyproject.toml

@@ -20,6 +20,8 @@ dependencies = [
     "torch",
     "accelerate==1.6.0",
     "transformers==4.51.3", # we develop this project on transformers==4.51.3, later version may not be compatible
+    "datasets==3.5.0", # avoid issues with torchcodec
+    "peft", # for finetuning
     "llvmlite>=0.40.0",
     "numba>=0.57.0",
     "diffusers",

vibevoice/finetune/__init__.py

@@ -0,0 +1 @@
+# VibeVoice finetuning module

vibevoice/finetune/data_vibevoice.py

@@ -0,0 +1,362 @@
+import math
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
+
+import numpy as np
+import torch
+import warnings
+import random
+
+try:
+    import librosa  # type: ignore
+except Exception:  # pragma: no cover
+    librosa = None  # Fallback: user must install librosa when using local audio paths
+
+try:
+    import resampy  # type: ignore
+except Exception:  # pragma: no cover
+    resampy = None
+
+
+def _resample_if_needed(wav: np.ndarray, orig_sr: int, target_sr: int) -> np.ndarray:
+    if orig_sr == target_sr:
+        return wav.astype(np.float32, copy=False)
+    if resampy is not None:
+        return resampy.resample(wav.astype(np.float32), orig_sr, target_sr)
+    if librosa is not None:
+        return librosa.resample(y=wav.astype(np.float32), orig_sr=orig_sr, target_sr=target_sr)
+    warnings.warn(
+        "No resampler available; treating audio as target_sr without resampling. Install resampy or librosa.",
+        RuntimeWarning,
+    )
+    return wav.astype(np.float32, copy=False)
+
+
+# Lightweight HF-style dataset wrapper (optional). Trainer can also pass raw HF datasets directly.
+class VibeVoiceDataset:
+    def __init__(
+        self,
+        dataset: Any,
+        text_column: str = "text",
+        audio_column: str = "audio",
+        voice_prompts_column: Optional[str] = "voice_prompts",
+    ) -> None:
+        self.dataset = dataset
+        self.text_column = text_column
+        self.audio_column = audio_column
+        self.voice_prompts_column = voice_prompts_column
+
+    def __len__(self) -> int:
+        return len(self.dataset)
+
+    def __getitem__(self, idx: int) -> Dict[str, Any]:
+        item = self.dataset[idx]
+        data: Dict[str, Any] = {}
+        data["text"] = item[self.text_column]
+        data["audio"] = item[self.audio_column]
+
+        user_provided_prompt = None
+        if self.voice_prompts_column and self.voice_prompts_column in item:
+            user_provided_prompt = item[self.voice_prompts_column]
+
+        if user_provided_prompt:
+            # A prompt was provided in the dataset, so we use it.
+            if not isinstance(user_provided_prompt, list):
+                data["voice_prompts"] = [user_provided_prompt]
+            else:
+                data["voice_prompts"] = user_provided_prompt
+        else:
+            # FALLBACK: No prompt provided, so we auto-generate one from the target audio.
+            try:
+                target_sr = 24000
+                wav_array = _load_audio_to_24k(item[self.audio_column], target_sr=target_sr)
+                audio_len_seconds = len(wav_array) / target_sr
+
+                min_len_sec = min(5.0, audio_len_seconds / 4.0)
+                max_len_sec = min(15.0, audio_len_seconds / 2.0)
+                
+                if min_len_sec > max_len_sec:
+                    min_len_sec = max_len_sec
+                max_len_sec = min(max_len_sec, audio_len_seconds)
+
+                if max_len_sec > 0.1:
+                    prompt_len_sec = random.uniform(min_len_sec, max_len_sec)
+                    prompt_len_samples = int(prompt_len_sec * target_sr)
+
+                    max_start_sample = len(wav_array) - prompt_len_samples
+                    start_sample = random.randint(0, max_start_sample)
+                    
+                    prompt_crop = wav_array[start_sample : start_sample + prompt_len_samples]
+                    
+                    data["voice_prompts"] = [prompt_crop]
+                else:
+                    data["voice_prompts"] = None
+
+            except Exception as e:
+                warnings.warn(f"Could not create voice prompt for item {idx}: {e}")
+                data["voice_prompts"] = None            
+        return data
+
+
+def _load_audio_to_24k(audio: Union[str, np.ndarray, torch.Tensor, Dict[str, Any]], *, target_sr: int = 24000) -> np.ndarray:
+    if isinstance(audio, np.ndarray):
+        return audio.astype(np.float32)
+    if isinstance(audio, torch.Tensor):
+        return audio.detach().cpu().float().numpy()
+    if isinstance(audio, str):
+        if librosa is None:
+            raise RuntimeError("librosa is required to load audio file paths. Please pip install librosa.")
+        wav, sr = librosa.load(audio, sr=None, mono=True)
+        wav = _resample_if_needed(wav, int(sr), target_sr)
+        return wav
+    if isinstance(audio, dict) and "array" in audio and "sampling_rate" in audio:
+        arr = np.asarray(audio["array"], dtype=np.float32)
+        sr = int(audio["sampling_rate"])
+        arr = _resample_if_needed(arr, sr, target_sr)
+        return arr
+    raise ValueError(f"Unsupported audio type: {type(audio)}")
+
+
+@dataclass
+class VibeVoiceCollator:
+    processor: Any  # VibeVoiceProcessor
+    max_length: Optional[int] = None
+    speech_compress_ratio: int = 3200
+    semantic_vae_dim: int = 128
+    compute_semantics: bool = False
+    debug_checks: bool = False
+
+    text_field: str = "text"
+    audio_field: str = "audio"
+    voice_prompts_field: str = "voice_prompts"
+    voice_prompt_drop_rate: float = 0.0
+
+    def __call__(self, features: Sequence[Dict[str, Any]]) -> Dict[str, Any]:
+        batch_size = len(features)
+
+        sample_input_ids: List[List[int]] = []
+        sample_attention_masks: List[List[int]] = []
+        sample_acoustic_input_masks: List[List[bool]] = []
+        sample_acoustic_loss_masks: List[List[bool]] = []
+
+        all_speech_waveforms: List[np.ndarray] = []
+        all_speech_latent_lengths: List[int] = []
+        per_segment_is_target: List[bool] = []
+
+        for ex in features:
+            text: str = ex.get(self.text_field, "")
+            voice_prompts: Optional[List[Union[str, np.ndarray, torch.Tensor]]] = ex.get(self.voice_prompts_field)
+            target_audio: Union[str, np.ndarray, torch.Tensor, Dict[str, Any]] = ex.get(self.audio_field)
+
+            # Clamp drop rate for safety
+            _drop_rate = self.voice_prompt_drop_rate
+            if _drop_rate < 0.0:
+                _drop_rate = 0.0
+            elif _drop_rate > 1.0:
+                _drop_rate = 1.0
+
+            proc = self.processor(
+                text=[text],
+                voice_samples=[voice_prompts] if voice_prompts is not None and random.random() >= _drop_rate else None,
+                padding=False,
+                truncation=False,
+                max_length=self.max_length,
+                return_tensors="pt",
+            )
+
+            ids = proc["input_ids"][0].tolist()
+            attn = proc.get("attention_mask", torch.ones_like(proc["input_ids"]))[0].tolist()
+            speech_input_mask = proc.get("speech_input_mask")
+            if speech_input_mask is None:
+                speech_input_mask = torch.zeros_like(proc["input_ids"], dtype=torch.bool)
+            speech_input_mask_list = speech_input_mask[0].tolist()
+
+            wav_target = _load_audio_to_24k(target_audio, target_sr=24000)
+            # Prefer exact frame count from acoustic tokenizer if available; fallback to compress ratio
+            target_latent_len = None
+            try:
+                acoustic_tok = getattr(self.processor, "acoustic_tokenizer", None)
+                if acoustic_tok is not None and hasattr(acoustic_tok, "encode"):
+                    enc_out = acoustic_tok.encode(wav_target)
+                    # Normalize various possible return formats to get time dimension
+                    T = None
+                    try:
+                        # Direct array-like with shape (T, D) or (T,)
+                        if hasattr(enc_out, "shape") and len(getattr(enc_out, "shape", [])) >= 1:
+                            T = int(enc_out.shape[0])
+                        else:
+                            # Nested lists/tuples or ModelOutput-like
+                            cand = enc_out
+                            # Drill down a couple of levels safely
+                            for _ in range(2):
+                                if isinstance(cand, (list, tuple)) and len(cand) > 0:
+                                    cand = cand[0]
+                            if hasattr(cand, "shape") and len(getattr(cand, "shape", [])) >= 1:
+                                T = int(cand.shape[0])
+                    except Exception:
+                        T = None
+                    if T is not None and T > 0:
+                        target_latent_len = T
+            except Exception:
+                target_latent_len = None
+            if target_latent_len is None:
+                target_latent_len = max(1, int(math.ceil(len(wav_target) / float(self.speech_compress_ratio))))
+
+            speech_diff_id = self.processor.tokenizer.speech_diffusion_id
+            target_placeholders = [speech_diff_id] * target_latent_len
+
+            ids_extended = ids + target_placeholders
+            attn_extended = attn + [1] * target_latent_len
+
+            acoustic_input_mask = speech_input_mask_list + [True] * target_latent_len
+            acoustic_loss_mask = ([False] * len(speech_input_mask_list)) + [True] * target_latent_len
+
+            speech_end_id = self.processor.tokenizer.speech_end_id
+            ids_extended.append(speech_end_id)
+            attn_extended.append(1)
+            acoustic_input_mask.append(False)
+            acoustic_loss_mask.append(False)
+
+            if self.max_length is not None and len(ids_extended) > self.max_length:
+                cut = len(ids_extended) - int(self.max_length)
+                leading_non_acoustic = 0
+                for v in acoustic_input_mask:
+                    if v:
+                        break
+                    leading_non_acoustic += 1
+                if cut > leading_non_acoustic:
+                    raise ValueError(
+                        f"--max_length={self.max_length} would truncate into acoustic tokens. "
+                        f"Needed cut={cut}, but only {leading_non_acoustic} leading non-acoustic tokens available. "
+                        "Increase max_length or shorten text/voice-prompt preamble."
+                    )
+                ids_extended = ids_extended[cut:]
+                attn_extended = attn_extended[cut:]
+                acoustic_input_mask = acoustic_input_mask[cut:]
+                acoustic_loss_mask = acoustic_loss_mask[cut:]
+
+            sample_input_ids.append(ids_extended)
+            sample_attention_masks.append(attn_extended)
+            sample_acoustic_input_masks.append(acoustic_input_mask)
+            sample_acoustic_loss_masks.append(acoustic_loss_mask)
+
+            voice_speeches = []
+            voice_latent_lengths = []
+            if proc.get("speech_tensors") is not None:
+                voice_np = proc["speech_tensors"].cpu().numpy()
+                voice_masks = proc["speech_masks"].cpu().numpy().astype(bool)
+                for seg_idx in range(voice_np.shape[0]):
+                    voice_speeches.append(voice_np[seg_idx])
+                    voice_latent_lengths.append(int(voice_masks[seg_idx].sum()))
+
+            all_speech_waveforms.extend(voice_speeches)
+            all_speech_latent_lengths.extend(voice_latent_lengths)
+            per_segment_is_target.extend([False] * len(voice_speeches))
+
+            all_speech_waveforms.append(wav_target)
+            all_speech_latent_lengths.append(target_latent_len)
+            per_segment_is_target.append(True)
+
+        max_seq_len = max(len(x) for x in sample_input_ids)
+        padded_input_ids = []
+        padded_attention_masks = []
+        padded_acoustic_input_masks = []
+        padded_acoustic_loss_masks = []
+        tok = self.processor.tokenizer
+        pad_token_id = getattr(tok, "pad_token_id", None)
+        if pad_token_id is None or pad_token_id < 0:
+            pad_token_id = getattr(tok, "eos_token_id", None)
+            if pad_token_id is None or pad_token_id < 0:
+                raise ValueError(
+                    "Tokenizer has no pad_token_id or eos_token_id; please set one or pass a valid pad id."
+                )
+        for ids, attn, ain_mask, aloss_mask in zip(
+            sample_input_ids, sample_attention_masks, sample_acoustic_input_masks, sample_acoustic_loss_masks
+        ):
+            pad_len = max_seq_len - len(ids)
+            padded_input_ids.append(ids + [pad_token_id] * pad_len)
+            padded_attention_masks.append(attn + [0] * pad_len)
+            padded_acoustic_input_masks.append(ain_mask + [False] * pad_len)
+            padded_acoustic_loss_masks.append(aloss_mask + [False] * pad_len)
+
+        input_ids_tensor = torch.tensor(padded_input_ids, dtype=torch.long)
+        attention_mask_tensor = torch.tensor(padded_attention_masks, dtype=torch.long)
+        acoustic_input_mask_tensor = torch.tensor(padded_acoustic_input_masks, dtype=torch.bool)
+        acoustic_loss_mask_tensor = torch.tensor(padded_acoustic_loss_masks, dtype=torch.bool)
+
+        if all_speech_waveforms:
+            max_wave_len = max(w.shape[0] for w in all_speech_waveforms)
+            padded_speeches = np.zeros((len(all_speech_waveforms), max_wave_len), dtype=np.float32)
+            for i, w in enumerate(all_speech_waveforms):
+                L = w.shape[0]
+                padded_speeches[i, :L] = w
+
+            max_latent_len = max(all_speech_latent_lengths) if all_speech_latent_lengths else 1
+            speech_masks_np = np.zeros((len(all_speech_waveforms), max_latent_len), dtype=np.bool_)
+            for i, L_lat in enumerate(all_speech_latent_lengths):
+                speech_masks_np[i, :L_lat] = True
+
+            speech_tensors_tensor = torch.tensor(padded_speeches, dtype=torch.float32)
+            speech_masks_tensor = torch.tensor(speech_masks_np, dtype=torch.bool)
+
+            speeches_loss_input_np = np.zeros_like(speech_masks_np, dtype=np.bool_)
+            for i, is_target in enumerate(per_segment_is_target):
+                if is_target:
+                    speeches_loss_input_np[i] = speech_masks_np[i]
+            speeches_loss_input_tensor = torch.tensor(speeches_loss_input_np, dtype=torch.bool)
+
+            # Semantic features
+            if self.compute_semantics and hasattr(self.processor, "semantic_tokenizer") and self.processor.semantic_tokenizer is not None:
+                sem_feats: List[np.ndarray] = []
+                for w in all_speech_waveforms:
+                    try:
+                        # Expect [T, D]  where T ≈ ceil(len(w)/compress_ratio)
+                        sem = self.processor.semantic_tokenizer.encode(w)
+                        sem = np.asarray(sem, dtype=np.float32)
+                    except Exception:
+                        sem = np.zeros((0, self.semantic_vae_dim), dtype=np.float32)
+                    if sem.ndim != 2:
+                        raise RuntimeError(f"Semantic tokenizer returned unexpected shape {sem.shape}. Expect [T, D].")
+                    L = sem.shape[0]
+                    D = sem.shape[1]
+                    if D != self.semantic_vae_dim:
+                        if D < self.semantic_vae_dim:
+                            pad_d = np.zeros((L, self.semantic_vae_dim - D), dtype=np.float32)
+                            sem = np.concatenate([sem, pad_d], axis=1)
+                        else:
+                            sem = sem[:, : self.semantic_vae_dim]
+                    if L < max_latent_len:
+                        pad = np.zeros((max_latent_len - L, self.semantic_vae_dim), dtype=np.float32)
+                        sem = np.concatenate([sem, pad], axis=0)
+                    elif L > max_latent_len:
+                        sem = sem[:max_latent_len]
+                    sem_feats.append(sem.astype(np.float32))
+                speech_semantic_tensors = torch.tensor(np.stack(sem_feats, axis=0), dtype=torch.float32)
+            else:
+                # Semantic tokenizer unavailable while semantics are required for training.
+                # Raise to avoid silently degrading alignment with zeroed features.
+                raise RuntimeError(
+                    "Semantic features are required but could not be computed. "
+                    "Ensure processor.semantic_tokenizer is available or precompute and provide features."
+                )
+        else:
+            speech_tensors_tensor = None
+            speech_masks_tensor = None
+            speeches_loss_input_tensor = None
+            speech_semantic_tensors = None  # No segments in batch
+
+        if self.debug_checks:
+            assert (input_ids_tensor >= 0).all(), "input_ids contains negative indices"
+            if speech_tensors_tensor is not None:
+                assert speech_tensors_tensor.dim() == 2, "Expected speech_tensors 2D [segments, samples]"
+
+        return {
+            "input_ids": input_ids_tensor,
+            "attention_mask": attention_mask_tensor,
+            "speech_tensors": speech_tensors_tensor,
+            "speech_masks": speech_masks_tensor,
+            "speech_semantic_tensors": speech_semantic_tensors,
+            "acoustic_input_mask": acoustic_input_mask_tensor,
+            "acoustic_loss_mask": acoustic_loss_mask_tensor,
+            "speeches_loss_input": speeches_loss_input_tensor,
+        }

vibevoice/finetune/train_vibevoice.py

@@ -0,0 +1,902 @@
+# train_vibevoice_lora.py
+import logging
+import os
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from datasets import load_dataset, DatasetDict, VerificationMode
+
+from transformers import (
+    HfArgumentParser,
+    Trainer,
+    set_seed,
+    TrainerCallback,
+)
+from transformers import TrainingArguments as HfTrainingArguments
+
+from peft import LoraConfig, get_peft_model, TaskType
+
+from vibevoice.modular.modeling_vibevoice import VibeVoiceForConditionalGeneration
+from vibevoice.modular.configuration_vibevoice import VibeVoiceConfig
+from vibevoice.processor.vibevoice_processor import VibeVoiceProcessor
+
+from vibevoice.finetune.data_vibevoice import VibeVoiceDataset, VibeVoiceCollator
+
+logger = logging.getLogger(__name__)
+
+# ================== SAMPLE CALLBACK UTILS ==================
+
+import copy
+import torch
+from transformers import TrainerCallback
+
+class EmaCallback(TrainerCallback):
+    def __init__(self, attr_path="model.prediction_head", decay=0.999, device="cpu"):
+        """
+        attr_path: where the head lives under self.model (Trainer wraps your VibeVoiceForConditionalGeneration)
+        decay:     EMA decay (0.999 ~ stable, 0.9999 ~ very smooth, slower to adapt)
+        """
+        self.attr_path = attr_path
+        self.decay = float(decay)
+        self.device = torch.device(device)
+        self.shadow = None
+        self._orig = None  # store non-EMA weights when we swap
+
+    def _get_module(self, model):
+        # Resolve dotted path like "model.prediction_head"
+        mod = model
+        for name in self.attr_path.split('.'):
+            mod = getattr(mod, name)
+        return mod
+
+    def on_train_begin(self, args, state, control, model=None, **kwargs):
+        head = self._get_module(model)
+        self.shadow = {k: p.detach().to(self.device).clone()
+                       for k, p in head.state_dict().items()}
+
+    def on_step_end(self, args, state, control, model=None, **kwargs):
+        if self.shadow is None: return
+        head = self._get_module(model)
+        with torch.no_grad():
+            for k, v in head.state_dict().items():
+                self.shadow[k].mul_(self.decay).add_(v.detach().to(self.device), alpha=(1.0 - self.decay))
+
+    # ---- Swap helpers ----
+    def _swap_in_ema(self, model):
+        head = self._get_module(model)
+        self._orig = copy.deepcopy(head.state_dict())
+        head.load_state_dict(self.shadow, strict=False)
+
+    def _swap_back(self, model):
+        if self._orig is None: return
+        head = self._get_module(model)
+        head.load_state_dict(self._orig, strict=False)
+        self._orig = None
+
+    def on_evaluate(self, args, state, control, model=None, **kwargs):
+        # use EMA during eval
+        self._swap_in_ema(model)
+
+    def on_evaluate_end(self, args, state, control, model=None, **kwargs):
+        self._swap_back(model)
+
+    def on_save(self, args, state, control, model=None, **kwargs):
+        # temporarily swap to EMA, let Trainer save, then swap back
+        self._swap_in_ema(model)
+
+    def on_save_end(self, args, state, control, model=None, **kwargs):
+        self._swap_back(model)
+
+    def on_train_end(self, args, state, control, model=None, **kwargs):
+        # final checkpoint: persist EMA
+        self._swap_in_ema(model)
+
+
+@dataclass
+class ModelArguments:
+    model_name_or_path: Optional[str] = field(
+        default=None, metadata={"help": "Path to VibeVoice base model with config.json"}
+    )
+    processor_name_or_path: Optional[str] = field(
+        default=None, metadata={"help": "Path to processor dir (preprocessor_config.json). Defaults to model path."}
+    )
+    cache_dir: Optional[str] = field(default=None)
+    freeze_acoustic_tokenizer: bool = field(default=True)
+    freeze_semantic_tokenizer: bool = field(default=True)
+    lora_r: int = field(default=8)
+    lora_alpha: int = field(default=32)
+    lora_dropout: float = field(default=0.05)
+    lora_target_modules: str = field(
+        default="q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj",
+        metadata={"help": "Comma-separated list of target module names in the LLM blocks"},
+    )
+    lora_wrap_diffusion_head: bool = field(default=False, metadata={"help": "Wrap diffusion head with PEFT LoRA"})
+    train_diffusion_head: bool = field(default=False, metadata={"help": "Train diffusion prediction head (full fine-tune)"})
+    train_connectors: bool = field(default=False, metadata={"help": "Train acoustic/semantic connectors (full fine-tune)"})
+    layers_to_freeze: Optional[str] = field(
+        default=None, 
+        metadata={"help": "Comma-separated indices of diffusion head layers to freeze (e.g., '0,1,5,7,8')."}
+    )
+
+@dataclass
+class DataArguments:
+    dataset_name: Optional[str] = field(default=None, metadata={"help": "HF dataset name or 'json' with --train_jsonl for local files"})
+    dataset_config_name: Optional[str] = field(default=None)
+    train_split_name: str = field(default="train")
+    eval_split_name: Optional[str] = field(default="validation")
+    text_column_name: str = field(default="text")
+    audio_column_name: str = field(default="audio")
+    voice_prompts_column_name: Optional[str] = field(default="voice_prompts")
+    eval_split_size: float = field(default=0.0)
+    ignore_verifications: bool = field(default=False)
+    max_length: Optional[int] = field(default=None)
+    train_jsonl: Optional[str] = field(default=None, metadata={"help": "Path to local train JSONL with {text, audio, [voice_prompts]}"})
+    validation_jsonl: Optional[str] = field(default=None, metadata={"help": "Optional path to local validation JSONL"})
+    voice_prompt_drop_rate: float = field(
+        default=0.0,
+        metadata={"help": "Probability to drop conditioning voice prompt during training (0.0 keep always, 1.0 drop always)."},
+    )
+
+@dataclass
+class CustomTrainingArguments(HfTrainingArguments):
+    ddpm_batch_mul: int = field(default=1)
+    ce_loss_weight: float = field(default=1.0)
+    diffusion_loss_weight: float = field(default=1.0)
+    debug_ce_details: bool = field(default=False)
+    debug_ce_topk: int = field(default=5)
+    debug_ce_max_examples: int = field(default=1)
+    debug_ce_every_n_steps: int = field(default=200)
+    gradient_clipping: bool = field(
+        default=False,
+        metadata={"help": "Enable gradient clipping using max_grad_norm (set via --max_grad_norm, default 1.0). When False, disables clipping by forcing max_grad_norm=0.0."},
+    )
+    debug_save: bool = field(
+        default=False,
+        metadata={"help": "If set, saves model components BEFORE training starts, into output_dir/debug_initial."},
+    )
+
+def build_lora_config(args: ModelArguments) -> LoraConfig:
+    target_modules = [s.strip() for s in args.lora_target_modules.split(",") if s.strip()]
+    return LoraConfig(
+        r=args.lora_r,
+        lora_alpha=args.lora_alpha,
+        lora_dropout=args.lora_dropout,
+        bias="none",
+        task_type=TaskType.CAUSAL_LM,
+        target_modules=target_modules,
+    )
+
+def build_head_lora_config(args: ModelArguments) -> LoraConfig:
+    target_modules = ["noisy_images_proj","cond_proj","gate_proj","up_proj","down_proj","linear"]
+    return LoraConfig(
+        r=args.lora_r,
+        lora_alpha=args.lora_alpha,
+        lora_dropout=args.lora_dropout,
+        bias="none",
+        task_type=TaskType.FEATURE_EXTRACTION,
+        target_modules=target_modules,
+    )
+
+def mask_for_ce(labels: torch.Tensor, attention_mask: torch.Tensor, acoustic_input_mask: torch.Tensor, pad_id: int = -100) -> torch.Tensor:
+    shifted = labels[:, 1:].contiguous()
+    base_mask = attention_mask[:, 1:].contiguous().eq(1) if (attention_mask is not None and attention_mask.numel() > 0) else torch.ones_like(shifted, dtype=torch.bool)
+    label_is_acoustic = acoustic_input_mask[:, 1:].contiguous()
+    final_mask = base_mask & (~label_is_acoustic)
+    out = shifted.clone()
+    out[~final_mask] = pad_id
+    return out
+
+def _patch_acoustic_encode_for_legacy_indexing(model_obj, logger_):
+    try:
+        acoustic = getattr(getattr(model_obj, "model", model_obj), "acoustic_tokenizer", None)
+        if acoustic is None or not hasattr(acoustic, "encode"):
+            logger_.warning("No acoustic_tokenizer.encode() found to patch.")
+            return
+        base_encode = acoustic.encode
+        def encode_wrapped(*args, **kwargs):
+            out = base_encode(*args, **kwargs)
+            try:
+                _ = out[0][0]
+                return out
+            except Exception:
+                pass
+            if isinstance(out, dict):
+                for k in ("frames", "codes", "tokens", "latents", "hidden_states"):
+                    if k in out:
+                        return [[out[k]]]
+                if len(out) > 0:
+                    return [[next(iter(out.values()))]]
+            for attr in ("frames", "codes", "tokens", "latents", "hidden_states"):
+                if hasattr(out, attr):
+                    return [[getattr(out, attr)]]
+            try:
+                if isinstance(out, torch.Tensor):
+                    return [[out]]
+            except Exception:
+                pass
+            return [[out]]
+        acoustic.encode = encode_wrapped
+        logger_.info("Patched acoustic_tokenizer.encode() to return [[...]] for legacy indexing.")
+    except Exception as e:
+        logger_.warning(f"Failed to patch acoustic_tokenizer.encode(): {e}")
+
+def main() -> None:
+    parser = HfArgumentParser((ModelArguments, DataArguments, CustomTrainingArguments))
+    model_args, data_args, training_args = parser.parse_args_into_dataclasses()
+
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+        datefmt="%m/%d/%Y %H:%M:%S",
+        level=logging.INFO if training_args.local_rank in [-1, 0] else logging.WARN,
+    )
+    logger.info("Training/evaluation parameters %s", training_args)
+    set_seed(training_args.seed)
+
+    # Configure gradient clipping
+    if not getattr(training_args, "gradient_clipping", False):
+        if hasattr(training_args, "max_grad_norm"):
+            training_args.max_grad_norm = 0.0
+            logger.info("Gradient clipping disabled (set max_grad_norm=0.0). Use --gradient_clipping to enable.")
+    else:
+        if (not hasattr(training_args, "max_grad_norm")) or training_args.max_grad_norm is None or training_args.max_grad_norm <= 0:
+            training_args.max_grad_norm = 1.0
+        logger.info(f"Gradient clipping enabled: max_grad_norm={training_args.max_grad_norm}")
+
+    # Load processor
+    processor_path = model_args.processor_name_or_path or model_args.model_name_or_path
+    if processor_path is None:
+        raise ValueError("--model_name_or_path (or --processor_name_or_path) must be provided")
+    processor: VibeVoiceProcessor = VibeVoiceProcessor.from_pretrained(processor_path)
+
+    # Required special tokens
+    tok = processor.tokenizer
+    for required in ["speech_start_id", "speech_diffusion_id", "speech_end_id"]:
+        if not hasattr(tok, required) or getattr(tok, required) is None:
+            raise RuntimeError(f"Tokenizer missing required special id: {required}")
+
+    # Load model
+    if model_args.model_name_or_path is None:
+        raise ValueError("--model_name_or_path is required to load VibeVoice base model")
+    dtype = torch.float32
+    if training_args.bf16:
+        dtype = torch.bfloat16
+    elif getattr(training_args, "fp16", False):
+        dtype = torch.float16
+    model = VibeVoiceForConditionalGeneration.from_pretrained(
+        model_args.model_name_or_path,
+        torch_dtype=dtype,
+    )
+    _patch_acoustic_encode_for_legacy_indexing(model, logger)
+    processor.semantic_tokenizer = getattr(model.model, "semantic_tokenizer", None)
+
+    # Diagnostics: LM head tie
+    try:
+        in_emb_mod = model.get_input_embeddings()
+        out_emb_mod = model.get_output_embeddings()
+        in_w = getattr(in_emb_mod, "weight", None)
+        out_w = getattr(out_emb_mod, "weight", None)
+        shared_ptr = bool(in_w is not None and out_w is not None and in_w.data_ptr() == out_w.data_ptr())
+        values_equal = False
+        if in_w is not None and out_w is not None and in_w.shape == out_w.shape:
+            try:
+                values_equal = bool(torch.allclose(in_w, out_w))
+            except Exception:
+                values_equal = False
+        try:
+            tie_cfg = getattr(getattr(model.config, "decoder_config", model.config), "tie_word_embeddings", None)
+        except Exception:
+            tie_cfg = getattr(model.config, "tie_word_embeddings", None)
+        logger.info(f"LM head diagnostics -> shared_params={shared_ptr}, values_equal={values_equal}, tie_word_embeddings={tie_cfg}")
+        if out_w is not None:
+            logger.info(f"LM head requires_grad before freeze: {bool(out_w.requires_grad)}")
+    except Exception as e:
+        logger.warning(f"LM head tie diagnostics failed: {e}")
+
+    # Hard-tie LM head
+    try:
+        emb_module = model.get_input_embeddings()
+        head_module = model.get_output_embeddings()
+        if hasattr(emb_module, "weight") and hasattr(head_module, "weight"):
+            if emb_module.weight.shape == head_module.weight.shape and emb_module.weight.data_ptr() != head_module.weight.data_ptr():
+                with torch.no_grad():
+                    head_module.weight = emb_module.weight
+                logger.info("Force-tied LM head weight to input embeddings (pointer share).")
+    except Exception as e:
+        logger.warning(f"Force-tie of LM head failed: {e}")
+
+    # Validate special IDs (info logs only)
+    try:
+        special_names = ["speech_start_id", "speech_diffusion_id", "speech_end_id"]
+        try:
+            vocab_size = int(getattr(model.config.decoder_config, "vocab_size", 0))
+        except Exception:
+            vocab_size = 0
+        in_emb_mod = model.get_input_embeddings()
+        out_emb_mod = model.get_output_embeddings()
+        in_w = getattr(in_emb_mod, "weight", None)
+        out_w = getattr(out_emb_mod, "weight", None)
+        for name in special_names:
+            val = getattr(tok, name, None)
+            exists = (val is not None)
+            in_range = (exists and isinstance(val, int) and 0 <= val < vocab_size)
+            equal_row = None
+            if in_range and in_w is not None and out_w is not None and in_w.shape == out_w.shape and in_w.size(0) > val:
+                try:
+                    equal_row = bool(torch.allclose(in_w[val], out_w[val]))
+                except Exception:
+                    equal_row = False
+            decoded_str = None
+            if exists and isinstance(val, int):
+                try:
+                    decoded_str = tok.decode([val])
+                except Exception:
+                    try:
+                        decoded_str = tok.convert_ids_to_tokens(val)
+                    except Exception:
+                        decoded_str = "<decode_failed>"
+            logger.info(f"Special token check -> {name}={val}, decoded='{decoded_str}', exists={exists}, in_vocab_range={in_range}, emb_vs_head_row_equal={equal_row}")
+    except Exception as e:
+        logger.warning(f"Special token ID/row validation failed: {e}")
+
+    # Quick tokenizer diagnostics (optional)
+    try:
+        logger.info("=== TOKENIZER DIAGNOSTICS ===")
+        logger.info(f"Tokenizer class: {type(tok).__name__}")
+        logger.info(f"Tokenizer vocab_size: {tok.vocab_size}")
+        # tiny CE smoke test
+        with torch.no_grad():
+            simple_text = "The cat sat on the mat."
+            simple_ids = torch.tensor([tok.encode(simple_text, add_special_tokens=True)], device=model.device)
+            simple_mask = torch.ones_like(simple_ids)
+            x = model.get_input_embeddings()(simple_ids)
+            outputs = model.model(inputs_embeds=x, attention_mask=simple_mask, return_dict=True)
+            logits = model.lm_head(outputs.last_hidden_state)
+            shift_logits = logits[:, :-1, :].contiguous()
+            shift_labels = simple_ids[:, 1:].contiguous()
+            ce_loss = F.cross_entropy(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1), reduction='mean')
+            logger.info(f"Simple text CE loss: {ce_loss.item():.4f}")
+    except Exception as e:
+        logger.warning(f"Tokenizer diagnostics failed: {e}")
+
+    # Disable cache during training
+    if hasattr(model.config, "use_cache") and training_args.do_train:
+        model.config.use_cache = False
+
+    # Freeze tokenizers
+    if model_args.freeze_acoustic_tokenizer and hasattr(model.model, "acoustic_tokenizer"):
+        for p in model.model.acoustic_tokenizer.parameters():
+            p.requires_grad = False
+    if model_args.freeze_semantic_tokenizer and hasattr(model.model, "semantic_tokenizer"):
+        for p in model.model.semantic_tokenizer.parameters():
+            p.requires_grad = False
+
+    # LoRA wrap LLM (optional)
+    lora_cfg = build_lora_config(model_args)
+    tm_lower = [s.strip().lower() for s in model_args.lora_target_modules.split(",") if s.strip()]
+    skip_lm_lora = (len(tm_lower) == 0) or all(t in ("none", "off", "disable", "disabled") for t in tm_lower)
+    if not skip_lm_lora:
+        model.model.language_model = get_peft_model(model.model.language_model, lora_cfg)
+    else:
+        logger.info("Skipping LLM LoRA wrapping (lora_target_modules indicates none).")
+
+    try:
+        model.tie_weights()
+    except Exception:
+        pass
+
+    # Freeze all then enable trainable subsets
+    for _, p in model.named_parameters():
+        p.requires_grad = False
+
+    try:
+        for n, p in model.model.language_model.named_parameters():
+            if "lora_A" in n or "lora_B" in n:
+                p.requires_grad = True
+    except Exception:
+        logger.warning("Could not re-enable LoRA params on language_model.")
+
+    # Diffusion head LoRA wrapping (optional)
+    if getattr(model_args, "lora_wrap_diffusion_head", False) and hasattr(model.model, "prediction_head"):
+        class _HeadForwardShim(nn.Module):
+            def __init__(self, base: nn.Module): super().__init__(); self.base = base
+            def forward(self, *args, **kwargs):
+                if len(args) >= 3:
+                    noisy_images, timesteps, condition = args[:3]
+                else:
+                    noisy_images = kwargs.get("noisy_images")
+                    timesteps = kwargs.get("timesteps")
+                    condition = kwargs.get("condition")
+                return self.base(noisy_images, timesteps, condition)
+        try:
+            shim = _HeadForwardShim(model.model.prediction_head)
+            model.model.prediction_head = get_peft_model(shim, build_head_lora_config(model_args))
+            for n, p in model.model.prediction_head.named_parameters():
+                if "lora_A" in n or "lora_B" in n:
+                    p.requires_grad = True
+        except Exception as e:
+            logger.warning(f"Could not LoRA-wrap diffusion head: {e}")
+
+    # Train full diffusion head (optional)
+    if getattr(model_args, "train_diffusion_head", False) and hasattr(model.model, "prediction_head"):
+        for p in model.model.prediction_head.parameters():
+            p.requires_grad = True
+
+    # Freeze diffusion head layers (optional)
+    if model_args.layers_to_freeze is not None and hasattr(model.model, "prediction_head"):
+        head_params = list(model.model.prediction_head.named_parameters())
+        try:
+            indices_to_freeze = {int(x.strip()) for x in model_args.layers_to_freeze.split(',') if x.strip()}
+            frozen_count = 0
+            for i, (name, param) in enumerate(head_params):
+                if i in indices_to_freeze:
+                    param.requires_grad = False
+                    frozen_count += 1
+                    logger.info(f"Froze layer [{i}]: {name}")
+            logger.info(f"Successfully froze {frozen_count} parameter groups in the diffusion head.")
+        except Exception as e:
+            logger.error(f"Could not parse --layers_to_freeze: {e}")
+            raise
+    
+    # Connectors
+    if getattr(model_args, "train_connectors", False):
+        if hasattr(model.model, "acoustic_connector"):
+            for p in model.model.acoustic_connector.parameters():
+                p.requires_grad = True
+        if hasattr(model.model, "semantic_connector"):
+            for p in model.model.semantic_connector.parameters():
+                p.requires_grad = True
+    else:
+        if hasattr(model.model, "acoustic_connector"):
+            for p in model.model.acoustic_connector.parameters():
+                p.requires_grad = False
+        if hasattr(model.model, "semantic_connector"):
+            for p in model.model.semantic_connector.parameters():
+                p.requires_grad = False
+
+    # Freeze embedding + head
+    try:
+        emb = model.get_input_embeddings()
+        if hasattr(emb, "weight"):
+            emb.weight.requires_grad_(False)
+        head = model.get_output_embeddings()
+        if head is not None and hasattr(head, "weight"):
+            head.weight.requires_grad_(False)
+    except Exception:
+        pass
+
+    # Diagnostics
+    def _sum_params(named_iter):
+        return sum(p.numel() for _, p in named_iter if p.requires_grad)
+    try:
+        lm_lora = _sum_params(model.model.language_model.named_parameters()) if hasattr(model.model, "language_model") else 0
+        pred_head_train = _sum_params(model.model.prediction_head.named_parameters()) if hasattr(model.model, "prediction_head") else 0
+        ac_conn_train = _sum_params(model.model.acoustic_connector.named_parameters()) if hasattr(model.model, "acoustic_connector") else 0
+        se_conn_train = _sum_params(model.model.semantic_connector.named_parameters()) if hasattr(model.model, "semantic_connector") else 0
+        total_trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+        logger.info(f"Trainable by block -> LLM-LoRA: {lm_lora:,} | diff_head: {pred_head_train:,} | ac_conn: {ac_conn_train:,} | se_conn: {se_conn_train:,}")
+        logger.info("TOTAL trainable: %s", f"{total_trainable:,}")
+    except Exception:
+        pass
+
+    # Datasets
+    verification_mode = VerificationMode.NO_CHECKS if data_args.ignore_verifications else VerificationMode.BASIC_CHECKS
+    if data_args.train_jsonl is not None:
+        data_files: Dict[str, str] = {"train": data_args.train_jsonl}
+        if data_args.validation_jsonl is not None:
+            data_files["validation"] = data_args.validation_jsonl
+        raw = load_dataset("json", data_files=data_files, verification_mode=verification_mode, cache_dir=model_args.cache_dir)
+    else:
+        if data_args.dataset_name is None:
+            raise ValueError("Provide --dataset_name (HF datasets) or use --train_jsonl/--validation_jsonl for local files.")
+        raw = load_dataset(
+            data_args.dataset_name,
+            data_args.dataset_config_name,
+            verification_mode=verification_mode,
+            cache_dir=model_args.cache_dir,
+        )
+    train_ds = raw[data_args.train_split_name]
+    eval_ds = None
+    if training_args.do_eval:
+        if data_args.eval_split_name and data_args.eval_split_name in raw:
+            eval_ds = raw[data_args.eval_split_name]
+        elif data_args.eval_split_size and data_args.eval_split_size > 0 and len(train_ds) > 1:
+            split = train_ds.train_test_split(test_size=data_args.eval_split_size, seed=training_args.seed)
+            train_ds, eval_ds = split["train"], split["test"]
+
+    train_dataset = VibeVoiceDataset(
+        train_ds,
+        text_column=data_args.text_column_name,
+        audio_column=data_args.audio_column_name,
+        voice_prompts_column=data_args.voice_prompts_column_name,
+    )
+    eval_dataset = None
+    if eval_ds is not None:
+        eval_dataset = VibeVoiceDataset(
+            eval_ds,
+            text_column=data_args.text_column_name,
+            audio_column=data_args.audio_column_name,
+            voice_prompts_column=data_args.voice_prompts_column_name,
+        )
+
+    # Ratios/dims from processor+model
+    speech_compress_ratio = getattr(processor, "speech_tok_compress_ratio", 3200)
+    semantic_dim = getattr(model.config, "semantic_vae_dim", None)
+    if semantic_dim is None:
+        try:
+            semantic_dim = int(getattr(model.config.semantic_tokenizer_config, "vae_dim", 128))
+        except Exception:
+            semantic_dim = 128
+
+    compute_semantics_flag = hasattr(processor, "semantic_tokenizer") and processor.semantic_tokenizer is not None
+
+    data_collator = VibeVoiceCollator(
+        processor=processor,
+        max_length=data_args.max_length,
+        speech_compress_ratio=speech_compress_ratio,
+        semantic_vae_dim=semantic_dim,
+        compute_semantics=compute_semantics_flag,
+        debug_checks=False,
+        voice_prompt_drop_rate=data_args.voice_prompt_drop_rate,
+    )
+
+    class LoRADebugCallback(TrainerCallback):
+        def __init__(self, log_every_n_steps: int = 50):
+            self.log_every_n_steps = max(1, int(log_every_n_steps))
+            self.prev_param_norms: Dict[str, float] = {}
+            self.lora_param_names: List[str] = []
+
+        def on_train_begin(self, args, state, control, model=None, **kwargs):
+            try:
+                if model is None:
+                    return
+                named: Dict[str, torch.nn.Parameter] = dict(model.named_parameters())
+                self.lora_param_names = [n for n in named.keys() if ("lora_A" in n or "lora_B" in n)]
+                for n in self.lora_param_names:
+                    p = named[n]
+                    self.prev_param_norms[n] = float(p.data.norm().item())
+                total = len(self.lora_param_names)
+                req_grad = sum(1 for n in self.lora_param_names if named[n].requires_grad)
+                num_A = sum(1 for n in self.lora_param_names if "lora_A" in n)
+                num_B = sum(1 for n in self.lora_param_names if "lora_B" in n)
+                zero_B = sum(1 for n in self.lora_param_names if ("lora_B" in n and float(named[n].data.norm().item()) == 0.0))
+                logger.info(f"LoRA debug: found {total} LoRA params (A={num_A}, B={num_B}); trainable={req_grad}. Initial lora_B_zero={zero_B}.")
+                if total == 0:
+                    logger.warning("LoRA debug: No LoRA parameters found. Check lora_target_modules.")
+                if req_grad != total:
+                    logger.warning("LoRA debug: Some LoRA params are frozen. They should be trainable.")
+            except Exception as e:
+                logger.warning(f"LoRA debug (on_train_begin) failed: {e}")
+
+        def on_step_end(self, args, state, control, model=None, **kwargs):
+            try:
+                if model is None or len(self.lora_param_names) == 0:
+                    return
+                step = int(getattr(state, "global_step", 0) or 0)
+                if step % self.log_every_n_steps != 0 and step != 1:
+                    return
+                named: Dict[str, torch.nn.Parameter] = dict(model.named_parameters())
+                changed_A = 0
+                changed_B = 0
+                zero_B = 0
+                eps = 1e-12
+                for n in self.lora_param_names:
+                    p = named.get(n, None)
+                    if p is None:
+                        continue
+                    prev = self.prev_param_norms.get(n, 0.0)
+                    curr = float(p.data.norm().item())
+                    if "lora_A" in n and abs(curr - prev) > eps:
+                        changed_A += 1
+                    if "lora_B" in n:
+                        if abs(curr - prev) > eps:
+                            changed_B += 1
+                        if curr == 0.0:
+                            zero_B += 1
+                    self.prev_param_norms[n] = curr
+                total_A = sum(1 for n in self.lora_param_names if "lora_A" in n)
+                total_B = sum(1 for n in self.lora_param_names if "lora_B" in n)
+                logger.info(f"LoRA debug step {step}: changed A {changed_A}/{total_A}, changed B {changed_B}/{total_B}, lora_B_zero_now={zero_B}.")
+            except Exception as e:
+                logger.warning(f"LoRA debug (on_step_end) failed: {e}")
+
+    class VibeVoiceTrainer(Trainer):
+        def compute_loss(self, model: VibeVoiceForConditionalGeneration, inputs: Dict[str, Any], return_outputs=False, num_items_in_batch: Optional[int] = None):
+            labels = inputs.get("input_ids")
+            attention_mask = inputs.get("attention_mask")
+            acoustic_input_mask = inputs.get("acoustic_input_mask")
+
+            # Ensure semantic tensors exist and have correct dtype/device
+            sem = inputs.get("speech_semantic_tensors", None)
+            try:
+                target_dtype = next(model.model.semantic_connector.parameters()).dtype
+            except Exception:
+                target_dtype = model.get_input_embeddings().weight.dtype
+
+            if sem is None:
+                sm = inputs.get("speech_masks")
+                if sm is not None:
+                    zeros = torch.zeros(
+                        sm.size(0), sm.size(1),
+                        getattr(model.config, "semantic_vae_dim", 128),
+                        dtype=target_dtype,
+                        device=sm.device,
+                    )
+                    inputs["speech_semantic_tensors"] = zeros
+            else:
+                if isinstance(sem, torch.Tensor):
+                    inputs["speech_semantic_tensors"] = sem.to(dtype=target_dtype)
+
+            outputs = model(
+                input_ids=inputs.get("input_ids"),
+                attention_mask=attention_mask,
+                speech_tensors=inputs.get("speech_tensors"),
+                speech_masks=inputs.get("speech_masks"),
+                speech_semantic_tensors=inputs.get("speech_semantic_tensors"),
+                acoustic_input_mask=acoustic_input_mask,
+                acoustic_loss_mask=inputs.get("acoustic_loss_mask"),
+                speeches_loss_input=inputs.get("speeches_loss_input"),
+                ddpm_batch_mul=training_args.ddpm_batch_mul,
+            )
+
+            # Invariants: token/latent selection equality across views (warn, don't assert)
+            try:
+                al_mask = inputs.get("acoustic_loss_mask")
+                sp_masks = inputs.get("speech_masks")
+                sp_loss_sel = inputs.get("speeches_loss_input")
+                num_tok_total = int(acoustic_input_mask.sum().item()) if acoustic_input_mask is not None else 0
+                num_tok_loss = int(al_mask.sum().item()) if al_mask is not None else 0
+                num_lat_total = int(sp_masks.sum().item()) if sp_masks is not None else 0
+                num_lat_loss = int(((sp_loss_sel & sp_masks).sum().item())) if (sp_loss_sel is not None and sp_masks is not None) else 0
+                self.log({
+                    "debug/num_tok_total": float(num_tok_total),
+                    "debug/num_tok_loss": float(num_tok_loss),
+                    "debug/num_lat_total": float(num_lat_total),
+                    "debug/num_lat_loss": float(num_lat_loss),
+                })
+                if sp_loss_sel is not None and sp_masks is not None and al_mask is not None:
+                    if num_tok_loss != num_lat_loss:
+                        logger.warning(f"Loss selection mismatch: acoustic_loss_mask={num_tok_loss} vs speeches_loss_input={num_lat_loss}")
+            except Exception:
+                pass
+
+            # CE Loss
+            logits = outputs.logits
+            ce_labels = mask_for_ce(labels, attention_mask, acoustic_input_mask, pad_id=-100)
+            shift_logits = logits[:, :-1, :].contiguous()
+            loss_fct = nn.CrossEntropyLoss(ignore_index=-100)
+            ce_loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), ce_labels.view(-1))
+
+            # Optional CE diagnostics
+            try:
+                self._debug_ce(shift_logits, ce_labels, attention_mask, acoustic_input_mask)
+            except Exception as e:
+                logger.warning(f"Failed invoking CE debug: {e}")
+
+            # Diffusion loss
+            diffusion_loss = outputs.diffusion_loss if outputs.diffusion_loss is not None else torch.tensor(0.0, device=ce_loss.device)
+            total = training_args.ce_loss_weight * ce_loss + training_args.diffusion_loss_weight * diffusion_loss
+
+            # Logs
+            try:
+                prefix = "train" if model.training else "eval"
+                self.log({
+                    f"{prefix}/ce_loss": ce_loss.detach().item(),
+                    f"{prefix}/diffusion_loss": diffusion_loss.detach().item() if isinstance(diffusion_loss, torch.Tensor) else float(diffusion_loss),
+                })
+                if hasattr(self, "optimizer") and self.optimizer is not None and len(self.optimizer.param_groups) > 0:
+                    lr_val = self.optimizer.param_groups[0].get("lr", None)
+                    if lr_val is not None:
+                        self.log({"train/learning_rate_real": float(lr_val)})
+            except Exception:
+                pass
+
+            return (total, outputs) if return_outputs else total
+
+        def _debug_ce(self, shift_logits: torch.Tensor, ce_labels: torch.Tensor, attention_mask: Optional[torch.Tensor], acoustic_input_mask: Optional[torch.Tensor]):
+            try:
+                if not getattr(training_args, "debug_ce_details", False):
+                    return
+                step = int(getattr(self.state, "global_step", 0) or 0)
+                every_n = max(1, int(getattr(training_args, "debug_ce_every_n_steps", 200) or 200))
+                if not (step <= 1 or (step % every_n == 0)):
+                    return
+
+                with torch.no_grad():
+                    vocab = shift_logits.size(-1)
+                    per_token_loss = F.cross_entropy(
+                        shift_logits.view(-1, vocab),
+                        ce_labels.view(-1),
+                        reduction="none",
+                        ignore_index=-100,
+                    ).view_as(ce_labels)
+
+                    valid_mask = ce_labels.ne(-100)
+                    num_valid = int(valid_mask.sum().item())
+                    avg_loss = float((per_token_loss[valid_mask].mean().item())) if num_valid > 0 else float("nan")
+
+                    per_ex_avgs = []
+                    max_examples = max(1, int(getattr(training_args, "debug_ce_max_examples", 1) or 1))
+                    B = ce_labels.size(0)
+                    for b in range(min(B, max_examples)):
+                        vb = valid_mask[b]
+                        if int(vb.sum().item()) > 0:
+                            per_ex_avgs.append(float(per_token_loss[b][vb].mean().item()))
+                        else:
+                            per_ex_avgs.append(float("nan"))
+                    logger.info(f"CE debug: tokens_in_loss={num_valid}, avg_loss={avg_loss:.4f}, per_example_avgs={[round(x,4) if x==x else None for x in per_ex_avgs]}")
+            except Exception as e:
+                logger.warning(f"CE detailed debug failed: {e}")
+
+        # --------- CRITICAL SAVE OVERRIDES: also dump FULL head/connectors for inference ---------
+  
+
+        def _save(self, output_dir: Optional[str] = None, state_dict=None) -> None:
+            try:
+                target_dir = output_dir or self.args.output_dir
+                lora_out = os.path.join(target_dir, "lora")
+                os.makedirs(lora_out, exist_ok=True)
+    
+                # --- LLM PEFT adapters (if LoRA-wrapped) ---
+                language_model = getattr(self.model.model, "language_model", None)
+                if hasattr(language_model, "save_pretrained"):
+                    language_model.save_pretrained(lora_out)
+    
+                # --- Diffusion head PEFT adapters (if LoRA-wrapped) ---
+                pred_head = getattr(self.model.model, "prediction_head", None)
+                if hasattr(pred_head, "save_pretrained"):
+                    ph_dir = os.path.join(lora_out, "diffusion_head")
+                    os.makedirs(ph_dir, exist_ok=True)
+                    pred_head.save_pretrained(ph_dir)
+    
+                # --- ALWAYS save FULL diffusion head state_dict for fallback ---
+                if pred_head is not None and hasattr(pred_head, "state_dict"):
+                    sd = pred_head.state_dict()
+                    torch.save(sd, os.path.join(lora_out, "diffusion_head_full.bin"))
+                    ph_dir = os.path.join(lora_out, "diffusion_head")
+                    os.makedirs(ph_dir, exist_ok=True)
+                    torch.save(sd, os.path.join(ph_dir, "diffusion_head_full.bin"))
+    
+                # --- Connectors (plain state_dicts) ---
+                ac = getattr(self.model.model, "acoustic_connector", None)
+                if ac is not None:
+                    ac_dir = os.path.join(lora_out, "acoustic_connector")
+                    os.makedirs(ac_dir, exist_ok=True)
+                    torch.save(ac.state_dict(), os.path.join(ac_dir, "pytorch_model.bin"))
+    
+                se = getattr(self.model.model, "semantic_connector", None)
+                if se is not None:
+                    se_dir = os.path.join(lora_out, "semantic_connector")
+                    os.makedirs(se_dir, exist_ok=True)
+                    torch.save(se.state_dict(), os.path.join(se_dir, "pytorch_model.bin"))
+    
+            except Exception as e:
+                logger.warning(f"Failed to save LoRA assets: {e}")
+
+
+    # ------------- Build the Trainer -------------
+
+    # Resolve which adapters to apply in samples
+
+    ema_cb = EmaCallback(attr_path="model.prediction_head", decay=0.999, device="cpu")
+
+    trainer = VibeVoiceTrainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        data_collator=data_collator,
+        callbacks=[ema_cb, LoRADebugCallback(log_every_n_steps=(int(getattr(training_args, "logging_steps", 50) or 50)))],
+    )
+
+    # Optional debug pre-training save
+    if getattr(training_args, "debug_save", False):
+        try:
+            debug_dir = os.path.join(training_args.output_dir, "debug_initial")
+            lora_out = os.path.join(debug_dir, "lora")
+            os.makedirs(lora_out, exist_ok=True)
+            logger.info(f"[debug_save] Saving initial (pre-training) model components to: {debug_dir}")
+            # language model adapters / base
+            try:
+                if hasattr(model.model.language_model, "save_pretrained"):
+                    model.model.language_model.save_pretrained(lora_out)
+            except Exception as e_lm:
+                logger.warning(f"[debug_save] Failed to save language_model: {e_lm}")
+            # diffusion head
+            try:
+                if hasattr(model.model, "prediction_head") and hasattr(model.model.prediction_head, "save_pretrained"):
+                    model.model.prediction_head.save_pretrained(os.path.join(lora_out, "diffusion_head"))
+            except Exception as e_head:
+                logger.warning(f"[debug_save] Failed to save prediction_head: {e_head}")
+            # NEW: full diffusion head state_dict as fallback
+            try:
+                ph = getattr(model.model, "prediction_head", None)
+                if ph is not None and hasattr(ph, "state_dict"):
+                    sd = ph.state_dict()
+                    torch.save(sd, os.path.join(lora_out, "diffusion_head_full.bin"))
+                    os.makedirs(os.path.join(lora_out, "diffusion_head"), exist_ok=True)
+                    torch.save(sd, os.path.join(lora_out, "diffusion_head", "diffusion_head_full.bin"))
+            except Exception as e:
+                logger.warning(f"[debug_save] Failed to save FULL diffusion head: {e}")
+            # connectors
+            try:
+                ac_conn = getattr(model.model, "acoustic_connector", None)
+                if ac_conn is not None:
+                    ac_dir = os.path.join(lora_out, "acoustic_connector")
+                    os.makedirs(ac_dir, exist_ok=True)
+                    torch.save(ac_conn.state_dict(), os.path.join(ac_dir, "pytorch_model.bin"))
+            except Exception as e_ac:
+                logger.warning(f"[debug_save] Failed to save acoustic_connector: {e_ac}")
+            try:
+                se_conn = getattr(model.model, "semantic_connector", None)
+                if se_conn is not None:
+                    se_dir = os.path.join(lora_out, "semantic_connector")
+                    os.makedirs(se_dir, exist_ok=True)
+                    torch.save(se_conn.state_dict(), os.path.join(se_dir, "pytorch_model.bin"))
+            except Exception as e_se:
+                logger.warning(f"[debug_save] Failed to save semantic_connector: {e_se}")
+        except Exception as e:
+            logger.warning(f"[debug_save] Unexpected failure saving initial components: {e}")
+
+    if getattr(training_args, "gradient_checkpointing", False):
+        try:
+            model.gradient_checkpointing_enable()
+        except Exception:
+            logger.warning("Failed to enable gradient checkpointing on the model.")
+
+    if training_args.do_train:
+        trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint)
+    
+        lora_out = os.path.join(training_args.output_dir, "lora")
+        os.makedirs(lora_out, exist_ok=True)
+    
+        # LLM PEFT (if any)
+        lm = getattr(model.model, "language_model", None)
+        if hasattr(lm, "save_pretrained"):
+            lm.save_pretrained(lora_out)
+    
+        # Diffusion head PEFT (if any)
+        ph = getattr(model.model, "prediction_head", None)
+        if hasattr(ph, "save_pretrained"):
+            ph_dir = os.path.join(lora_out, "diffusion_head")
+            os.makedirs(ph_dir, exist_ok=True)
+            ph.save_pretrained(ph_dir)
+    
+        # ALWAYS: full diffusion head state_dict fallback
+        try:
+            if ph is not None and hasattr(ph, "state_dict"):
+                sd = ph.state_dict()
+                torch.save(sd, os.path.join(lora_out, "diffusion_head_full.bin"))
+                ph_dir = os.path.join(lora_out, "diffusion_head")
+                os.makedirs(ph_dir, exist_ok=True)
+                torch.save(sd, os.path.join(ph_dir, "diffusion_head_full.bin"))
+        except Exception as e:
+            logger.warning(f"Failed to save FULL diffusion head at end: {e}")
+    
+        # Connectors (if trained)
+        try:
+            ac = getattr(model.model, "acoustic_connector", None)
+            if ac is not None:
+                ac_dir = os.path.join(lora_out, "acoustic_connector")
+                os.makedirs(ac_dir, exist_ok=True)
+                torch.save(ac.state_dict(), os.path.join(ac_dir, "pytorch_model.bin"))
+        except Exception as e:
+            logger.warning(f"Failed to save acoustic_connector: {e}")
+    
+        try:
+            se = getattr(model.model, "semantic_connector", None)
+            if se is not None:
+                se_dir = os.path.join(lora_out, "semantic_connector")
+                os.makedirs(se_dir, exist_ok=True)
+                torch.save(se.state_dict(), os.path.join(se_dir, "pytorch_model.bin"))
+        except Exception as e:
+            logger.warning(f"Failed to save semantic_connector: {e}")
+
+    if training_args.do_eval and eval_dataset is not None:
+        trainer.evaluate()
+
+
+if __name__ == "__main__":
+    main()

vibevoice/modular/modeling_vibevoice.py

@@ -368,11 +368,18 @@ class VibeVoiceForConditionalGeneration(VibeVoicePreTrainedModel):
                 )
             if speech_tensors is not None:
                 if semantic_speech_all_connect_features is not None:
-                    x[acoustic_input_mask] = speech_all_connect_features[speech_masks] + semantic_speech_all_connect_features[speech_masks]
+                    x[acoustic_input_mask] = (
+                        speech_all_connect_features[speech_masks]
+                        + semantic_speech_all_connect_features[speech_masks]
+                    )
                 else:
                     x[acoustic_input_mask] = speech_all_connect_features[speech_masks]
-                speech_features = speech_all_features[speeches_loss_input.unsqueeze(-1) & speech_masks] # only part audio need diffuse
+
-                speech_connect_features = speech_all_connect_features[speeches_loss_input.unsqueeze(-1) & speech_masks]
+                # Select only the target segments' latents for diffusion loss.
+                # Both masks are [num_segments, max_latent_len]; using 2D mask on [B,T,D] selects [N_true, D].
+                target_latent_mask = speeches_loss_input & speech_masks
+                speech_features = speech_all_features[target_latent_mask]
+                speech_connect_features = speech_all_connect_features[target_latent_mask]
         else:
             speech_features, speech_connect_features = self.forward_speech_features(
                     speech_tensors=speech_tensors.type_as(x) if speech_tensors is not None else None,

vibevoice/processor/vibevoice_processor.py

@@ -426,6 +426,14 @@ class VibeVoiceProcessor:
             if isinstance(speaker_audio, str):
                 # Load audio from file
                 wav = self.audio_processor._load_audio_from_path(speaker_audio)
+            elif isinstance(speaker_audio, dict):
+                # Handle dict format with 'array' or 'audio' key
+                if 'array' in speaker_audio:
+                    wav = np.array(speaker_audio['array'], dtype=np.float32)
+                elif 'audio' in speaker_audio:
+                    wav = np.array(speaker_audio['audio'], dtype=np.float32)
+                else:
+                    raise ValueError(f"Dictionary audio input must have 'array' or 'audio' key, got: {speaker_audio.keys()}")
             else:
                 wav = np.array(speaker_audio, dtype=np.float32)