Improve Inference Code

.gitattributes

@@ -35,3 +35,8 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
 *.jpg filter=lfs diff=lfs merge=lfs -text
 *.png filter=lfs diff=lfs merge=lfs -text
+fp32.wav filter=lfs diff=lfs merge=lfs -text
+int4.wav filter=lfs diff=lfs merge=lfs -text
+samples/voices/en-Alice_woman.wav filter=lfs diff=lfs merge=lfs -text
+tts_out.wav filter=lfs diff=lfs merge=lfs -text
+why1.wav filter=lfs diff=lfs merge=lfs -text

HANDOFF.md

@@ -0,0 +1,191 @@
+# VibeVoice → ONNX — Engineer's Handoff Note
+
+_Written by the departing engineer. Everything here was learned the hard way; read it before
+touching the code. Companion files: `STATUS.md` (running log), `.claude/skills/vibevoice-onnx/`
+(the operational skill), and the checklist at the bottom of this note._
+
+---
+
+## 1. The one-paragraph mental model
+
+The VibeVoice family is **one Qwen2 LLM backbone + audio VAE tokenizers + small glue MLPs**, in
+four checkpoints that *look* alike but differ in critical details. We do NOT export the composite
+model. We decompose each checkpoint into standalone ONNX sub-parts — LLM decoder via
+onnxruntime-genai **ModelBuilder** (int4/fp16, KV cache), everything else via **Olive**
+(dynamo exporter) — and re-compose them at inference time in plain Python + onnxruntime
+(`common.py`). Every exported part is parity-checked against PyTorch (cosine ~1.0) before we
+trust it. That decomposition + parity discipline is the whole method.
+
+## 2. The four checkpoints — what actually differs (memorize this table)
+
+| Key | HF layout | LLM | lm_head? | Acoustic tokenizer naming | Loads via |
+|---|---|---|---|---|---|
+| `1.5b` | codes/ (`vibevoice`) | Qwen2.5-1.5B, 28L | **No** (head = diffusion) | `encoder.downsample_layers.*` | `codes/` |
+| `asr` | codes/ (`vibevoice`, arch `VibeVoiceForASRTraining`) | Qwen2.5-**7B** | **Yes** (top-level `lm_head.weight`) | same as 1.5b | `codes/` |
+| `asr-hf` | transformers-native (`vibevoice_asr`) | Qwen2.5-**7B** | **Yes** (`language_model.lm_head.*`) | `acoustic_tokenizer_encoder.conv_layers.*` | transformers |
+| `realtime` | codes/ (`vibevoice_streaming`) | Qwen2.5-0.5B, **20L (config says 24!)** | No | `decoder.stages.*` — **decoder-only, no encoder shipped** | `codes/` |
+
+Rules that fall out of this:
+- **Never assume tokenizer classes are interchangeable across checkpoints.** Same class name,
+  three different weight namings. Always load with `strict=False` and **assert 0 missing /
+  0 unexpected** — that assert is the tripwire that caught every mismatch.
+- **TTS backbones** build with `exclude_embeds + exclude_lm_head` (inputs_embeds → hidden_states);
+  **ASR decoders keep lm_head** (→ logits). `optimize.py`'s `MODELS` registry encodes this.
+- `asr` vs `asr-hf` are **different models** with different key layouts — resist merging them.
+- Realtime has **TWO embed tables** (`language_model` + `tts_language_model`) that differ.
+  `common.EMBED_KEY` maps each key to the exact table matching the exported decoder. A
+  first-match heuristic here silently produced garbage audio once. Don't reintroduce it.
+
+## 3. Traps that cost real time (do not rediscover these)
+
+1. **`import vibevoice` collides with transformers** (both register the same model_type) and the
+   package `__init__` pulls diffusers + a qwen2-tokenizer module renamed in transformers ≥5.10.
+   Fix: `_codes_import()` in `user_script.py` — inject empty namespace packages so only the one
+   target module is imported, and shim `Auto*.register` to swallow duplicate registrations.
+2. **dynamo ignores `dynamic_axes`.** Use `dynamic_shapes` in io_configs or the dim is baked
+   static. Also: the frame axis of `latents[B, frames, 64]` is **dim 1** (dim 2 is vae_dim) —
+   we shipped that off-by-one once.
+3. **Timesteps must be float32.** The diffusion head's `TimestepEmbedder` casts its sinusoidal
+   embedding back to `t.dtype` before a float MLP; int64 timesteps crash the matmul.
+4. **CFG needs ONE sample.** Conditional and unconditional eps must come from the SAME noisy
+   latent (duplicate the sample for the head call only, each step). Two independent samples =
+   guidance mixed with a noise difference = degraded audio, no error.
+5. **int4 is fine for token generation, lossy for raw hidden states** (first-token hidden cos
+   ~0.83 vs fp32 across 28 layers of 4-bit RTN). The TTS diffusion head is conditioned on hidden
+   states → **use fp16 builds for audio quality**; int4 for footprint experiments only.
+6. **7B/8B int4 ModelBuilder serialize needs ~16–20 GB free RAM** (chandra-8B and both ASR 7Bs
+   hit this). Extraction is separate and cheap — stream shards, `del` after each. Subprocess
+   isolation does NOT help a single large serialize; you need physical RAM.
+7. **Parity metric footnote:** cosine collapses on near-silent outputs (decoder fed random
+   latents ≈ silence → cosine noise-dominated while max|Δ| ≈ 1e-10). Pass criterion is
+   `cos ≥ 0.99 OR max|Δ| tiny` — that's deliberate, not sloppy.
+8. **Windows consoles are cp1252.** Olive logs emoji → `UnicodeEncodeError`. Every entrypoint
+   reconfigures stdout/stderr to UTF-8 at import. Keep that block.
+9. **`argparse nargs='*'` + trailing positional don't mix** — `--components` is comma-separated
+   for that reason.
+10. **VibeVoice repos ship no tokenizer/processor.** We fetch the plain Qwen2.5 tokenizer; exact
+    prompt layout and voice-cloning need the original processor assets (see codes/ `processor/`).
+11. **fp16 conversion needs an `op_block_list` for shape/resample ops.** `OnnxFloatToFloat16` on the
+    VAE/conv graphs will emit an invalid graph if it touches `ConstantOfShape` (its fp16 output hits a
+    float32 consumer → `Type Error … does not match expected type (tensor(float))`), and `ConvTranspose`
+    /`Resize`/`Range` gain nothing from fp16. `build_olive` blocks `["ConstantOfShape","ConvTranspose",
+    "Resize","Range"]`. (The standalone `acoustic` encoder is where this first bit — the codes/-layout
+    1.5b acoustic encoder happened not to.)
+12. **int4 only quantizes the LLM.** Keys with no LLM (`acoustic`) or the audio/VAE/DiT components have
+    no MatMulNBits-quantizable weights, so int4 is a no-op that silently re-emits fp32. `build_model`
+    warns and downgrades `int4→fp32` for LLM-less keys so you don't ship a misnamed fp32 copy.
+13. **Eval must feed each ONNX its DECLARED input dtype.** fp16 builds expect `float16`; feeding fp32
+    dummies → `Unexpected input data type`. `eval.py` (`parity_component` + `whole_pipeline._feed`) casts
+    per graph; float32-pinned inputs (diffusion `timesteps`) stay fp32. Without this, fp16 builds falsely
+    "fail" eval though the graphs are fine.
+14. **The VibeVoice source is VENDORED at `VibeVoice/vibevoice/` (~591 KB, MIT — see `VIBEVOICE_LICENSE`).**
+    NOT a submodule, NOT a pip/git dependency. `_vibevoice_dir()` (user_script + common) returns that
+    directory and raises if it's missing. The former `codes/` submodule, the `vibevoice-repo/` workspace
+    clone (~263 MB each), the `[submodule]` in `.gitmodules`, the `vibevoice @ git+…` line in the onnx
+    pyprojects, and the root-pyproject `vibevoice` dep + `vibevoice-repo` workspace were ALL removed so
+    the tree is self-contained and uploadable. Imports still go through the isolated-import shim
+    (`_codes_import` / the namespace-injection in `make_scheduler`/`_load_vv_processor`) — see trap 1;
+    never `import vibevoice`. Only the required subtree is vendored (modular/, processor/, schedule/,
+    scripts/, configs/ — no demo/, finetuning-asr/, vllm_plugin/).
+15. **The 1.5B has a learned EOS — do NOT ship a fixed frame budget.** `lm_head` is TIED to
+    `embed_tokens` (`_tied_weights_keys=["lm_head.weight"]`), so logits = `hidden @ embedᵀ` with NO
+    lm_head export. VibeVoice reuses vision tokens for speech: stop when the model predicts speech-end
+    (`<|vision_end|>`) or `<|endoftext|>`. `inference.py` computes this each frame via
+    `common.load_embed_matrix()` and breaks. This killed the "jargon tail" (a real sonnet stopped at
+    frame 101 of a 172 cap → 13.5 s not 22.9 s). `--max-frames` is now just a safety cap. Upstream
+    ships no non-streaming 1.5B `generate` (codes/ @ 303b283 = training forward + streaming only);
+    single-shot AR + this EOS is our reconstruction and it works — chunking regressed it, do not reintroduce.
+    (Upstream check was against the vendored `vibevoice/modular/modeling_vibevoice.py` @ upstream 303b283.)
+
+## 4. Architecture of our code (7 files, one direction of dependency)
+
+```
+optimize.py      MODELS registry (per-key: extract fn, exclude flags, olive component specs)
+   │             + ensure_checkpoint (auto-download) + detect_model_type + output layout
+   ▼
+user_script.py   ALL loaders. codes/ isolated-import shim; per-checkpoint weight collectors;
+                 4× extract_qwen2_* (standalone HF dirs for ModelBuilder); io_configs; dummies.
+   ▼
+common.py        Inference primitives: OnnxLLM (manual KV-cache driver over the genai graph,
+                 hidden_states OR logits), DiffusionSampler (DPM + CFG), audio io, EMBED_KEY,
+                 resolve_from_path (parses onnx/{key}/{device}_{precision}).
+   ▼
+inference.py / inference_asr.py / inference_realtime.py   thin drivers (one positional: built dir)
+eval.py          parity (A) + whole-pipeline TTS checks (B); reuses the registry.
+```
+
+Conventions: output layout `onnx/{key}/{device}_{precision}`; model keyword is always the FINAL
+positional; drivers take ONLY the built dir and derive model_id/device/precision from the path;
+`codes/` = unmodified `github.com/microsoft/VibeVoice` @ `303b283` (submodule-able, imported
+read-only by path — never pip-install it, never edit it).
+
+## 5. What is done vs pending
+
+**Done & parity-verified (cos ~1.0):** 19 sub-models — 1.5b (7/7 incl. int4 LLM), realtime (4/4),
+asr front-end (5), asr-hf front-end (3). End-to-end TTS runs for 1.5b + realtime.
+
+**Pending:** (a) 7B ASR LLM builds — RAM-bound, run on a big-memory machine; front-end +
+`inference_asr.py` are ready and degrade cleanly without it. (b) fp16 builds for TTS quality.
+(c) Encoders bake audio length at 24000 samples (switch io_configs to `dynamic_shapes` to lift).
+(d) Learned EOS (1.5b lm_head / realtime `tts_eos_classifier`) not exported → `--max-frames`
+stop. (e) Cleanup: fold the 4 extractors into one helper; dedupe the ~9 shard-collect loops;
+delete dead stub classes; factor the duplicated driver frame-loop into `common.py`.
+
+**Environment:** use the PROJECT env (`uv run` from the repo; deps in root `pyproject.toml`).
+Do NOT add PEP-723 inline headers — they bypass the curated venv. CUDA needs GPU builds of
+onnxruntime/genai present in the env. `.gitmodules` needs fixing: root-level,
+`path = VibeVoice/codes`, url `https://github.com/microsoft/VibeVoice`, pin `303b283`.
+
+---
+
+# THE CHECKLIST — adding a new VibeVoice-family checkpoint (or re-running everything)
+
+## A. Recon (30 min — do not skip)
+- [ ] Dump `config.json`: `model_type`, `architectures`, sub-configs (`decoder_config` /
+      `text_config`, `*_tokenizer_config`, `diffusion_head_config` — note `hidden_size`!).
+- [ ] Dump weight groups: `Counter('.'.join(k.split('.')[:2]) for k in index/weight_map)`.
+      Identify: LLM prefix, lm_head (present? where?), tokenizer groups + their naming style,
+      connectors, heads. Compare against the table in §2.
+- [ ] Decide codes/ vs transformers-native per component by **instantiating and loading with
+      strict=False; require 0 missing / 0 unexpected**. Try both if unsure.
+- [ ] Check actual layer count vs config (`realtime` lied: 20 real vs 24 configured).
+
+## B. Wire loaders (`user_script.py`)
+- [ ] Weight collector per component (prefix-strip; stream shards for >2 GB; assert 0/0).
+- [ ] Extractor for the LLM → standalone Qwen2 dir (correct config sub-key, tokenizer id,
+      lm_head kept iff the model emits text; override layer count if config lies).
+- [ ] io_config: `dynamic_shapes` (NOT `dynamic_axes`), frame axis = dim 1, float32 timesteps.
+- [ ] Register in `optimize.py` `MODELS` (+ `HF_REPO`, `common.EMBED_KEY`, `common.MODEL_IDS`).
+
+## C. Build & verify (never skip verify)
+- [ ] `uv run optimize.py --device cuda --precision fp16 <key>` (int4 only for the LLM footprint;
+      `--exclude-llm` if RAM-bound). Output: `onnx/<key>/<device>_<precision>/`.
+- [ ] `uv run eval.py <key>` → every component PASS; whole-pipeline codec round-trip
+      corr > 0.99 / SNR > +15 dB on the tone test.
+- [ ] LLM structural check: MatMulNBits/GQA counts = layer count; inputs = inputs_embeds (TTS)
+      or logits output (ASR). KV-driver sanity: incremental step == full prefill (cos 1.0).
+- [ ] Smoke the matching driver end-to-end (`--max-frames 8` is enough to prove the loop).
+
+## D. Ship
+- [ ] Update `STATUS.md` matrix row + a dated section (what, parity numbers, gotchas hit).
+- [ ] No scratch files (`_*.py`, `*.log`), no stale `qwen2_*_standalone/` (15 GB each!), no
+      old-layout outputs. Disk-full here corrupts builds mid-serialize.
+- [ ] If it's a new trap, add it to §3 of this note. That's the contract.
+
+# THE 5-MINUTE OPERATOR CARD (for someone who just wants to run it)
+
+```bash
+# build (downloads checkpoint if absent; output → onnx/<key>/<device>_<precision>)
+uv run optimize.py --device cuda --precision fp16 1.5b        # or asr | asr-hf | realtime | all
+uv run optimize.py --exclude-llm asr                          # front-end only (low RAM)
+
+# verify
+uv run eval.py 1.5b
+
+# run
+uv run inference.py          --text "Hello world." onnx/1.5b/cuda_fp16
+uv run inference_realtime.py --text "Hi."          onnx/realtime/cuda_fp16
+uv run inference_asr.py      --audio speech.wav    onnx/asr-hf/cuda_fp16
+```
+Drivers print `model_id / device / precision` derived from the path and refuse the wrong model
+type. If ASR says the 7B LLM isn't built, build it on a ≥32 GB-free-RAM machine.

STATUS.md

@@ -0,0 +1,286 @@
+# VibeVoice-1.5B — Current Status
+
+> **Handoff:** read [HANDOFF.md](HANDOFF.md) — mental model, per-checkpoint table, the 10 traps, the new-checkpoint checklist, and the 5-minute operator card. Skill: `.claude/skills/vibevoice-onnx/`.
+
+_Updated: 2026-07-17_
+
+## 2026-07-17 — 1.5b rebuilt & re-verified on CPU (after checkpoints re-downloaded)
+
+Checkpoints re-downloaded: **1.5B** (`model/`, 3/3 shards ✅) and **realtime** (`realtime/`, ✅)
+complete; **asr** (3/8) and **asr-hf** (4/8) still downloading; `acoustic/` config-only. This box
+is **CPU-only (no NVIDIA)**, so builds use `--device cpu` (LLM int4 — genai has no fp16-CPU).
+
+Rebuilt all 7 **1.5b** sub-parts → `onnx/1.5b/cpu_int4/` (5.2 GB: llm_decoder int4 + acoustic
+enc/dec + semantic enc + diffusion_head + acoustic/semantic connectors). `eval.py 1.5b` = **9/9
+PASS** (llm: 140 MatMulNBits / 28 GQA / inputs_embeds; encoders/head/connectors cos 1.0;
+acoustic_decoder cos 0.935 but max|Δ| 3.5e-10 = near-silent case; codec round-trip corr +0.996 /
+SNR +19 dB; tts chain smoke OK). Driver smoke: `inference.py … onnx/1.5b/cpu_int4` → 2.13 s WAV.
+
+Rebuilt all 4 **realtime** sub-parts → `onnx/realtime/cpu_int4/` (1.7 GB: llm_decoder int4 [20
+layers, decoder-only] + acoustic_decoder + diffusion_head + acoustic_connector). `eval.py realtime`
+= **5/5 PASS** (llm 100 MatMulNBits / 20 GQA / inputs_embeds; diffusion_head + connector cos 1.0;
+acoustic_decoder cos 0.893 but max|Δ| 3.07e-10 near-silent; tts chain smoke H=896 OK). Driver smoke:
+`inference_realtime.py … onnx/realtime/cpu_int4` → 2.13 s WAV. Scratch `qwen2_*_standalone/` removed.
+
+Built + verified both **7B front-ends** on CPU (`--exclude-llm`; the 7B LLM int4 serialize is still
+RAM-bound → needs a ≥32 GB-free box):
+- **asr-hf** → `onnx/asr-hf/cpu_int4/` (acoustic_encoder + semantic_encoder + multi_modal_projector).
+  `eval.py asr-hf` = **4/4 PASS** (encoders + projector cos 1.0; asr fusion chain fused=(1,7,3584);
+  LLM/codec/tts-chain SKIP as expected without the LLM).
+- **asr** → `onnx/asr/cpu_int4/` (acoustic enc/dec + semantic enc + acoustic/semantic connectors).
+  `eval.py asr` = **6/6 PASS** (encoders/connectors cos 1.0; acoustic_decoder cos 0.935 maxd 3.9e-10
+  near-silent; codec round-trip corr +0.996 / SNR +19 dB).
+
+Added **`acoustic`** registry key — the standalone `vibevoice_acoustic_tokenizer` checkpoint
+(`acoustic/`, 1.3 GB), transformers-native (`AutoModel`, no `codes/`). New loaders
+`get_acoustic_std_{encoder,decoder}_model` (encode→`.latents`, decode→`.sample`); `all_components`
+handles the no-LLM case; `detect_model_type` maps `vibevoice_acoustic_tokenizer`→`acoustic`. Built
+→ `onnx/acoustic/cpu_fp32/`; `eval.py --precision fp32 acoustic` = **3/3 PASS** (encoder cos 0.999,
+decoder cos 0.939 near-silent, codec round-trip corr +0.999 / SNR +22.9 dB).
+
+All checkpoints downloaded: 1.5b, realtime, asr (17 GB), asr-hf (16 GB), standalone `acoustic` (1.3 GB).
+
+`.gitmodules` fixed: `codes` submodule URL corrected `huggingface/vibevoice` → `microsoft/VibeVoice`
+(pin 303b283) — the old URL left `codes/` empty on fresh clones → `ModuleNotFoundError:
+vibevoice.modular.modular_vibevoice_tokenizer` in every codes/-backed loader.
+
+Audio-quality caveat stands for the TTS builds — int4 hidden-state conditioning is lossy; fp16 would
+need a GPU (genai has no fp16-on-CPU). Remaining: the three **7B ASR LLM decoders** (asr, asr-hf +
+lm_heads) — build on a big-RAM machine; front-ends are done and `inference_asr.py` degrades cleanly.
+
+### Later 2026-07-17 — fp16 eval fix, acoustic fp16/int4, git-dependency option
+
+- **eval.py made dtype-aware** (`parity_component` + `whole_pipeline._feed`): feed each ONNX its
+  DECLARED input dtype (fp16 graphs want float16; timesteps stay fp32). Fixes false "failures" on fp16
+  builds. `1.5b/cpu_fp16` now evals **9/9** (it was fine all along — the harness fed fp32). HANDOFF trap 13.
+- **1.5b CPU variants:** `cpu_int4` 9/9, `cpu_fp32` 9/9, `cpu_fp16` 9/9. `gpu_*` can't be built/run here
+  (CPU-only box; `torch +cpu`, no CUDAExecutionProvider). The user-supplied fp16-on-CPU **LLM** is genai's
+  unfused GQA×0 build needing `position_ids` (not driver-compatible) — a GPU fp16 build avoids that.
+- **acoustic:** `fp32` 3/3 ✅; `fp16` rebuilt with `op_block_list=[ConstantOfShape,ConvTranspose,Resize,
+  Range]` → 3/3 ✅ (ConstantOfShape had emitted fp16 into a float32 consumer — HANDOFF trap 11); `int4`
+  removed — it's a no-op on a conv VAE, so `build_model` now **warns + downgrades int4→fp32** for
+  LLM-less keys (HANDOFF trap 12).
+- **`optimize.py` acoustic key** finalized: `HF_REPO["acoustic"]=microsoft/VibeVoice-AcousticTokenizer`,
+  listed in `--help`/`--list`, transformers-native loaders (no `codes/`).
+- **2026-07-22 — VibeVoice source VENDORED, submodule + git-dep removed (uploadable).** The required
+  upstream subset (modular/, processor/, schedule/, scripts/, configs/ — ~591 KB, MIT/`VIBEVOICE_LICENSE`)
+  now ships as `VibeVoice/vibevoice/` and (copied) `onnx/1.5b/vibevoice/`. `_vibevoice_dir()` returns that
+  tree and raises if absent — no pip/git fallback. Removed: `codes/` submodule, `vibevoice-repo/` clone
+  (~263 MB each), `.gitmodules`, the `vibevoice @ git+…` line in all three onnx pyprojects, and the
+  root-pyproject `vibevoice` dep + `vibevoice-repo` workspace. Imports still use the isolated shim
+  (trap 1). Verified: modular/schedule/processor all resolve from the vendored tree with codes/ deleted;
+  the `onnx/1.5b/` package resolves its own copy. HANDOFF trap 14.
+
+## 2026-07-17 — faithful 1.5B TTS voice-cloning inference built (from source)
+
+No upstream 1.5B TTS `generate` exists (not in transformers — only `vibevoice_asr`/`vibevoice_acoustic_tokenizer`; not in the HF repo; `codes/` has only the training `forward`). Reconstructed it on the ONNX sub-parts, replicating from source:
+- **Dynamic acoustic encoder** — fixed `dynamic_axes`→`dynamic_shapes` (trap #2); now emits `samples/3200` frames (was baked 24000 → 7 vs 7.5 drift), aligning with the processor's `speech_tok_compress_ratio=3200`. Also fixes the codec round-trip length drift.
+- **Voice+prompt prefill** (`common.voice_prompt_embeds`) — runs `codes/` `VibeVoiceProcessor` (with trap-#1 shims: qwen2-fast alias + empty namespaces) → `input_ids` + `speech_input_mask` (70) + reference `speech_tensors`; acoustic-encodes the (24 kHz) voice, applies checkpoint `speech_scaling_factor`/`bias` (0.196/−0.049), connects, scatters the voice embeds into the masked positions. Replicates `forward_speech_features` (acoustic-only; TTS has no semantic tensors).
+- **CFG negative** — parallel `<|image_pad|>` (id 151655) LLM context, prefilled + stepped alongside the positive one (from the streaming `generate`; NOT the old zero-hidden). `inference.py --voice` (default `samples/voices/en-Alice_woman.wav`).
+
+Signal-level result (can't audition here): output went unconditioned→voiced→sharper as each piece landed — RMS 0.004→0.029, ZCR 0.110→0.052, centroid 1721→1065 Hz, sub-4 kHz energy 0.89→0.97. Intelligibility not verified (needs listening); remaining is quality/tuning, not missing machinery.
+
+Notes: the voice path needs the **dynamic** acoustic_encoder (re-exported into `cpu_fp32` + `cpu_int4`). Inference now pulls `codes/` (processor + scheduler) — no longer onnxruntime-only, inherent to VibeVoice's prompt format.
+
+### Working + tuned (later 2026-07-17)
+Confirmed by listening — the pipeline produces **intelligible human speech** in the reference voice. Tuning applied:
+- **fp32, not int4** — int4's raw-hidden quantization adds audible "hiss/background" (trap #5); fp32 is clear. (The `*_bgm.wav` sample voices carry real background music; `en-Alice_woman.wav` is clean.)
+- **CFG negative = static `<|image_pad|>`** (prefill once, reuse) — cfg 1.3 preferred over 1.0; keeps voice while staying **single-session** (fp32 ~5 GB, not ~10 GB).
+- **Auto frame-length** (`--max-frames 0`): ~5 frames/word (no EOS classifier is exported → fixed budget, not auto-stop). Prevents truncation of long text.
+- **Sentence chunking**: split on `[.!?]`, generate each chunk short, concat (150 ms gaps) → keeps each generation in the stable short regime.
+
+**Resolved — the generate works** (single-shot; earlier "chunking" was a regression and is removed). Confirmed intelligible on prose AND a Shakespeare sonnet with the reference voice. Final recipe:
+- **Single-shot** generation (whole prompt in one context) — chunking split phrasing and made some inputs noise; reverted.
+- **Learned EOS (the key fix)** — the 1.5B `lm_head` is **tied to `embed_tokens`**, so `logits = hidden @ embedᵀ` needs no lm_head export. VibeVoice reuses vision tokens for speech; generation **stops when lm_head predicts `<|vision_end|>` (speech-end) or `<|endoftext|>` (EOS)**. This removed the "jargon tail" at the source (e.g. sonnet stopped at frame 101 of a 172 budget → 13.5 s not 22.9 s). `--max-frames` is now just a safety cap. **So "no EOS exported" is NOT a limitation** — it's recovered from the tied weights.
+- **fp32, not int4** for clarity (int4 hidden-quant → hiss). **static `<|image_pad|>` CFG negative**, cfg 1.3. Voice-conditioned prefill (processor + acoustic splice + scale/bias).
+
+**Remaining minor** (1.5B model quality, not pipeline): approximate timbre cloning (gender not always captured); occasional pronunciation quirks ("boy"→"bow") and weak digit reading; fp32 CPU ~2–3 s/frame (GPU `gpu_fp16` far faster). The 7B model would improve fidelity.
+
+## Status matrix (model × component)
+
+Legend: ✅ converted & parity-verified · ⚠️ pending / partial / memory-bound · ❌ not started or blocked · — not applicable
+
+| Model | LLM decoder | Acoustic enc | Acoustic dec | Semantic tok | Diffusion head | Connectors / projector |
+|---|:---:|:---:|:---:|:---:|:---:|:---:|
+| **VibeVoice-1.5B** (TTS) | ✅ int4 | ✅ cos 1.0 | ✅ cos 1.0 | ✅ cos 1.0 (enc) | ✅ cos 1.0 | ✅ cos 1.0 (ac+sem) |
+| **VibeVoice-ASR-HF** (7B) | ⚠️ extracted, int4 OOM | ✅ cos 1.0 | — | ✅ cos 1.0 | — | ✅ cos 1.0 (mm_projector) |
+| **VibeVoice-ASR** (7B) | ⚠️ 7B int4 OOM | ✅ cos 1.0 | ✅ cos 1.0² | ✅ cos 1.0 | — none | ✅ cos 1.0 (ac+sem) |
+| **Realtime-0.5B** | ✅ int4 (tts) | — (none shipped) | ✅ cos 1.0 | — none | ✅ cos 1.0 | ✅ cos 1.0 (acoustic) |
+
+Encoder note: all acoustic/semantic **encoders** (1.5B, ASR, ASR-HF) currently bake the audio `samples` dim (24000 = 1 s @ 24 kHz) — the shared io_config uses `dynamic_axes`, which the dynamo exporter fixes. Frames/latents outputs are correct; a cross-cutting `dynamic_shapes` switch would make audio length variable (not yet applied).
+
+Notes: the three **7B LLMs** (ASR-HF / ASR text decoders + their lm_heads) all hit the int4-serialize OOM wall at current free RAM (peak ~16–20 GB, same as chandra-8B) — extraction paths ready, re-run after freeing RAM. Realtime **acoustic enc** = none shipped (decoder-only checkpoint); Realtime **semantic tok** = none in that checkpoint. ASR acoustic dec is off the audio→text path (see ²).
+
+## Architecture (confirmed from downloaded checkpoint)
+`VibeVoiceForConditionalGeneration` (`model_type: vibevoice`, auto_map null → needs the custom
+`vibevoice` package; not in transformers). Weight groups (model.safetensors.index.json):
+- `model.language_model.*` (338) — **Qwen2.5-1.5B backbone** (1536/28L/12h/2kv, q/k/v bias,
+  tied embeddings, **no lm_head** — the head is the diffusion prediction_head).
+- `model.acoustic_tokenizer.*` (552) — acoustic VAE/codec.
+- `model.semantic_tokenizer.*` (276) — semantic tokenizer.
+- `model.prediction_head.*` (26) — DiT-style diffusion denoiser (adaLN + ffn).
+- `model.acoustic_connector.*` / `model.semantic_connector.*` (5 each) — projection MLPs.
+
+## Done ✅
+- ✅ Model downloaded to `model/` (3 shards). Tokenizer NOT in repo → fetched Qwen2.5-1.5B's.
+- ✅ **`llm_decoder` BUILT (cpu int4)** → `cpu_int4/models/llm_decoder.onnx` (+857 MB data),
+  `genai_config.json`, tokenizer. Pipeline: `user_script.extract_qwen2_standalone` remaps
+  `model.language_model.*` → a standalone Qwen2ForCausalLM dir; `optimize.py` runs ModelBuilder
+  INT4 with **exclude_embeds + exclude_lm_head** → `inputs_embeds → hidden_states` + KV cache.
+  Verified: 140 MatMulNBits (int4) + 28 GroupQueryAttention; genai_config wires inputs_embeds→hidden.
+- Build: `python optimize.py --skip-download --components llm` (or no args to also download).
+
+## Pending ⏳ (need the custom `vibevoice` package)
+- `acoustic_tokenizer`, `semantic_tokenizer`, `diffusion_head`, connectors → Olive.
+  `auto_map` is null and the classes aren't in transformers, so loading the custom modules for
+  export needs `pip install vibevoice` (or the github source). `user_script.py` has stub loaders
+  that raise a clear message until that's wired.
+- Full TTS inference (`inference.py`) + `eval.py`: text → embed+connectors → llm_decoder → hidden
+  → diffusion head (iterative denoise) → acoustic latents → acoustic decoder → waveform. The
+  diffusion sampling loop + acoustic VAE are the novelty/risk (expect custom-op / fixed-shape work).
+
+## Files
+`optimize.py` (LLM build), `user_script.py` (Qwen2 extraction + audio stubs), scaffold (info.yml,
+README, requirements). Target: CPU INT4.
+
+## ✅ AcousticTokenizer (from VibeVoice-1.5B) — CONVERTED & verified
+Handled per-checkpoint (acoustic tokenizers differ across the family: 1.5B=downsample_layers,
+Realtime=stages/head, ASR-HF=conv_layers). For **1.5B**, the vendored `codes/`
+`VibeVoiceAcousticTokenizerModel` matches its weights EXACTLY (552 tensors, 0 missing).
+Loaded via an isolated import (bypasses the package __init__ that pulls diffusers + a
+transformers-5.10.2-renamed qwen2 tokenizer) + an Auto*.register shim (coexists with
+transformers' built-in vibevoice_acoustic_tokenizer). Exported both halves via Olive (fp32, dynamo):
+- `acoustic_encoder.onnx` — `audio[B,1,T] → latents[B,8,64]` (VAE `.mean`)
+- `acoustic_decoder.onnx` — `latents[B,8,64] → audio[B,L]`
+**Parity vs PyTorch: cosine 1.0** (encoder max|Δ| ~1e-4, decoder ~5e-6). `user_script.py`
+(`_load_acoustic` + get_acoustic_encoder/decoder_*), `optimize.py --components acoustic_encoder acoustic_decoder`.
+(Each file ~1.37 GB — the codec is a large VAE; both currently carry full weights.)
+
+## Deps note
+`codes/` needs `diffusers` (installed via uv) for its diffusion scheduler; the acoustic tokenizer
+itself is imported in isolation and does NOT need it. transformers 5.10.2 has native
+`vibevoice_acoustic_tokenizer` + `vibevoice_asr` (different weight-namings than 1.5B — see above).
+
+## Next
+ASR (note: `VibeVoice-ASR` and `VibeVoice-ASR-HF` are DIFFERENT models — HF = transformers-native
+8B `vibevoice_asr`), then Realtime-0.5B (vibevoice_streaming, via codes/).
+
+## ✅ ASR-HF acoustic encoder — CONVERTED & verified (transformers-native)
+`VibeVoice-ASR-HF` (`vibevoice_asr`, transformers-native, 8 shards / 16.7 GB, LLM = Qwen2.5-**7B**)
+is DISTINCT from `VibeVoice-ASR` (`VibeVoiceForASRTraining`, codes/-format, config-only here).
+Its acoustic encoder is a DIFFERENT arch than 1.5B (`acoustic_tokenizer_encoder.conv_layers.*`) and
+loads NATIVELY via transformers `VibeVoiceAcousticTokenizerEncoderModel` — no codes/ / shim.
+Loaded ONLY the `acoustic_tokenizer_encoder.*` weights (276, 0 missing) — not the 7B LLM — so it
+fits in memory. Exported (Olive fp32, dynamo) → `asr-hf/cpu_int4/models/acoustic_encoder.onnx`
+(`audio[B,1,T] → latents[B,7,64]`). **Parity vs PyTorch: cosine 1.0**, max|Δ| ~1e-4.
+`user_script.get_asrhf_acoustic_encoder_*`.
+
+## ✅ ASR-HF semantic encoder + multi_modal_projector — CONVERTED & verified
+- **semantic_encoder** — SAME transformers-native class as the acoustic encoder
+  (`vibevoice_acoustic_tokenizer_encoder`), just a different config + weight prefix
+  (`semantic_tokenizer_encoder.*`, hidden 128). `audio[B,1,T] → latents[B,frames,128]`.
+  **cos 1.0**, max|Δ| ~4e-5. `user_script.get_asrhf_semantic_encoder_*` (shared `_load_asrhf_encoder`).
+- **multi_modal_projector** (`VibeVoiceAsrMultiModalProjector`) — fuses
+  acoustic[B,T,64] + semantic[B,T,128] → LLM features[B,T,3584] (two 2-layer MLPs + RMSNorm, summed).
+  Exported with dynamo **dynamic_shapes** (frames dynamic — verified T=8/20/37). **cos 1.0**, max|Δ| ~7e-6.
+  `→ asr-hf/cpu_int4/models/{semantic_encoder,multi_modal_projector}.onnx`.
+
+## ASR-HF LLM — extraction DONE, int4 serialize OOM-blocked
+- **Extraction ✅**: `user_script.extract_qwen2_asrhf` streams shards → standalone `Qwen2ForCausalLM`
+  dir `qwen2_asrhf_standalone/` (**WITH lm_head** — ASR emits text; `language_model.model.*`→`model.*`,
+  `language_model.lm_head.*`→`lm_head.*`; config = text_config = Qwen2.5-7B: 3584/28L/28h/4kv/vocab
+  152064, untied). 15.2 GB `model.safetensors` + Qwen2.5-7B tokenizer written. Ready to build.
+- **int4 build ⚠️ OOM**: `create_model(..., "int4", "cpu")` (WITH embeds+lm_head) crashed mid-serialize
+  at ~2–3 GB free RAM (peak need ~16–20 GB) — SAME wall as chandra-8B. No `llm_decoder.onnx` produced.
+  Run `_build_asrhf_llm.py` again after freeing RAM (close Visual Studio ~3 GB + WSL) — extraction is
+  cached so it resumes straight to ModelBuilder. Subprocess isolation can't shrink a single 7B serialize.
+
+## ✅ Realtime-0.5B (vibevoice_streaming) — LLM + acoustic decoder CONVERTED & verified
+Streaming TTS checkpoint (`model.safetensors`, single file). Key groups: `tts_language_model.*`
+(Qwen2.5-0.5B backbone, **20 layers**, no lm_head), `acoustic_tokenizer.*` (**DECODER-ONLY**, 276 —
+no encoder shipped: inference only decodes generated latents), `language_model.*` (4-layer base),
+`prediction_head.*` (diffusion), `acoustic_connector.*`, `tts_eos_classifier.*`.
+- **tts llm_decoder ✅ int4** → `realtime/cpu_int4/models/llm_decoder.onnx` (+192 MB data). Built
+  0.5B in-memory (no OOM). `extract_qwen2_realtime` remaps `tts_language_model.*`→standalone Qwen2
+  (config = decoder_config with num_hidden_layers overridden to actual 20), ModelBuilder int4
+  **exclude_embeds+exclude_lm_head** → inputs_embeds→hidden. Verified 100 MatMulNBits + 20 GQA.
+- **acoustic_decoder ✅** → `realtime/cpu_int4/models/acoustic_decoder.onnx` (1.38 GB). codes/ class
+  matches EXACTLY (decoder 0 missing/0 unexpected; stages/head naming — NOT transformers-native
+  conv_layers/convtr). Decoder-only load (`decoder.*` weights, encoder=None). `latents[B,T,64] →
+  audio[B,samples]`, dynamo **dynamic_shapes** (frames dynamic — verified T=10/25/50). **cos 1.0**,
+  max|Δ| ~4e-6. `user_script.get_realtime_acoustic_decoder_*`, `_codes_tokenizer` (shared codes/ import).
+- **Acoustic ENCODER: none** — not in the streaming checkpoint (nothing to convert).
+
+## Realtime pending
+- prediction_head (diffusion denoiser) + acoustic_connector + tts_eos_classifier + 4-layer base
+  language_model → the streaming generation loop (text → tts backbone → diffusion → latents →
+  acoustic decoder → audio). Same diffusion-sampling novelty/risk as the 1.5B TTS head.
+
+## ✅ Diffusion prediction_head + speech connectors (1.5B & Realtime) — CONVERTED & verified
+Shared `codes/` classes (`_codes_import` isolated-import helper). Exported via Olive (fp32, dynamo).
+- **diffusion_head** (`VibeVoiceDiffusionHead`) — ONE DDPM denoise step:
+  `(noisy_images[B,64], timesteps[B] FLOAT, condition[B,H]) → pred[B,64]`. The ~20-step sampling
+  loop stays in the pipeline; ONNX = one step. `H`=1536 (1.5B) / 896 (Realtime). **Batch dynamic**
+  (verified B=4/7). **cos 1.0**, max|Δ| ~1e-6. Note: timesteps MUST be float32 — `TimestepEmbedder`
+  casts its sinusoidal embedding back to `t.dtype` before the float MLP.
+- **connectors** (`SpeechConnector`: fc1→RMSNorm→fc2) — project VAE latents into LLM hidden space,
+  frames dynamic. 1.5B: `acoustic_connector` 64→1536 + `semantic_connector` 128→1536; Realtime:
+  `acoustic_connector` 64→896. All **cos 1.0**, max|Δ| ≤5e-7.
+- `user_script.get_diffusion_head_* / get_{acoustic,semantic}_connector_*`. Head dummy hidden via
+  env `VV_HEAD_HIDDEN` (1536 default; set 896 for Realtime).
+
+## Sub-model tally: 13 ONNX parts converted & parity-verified (cos 1.0)
+- **1.5B** (6): llm_decoder(int4), acoustic_encoder, acoustic_decoder, diffusion_head,
+  acoustic_connector, semantic_connector.  Pending: semantic_tokenizer encoder; end-to-end pipeline.
+- **ASR-HF** (3): acoustic_encoder, semantic_encoder, multi_modal_projector. LLM extracted (int4 OOM).
+- **Realtime** (4): llm_decoder(int4), acoustic_decoder, diffusion_head, acoustic_connector.
+Remaining to reach end-to-end TTS: wire the sampling pipeline (text → embed+connector → llm_decoder →
+hidden → DDPM loop over diffusion_head → acoustic latents → acoustic_decoder → waveform) + eval.
+
+² ASR acoustic_decoder frames baked at 8 (shared 1.5B io_config uses dynamic_axes, which dynamo
+ignores) — harmless: the decoder is NOT on the ASR audio→text path. Encoders + connectors are frames-dynamic.
+
+## ✅ VibeVoice-ASR (weights now downloaded) — audio front-end CONVERTED & verified
+`VibeVoiceForASRTraining` (`model_type vibevoice`, auto_map null → codes/), 8 shards. Same codes/
+family as 1.5B TTS, but: LLM = Qwen2.5-**7B** (decoder_config 3584/28L) **WITH** top-level
+`lm_head.weight` (audio→text), and **NO prediction_head** (no audio generation). Weight groups:
+acoustic_tokenizer (552, enc+dec), semantic_tokenizer (276, ENCODE-only), language_model (338) +
+lm_head, acoustic/semantic connectors (5 each). All front-end pieces load via the EXISTING codes/
+loaders (matched exactly: acoustic 552/0-miss, semantic 276/0-miss) — no per-checkpoint divergence
+here (unlike the ASR-HF transformers-native tokenizers). Exported (Olive fp32, dynamo) →
+`asr/cpu_int4/models/`:
+- **acoustic_encoder** `audio[B,1,T]→latents[B,f,64]` (VAE .mean) — cos 1.0
+- **acoustic_decoder** `latents[B,8,64]→audio` — cos 1.0 (frames fixed 8, off-path — see ²)
+- **semantic_encoder** (`VibeVoiceSemanticTokenizerModel`, encode-only) `audio→latents[B,f,128]` — cos 1.0
+- **acoustic_connector** 64→3584, **semantic_connector** 128→3584 — cos 1.0, frames-dynamic
+`user_script.get_semantic_tokenizer_encoder_*` (new) + reused `_load_acoustic`/`_load_connector`.
+
+## VibeVoice-ASR LLM — extraction ready, int4 serialize OOM (same 7B wall)
+`extract_qwen2_asr` remaps `model.language_model.*`→`model.*` + top-level `lm_head.weight` (differs
+from ASR-HF's `language_model.model.*`/`language_model.lm_head.*` layout), config = decoder_config
+(Qwen2.5-7B), fetches Qwen2.5-7B tok. NOT run yet (avoids a 2nd 15 GB standalone dir on disk); the
+int4 ModelBuilder serialize would OOM at current free RAM exactly like ASR-HF. Run after freeing RAM.
+
+## Updated tally: 19 ONNX parts converted & parity-verified (cos ~1.0)
+1.5B(7) + ASR-HF(3) + Realtime(4) + **ASR(5)**. Remaining: the three 7B
+LLMs (ASR-HF/ASR extracted-or-ready but int4 OOM-bound); end-to-end pipelines.
+
+## Inference (ONNX) — inference_common.py + 3 drivers
+`inference_common.py` (shared): resolve() (reuses optimize registry), **OnnxLLM** (drives the
+genai llm_decoder.onnx directly via ORT — 28-layer KV cache + growing attn mask; verified
+KV-incremental==full cos 1.0; handles hidden_states/logits outputs), OnnxOp, **DiffusionSampler**
+(diffusion_head.onnx + codes/ DPM scheduler + CFG), audio load/normalize(-25dBFS)/save,
+embed_tokens lookup, load_scaling, acoustic_decode_to_wav.
+- **inference.py** — VibeVoice-1.5B TTS: text→embed→prefill→[hidden→diffuse→latent→connector→step]→
+  decode all latents→wav. Verified: 8 frames → 1.07 s wav.
+- **inference_realtime.py** — Realtime-0.5B streaming TTS: same loop, per-frame decode (dynamic
+  decoder) streamed. Verified: 6 frames → 0.80 s wav.
+- **inference_asr.py** — asr/asr-hf audio→text: encoders→(connectors|projector)→fuse at speech-pad
+  positions→prefill(lm_head→logits)→greedy decode. Verified: front-end → fused [N,3584]; degrades
+  cleanly when the 7B llm_decoder.onnx isn't built.
+Caveats (documented in files): int4 is lossy in raw-hidden space (first-token hidden cos ~0.83 vs
+fp32) — prefer fp16 for audio fidelity; learned EOS (1.5b lm_head / realtime tts_eos_classifier)
+isn't in the sub-part set → stops at --max-frames; exact prompt layout/voice-clone needs the
+original VibeVoice processor/tokenizer assets (repos ship none → Qwen2.5 tokenizer used). Fixed a
+real bug: acoustic_decoder io_config marked the vae_dim (dim 2) as "frames" instead of dim 1 →
+frames stayed static; switched to dynamic_shapes on dim 1 and re-exported.

VIBEVOICE_LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Microsoft
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

common.py

@@ -10,7 +10,7 @@ same primitives, factored here:
                            the ASR decoder. One prefill(embeds)->hidden + step(embeds)->hidden loop.
   * OnnxOp               — thin wrapper over a single-file ONNX session (encoders/decoder/
                            connector/diffusion_head/projector).
-  * DiffusionSampler     — the DDPM/DPM denoise loop (diffusion_head.onnx + codes/ scheduler + CFG).
+  * DiffusionSampler     — the DDPM/DPM denoise loop (diffusion_head.onnx + vendored vibevoice/ scheduler + CFG).
   * audio load / normalize (-25 dBFS) / save  — 24 kHz mono (VibeVoice standard).
   * load_tokenizer, load_scaling — Qwen2.5 tokenizer + the stored speech scaling/bias factors.
 
@@ -162,19 +162,13 @@ def load_scaling(src):
     return scale, bias
 
 
-# --------------------------------------------------------------------------- codes/ scheduler
+# -------------------------------------------------------------------- vendored vibevoice scheduler
 def _vibevoice_dir():
-    """vibevoice source: local codes/ submodule, else the pip-installed `vibevoice` git dependency."""
-    local = HERE / "codes" / "vibevoice"
+    """The vendored vibevoice source tree shipped alongside this code (VibeVoice/vibevoice/)."""
+    local = HERE / "vibevoice"
     if local.is_dir():
         return str(local)
-    import importlib.util
-    spec = importlib.util.find_spec("vibevoice")
-    if spec and spec.submodule_search_locations:
-        return list(spec.submodule_search_locations)[0]
-    raise ModuleNotFoundError(
-        "vibevoice source not found — populate codes/ or install "
-        "'vibevoice @ git+https://github.com/microsoft/VibeVoice.git@303b283'")
+    raise ModuleNotFoundError(f"vendored vibevoice source missing at {local}")
 
 
 def make_scheduler(diffusion_cfg):
@@ -252,11 +246,17 @@ class OnnxLLM:
         outs = [o.name for o in self.sess.get_outputs()]
         # TTS backbone emits 'hidden_states'; an ASR decoder that kept lm_head emits 'logits'.
         self.out_name = "logits" if "logits" in outs else "hidden_states"
+        # Fused GQA builds compute positions internally; the unfused fallback (fp16-on-CPU / fp32-on-CUDA)
+        # declares a 'position_ids' input we must feed. Match the graph's float dtype (fp16 build wants
+        # float16 embeds + KV cache), else ORT rejects the feed.
+        self.wants_pos = "position_ids" in self.in_names
+        emb_t = next(i.type for i in self.sess.get_inputs() if i.name == "inputs_embeds")
+        self.float_dt = np.float16 if "float16" in emb_t else np.float32
         self.hidden = None
         self._reset()
 
     def _reset(self, batch=1):
-        z = np.zeros((batch, self.kv_heads, 0, self.head_dim), dtype=np.float32)
+        z = np.zeros((batch, self.kv_heads, 0, self.head_dim), dtype=self.float_dt)
         self.past = {}
         for i in range(self.n_layers):
             self.past[f"past_key_values.{i}.key"] = z.copy()
@@ -265,11 +265,14 @@ class OnnxLLM:
         self.batch = batch
 
     def _run(self, embeds):
-        embeds = np.asarray(embeds, dtype=np.float32)
+        embeds = np.asarray(embeds, dtype=self.float_dt)
         b, s, _ = embeds.shape
+        prev = self.total
         self.total += s
         feed = {"inputs_embeds": embeds,
                 "attention_mask": np.ones((b, self.total), dtype=np.int64)}
+        if self.wants_pos:  # positions of the newly-appended tokens
+            feed["position_ids"] = np.arange(prev, prev + s, dtype=np.int64)[None, :].repeat(b, 0)
         feed.update(self.past)
         outs = self.sess.run(None, feed)
         names = [o.name for o in self.sess.get_outputs()]
@@ -316,6 +319,72 @@ def embed_tokens(src, token_ids, model_key):
     raise RuntimeError(f"{key} not found in {src}")
 
 
+def _load_vv_processor(src):
+    """Load the VibeVoiceProcessor from the vendored vibevoice/ source with the trap-#1 shims: alias
+    the pre-5.10 qwen2 fast tokenizer path, inject empty vibevoice namespaces (no package __init__),
+    swallow Auto*.register collisions (see _vibevoice_dir — vendored VibeVoice/vibevoice/)."""
+    import os, types, importlib
+    base = _vibevoice_dir()                          # VibeVoice/vibevoice/ (vendored source)
+    root = os.path.dirname(base)
+    import transformers as _tf
+    qm = types.ModuleType("transformers.models.qwen2.tokenization_qwen2_fast")
+    qm.Qwen2TokenizerFast = _tf.Qwen2TokenizerFast
+    sys.modules.setdefault("transformers.models.qwen2.tokenization_qwen2_fast", qm)
+    for name, sub in [("vibevoice", ""), ("vibevoice.processor", "processor"), ("vibevoice.modular", "modular")]:
+        if name not in sys.modules:
+            m = types.ModuleType(name); m.__path__ = [os.path.join(base, sub) if sub else base]
+            sys.modules[name] = m
+    from transformers import AutoConfig, AutoModel
+    for cls in (AutoConfig, AutoModel):
+        r = cls.register
+        def _safe(*a, __r=r, **k):
+            try: __r(*a, **k)
+            except Exception: pass
+        cls.register = staticmethod(_safe)
+    VP = importlib.import_module("vibevoice.processor.vibevoice_processor").VibeVoiceProcessor
+    return VP.from_pretrained(str(src))
+
+
+def voice_prompt_embeds(src, onnx_dir, text, voice_path, scale, bias, device="cpu", hop=3200):
+    """Build the TTS prefill embeddings the way VibeVoice's forward_speech_features does, but on ONNX:
+      processor(text, voice) -> input_ids + speech_input_mask + reference speech_tensors
+      voice -> acoustic_encoder.onnx -> latents -> (latents+bias)*scale -> acoustic_connector.onnx
+      inputs_embeds = embed_tokens(input_ids); scatter the voice embeds into the masked positions.
+    Returns (inputs_embeds [1,S,H] float32, input_ids [1,S])."""
+    proc = _load_vv_processor(src)
+    if not text.strip().lower().startswith("speaker"):
+        text = "Speaker 1: " + text
+    wav = load_audio(voice_path, SR)                                  # mono @ 24 kHz
+    out = proc(text=[text], voice_samples=[[wav]], return_tensors="pt")
+    input_ids = np.asarray(out["input_ids"])                         # [1,S]
+    mask = np.asarray(out["speech_input_mask"]).astype(bool)         # [1,S]
+    x = embed_tokens(src, input_ids[0], "1.5b").astype(np.float32)   # [1,S,H]
+    n_fr = int(mask.sum())
+    if n_fr:
+        v = np.asarray(out["speech_tensors"], np.float32).ravel()   # reference samples @ 24 kHz
+        pad = n_fr * hop                                            # pad to exactly n_fr frames
+        v = np.pad(v, (0, max(0, pad - len(v))))[:pad][None, None]  # [1,1,n_fr*hop]
+        lat = OnnxOp(onnx_dir / "acoustic_encoder.onnx", device).run(audio=v)  # [1,n_fr,64]
+        audio_feat = (lat.astype(np.float32) + bias) * scale        # checkpoint speech scale/bias
+        conn = OnnxOp(onnx_dir / "acoustic_connector.onnx", device)
+        vemb = conn.run(**{conn.inames[0]: audio_feat})             # [1,n_fr,H]
+        x[0][mask[0]] = vemb[0]
+    return x, input_ids
+
+
+def load_embed_matrix(src, model_key):
+    """Full token-embedding matrix [vocab, H] for `model_key`. The 1.5B lm_head is TIED to these
+    embeddings, so lm_head logits = hidden @ matrix.T — lets us detect the learned speech-end/EOS
+    token without exporting lm_head (it was excluded; head is the diffusion head)."""
+    from safetensors.torch import load_file
+    key = EMBED_KEY.get(model_key)
+    for sf in glob.glob(str(Path(src) / "*.safetensors")):
+        d = load_file(sf)
+        if key in d:
+            return d[key].float().numpy().astype(np.float32)     # [vocab, H]
+    raise RuntimeError(f"{key} not found in {src}")
+
+
 def acoustic_decode_to_wav(dec_op: OnnxOp, latents, scale, bias):
     """latents [B,T,64] (LLM-space) -> waveform. Applies /scale - bias then acoustic_decoder."""
     lat = (np.asarray(latents, dtype=np.float32) / (scale if scale else 1.0)) - bias

cpu_fp32/acoustic_encoder.onnx

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:a6d5a0f7cfd27e06e39a6f11bcbc1abba63245e9e6b256b407e2d43bd67396ad
-size 1375985271
+oid sha256:7223ac37cbf144296dbbcd3298e18cb1a30d44b4d5689d84892fe3bb4458f61c
+size 1376385625

cpu_int4/acoustic_encoder.onnx

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:a6d5a0f7cfd27e06e39a6f11bcbc1abba63245e9e6b256b407e2d43bd67396ad
-size 1375985271
+oid sha256:7223ac37cbf144296dbbcd3298e18cb1a30d44b4d5689d84892fe3bb4458f61c
+size 1376385625

fp32.wav

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

inference.py

@@ -1,96 +1,130 @@
-"""VibeVoice-1.5B TTS inference over the exported ONNX sub-parts.
-
-Pipeline (reference algorithm, driven on ONNX via common):
-  text → tokenize → embed_tokens → OnnxLLM.prefill → then per acoustic frame:
-    last hidden state → DiffusionSampler (diffusion_head.onnx, DDPM/CFG) → acoustic latent
-    → acoustic_decoder.onnx → waveform chunk
-    → acoustic_connector.onnx (latent → LLM-space embed) → OnnxLLM.step (feeds next frame)
-  concatenate chunks → 24 kHz wav.
-
-Usage (prefer uv run) — the ONLY model input is the built ONNX dir; model_id/device/precision are
-derived from the path and printed:
-  uv run inference.py --text "Hello world." --out out.wav model/cpu_int4/models
-  uv run inference.py --text "..." --max-frames 200 model/cuda_fp16/models
-
-Notes:
-  * The 1.5B decoder is exported WITHOUT lm_head (its head is the diffusion head), so there is
-    no vocab EOS token to detect here — generation stops at --max-frames. Faithful EOS needs the
-    lm_head / a stop-token classifier (not part of this sub-part decomposition).
-  * int4 is lossy in raw-hidden space (what conditions the diffusion head); build a fp16 dir
-    for higher-fidelity audio. Exact prompt layout/voice-cloning needs the original VibeVoice
-    processor assets (repos ship none — we use the Qwen2.5 tokenizer).
-"""
-import argparse
-import json
-import sys
-from pathlib import Path
-
-import numpy as np
-
-sys.path.insert(0, str(Path(__file__).parent))
-import common 
-
-
-def main():
-    ap = argparse.ArgumentParser(description="VibeVoice-1.5B TTS (ONNX)")
-    ap.add_argument("model_path", help="built ONNX dir, e.g. model/cpu_int4/models")
-    ap.add_argument("--text", required=True, help="text to synthesize")
-    ap.add_argument("--out", default="tts_out.wav", help="output wav path")
-    ap.add_argument("--max-frames", type=int, default=200, help="acoustic frames to generate")
-    ap.add_argument("--cfg-scale", type=float, default=1.3)
-    args = ap.parse_args()
-
-    key, src, onnx_dir, device, precision, model_id = common.resolve_from_path(args.model_path)
-    print(f"=== TTS | model_id={model_id}  device={device}  precision={precision} ===")
-    if key != "1.5b":
-        sys.exit(f"[error] {model_id} ({key}) is not the 1.5B TTS model — "
-                 f"use inference_asr.py (asr/asr-hf) or inference_realtime.py (realtime)")
-    for need in ("llm_decoder", "diffusion_head", "acoustic_decoder", "acoustic_connector"):
-        if not (onnx_dir / f"{need}.onnx").exists():
-            sys.exit(f"missing {need}.onnx in {onnx_dir} — build it: uv run optimize.py {key}")
-
-    cfg = json.loads((src / "config.json").read_text())
-    dcfg = cfg["diffusion_head_config"]
-    scale, bias = common.load_scaling(src)
-    print(f"    scale={scale:.4f} bias={bias:.4f}")
-
-    tok = common.load_tokenizer(onnx_dir, src)
-    llm = common.OnnxLLM(onnx_dir / "llm_decoder.onnx", device)
-    head = common.OnnxOp(onnx_dir / "diffusion_head.onnx", device)
-    dec = common.OnnxOp(onnx_dir / "acoustic_decoder.onnx", device)
-    conn = common.OnnxOp(onnx_dir / "acoustic_connector.onnx", device)
-    sampler = common.DiffusionSampler(head, dcfg, device)
-
-    # Prompt: embed the text tokens and prefill the backbone. (Voice-clone/speaker prompts
-    # would prepend reference-audio acoustic embeds here — needs the processor's token layout.)
-    ids = tok.encode(args.text)
-    embeds = common.embed_tokens(src, ids, key)                  # [1,S,H]
-    hidden = llm.prefill(embeds)                            # [1,S,H]
-    H = hidden.shape[-1]
-    neg = np.zeros(H, dtype=np.float32)                     # unconditional = zero hidden
-
-    # Autoregressive frame loop: each latent is decoded together at the end (the conv codec has a
-    # cross-frame receptive field, so one decode of the whole sequence avoids block-boundary seams).
-    latents = []
-    for f in range(args.max_frames):
-        cond = hidden[0, -1, :]                             # condition on last hidden
-        latent = sampler.sample(cond, neg, cfg_scale=args.cfg_scale, n_frames=1, seed=f)  # [1,64]
-        latents.append(latent[0])
-        nxt = conn.run(features=latent[None])              # [1,1,H] LLM-space embed
-        hidden = llm.step(nxt.astype(np.float32))          # advance one frame
-        if (f + 1) % 25 == 0:
-            print(f"  frame {f+1}/{args.max_frames}")
-
-    lat_seq = np.stack(latents)[None].astype(np.float32)   # [1,T,64]
-    exp = dec.sess.get_inputs()[0].shape[1]                # decoder's frame axis (int if static)
-    if isinstance(exp, int) and exp != lat_seq.shape[1]:
-        sys.exit(f"acoustic_decoder expects {exp} frames but got {lat_seq.shape[1]} — this decoder "
-                 f"was built with a static frame count; re-export with dynamic frames: "
-                 f"uv run optimize.py --components acoustic_decoder {key}")
-    audio = np.asarray(common.acoustic_decode_to_wav(dec, lat_seq, scale, bias)).ravel()
-    common.save_wav(args.out, audio)
-    print(f"wrote {args.out}  ({len(audio)} samples, {len(audio)/common.SR:.2f}s, {len(latents)} frames)")
-
-
-if __name__ == "__main__":
-    main()
+"""VibeVoice-1.5B TTS inference over the exported ONNX sub-parts.
+
+Pipeline (reference algorithm, driven on ONNX via common):
+  text → tokenize → embed_tokens → OnnxLLM.prefill → then per acoustic frame:
+    last hidden state → DiffusionSampler (diffusion_head.onnx, DDPM/CFG) → acoustic latent
+    → acoustic_decoder.onnx → waveform chunk
+    → acoustic_connector.onnx (latent → LLM-space embed) → OnnxLLM.step (feeds next frame)
+  concatenate chunks → 24 kHz wav.
+
+Voice-cloning: the VibeVoice processor (vendored vibevoice/) builds the system + Speaker prompt and marks speech
+placeholders; the reference voice is acoustic-encoded and spliced there (replicates
+forward_speech_features). CFG uses a static "<|image_pad|>" negative. Single-shot (whole prompt in
+one context) — chunking split phrasing and regressed some inputs.
+
+Learned EOS: the 1.5B lm_head is TIED to embed_tokens, so logits = hidden @ embedᵀ (no lm_head
+export needed). VibeVoice reuses vision tokens for speech, so generation stops when the model
+predicts speech-end (<|vision_end|>) or EOS (<|endoftext|>). --max-frames is just a safety cap.
+
+Usage — the ONLY model input is the built ONNX dir; model_id/device/precision derive from the path:
+  uv run inference.py --text "Hello world." --voice samples/voices/en-Alice_woman.wav onnx/1.5b/cpu_fp32
+  uv run inference.py --text "..." --cfg-scale 1.3 --max-frames 0 onnx/1.5b/cpu_fp32   # 0 = auto cap
+
+Notes:
+  * Use **fp32** — int4 is lossy in raw-hidden space (conditions the diffusion head) → audible hiss.
+  * timbre-cloning is approximate and digits read weakly (1.5B model quality). See STATUS.md.
+"""
+import argparse
+import json
+import sys
+from pathlib import Path
+
+import numpy as np
+
+sys.path.insert(0, str(Path(__file__).parent))
+import common 
+
+
+def main():
+    ap = argparse.ArgumentParser(description="VibeVoice-1.5B TTS (ONNX)")
+    ap.add_argument("model_path", help="built ONNX dir, e.g. model/cpu_int4/models")
+    ap.add_argument("--text", required=True, help="text to synthesize")
+    ap.add_argument("--out", default="tts_out.wav", help="output wav path")
+    ap.add_argument("--max-frames", type=int, default=0,
+                    help="safety cap on acoustic frames (7.5/sec); generation normally stops earlier "
+                         "via the learned EOS. 0 = auto-cap from text length (~5/word)")
+    ap.add_argument("--cfg-scale", type=float, default=1.3)
+    ap.add_argument("--voice", default="samples/voices/en-Alice_woman.wav",
+                    help="reference voice wav for cloning (VibeVoice is voice-conditioned); "
+                         "'' = text-only smoke-test (unconditioned, not intelligible)")
+    args = ap.parse_args()
+
+    key, src, onnx_dir, device, precision, model_id = common.resolve_from_path(args.model_path)
+    print(f"=== TTS | model_id={model_id}  device={device}  precision={precision} ===")
+    if key != "1.5b":
+        sys.exit(f"[error] {model_id} ({key}) is not the 1.5B TTS model — "
+                 f"use inference_asr.py (asr/asr-hf) or inference_realtime.py (realtime)")
+    for need in ("llm_decoder", "diffusion_head", "acoustic_decoder", "acoustic_connector"):
+        if not (onnx_dir / f"{need}.onnx").exists():
+            sys.exit(f"missing {need}.onnx in {onnx_dir} — build it: uv run optimize.py {key}")
+
+    cfg = json.loads((src / "config.json").read_text())
+    dcfg = cfg["diffusion_head_config"]
+    scale, bias = common.load_scaling(src)
+    print(f"    scale={scale:.4f} bias={bias:.4f}")
+
+    tok = common.load_tokenizer(onnx_dir, src)
+    llm = common.OnnxLLM(onnx_dir / "llm_decoder.onnx", device)
+    head = common.OnnxOp(onnx_dir / "diffusion_head.onnx", device)
+    dec = common.OnnxOp(onnx_dir / "acoustic_decoder.onnx", device)
+    conn = common.OnnxOp(onnx_dir / "acoustic_connector.onnx", device)
+    sampler = common.DiffusionSampler(head, dcfg, device)
+
+    # CFG negative (static "<|image_pad|>" — see below), computed once and reused for every chunk.
+    neg_vec = None
+    if args.cfg_scale != 1.0:
+        neg_id = tok.convert_tokens_to_ids("<|image_pad|>")
+        if neg_id is None or neg_id < 0:
+            print("    [warn] no <|image_pad|> token — CFG disabled")
+        else:
+            nh = llm.prefill(common.embed_tokens(src, [neg_id], key))
+            neg_vec = nh[0, -1, :].copy(); llm._reset()
+            print(f"    CFG negative: static '<|image_pad|>' (id={neg_id}), scale={args.cfg_scale}")
+
+    # Single-shot generation (chunking regressed some inputs — the whole prompt in one context is
+    # what works). Voice-conditioned prefill, then the AR diffusion frame loop.
+    from pathlib import Path as _P
+    if args.voice and _P(args.voice).exists():
+        print(f"    voice-cloning prompt from {args.voice}")
+        embeds, ids = common.voice_prompt_embeds(src, onnx_dir, args.text, args.voice, scale, bias, device)
+    else:
+        if args.voice:
+            print(f"    [warn] voice '{args.voice}' not found — text-only (unconditioned)")
+        embeds = common.embed_tokens(src, tok.encode(args.text), key)
+    llm._reset()
+    hidden = llm.prefill(embeds)
+    H = hidden.shape[-1]
+
+    # Safety cap only — the learned EOS below normally stops first. Auto-estimated from text (~5/word).
+    n_frames = args.max_frames if args.max_frames > 0 else min(400, 12 + len(args.text.split()) * 5)
+    if args.max_frames <= 0:
+        print(f"    frame cap={n_frames} ({len(args.text.split())} words, ~{n_frames/7.5:.1f}s) — EOS stops earlier")
+
+    # Learned EOS: lm_head is TIED to embed_tokens, so logits = hidden @ embedᵀ. VibeVoice reuses
+    # vision tokens for speech — stop when the model predicts speech-end (<|vision_end|>) or EOS
+    # (<|endoftext|>). This removes the no-EOS overshoot ("jargon tail") instead of guessing a budget.
+    emb_W = common.load_embed_matrix(src, key)                 # [vocab, H]
+    end_ids = {tok.convert_tokens_to_ids(t) for t in ("<|vision_end|>", "<|endoftext|>")}
+    end_ids = {i for i in end_ids if isinstance(i, int) and i >= 0}
+
+    latents = []
+    for f in range(n_frames):
+        cond = hidden[0, -1, :]
+        neg = neg_vec if neg_vec is not None else np.zeros(H, dtype=np.float32)
+        latent = sampler.sample(cond, neg, cfg_scale=args.cfg_scale, n_frames=1, seed=f)
+        latents.append(latent[0])
+        hidden = llm.step(conn.run(features=latent[None]).astype(np.float32))
+        if end_ids and int((hidden[0, -1] @ emb_W.T).argmax()) in end_ids:  # tied-lm_head EOS
+            print(f"  EOS predicted at frame {f+1}/{n_frames} — stop")
+            break
+        if (f + 1) % 25 == 0:
+            print(f"  frame {f+1}/{n_frames}")
+
+    lat_seq = np.stack(latents)[None].astype(np.float32)
+    audio = np.asarray(common.acoustic_decode_to_wav(dec, lat_seq, scale, bias)).ravel()
+    common.save_wav(args.out, audio)
+    print(f"wrote {args.out}  ({len(audio)} samples, {len(audio)/common.SR:.2f}s, {len(latents)} frames)")
+
+
+if __name__ == "__main__":
+    main()

inference_asr.py

@@ -1,123 +1,123 @@
-"""VibeVoice ASR inference (audio → text) over the exported ONNX sub-parts.
-
-Handles BOTH ASR checkpoints (choose via the positional model — key or path):
-  asr     (VibeVoice-ASR, codes/ family)  : acoustic_connector + semantic_connector, fused = sum
-  asr-hf  (VibeVoice-ASR-HF, transformers) : multi_modal_projector fuses acoustic+semantic
-
-Pipeline (reference algorithm on ONNX via common):
-  audio → acoustic_encoder + semantic_encoder → (connectors | projector) → speech features
-  → build [system + <speech_start> <speech_pad>*N <speech_end> + instruction] embeds, inject the
-    speech features at the <speech_pad> positions → OnnxLLM.prefill (decoder KEEPS lm_head → logits)
-  → greedy-decode tokens until EOS → tokenizer.decode → transcript.
-
-Usage (prefer uv run) — the ONLY model input is the built ONNX dir; model_id/device/precision are
-derived from the path and printed:
-  uv run inference_asr.py --audio speech.wav asr-hf/cpu_int4/models
-  uv run inference_asr.py --audio speech.wav asr/cuda_fp16/models
-
-Note: the ASR LLM is Qwen2.5-7B; its int4 serialize is RAM-bound (see STATUS.md), so
-llm_decoder.onnx may not be built. Without it this still runs the audio front-end and reports the
-fused-feature shape, then explains that the 7B decoder must be built to produce text.
-"""
-import argparse
-import json
-import sys
-from pathlib import Path
-
-import numpy as np
-
-sys.path.insert(0, str(Path(__file__).parent))
-import common as C
-
-SPEECH_TOKENS = ["<|speech_start|>", "<|speech_pad|>", "<|speech_end|>"]
-
-
-def _resolve_speech_ids(tok):
-    ids = {}
-    for t in SPEECH_TOKENS:
-        i = tok.convert_tokens_to_ids(t)
-        ids[t] = i if isinstance(i, int) and i >= 0 else None
-    return ids
-
-
-def encode_audio_features(key, src, onnx_dir, device, wav):
-    """Return fused speech features [N, H] and frame count N from the ONNX front-end."""
-    a_enc = C.OnnxOp(onnx_dir / "acoustic_encoder.onnx", device)
-    s_enc = C.OnnxOp(onnx_dir / "semantic_encoder.onnx", device)
-    x = wav[None, None, :].astype(np.float32)               # [1,1,T]
-    al = a_enc.run(audio=x)                                  # [1,fa,64]
-    sl = s_enc.run(audio=x)                                  # [1,fs,128]
-    n = min(al.shape[1], sl.shape[1])
-    al, sl = al[:, :n].astype(np.float32), sl[:, :n].astype(np.float32)
-    proj = C.component_path(onnx_dir, "multi_modal_projector")
-    if proj:                                                # asr-hf
-        fused = C.OnnxOp(proj, device).run(acoustic_latents=al, semantic_latents=sl)   # [1,n,H]
-    else:                                                   # asr: sum of the two connectors
-        ac = C.OnnxOp(onnx_dir / "acoustic_connector.onnx", device).run(features=al)
-        sc = C.OnnxOp(onnx_dir / "semantic_connector.onnx", device).run(features=sl)
-        fused = ac + sc
-    return fused[0].astype(np.float32), n                    # [n,H], n
-
-
-def main():
-    ap = argparse.ArgumentParser(description="VibeVoice ASR (ONNX)")
-    ap.add_argument("model_path", help="built ONNX dir, e.g. asr-hf/cpu_int4/models")
-    ap.add_argument("--audio", required=True, help="input wav")
-    ap.add_argument("--max-new-tokens", type=int, default=256)
-    ap.add_argument("--prompt", default="Please transcribe the audio.")
-    args = ap.parse_args()
-
-    key, src, onnx_dir, device, precision, model_id = C.resolve_from_path(args.model_path)
-    print(f"=== ASR | model_id={model_id}  device={device}  precision={precision} ===")
-    if key not in ("asr", "asr-hf"):
-        sys.exit(f"[error] {model_id} ({key}) is not an ASR model — "
-                 f"use inference.py (1.5b) or inference_realtime.py (realtime)")
-    for need in ("acoustic_encoder", "semantic_encoder"):
-        if not (onnx_dir / f"{need}.onnx").exists():
-            sys.exit(f"missing {need}.onnx in {onnx_dir} — build: uv run optimize.py {key}")
-
-    wav = C.normalize_audio(C.load_audio(args.audio))
-    feats, n = encode_audio_features(key, src, onnx_dir, device, wav)
-    H = feats.shape[-1]
-    print(f"  audio {len(wav)/C.SR:.2f}s → {n} speech frames → fused features [{n},{H}]")
-
-    llm_path = onnx_dir / "llm_decoder.onnx"
-    if not llm_path.exists():
-        print(f"\n[front-end OK] llm_decoder.onnx not built (Qwen2.5-7B int4 is RAM-bound — see STATUS.md).")
-        print(f"  Build it after freeing RAM: uv run optimize.py --components llm {key}")
-        print(f"  Then re-run for the transcript.")
-        return
-
-    tok = C.load_tokenizer(onnx_dir, src)
-    sid = _resolve_speech_ids(tok)
-    # Prompt: system + user(<speech_start> <speech_pad>*n <speech_end> + instruction).
-    pad = sid["<|speech_pad|>"]
-    if pad is None:
-        # tokenizer lacks the speech tokens (repo shipped none) — add them, then use a pad id.
-        tok.add_special_tokens({"additional_special_tokens": SPEECH_TOKENS})
-        sid = _resolve_speech_ids(tok); pad = sid["<|speech_pad|>"]
-    start = sid["<|speech_start|>"] if sid["<|speech_start|>"] is not None else pad
-    end = sid["<|speech_end|>"] if sid["<|speech_end|>"] is not None else pad
-    pre = tok.encode(f"{args.prompt}\n")
-    ids = pre + [start] + [pad] * n + [end]
-    embeds = C.embed_tokens(src, ids, key)                 # [1,S,H]
-    # inject at the n dedicated speech-pad slots (the [pad]*n block), not any incidental pad id
-    pad_pos = list(range(len(pre) + 1, len(pre) + 1 + n))
-    embeds[0, pad_pos, :] = feats[:len(pad_pos)]            # inject speech features
-
-    llm = C.OnnxLLM(llm_path, device)
-    logits = llm.prefill(embeds)                            # [1,S,V]
-    eos = getattr(tok, "eos_token_id", None)
-    out = []
-    for _ in range(args.max_new_tokens):
-        nxt = int(np.asarray(logits)[0, -1].argmax())
-        if eos is not None and nxt == eos:
-            break
-        out.append(nxt)
-        logits = llm.step(C.embed_tokens(src, [nxt], key))
-    text = tok.decode(out, skip_special_tokens=True)
-    print(f"\nTRANSCRIPT:\n{text}")
-
-
-if __name__ == "__main__":
-    main()
+"""VibeVoice ASR inference (audio → text) over the exported ONNX sub-parts.
+
+Handles BOTH ASR checkpoints (choose via the positional model — key or path):
+  asr     (VibeVoice-ASR, vibevoice/ family)  : acoustic_connector + semantic_connector, fused = sum
+  asr-hf  (VibeVoice-ASR-HF, transformers) : multi_modal_projector fuses acoustic+semantic
+
+Pipeline (reference algorithm on ONNX via common):
+  audio → acoustic_encoder + semantic_encoder → (connectors | projector) → speech features
+  → build [system + <speech_start> <speech_pad>*N <speech_end> + instruction] embeds, inject the
+    speech features at the <speech_pad> positions → OnnxLLM.prefill (decoder KEEPS lm_head → logits)
+  → greedy-decode tokens until EOS → tokenizer.decode → transcript.
+
+Usage (prefer uv run) — the ONLY model input is the built ONNX dir; model_id/device/precision are
+derived from the path and printed:
+  uv run inference_asr.py --audio speech.wav asr-hf/cpu_int4/models
+  uv run inference_asr.py --audio speech.wav asr/cuda_fp16/models
+
+Note: the ASR LLM is Qwen2.5-7B; its int4 serialize is RAM-bound (see STATUS.md), so
+llm_decoder.onnx may not be built. Without it this still runs the audio front-end and reports the
+fused-feature shape, then explains that the 7B decoder must be built to produce text.
+"""
+import argparse
+import json
+import sys
+from pathlib import Path
+
+import numpy as np
+
+sys.path.insert(0, str(Path(__file__).parent))
+import common as C
+
+SPEECH_TOKENS = ["<|speech_start|>", "<|speech_pad|>", "<|speech_end|>"]
+
+
+def _resolve_speech_ids(tok):
+    ids = {}
+    for t in SPEECH_TOKENS:
+        i = tok.convert_tokens_to_ids(t)
+        ids[t] = i if isinstance(i, int) and i >= 0 else None
+    return ids
+
+
+def encode_audio_features(key, src, onnx_dir, device, wav):
+    """Return fused speech features [N, H] and frame count N from the ONNX front-end."""
+    a_enc = C.OnnxOp(onnx_dir / "acoustic_encoder.onnx", device)
+    s_enc = C.OnnxOp(onnx_dir / "semantic_encoder.onnx", device)
+    x = wav[None, None, :].astype(np.float32)               # [1,1,T]
+    al = a_enc.run(audio=x)                                  # [1,fa,64]
+    sl = s_enc.run(audio=x)                                  # [1,fs,128]
+    n = min(al.shape[1], sl.shape[1])
+    al, sl = al[:, :n].astype(np.float32), sl[:, :n].astype(np.float32)
+    proj = C.component_path(onnx_dir, "multi_modal_projector")
+    if proj:                                                # asr-hf
+        fused = C.OnnxOp(proj, device).run(acoustic_latents=al, semantic_latents=sl)   # [1,n,H]
+    else:                                                   # asr: sum of the two connectors
+        ac = C.OnnxOp(onnx_dir / "acoustic_connector.onnx", device).run(features=al)
+        sc = C.OnnxOp(onnx_dir / "semantic_connector.onnx", device).run(features=sl)
+        fused = ac + sc
+    return fused[0].astype(np.float32), n                    # [n,H], n
+
+
+def main():
+    ap = argparse.ArgumentParser(description="VibeVoice ASR (ONNX)")
+    ap.add_argument("model_path", help="built ONNX dir, e.g. asr-hf/cpu_int4/models")
+    ap.add_argument("--audio", required=True, help="input wav")
+    ap.add_argument("--max-new-tokens", type=int, default=256)
+    ap.add_argument("--prompt", default="Please transcribe the audio.")
+    args = ap.parse_args()
+
+    key, src, onnx_dir, device, precision, model_id = C.resolve_from_path(args.model_path)
+    print(f"=== ASR | model_id={model_id}  device={device}  precision={precision} ===")
+    if key not in ("asr", "asr-hf"):
+        sys.exit(f"[error] {model_id} ({key}) is not an ASR model — "
+                 f"use inference.py (1.5b) or inference_realtime.py (realtime)")
+    for need in ("acoustic_encoder", "semantic_encoder"):
+        if not (onnx_dir / f"{need}.onnx").exists():
+            sys.exit(f"missing {need}.onnx in {onnx_dir} — build: uv run optimize.py {key}")
+
+    wav = C.normalize_audio(C.load_audio(args.audio))
+    feats, n = encode_audio_features(key, src, onnx_dir, device, wav)
+    H = feats.shape[-1]
+    print(f"  audio {len(wav)/C.SR:.2f}s → {n} speech frames → fused features [{n},{H}]")
+
+    llm_path = onnx_dir / "llm_decoder.onnx"
+    if not llm_path.exists():
+        print(f"\n[front-end OK] llm_decoder.onnx not built (Qwen2.5-7B int4 is RAM-bound — see STATUS.md).")
+        print(f"  Build it after freeing RAM: uv run optimize.py --components llm {key}")
+        print(f"  Then re-run for the transcript.")
+        return
+
+    tok = C.load_tokenizer(onnx_dir, src)
+    sid = _resolve_speech_ids(tok)
+    # Prompt: system + user(<speech_start> <speech_pad>*n <speech_end> + instruction).
+    pad = sid["<|speech_pad|>"]
+    if pad is None:
+        # tokenizer lacks the speech tokens (repo shipped none) — add them, then use a pad id.
+        tok.add_special_tokens({"additional_special_tokens": SPEECH_TOKENS})
+        sid = _resolve_speech_ids(tok); pad = sid["<|speech_pad|>"]
+    start = sid["<|speech_start|>"] if sid["<|speech_start|>"] is not None else pad
+    end = sid["<|speech_end|>"] if sid["<|speech_end|>"] is not None else pad
+    pre = tok.encode(f"{args.prompt}\n")
+    ids = pre + [start] + [pad] * n + [end]
+    embeds = C.embed_tokens(src, ids, key)                 # [1,S,H]
+    # inject at the n dedicated speech-pad slots (the [pad]*n block), not any incidental pad id
+    pad_pos = list(range(len(pre) + 1, len(pre) + 1 + n))
+    embeds[0, pad_pos, :] = feats[:len(pad_pos)]            # inject speech features
+
+    llm = C.OnnxLLM(llm_path, device)
+    logits = llm.prefill(embeds)                            # [1,S,V]
+    eos = getattr(tok, "eos_token_id", None)
+    out = []
+    for _ in range(args.max_new_tokens):
+        nxt = int(np.asarray(logits)[0, -1].argmax())
+        if eos is not None and nxt == eos:
+            break
+        out.append(nxt)
+        logits = llm.step(C.embed_tokens(src, [nxt], key))
+    text = tok.decode(out, skip_special_tokens=True)
+    print(f"\nTRANSCRIPT:\n{text}")
+
+
+if __name__ == "__main__":
+    main()

inference_tokenizer.py

@@ -0,0 +1,119 @@
+"""VibeVoice Acoustic Tokenizer — ONNX encode/decode driver (the shared 7.5 Hz speech codec used by
+the VibeVoice TTS and ASR models). onnxruntime + numpy only (+ soundfile/librosa for audio I/O).
+
+Flow (hop = 3200 samples/frame @ 24 kHz):
+
+  audio ─[acoustic_encoder]→ latents [1, frames, 64] ─[acoustic_decoder]→ audio
+         (encoder fixed at 24000 samples/call → 7 frames; long audio is encoded in 1 s chunks)
+
+Usage:
+  python inference_tokenizer.py --models-dir fp32 --input in.wav --output recon.wav   # round-trip
+  python inference_tokenizer.py --models-dir fp32 --input in.wav --latents-out z.npy --encode-only
+  python inference_tokenizer.py --models-dir fp32 --latents-in z.npy --output out.wav --decode-only
+"""
+import argparse
+from pathlib import Path
+
+import numpy as np
+
+SR = 24_000          # sampling rate
+CHUNK = 24_000       # encoder's baked input length (1 s)
+HOP = 3_200          # samples per latent frame
+VAE_DIM = 64
+
+
+def _session(path, provider):
+    import onnxruntime as ort
+    so = ort.SessionOptions(); so.log_severity_level = 3
+    return ort.InferenceSession(str(path), sess_options=so, providers=[provider])
+
+
+def _npdtype(sess):
+    return np.float16 if "float16" in sess.get_inputs()[0].type else np.float32
+
+
+def load_audio(path):
+    import soundfile as sf
+    wav, sr = sf.read(path, dtype="float32", always_2d=False)
+    if wav.ndim > 1:
+        wav = wav.mean(1)
+    if sr != SR:
+        import librosa
+        wav = librosa.resample(wav, orig_sr=sr, target_sr=SR)
+    return wav.astype(np.float32)
+
+
+def save_audio(path, wav):
+    import soundfile as sf
+    sf.write(path, np.asarray(wav, np.float32).ravel(), SR, subtype="PCM_16")
+
+
+def encode(enc, wav):
+    """wav [T] → latents [1, frames, 64], printing the encoder flow. Chunks into 1 s windows."""
+    npdt = _npdtype(enc)
+    n = int(np.ceil(len(wav) / CHUNK))
+    padded = np.pad(wav, (0, n * CHUNK - len(wav)))
+    dt = "fp16" if npdt == np.float16 else "fp32"
+    print(f"[encode] {len(wav)} samples ({len(wav)/SR:.2f}s) @ {SR} Hz  ->  {n} x {CHUNK} chunk(s) "
+          f"(padded to {len(padded)}), {dt}")
+    lat = []
+    for i in range(n):
+        chunk = padded[i * CHUNK:(i + 1) * CHUNK][None, None, :].astype(npdt)
+        z = enc.run(["latents"], {"audio": chunk})[0]
+        lat.append(z)
+        print(f"  encoder: audio[1,1,{CHUNK}]  ->  latents{list(z.shape)}   (chunk {i+1}/{n})")
+    latents = np.concatenate(lat, axis=1)
+    print(f"[encode] done: latents {list(latents.shape)}  (7.5 Hz, {VAE_DIM}-dim per frame)")
+    return latents
+
+
+def decode(dec, latents):
+    """latents [1, frames, 64] → audio [samples], printing the decoder flow."""
+    npdt = _npdtype(dec)
+    frames = latents.shape[1]
+    audio = dec.run(["audio"], {"latents": latents.astype(npdt)})[0].ravel()
+    print(f"[decode] decoder: latents[1,{frames},{VAE_DIM}]  ->  audio[1,1,{len(audio)}]  "
+          f"({len(audio)/SR:.2f}s, {HOP}/frame)")
+    return audio
+
+
+def main():
+    ap = argparse.ArgumentParser(description="VibeVoice Acoustic Tokenizer — ONNX encode/decode")
+    ap.add_argument("--models-dir", default=".", help="dir with acoustic_encoder/decoder.onnx (fp32|fp16)")
+    ap.add_argument("--input", "-i", default=None)
+    ap.add_argument("--output", "-o", default="reconstructed.wav")
+    ap.add_argument("--latents-out", default=None)
+    ap.add_argument("--latents-in", default=None)
+    ap.add_argument("--encode-only", action="store_true")
+    ap.add_argument("--decode-only", action="store_true")
+    ap.add_argument("--cuda", action="store_true")
+    args = ap.parse_args()
+    prov = "CUDAExecutionProvider" if args.cuda else "CPUExecutionProvider"
+    d = Path(args.models_dir)
+    print(f"=== VibeVoice Acoustic Tokenizer (ONNX) | models={d} | {prov} ===")
+
+    if args.decode_only:
+        latents = np.load(args.latents_in)
+        wav = decode(_session(d / "acoustic_decoder.onnx", prov), latents)
+        save_audio(args.output, wav)
+        print(f"saved -> {args.output}")
+        return
+
+    wav = load_audio(args.input)
+    latents = encode(_session(d / "acoustic_encoder.onnx", prov), wav)
+    if args.latents_out:
+        np.save(args.latents_out, latents); print(f"latents -> {args.latents_out}")
+    if args.encode_only:
+        return
+    recon = decode(_session(d / "acoustic_decoder.onnx", prov), latents)
+    save_audio(args.output, recon)
+    n = min(len(recon), len(wav))
+    corr = float(np.corrcoef(wav[:n], recon[:n])[0, 1]) if n > 1 else float("nan")
+    print(f"saved -> {args.output}   round-trip corr={corr:+.3f}")
+    if len(wav) > CHUNK:
+        print("  note: multi-chunk round-trip loses ~1600 samples/chunk at the fixed-encoder edge "
+              "(cumulative time offset lowers corr); per-1s-window parity is corr ~0.999.")
+
+
+if __name__ == "__main__":
+    main()

int4.wav

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

optimize.py

@@ -211,7 +211,11 @@ def build_olive(model: str, component: str, model_src: Path, target: Path, preci
     tmp = (target / f"_{component}_tmp").resolve()
     passes = {"c": {"type": "OnnxConversion", "use_dynamo_exporter": True}}
     if precision == "fp16":
-        passes["f16"] = {"type": "OnnxFloatToFloat16"}
+        # keep shape/resample ops in fp32 — converting them yields invalid graphs (e.g. the acoustic
+        # tokenizer's ConstantOfShape emitted fp16 where a float32 consumer expects it) and they gain
+        # nothing from fp16. Same class of block-list the Higgs encoders needed.
+        passes["f16"] = {"type": "OnnxFloatToFloat16",
+                         "op_block_list": ["ConstantOfShape", "ConvTranspose", "Resize", "Range"]}
     cfg = {
         "input_model": {"type": "PyTorchModel", "model_path": str(model_src),
                         "model_loader": loader, "model_script": str(HERE / "user_script.py"),
@@ -257,6 +261,17 @@ def build_model(model: str, model_src: Path, device: str, precision: str, compon
         comps = [c for c in comps if c != "llm"]
         print(f"[{model}] --exclude-llm: skipping the LLM decoder build")
 
+    # int4 only quantizes the LLM (MatMulNBits). Keys with no LLM (e.g. acoustic — a conv VAE) have
+    # no quantizable MatMul weights, so int4 is a no-op that just re-emits fp32. Warn + downgrade so
+    # the output isn't a misleadingly-named fp32 copy.
+    if precision == "int4" and "llm" not in comps:
+        print(f"[{model}] WARNING: int4 is a no-op here (no LLM / conv-VAE components have no "
+              f"quantizable MatMul weights) — building fp32 instead.")
+        precision = "fp32"
+        if not output_dir:
+            target = HERE / "onnx" / model / f"{device}_fp32"
+            target.mkdir(parents=True, exist_ok=True)
+
     print(f"\n=== {model}  src={model_src}  device={device}  precision={precision} ===")
     print(f"    target={target}  components={comps}")
     for comp in comps:

pyproject.toml

@@ -14,13 +14,11 @@ dependencies = [
     "safetensors>=0.4",
     "soundfile>=0.12",
     "librosa>=0.10",
-    # VibeVoice source — provides the DPMSolverMultistepScheduler + modular tokenizer/diffusion-head
-    # classes the drivers import (via the isolated-import shim). Replaces the codes/ git submodule.
-    # NOTE (HANDOFF trap 1): do NOT `import vibevoice` directly — its package __init__ collides with
-    # transformers' native `vibevoice` registration and pulls diffusers + a renamed qwen2 tokenizer.
-    # The recipe imports individual leaf modules by path; find_spec() locates this install without
-    # executing that __init__.
-    "vibevoice @ git+https://github.com/microsoft/VibeVoice.git@main",
+    # The VibeVoice source (DPMSolverMultistepScheduler + modular tokenizer/diffusion-head classes)
+    # is VENDORED — it ships as the `vibevoice/` directory alongside these drivers, so there is NO
+    # pip/git dependency on it. The drivers import individual leaf modules from that tree by path
+    # (isolated-import shim), never `import vibevoice` (HANDOFF trap 1: its package __init__ collides
+    # with transformers' native `vibevoice` registration and pulls diffusers + a renamed qwen2 tokenizer).
 ]
 
 # Build the ONNX sub-parts (needs the source checkpoint + onnxruntime-genai ModelBuilder):

samples/text_examples/1p_abs.txt

@@ -0,0 +1,3 @@
+Speaker 1: Generating long-form, multi-speaker conversational audio like podcasts poses significant challenges for traditional Text-to-Speech (TTS) systems, particularly in scalability, speaker consistency, and natural turn-taking. This report presents VibeVoice, a novel model designed to synthesize long-form speech with multiple speakers by employing the next-token diffusion framework, a unified method for modeling continuous data by autoregressively generating latent vectors via diffusion. 
+
+Speaker 1: A core component of our approach is the continuous speech tokenizers operating at an ultra-low frame rate of 7.5. This tokenizer effectively preserves audio fidelity while significantly boosting computational efficiency for processing long sequences. This enables VibeVoice to synthesize long-form speech for up to 90 minutes (in a 64K context window length) with up to 4 speakers, capturing the authentic conversational "vibe" and surpassing all known open-source and closed-source dialogue models (for example, Gemini 2.5 Pro Preview TTS). Code and checkpoint are available now. 

samples/voices/en-Alice_woman.wav

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

tts_out.wav

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

user_script.py

@@ -126,7 +126,7 @@ def extract_qwen2_asrhf(model_path: str, output_dir: str) -> str:
 
 
 # =============================================================================
-# Acoustic tokenizer (VAE codec) — vendored `codes/` matches VibeVoice-1.5B EXACTLY
+# Acoustic tokenizer (VAE codec) — the vendored vibevoice/ source matches VibeVoice-1.5B EXACTLY
 # (552 weights, 0 missing). We import ONLY the tokenizer module in isolation (the
 # package __init__ pulls the streaming/diffusion chain → diffusers + a qwen2-tokenizer
 # import that transformers 5.10.2 renamed), and shim Auto*.register so it coexists with
@@ -135,27 +135,21 @@ def extract_qwen2_asrhf(model_path: str, output_dir: str) -> str:
 # stages/head; ASR-HF: conv_layers) — this loader targets the 1.5B one.
 # =============================================================================
 import sys as _sys
-_CODES = str(Path(__file__).parent / "codes")
+# The required VibeVoice source is VENDORED here at VibeVoice/vibevoice/ (no submodule, no git
+# dependency) — see VIBEVOICE_LICENSE. We still import it in isolation (below) because its package
+# __init__ collides with transformers' native registration.
+_VENDORED = str(Path(__file__).parent / "vibevoice")
 
 
 def _vibevoice_dir():
-    """Locate the vibevoice source tree: prefer the local codes/ submodule, else the pip-installed
-    `vibevoice` package (git dependency). find_spec does NOT run the package __init__ (no collision)."""
-    local = os.path.join(_CODES, "vibevoice")
-    if os.path.isdir(local):
-        return local
-    import importlib.util
-    spec = importlib.util.find_spec("vibevoice")
-    if spec and spec.submodule_search_locations:
-        return list(spec.submodule_search_locations)[0]
-    raise ModuleNotFoundError(
-        "vibevoice source not found. Either populate the codes/ submodule\n"
-        "  (git submodule update --init --force VibeVoice/codes), or install the git dependency:\n"
-        "  uv pip install 'vibevoice @ git+https://github.com/microsoft/VibeVoice.git@303b283'")
+    """The vendored vibevoice source tree shipped alongside this code."""
+    if os.path.isdir(_VENDORED):
+        return _VENDORED
+    raise ModuleNotFoundError(f"vendored vibevoice source missing at {_VENDORED}")
 
 
 def _codes_import(submodule):
-    """Isolated import of a single codes/vibevoice/modular/<submodule> module. Shims
+    """Isolated import of a single vibevoice/modular/<submodule> module. Shims
     Auto*.register (coexist with transformers) and injects empty `vibevoice[.modular]`
     parent packages so the package __init__ (diffusers + renamed qwen2-tokenizer) never runs."""
     import types, importlib
@@ -166,18 +160,17 @@ def _codes_import(submodule):
             try: __r(*a, **k)
             except Exception: pass
         cls.register = staticmethod(_safe)
-    # status: make the vibevoice source dependency (and where it resolves from) visible in the log
+    # status: make the vendored vibevoice source (and where it resolves from) visible in the log
     base = _vibevoice_dir()
     target = os.path.join(base, "modular", submodule + ".py")
-    src = "codes/" if base.startswith(_CODES) else "pip pkg"
-    print(f"[codes] isolated-import vibevoice.modular.{submodule}  <-  {target}  ({src})")
+    print(f"[vibevoice] isolated-import vibevoice.modular.{submodule}  <-  {target}  (vendored)")
     if not os.path.exists(target):
-        raise ModuleNotFoundError(f"{target} missing in the resolved vibevoice source ({base}).")
+        raise ModuleNotFoundError(f"{target} missing in the vendored vibevoice source ({base}).")
     for name, sub in [("vibevoice", ""), ("vibevoice.modular", "modular")]:  # empty parent pkgs
         m = types.ModuleType(name); m.__path__ = [os.path.join(base, sub)]
         _sys.modules[name] = m
     mod = importlib.import_module("vibevoice.modular." + submodule)
-    print(f"[codes] loaded {submodule} OK")
+    print(f"[vibevoice] loaded {submodule} OK (vendored)")
     return mod
 
 
@@ -188,7 +181,7 @@ def _codes_tokenizer():
 
 
 def _load_acoustic(model_path):
-    """Load VibeVoice-1.5B's acoustic tokenizer (VAE) via vendored codes/, weights loaded."""
+    """Load VibeVoice-1.5B's acoustic tokenizer (VAE) via the vendored vibevoice/ source, weights loaded."""
     import glob
     from safetensors.torch import load_file
     tok, ACfg = _codes_tokenizer()
@@ -223,15 +216,17 @@ def get_acoustic_encoder_model(model_path=None):
 
 
 def get_acoustic_encoder_io_config(model=None):
+    # dynamo IGNORES dynamic_axes (trap #2) — use dynamic_shapes so the audio length is variable,
+    # else it bakes 24000 (which isn't a multiple of the 3200 hop → 7 vs 7.5 frame drift that
+    # breaks alignment with the processor's speech-token count). frames = samples / 3200.
     return {"input_names": ["audio"], "output_names": ["latents"],
-            "input_shapes": [[1, 1, 24000]], "input_types": ["float32"],
-            "dynamic_axes": {"audio": {0: "batch", 2: "samples"},
-                             "latents": {0: "batch", 2: "frames"}}}
+            "input_shapes": [[1, 1, 25600]], "input_types": ["float32"],
+            "dynamic_shapes": {"audio": {0: "batch", 2: "samples"}}}
 
 
 def get_acoustic_encoder_dummy_inputs(model=None):
-    import torch
-    return {"audio": torch.randn(1, 1, 24000, dtype=torch.float32)}
+    import torch     # 25600 = 8 * 3200 hop → 8 frames (3200-aligned trace sample)
+    return {"audio": torch.randn(1, 1, 25600, dtype=torch.float32)}
 
 
 def get_acoustic_decoder_model(model_path=None):
@@ -298,7 +293,7 @@ def get_acoustic_std_decoder_model(model_path=None):
 
 # =============================================================================
 # ASR-HF acoustic encoder — transformers-NATIVE (VibeVoiceAcousticTokenizerEncoderModel),
-# different arch than 1.5B (conv_layers). No codes/ / shim needed. Loads only the
+# different arch than 1.5B (conv_layers). No vibevoice/ shim needed. Loads only the
 # `acoustic_tokenizer_encoder.*` weights (not the 7B LLM) so it fits in memory.
 # =============================================================================
 
@@ -441,12 +436,12 @@ def get_asrhf_projector_dummy_inputs(model=None):
 
 
 # =============================================================================
-# Realtime-0.5B (`vibevoice_streaming`, auto_map null → codes/) — a streaming TTS
+# Realtime-0.5B (`vibevoice_streaming`, auto_map null → vibevoice/) — a streaming TTS
 # checkpoint. Key groups: model.tts_language_model.* (Qwen2.5-0.5B backbone, 20 layers,
 # no lm_head), model.acoustic_tokenizer.* (DECODER-ONLY, 276 — no encoder shipped, since
 # inference only DECODES generated latents → audio), model.language_model.* (4-layer base),
 # model.prediction_head.* (diffusion), model.acoustic_connector.*, tts_eos_classifier.*.
-# The acoustic decoder matches the codes/ class EXACTLY (stages/head naming; decoder 0
+# The acoustic decoder matches the vendored vibevoice/ class EXACTLY (stages/head naming; decoder 0
 # missing / 0 unexpected) — NOT transformers-native (conv_layers/convtr naming).
 # =============================================================================
 RT_TTS_LM_PREFIX = "model.tts_language_model."
@@ -495,8 +490,8 @@ def extract_qwen2_realtime(model_path: str, output_dir: str) -> str:
 
 
 def _load_realtime_acoustic_decoder(model_path):
-    """Realtime acoustic tokenizer (DECODER-ONLY) via codes/. Loads only `decoder.*` weights
-    (encoder absent from the checkpoint), drops the encoder module. codes/ matches exactly."""
+    """Realtime acoustic tokenizer (DECODER-ONLY) via the vendored vibevoice/ source. Loads only `decoder.*` weights
+    (encoder absent from the checkpoint), drops the encoder module. the vendored source matches exactly."""
     import glob
     from safetensors.torch import load_file
     tok, ACfg = _codes_tokenizer()
@@ -538,7 +533,7 @@ def get_realtime_acoustic_decoder_dummy_inputs(model=None):
 
 
 # =============================================================================
-# Diffusion prediction_head + speech connectors (shared 1.5B / Realtime, via codes/).
+# Diffusion prediction_head + speech connectors (shared 1.5B / Realtime, via the vendored vibevoice/ source).
 #   diffusion_head: ONE denoise step (noisy_images[B,64], timesteps[B], condition[B,H]) → pred[B,64].
 #                   The ~20-step DDPM sampling loop stays in the pipeline; ONNX = one step.
 #   connector: SpeechConnector fc1(in→H) → RMSNorm(H) → fc2(H→H). 1.5B: acoustic 64→1536,
@@ -653,16 +648,16 @@ def get_semantic_connector_dummy_inputs(model=None):
 
 
 # =============================================================================
-# VibeVoice-ASR (`vibevoice`, VibeVoiceForASRTraining, codes/) — an audio→text ASR model:
-# same codes/ family as 1.5B TTS but the LLM is Qwen2.5-7B WITH lm_head (generates text) and
+# VibeVoice-ASR (`vibevoice`, VibeVoiceForASRTraining, vendored vibevoice/) — an audio→text ASR model:
+# same vibevoice/ family as 1.5B TTS but the LLM is Qwen2.5-7B WITH lm_head (generates text) and
 # there is NO prediction_head (no audio generation). Front-end = full acoustic tokenizer (552,
 # enc+dec) + semantic tokenizer (276, ENCODE-only) + acoustic/semantic connectors — all load via
-# the existing codes/ loaders (`_load_acoustic`, `_load_connector`, `_load_semantic`).
+# the existing vibevoice/ loaders (`_load_acoustic`, `_load_connector`, `_load_semantic`).
 # Weight layout: model.language_model.* (338) + top-level lm_head.weight (unlike ASR-HF).
 # =============================================================================
 
 def _load_semantic(model_path):
-    """Semantic tokenizer (ENCODE-only, deterministic latent = encode().mean) via codes/."""
+    """Semantic tokenizer (ENCODE-only, deterministic latent = encode().mean) via the vendored vibevoice/ source."""
     import glob
     from safetensors.torch import load_file
     tok = _codes_import("modular_vibevoice_tokenizer")

vibevoice/__init__.py

@@ -0,0 +1,16 @@
+# vibevoice/__init__.py
+from vibevoice.modular import (
+    VibeVoiceStreamingForConditionalGenerationInference,
+    VibeVoiceStreamingConfig,
+)
+from vibevoice.processor import (
+    VibeVoiceStreamingProcessor,
+    VibeVoiceTokenizerProcessor,
+)
+
+__all__ = [
+    "VibeVoiceStreamingForConditionalGenerationInference",
+    "VibeVoiceStreamingConfig",
+    "VibeVoiceStreamingProcessor",
+    "VibeVoiceTokenizerProcessor",
+]

vibevoice/configs/qwen2.5_1.5b_64k.json

@@ -0,0 +1,112 @@
+{
+  "_attn_implementation_autoset": true,
+  "acoustic_vae_dim": 64,
+  "acoustic_tokenizer_config": {
+    "causal": true,
+    "channels": 1,
+    "conv_bias": true,
+    "conv_norm": "none",
+    "corpus_normalize": 0.0,
+    "decoder_depths": null,
+    "decoder_n_filters": 32,
+    "decoder_ratios": [
+      8,
+      5,
+      5,
+      4,
+      2,
+      2
+    ],
+    "disable_last_norm": true,
+    "encoder_depths": "3-3-3-3-3-3-8",
+    "encoder_n_filters": 32,
+    "encoder_ratios": [
+      8,
+      5,
+      5,
+      4,
+      2,
+      2
+    ],
+    "fix_std": 0.5,
+    "layer_scale_init_value": 1e-06,
+    "layernorm": "RMSNorm",
+    "layernorm_elementwise_affine": true,
+    "layernorm_eps": 1e-05,
+    "mixer_layer": "depthwise_conv",
+    "model_type": "vibepod_acoustic_tokenizer",
+    "pad_mode": "constant",
+    "std_dist_type": "gaussian",
+    "vae_dim": 64,
+    "weight_init_value": 0.01
+  },
+  "decoder_config": {
+    "attention_dropout": 0.0,
+    "hidden_act": "silu",
+    "hidden_size": 1536,
+    "initializer_range": 0.02,
+    "intermediate_size": 8960,
+    "max_position_embeddings": 65536,
+    "max_window_layers": 28,
+    "model_type": "qwen2",
+    "num_attention_heads": 12,
+    "num_hidden_layers": 28,
+    "num_key_value_heads": 2,
+    "rms_norm_eps": 1e-06,
+    "rope_scaling": null,
+    "rope_theta": 1000000.0,
+    "sliding_window": null,
+    "tie_word_embeddings": true,
+    "torch_dtype": "bfloat16",
+    "use_cache": true,
+    "use_sliding_window": false,
+    "vocab_size": 151936
+  },
+  "diffusion_head_config": {
+    "ddpm_batch_mul": 4,
+    "ddpm_beta_schedule": "cosine",
+    "ddpm_num_inference_steps": 20,
+    "ddpm_num_steps": 1000,
+    "diffusion_type": "ddpm",
+    "head_ffn_ratio": 3.0,
+    "head_layers": 4,
+    "hidden_size": 1536,
+    "latent_size": 64,
+    "model_type": "vibepod_diffusion_head",
+    "prediction_type": "v_prediction",
+    "rms_norm_eps": 1e-05,
+    "speech_vae_dim": 64
+  },
+  "model_type": "vibepod",
+  "semantic_tokenizer_config": {
+    "causal": true,
+    "channels": 1,
+    "conv_bias": true,
+    "conv_norm": "none",
+    "corpus_normalize": 0.0,
+    "disable_last_norm": true,
+    "encoder_depths": "3-3-3-3-3-3-8",
+    "encoder_n_filters": 32,
+    "encoder_ratios": [
+      8,
+      5,
+      5,
+      4,
+      2,
+      2
+    ],
+    "fix_std": 0,
+    "layer_scale_init_value": 1e-06,
+    "layernorm": "RMSNorm",
+    "layernorm_elementwise_affine": true,
+    "layernorm_eps": 1e-05,
+    "mixer_layer": "depthwise_conv",
+    "model_type": "vibepod_semantic_tokenizer",
+    "pad_mode": "constant",
+    "std_dist_type": "none",
+    "vae_dim": 128,
+    "weight_init_value": 0.01
+  },
+  "semantic_vae_dim": 128,
+  "torch_dtype": "bfloat16"
+}

vibevoice/configs/qwen2.5_7b_32k.json

@@ -0,0 +1,113 @@
+{
+  "_attn_implementation_autoset": true,
+  "acoustic_vae_dim": 64,
+  "acoustic_tokenizer_config": {
+    "causal": true,
+    "channels": 1,
+    "conv_bias": true,
+    "conv_norm": "none",
+    "corpus_normalize": 0.0,
+    "decoder_depths": null,
+    "decoder_n_filters": 32,
+    "decoder_ratios": [
+      8,
+      5,
+      5,
+      4,
+      2,
+      2
+    ],
+    "disable_last_norm": true,
+    "encoder_depths": "3-3-3-3-3-3-8",
+    "encoder_n_filters": 32,
+    "encoder_ratios": [
+      8,
+      5,
+      5,
+      4,
+      2,
+      2
+    ],
+    "fix_std": 0.5,
+    "layer_scale_init_value": 1e-06,
+    "layernorm": "RMSNorm",
+    "layernorm_elementwise_affine": true,
+    "layernorm_eps": 1e-05,
+    "mixer_layer": "depthwise_conv",
+    "model_type": "vibepod_acoustic_tokenizer",
+    "pad_mode": "constant",
+    "std_dist_type": "gaussian",
+    "vae_dim": 64,
+    "weight_init_value": 0.01
+  },
+  "decoder_config": {
+    "attention_dropout": 0.0,
+    "hidden_act": "silu",
+    "hidden_size": 3584,
+    "initializer_range": 0.02,
+    "intermediate_size": 18944,
+    "max_position_embeddings": 32768,
+    "max_window_layers": 28,
+    "model_type": "qwen2",
+    "num_attention_heads": 28,
+    "num_hidden_layers": 28,
+    "num_key_value_heads": 4,
+    "rms_norm_eps": 1e-06,
+    "rope_theta": 1000000.0,
+    "sliding_window": null,
+    "tie_word_embeddings": false,
+    "torch_dtype": "bfloat16",
+    "transformers_version": "4.40.1",
+    "use_cache": true,
+    "use_mrope": false,
+    "use_sliding_window": false,
+    "vocab_size": 152064
+  },
+  "diffusion_head_config": {
+    "ddpm_batch_mul": 4,
+    "ddpm_beta_schedule": "cosine",
+    "ddpm_num_inference_steps": 20,
+    "ddpm_num_steps": 1000,
+    "diffusion_type": "ddpm",
+    "head_ffn_ratio": 3.0,
+    "head_layers": 4,
+    "hidden_size": 3584,
+    "latent_size": 64,
+    "model_type": "vibepod_diffusion_head",
+    "prediction_type": "v_prediction",
+    "rms_norm_eps": 1e-05,
+    "speech_vae_dim": 64
+  },
+  "model_type": "vibepod",
+  "semantic_tokenizer_config": {
+    "causal": true,
+    "channels": 1,
+    "conv_bias": true,
+    "conv_norm": "none",
+    "corpus_normalize": 0.0,
+    "disable_last_norm": true,
+    "encoder_depths": "3-3-3-3-3-3-8",
+    "encoder_n_filters": 32,
+    "encoder_ratios": [
+      8,
+      5,
+      5,
+      4,
+      2,
+      2
+    ],
+    "fix_std": 0,
+    "layer_scale_init_value": 1e-06,
+    "layernorm": "RMSNorm",
+    "layernorm_elementwise_affine": true,
+    "layernorm_eps": 1e-05,
+    "mixer_layer": "depthwise_conv",
+    "model_type": "vibepod_semantic_tokenizer",
+    "pad_mode": "constant",
+    "std_dist_type": "none",
+    "vae_dim": 128,
+    "weight_init_value": 0.01
+  },
+  "semantic_vae_dim": 128,
+  "torch_dtype": "bfloat16"
+}

vibevoice/modular/__init__.py

@@ -0,0 +1,14 @@
+# vibevoice/modular/__init__.py
+from .modeling_vibevoice_streaming_inference import VibeVoiceStreamingForConditionalGenerationInference
+from .configuration_vibevoice_streaming import VibeVoiceStreamingConfig
+from .modeling_vibevoice_streaming import VibeVoiceStreamingModel, VibeVoiceStreamingPreTrainedModel
+from .streamer import AudioStreamer, AsyncAudioStreamer
+
+__all__ = [
+    "VibeVoiceStreamingForConditionalGenerationInference",
+    "VibeVoiceStreamingConfig",
+    "VibeVoiceStreamingModel",
+    "VibeVoiceStreamingPreTrainedModel",
+    "AudioStreamer",
+    "AsyncAudioStreamer",
+]

vibevoice/modular/configuration_vibevoice.py

@@ -0,0 +1,406 @@
+""" VibeVoice_AcousticTokenizer model configuration"""
+
+from typing import Dict, List, Optional, Tuple
+
+import torch
+from transformers.configuration_utils import PretrainedConfig 
+from transformers.utils import logging
+
+from transformers.models.qwen2.configuration_qwen2 import Qwen2Config
+
+logger = logging.get_logger(__name__)
+
+
+def _convert_dtype_to_string(config_dict: dict) -> dict:
+    """
+    Convert torch.dtype objects to their string representation for JSON serialization.
+    
+    This fixes the "Object of type dtype is not JSON serializable" error that occurs
+    when transformers tries to log/serialize the config with torch_dtype as a torch.dtype object.
+    
+    See: https://github.com/microsoft/VibeVoice/issues/199
+    """
+    if "torch_dtype" in config_dict and config_dict["torch_dtype"] is not None:
+        dtype = config_dict["torch_dtype"]
+        if isinstance(dtype, torch.dtype):
+            # Convert torch.dtype to string (e.g., torch.bfloat16 -> "bfloat16")
+            config_dict["torch_dtype"] = str(dtype).replace("torch.", "")
+    return config_dict
+
+
+class VibeVoiceAcousticTokenizerConfig(PretrainedConfig):
+    model_type = "vibevoice_acoustic_tokenizer"
+
+    def __init__(
+        self,
+        channels: int = 1,
+        corpus_normalize: float = 0.0,
+        causal: bool = True,
+        vae_dim: int = 64,
+        fix_std: float = 0.5,
+        std_dist_type: str = 'gaussian',
+        # common 
+        mixer_layer: str = 'depthwise_conv',
+        conv_norm: str = 'none',
+        pad_mode: str = 'constant',
+        disable_last_norm: bool = True,
+        layernorm: str = 'RMSNorm',
+        layernorm_eps: float = 1e-5,
+        layernorm_elementwise_affine: bool = True,
+        conv_bias: bool = True,
+        layer_scale_init_value: float = 1e-6,
+        weight_init_value: float = 1e-2,
+        # encoder specific
+        encoder_n_filters: int = 32,
+        encoder_ratios: Optional[List[int]] = [8,5,5,4,2,2],
+        encoder_depths: str = "3-3-3-3-3-3-8",
+        # decoder specific
+        decoder_n_filters: int = 32,
+        decoder_ratios: Optional[List[int]] = None, # if None, same as encoder
+        decoder_depths: Optional[str] = None,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.channels = channels
+        self.corpus_normalize = corpus_normalize
+        self.causal = causal
+        self.vae_dim = vae_dim
+        self.fix_std = fix_std
+        self.std_dist_type = std_dist_type
+        
+        # common parameters
+        self.conv_norm = conv_norm
+        self.pad_mode = pad_mode
+        self.layernorm_eps = layernorm_eps
+        self.disable_last_norm = disable_last_norm
+        self.layernorm = layernorm
+        self.layernorm_elementwise_affine = layernorm_elementwise_affine
+        self.conv_bias = conv_bias
+        self.layer_scale_init_value = layer_scale_init_value
+        self.weight_init_value = weight_init_value
+        self.mixer_layer = mixer_layer
+
+        # encoder specific parameters
+        self.encoder_n_filters = encoder_n_filters
+        self.encoder_ratios = encoder_ratios
+        self.encoder_depths = encoder_depths
+        
+        # decoder specific parameters
+        self.decoder_ratios = decoder_ratios if decoder_ratios is not None else encoder_ratios
+        self.decoder_n_filters = decoder_n_filters
+        self.decoder_depths = decoder_depths
+
+
+class VibeVoiceSemanticTokenizerConfig(PretrainedConfig):
+    model_type = "vibevoice_semantic_tokenizer"
+    
+    def __init__(
+        self,
+        channels: int = 1,
+        corpus_normalize: float = 0.0,
+        causal: bool = True,
+        vae_dim: int = 64,
+        fix_std: float = 0,
+        std_dist_type: str = 'none',
+        # common 
+        mixer_layer: str = 'depthwise_conv',
+        conv_norm: str = 'none',
+        pad_mode: str = 'constant',
+        disable_last_norm: bool = True,
+        layernorm: str = 'RMSNorm',
+        layernorm_eps: float = 1e-5,
+        layernorm_elementwise_affine: bool = True,
+        conv_bias: bool = True,
+        layer_scale_init_value: float = 1e-6,
+        weight_init_value: float = 1e-2,
+        # encoder specific
+        encoder_n_filters: int = 32,
+        encoder_ratios: Optional[List[int]] = [8,5,5,4,2,2],
+        encoder_depths: str = "3-3-3-3-3-3-8",
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.channels = channels
+        self.corpus_normalize = corpus_normalize
+        self.causal = causal
+        self.vae_dim = vae_dim
+        self.fix_std = fix_std
+        self.std_dist_type = std_dist_type
+        
+        # common parameters
+        self.conv_norm = conv_norm
+        self.pad_mode = pad_mode
+        self.layernorm_eps = layernorm_eps
+        self.disable_last_norm = disable_last_norm
+        self.layernorm = layernorm
+        self.layernorm_elementwise_affine = layernorm_elementwise_affine
+        self.conv_bias = conv_bias
+        self.layer_scale_init_value = layer_scale_init_value
+        self.weight_init_value = weight_init_value
+        self.mixer_layer = mixer_layer
+
+        # encoder specific parameters
+        self.encoder_n_filters = encoder_n_filters
+        self.encoder_ratios = encoder_ratios
+        self.encoder_depths = encoder_depths
+        
+
+class VibeVoiceDiffusionHeadConfig(PretrainedConfig):
+    model_type = "vibevoice_diffusion_head"
+
+    def __init__(
+        self,
+        hidden_size=768,
+        head_layers=4,
+        head_ffn_ratio=3.0,
+        rms_norm_eps=1e-5,
+        latent_size=64,
+        speech_vae_dim=None,
+        prediction_type="v_prediction",
+        diffusion_type="ddpm",
+        ddpm_num_steps=1000,
+        ddpm_num_inference_steps=20,
+        ddpm_beta_schedule="cosine",
+        ddpm_batch_mul=4,
+        **kwargs
+    ):
+        self.hidden_size = hidden_size
+        self.head_layers = head_layers
+        self.head_ffn_ratio = head_ffn_ratio
+        self.rms_norm_eps = rms_norm_eps
+        self.latent_size = latent_size
+        self.speech_vae_dim = speech_vae_dim
+        self.prediction_type = prediction_type
+        self.diffusion_type = diffusion_type
+        self.ddpm_num_steps = ddpm_num_steps
+        self.ddpm_num_inference_steps = ddpm_num_inference_steps
+        self.ddpm_beta_schedule = ddpm_beta_schedule
+        self.ddpm_batch_mul = ddpm_batch_mul
+        
+        super().__init__(**kwargs)
+
+class VibeVoiceConfig(PretrainedConfig):
+    model_type = "vibevoice"
+    is_composition = True
+    sub_configs = {
+        "acoustic_tokenizer_config": VibeVoiceAcousticTokenizerConfig, 
+        "semantic_tokenizer_config": VibeVoiceSemanticTokenizerConfig,
+        "decoder_config": Qwen2Config,
+        "diffusion_head_config": VibeVoiceDiffusionHeadConfig,
+    }
+    # keys_to_ignore_at_inference = ["past_key_values"]
+    # Default tensor parallel plan for base model `Qwen2`
+    base_model_tp_plan = {
+        "layers.*.self_attn.q_proj": "colwise",
+        "layers.*.self_attn.k_proj": "colwise",
+        "layers.*.self_attn.v_proj": "colwise",
+        "layers.*.self_attn.o_proj": "rowwise",
+        "layers.*.mlp.gate_proj": "colwise",
+        "layers.*.mlp.up_proj": "colwise",
+        "layers.*.mlp.down_proj": "rowwise",
+    }
+    
+    def __init__(
+        self,
+        acoustic_tokenizer_config=None,
+        semantic_tokenizer_config=None,
+        decoder_config=None,
+        diffusion_head_config=None,
+        **kwargs
+    ):
+
+        # kwargs["_attn_implementation"] = "flash_attention_2"
+        kwargs["_attn_implementation_autoset"] = False 
+
+        if acoustic_tokenizer_config is None:
+            self.acoustic_tokenizer_config = self.sub_configs["acoustic_tokenizer_config"]()
+        elif isinstance(acoustic_tokenizer_config, dict):
+            acoustic_tokenizer_config["model_type"] = "vibevoice_acoustic_tokenizer"
+            self.acoustic_tokenizer_config = self.sub_configs["acoustic_tokenizer_config"](**acoustic_tokenizer_config)
+        elif isinstance(acoustic_tokenizer_config, VibeVoiceAcousticTokenizerConfig):
+            # If an instance of the config class is provided
+            self.acoustic_tokenizer_config = acoustic_tokenizer_config
+
+        if semantic_tokenizer_config is None:
+            self.semantic_tokenizer_config = self.sub_configs["semantic_tokenizer_config"]()
+        elif isinstance(semantic_tokenizer_config, dict):
+            semantic_tokenizer_config["model_type"] = "vibevoice_semantic_tokenizer"
+            self.semantic_tokenizer_config = self.sub_configs["semantic_tokenizer_config"](**semantic_tokenizer_config)
+        elif isinstance(semantic_tokenizer_config, VibeVoiceSemanticTokenizerConfig):
+            # If an instance of the config class is provided
+            self.semantic_tokenizer_config = semantic_tokenizer_config
+
+        if decoder_config is None:
+            self.decoder_config = self.sub_configs["decoder_config"]()
+        elif isinstance(decoder_config, dict):
+            # If a dictionary is provided, instantiate the config class with it
+            # self.decoder_config = self.sub_configs["decoder_config"](**decoder_config)
+            if decoder_config.get("model_type", '') == "qwen2":
+                self.decoder_config = Qwen2Config(**decoder_config)
+            else:
+                raise ValueError(f"Unsupported decoder model type: {decoder_config.get('model_type', '')}")
+        elif isinstance(decoder_config, (Qwen2Config,)):
+            # If an instance of the config class is provided
+            self.decoder_config = decoder_config
+
+        if diffusion_head_config is None:
+            self.diffusion_head_config = self.sub_configs["diffusion_head_config"]()
+        elif isinstance(diffusion_head_config, dict):
+            diffusion_head_config["model_type"] = "vibevoice_diffusion_head"
+            self.diffusion_head_config = self.sub_configs["diffusion_head_config"](**diffusion_head_config)
+        elif isinstance(diffusion_head_config, VibeVoiceDiffusionHeadConfig):
+            # If an instance of the config class is provided
+            self.diffusion_head_config = diffusion_head_config
+
+        # other parameters
+        self.acoustic_vae_dim = getattr(self.acoustic_tokenizer_config, 'vae_dim', 64)
+        self.semantic_vae_dim = getattr(self.semantic_tokenizer_config, 'vae_dim', 128)
+
+        super().__init__(**kwargs)
+
+    def get_text_config(self, decoder=False):
+        """
+        Returns the text config for this model.
+        
+        vLLM uses this method to get the text configuration from multimodal models.
+        This allows vLLM to correctly determine hidden_size, num_attention_heads,
+        and other properties needed for memory profiling and model execution.
+        
+        For VibeVoice, the "text config" is the decoder_config (Qwen2Config).
+        
+        Args:
+            decoder: If True, return the decoder config (for encoder-decoder models).
+                    For VibeVoice, this is always the decoder_config.
+        
+        Returns:
+            The decoder configuration (Qwen2Config) which contains hidden_size, etc.
+        """
+        return self.decoder_config
+
+    def to_dict(self):
+        """
+        Override to_dict to handle torch.dtype serialization.
+        
+        Fixes: https://github.com/microsoft/VibeVoice/issues/199
+        """
+        output = super().to_dict()
+        return _convert_dtype_to_string(output)
+
+class VibeVoiceASRConfig(PretrainedConfig):
+    model_type = "vibevoice"
+    is_composition = True
+    sub_configs = {
+        "acoustic_tokenizer_config": VibeVoiceAcousticTokenizerConfig, 
+        "semantic_tokenizer_config": VibeVoiceSemanticTokenizerConfig,
+        "decoder_config": Qwen2Config,
+    }
+    # keys_to_ignore_at_inference = ["past_key_values"]
+    # Default tensor parallel plan for base model `Qwen2`
+    base_model_tp_plan = {
+        "layers.*.self_attn.q_proj": "colwise",
+        "layers.*.self_attn.k_proj": "colwise",
+        "layers.*.self_attn.v_proj": "colwise",
+        "layers.*.self_attn.o_proj": "rowwise",
+        "layers.*.mlp.gate_proj": "colwise",
+        "layers.*.mlp.up_proj": "colwise",
+        "layers.*.mlp.down_proj": "rowwise",
+    }
+    
+    def __init__(
+        self,
+        acoustic_tokenizer_config=None,
+        semantic_tokenizer_config=None,
+        decoder_config=None,
+        **kwargs
+    ):
+
+        # kwargs["_attn_implementation"] = "flash_attention_2"
+        kwargs["_attn_implementation_autoset"] = False 
+
+        if acoustic_tokenizer_config is None:
+            self.acoustic_tokenizer_config = self.sub_configs["acoustic_tokenizer_config"]()
+        elif isinstance(acoustic_tokenizer_config, dict):
+            acoustic_tokenizer_config["model_type"] = "vibevoice_acoustic_tokenizer"
+            self.acoustic_tokenizer_config = self.sub_configs["acoustic_tokenizer_config"](**acoustic_tokenizer_config)
+        elif isinstance(acoustic_tokenizer_config, VibeVoiceAcousticTokenizerConfig):
+            # If an instance of the config class is provided
+            self.acoustic_tokenizer_config = acoustic_tokenizer_config
+
+        if semantic_tokenizer_config is None:
+            self.semantic_tokenizer_config = self.sub_configs["semantic_tokenizer_config"]()
+        elif isinstance(semantic_tokenizer_config, dict):
+            semantic_tokenizer_config["model_type"] = "vibevoice_semantic_tokenizer"
+            self.semantic_tokenizer_config = self.sub_configs["semantic_tokenizer_config"](**semantic_tokenizer_config)
+        elif isinstance(semantic_tokenizer_config, VibeVoiceSemanticTokenizerConfig):
+            # If an instance of the config class is provided
+            self.semantic_tokenizer_config = semantic_tokenizer_config
+
+        if decoder_config is None:
+            self.decoder_config = self.sub_configs["decoder_config"]()
+        elif isinstance(decoder_config, dict):
+            # If a dictionary is provided, instantiate the config class with it
+            # self.decoder_config = self.sub_configs["decoder_config"](**decoder_config)
+            if decoder_config.get("model_type", '') == "qwen2":
+                self.decoder_config = Qwen2Config(**decoder_config)
+            else:
+                raise ValueError(f"Unsupported decoder model type: {decoder_config.get('model_type', '')}")
+        elif isinstance(decoder_config, Qwen2Config):
+            # If an instance of the config class is provided
+            self.decoder_config = decoder_config
+
+        # other parameters
+        self.acoustic_vae_dim = getattr(self.acoustic_tokenizer_config, 'vae_dim', 64)
+        self.semantic_vae_dim = getattr(self.semantic_tokenizer_config, 'vae_dim', 128)
+
+        super().__init__(**kwargs)
+
+    def to_dict(self):
+        """
+        Override to_dict to handle torch.dtype serialization.
+        
+        Fixes: https://github.com/microsoft/VibeVoice/issues/199
+        """
+        output = super().to_dict()
+        return _convert_dtype_to_string(output)
+
+    def get_text_config(self, decoder: bool = False):
+        """Return the text (decoder) config for generation."""
+        return self.decoder_config
+
+    @property
+    def vocab_size(self):
+        """Return vocab_size from decoder config for generation compatibility."""
+        return self.decoder_config.vocab_size
+
+    @property
+    def num_attention_heads(self):
+        """Return num_attention_heads from decoder config for Ulysses SP compatibility."""
+        return self.decoder_config.num_attention_heads
+    
+    @property
+    def num_key_value_heads(self):
+        """Return num_key_value_heads from decoder config for Ulysses SP compatibility."""
+        return self.decoder_config.num_key_value_heads
+    
+    @property
+    def hidden_size(self):
+        """Return hidden_size from decoder config for model compatibility."""
+        return self.decoder_config.hidden_size
+    
+    @property
+    def num_hidden_layers(self):
+        """Return num_hidden_layers from decoder config for Ulysses SP compatibility."""
+        return self.decoder_config.num_hidden_layers
+    
+    @property
+    def head_dim(self):
+        """Return head_dim from decoder config for Ulysses SP compatibility."""
+        return getattr(self.decoder_config, 'head_dim', self.hidden_size // self.num_attention_heads)
+    
+__all__ = [
+    "VibeVoiceAcousticTokenizerConfig", 
+    "VibeVoiceSemanticTokenizerConfig", 
+    "VibeVoiceDiffusionHeadConfig", 
+    "VibeVoiceConfig",
+    "VibeVoiceASRConfig"
+]

vibevoice/modular/configuration_vibevoice_streaming.py

@@ -0,0 +1,104 @@
+""" VibeVoice Streaming model configuration"""
+
+import torch
+from transformers.configuration_utils import PretrainedConfig 
+from transformers.utils import logging
+
+from transformers.models.qwen2.configuration_qwen2 import Qwen2Config
+
+from .configuration_vibevoice import VibeVoiceAcousticTokenizerConfig, VibeVoiceDiffusionHeadConfig, _convert_dtype_to_string
+
+logger = logging.get_logger(__name__)
+
+
+class VibeVoiceStreamingConfig(PretrainedConfig):
+    model_type = "vibevoice_streaming"
+    is_composition = True
+    sub_configs = {
+        "acoustic_tokenizer_config": VibeVoiceAcousticTokenizerConfig,
+        "decoder_config": Qwen2Config,
+        "diffusion_head_config": VibeVoiceDiffusionHeadConfig,
+    }
+    # keys_to_ignore_at_inference = ["past_key_values"]
+    # Default tensor parallel plan for base model `Qwen2`
+    base_model_tp_plan = {
+        "layers.*.self_attn.q_proj": "colwise",
+        "layers.*.self_attn.k_proj": "colwise",
+        "layers.*.self_attn.v_proj": "colwise",
+        "layers.*.self_attn.o_proj": "rowwise",
+        "layers.*.mlp.gate_proj": "colwise",
+        "layers.*.mlp.up_proj": "colwise",
+        "layers.*.mlp.down_proj": "rowwise",
+    }
+    
+    def __init__(
+        self,
+        acoustic_tokenizer_config=None,
+        decoder_config=None,
+        diffusion_head_config=None,
+        tts_backbone_num_hidden_layers=20,
+        **kwargs
+    ):
+
+        # kwargs["_attn_implementation"] = "flash_attention_2"
+        kwargs["_attn_implementation_autoset"] = False 
+
+        if acoustic_tokenizer_config is None:
+            self.acoustic_tokenizer_config = self.sub_configs["acoustic_tokenizer_config"]()
+        elif isinstance(acoustic_tokenizer_config, dict):
+            acoustic_tokenizer_config["model_type"] = "vibevoice_acoustic_tokenizer"
+            self.acoustic_tokenizer_config = self.sub_configs["acoustic_tokenizer_config"](**acoustic_tokenizer_config)
+        elif isinstance(acoustic_tokenizer_config, VibeVoiceAcousticTokenizerConfig):
+            # If an instance of the config class is provided
+            self.acoustic_tokenizer_config = acoustic_tokenizer_config
+
+        if decoder_config is None:
+            self.decoder_config = self.sub_configs["decoder_config"]()
+        elif isinstance(decoder_config, dict):
+            # If a dictionary is provided, instantiate the config class with it
+            # self.decoder_config = self.sub_configs["decoder_config"](**decoder_config)
+            if decoder_config.get("model_type", '') == "qwen2":
+                self.decoder_config = Qwen2Config(**decoder_config)
+            else:
+                raise ValueError(f"Unsupported decoder model type: {decoder_config.get('model_type', '')}")
+        elif isinstance(decoder_config, (Qwen2Config,)):
+            # If an instance of the config class is provided
+            self.decoder_config = decoder_config
+
+        if diffusion_head_config is None:
+            self.diffusion_head_config = self.sub_configs["diffusion_head_config"]()
+        elif isinstance(diffusion_head_config, dict):
+            diffusion_head_config["model_type"] = "vibevoice_diffusion_head"
+            self.diffusion_head_config = self.sub_configs["diffusion_head_config"](**diffusion_head_config)
+        elif isinstance(diffusion_head_config, VibeVoiceDiffusionHeadConfig):
+            # If an instance of the config class is provided
+            self.diffusion_head_config = diffusion_head_config
+
+        # other parameters
+        self.acoustic_vae_dim = getattr(self.acoustic_tokenizer_config, 'vae_dim', 64)
+        # The decoder of the model is divided into two components. The lower Transformer layers are only used for encoding text, while the upper Transformer layers are used for encoding text and generating speech. `tts_backbone_num_hidden_layers` indicates the number of upper layers used for TTS.
+        self.tts_backbone_num_hidden_layers = tts_backbone_num_hidden_layers
+
+        super().__init__(**kwargs)
+
+    def get_text_config(self, decoder=False):
+        """Returns the decoder config (required for transformers >= 4.57 cache compatibility)."""
+        return self.decoder_config
+
+    @property
+    def num_hidden_layers(self):
+        """Proxy to decoder_config.num_hidden_layers (required for transformers >= 4.57)."""
+        return self.decoder_config.num_hidden_layers
+
+    def to_dict(self):
+        """
+        Override to_dict to handle torch.dtype serialization.
+        
+        Fixes: https://github.com/microsoft/VibeVoice/issues/199
+        """
+        output = super().to_dict()
+        return _convert_dtype_to_string(output)
+
+__all__ = [
+    "VibeVoiceStreamingConfig"
+]

vibevoice/modular/modeling_vibevoice.py

@@ -0,0 +1,496 @@
+# copied from https://github.com/vibevoice-community/VibeVoice/blob/main/vibevoice/modular/modeling_vibevoice.py
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Tuple, Union, Callable
+from tqdm import tqdm
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.distributed as dist
+
+from transformers.models.auto import AutoModel, AutoModelForCausalLM
+
+from transformers.activations import ACT2FN
+from transformers.modeling_outputs import CausalLMOutput, BaseModelOutputWithPast, ModelOutput
+from transformers.models.llama.modeling_llama import LlamaRMSNorm
+from transformers import modeling_utils
+from transformers.modeling_utils import PreTrainedModel
+from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
+from transformers.utils import logging
+
+
+from .modular_vibevoice_tokenizer import VibeVoiceTokenizerStreamingCache, VibeVoiceAcousticTokenizerModel, VibeVoiceSemanticTokenizerModel
+from .modular_vibevoice_diffusion_head import VibeVoiceDiffusionHead
+from vibevoice.schedule.dpm_solver import DPMSolverMultistepScheduler
+
+from .configuration_vibevoice import VibeVoiceConfig
+
+
+logger = logging.get_logger(__name__)
+
+if not hasattr(modeling_utils, "ALL_PARALLEL_STYLES") or modeling_utils.ALL_PARALLEL_STYLES is None:
+    modeling_utils.ALL_PARALLEL_STYLES = ["tp", "none", "colwise", "rowwise"]
+
+@dataclass
+class VibeVoiceCausalLMOutputWithPast(ModelOutput):
+    loss: Optional[torch.FloatTensor] = None
+    diffusion_loss: Optional[torch.FloatTensor] = None
+    speech_token_num: Optional[int] = None
+    logits: torch.FloatTensor = None
+    past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
+    hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None
+    attentions: Optional[Tuple[torch.FloatTensor, ...]] = None
+
+
+@dataclass
+class VibeVoiceGenerationOutput(ModelOutput):
+    """
+    Output type for VibeVoice generation.
+    
+    Args:
+        sequences (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+            The generated sequences. 
+        speech_outputs (`List[torch.FloatTensor]`, *optional*):
+            List of generated speech waveforms or latents for each speech segment.
+    """
+    sequences: torch.LongTensor = None
+    speech_outputs: Optional[List[torch.FloatTensor]] = None
+
+
+class SpeechConnector(nn.Module):
+    def __init__(self, input_dim, output_dim):
+        super().__init__()
+        self.fc1 = nn.Linear(input_dim, output_dim)
+        self.norm = LlamaRMSNorm(output_dim, eps=1e-6)
+        self.fc2 = nn.Linear(output_dim, output_dim)
+
+    def forward(self, features, **kwargs):    
+        x = self.fc1(features)
+        x = self.norm(x)
+        x = self.fc2(x)
+        return x
+
+
+# @auto_docstring
+class VibeVoicePreTrainedModel(PreTrainedModel):
+    config_class = VibeVoiceConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _skip_keys_device_placement = "past_key_values"
+    _supports_cache_class = True
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+    _supports_quantized_cache = True
+    _supports_static_cache = True
+    _supports_attention_backend = True
+
+    def _init_weights(self, module):
+        if isinstance(module, VibeVoiceDiffusionHead):
+            module.initialize_weights()
+            return
+
+        # Use the language model's initializer_range if available
+        if hasattr(self.config, 'language_model_config') and hasattr(self.config.language_model_config, 'initializer_range'):
+            std = self.config.language_model_config.initializer_range
+        elif hasattr(self.config, 'decoder_config') and hasattr(self.config.decoder_config, 'initializer_range'):
+            std = self.config.decoder_config.initializer_range
+        else:
+            std = 0.02  # Default value
+            
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.LayerNorm):
+            module.weight.data.fill_(1.0)
+            module.bias.data.zero_()
+
+# @auto_docstring
+class VibeVoiceModel(VibeVoicePreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        
+        if hasattr(config, 'torch_dtype') and config.torch_dtype is not None:
+            if isinstance(config.torch_dtype, str):
+                dtype = getattr(torch, config.torch_dtype)
+            else:
+                dtype = config.torch_dtype
+        else:
+            dtype = torch.float32
+        
+        # Initialize Qwen2 model for language modeling
+        lm_config = config.decoder_config 
+        self.language_model = AutoModel.from_config(lm_config)
+        
+        # Initialize speech components if needed
+        self.acoustic_tokenizer = AutoModel.from_config(config.acoustic_tokenizer_config).to(dtype)
+        self.semantic_tokenizer = AutoModel.from_config(config.semantic_tokenizer_config).to(dtype)
+
+        self.acoustic_connector = SpeechConnector(config.acoustic_vae_dim, lm_config.hidden_size).to(dtype)
+        self.semantic_connector = SpeechConnector(config.semantic_vae_dim, lm_config.hidden_size).to(dtype)
+        
+        # Register scaling factors as buffers - use 1D tensors for FSDP compatibility
+        self.register_buffer('speech_scaling_factor', torch.tensor(float('nan')))  
+        self.register_buffer('speech_bias_factor', torch.tensor(float('nan')))
+
+        # Initialize prediction head for speech generation
+        self.prediction_head = AutoModel.from_config(config.diffusion_head_config).to(dtype)
+
+        # Initialize noise scheduler
+        self.noise_scheduler = DPMSolverMultistepScheduler(
+            num_train_timesteps=config.diffusion_head_config.ddpm_num_steps,
+            beta_schedule=config.diffusion_head_config.ddpm_beta_schedule,
+            prediction_type=config.diffusion_head_config.prediction_type
+        )
+    
+    def get_input_embeddings(self):
+        if hasattr(self.language_model, 'embed_tokens'):
+            # If the language model has an embed_tokens attribute, return it
+            return self.language_model.embed_tokens
+        
+        for name, attr in self.language_model.fullmap.items(): # parallel by nnscaler, the name is changed
+            if attr.orig_name == 'embed_tokens.weight':
+                return getattr(self.language_model, name)
+        assert False, 'should not arrive here'
+
+    def set_input_embeddings(self, value):
+        self.language_model.embed_tokens = value
+    
+    def set_speech_tokenizers(self, acoustic_tokenizer=None, semantic_tokenizer=None):
+        """Set the speech tokenizers used for encoding and decoding speech."""
+        self.acoustic_tokenizer = acoustic_tokenizer
+        self.semantic_tokenizer = semantic_tokenizer
+        
+        # Reset the encoder to evaluation mode
+        if self.acoustic_tokenizer is not None:
+            self.acoustic_tokenizer.eval()
+            
+        if self.semantic_tokenizer is not None:
+            self.semantic_tokenizer.eval()
+    
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs,
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+        
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        
+        # Forward through language model
+        outputs = self.language_model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            cache_position=cache_position,
+            **kwargs,
+        )
+        
+        if not return_dict:
+            return outputs
+            
+        return BaseModelOutputWithPast(
+            last_hidden_state=outputs.last_hidden_state,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+class VibeVoiceForConditionalGeneration(VibeVoicePreTrainedModel):
+    _tied_weights_keys = ["lm_head.weight"]
+    _tp_plan = {"lm_head": "colwise_rep"}
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = VibeVoiceModel(config)
+        self.vocab_size = config.decoder_config.vocab_size
+        self.lm_head = nn.Linear(config.decoder_config.hidden_size, self.vocab_size, bias=False)
+
+        self.post_init()
+        
+    def get_input_embeddings(self):
+        return self.model.get_input_embeddings()
+
+    def set_input_embeddings(self, value):
+        self.model.set_input_embeddings(value)
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_decoder(self, decoder):
+        self.model.language_model = decoder
+
+    def get_decoder(self):
+        return self.model.language_model
+
+    def tie_weights(self):
+        """
+        Tie the weights between the input embeddings and the output embeddings.
+        """
+        if getattr(self.config.decoder_config, 'tie_word_embeddings', False):
+            # The standard PreTrainedModel method will handle the tying.
+            # It typically does a simple parameter object assignment, which is
+            # CORRECT to do BEFORE FSDP wraps the model.
+            output_embeddings = self.get_output_embeddings()
+            input_embeddings = self.get_input_embeddings()
+            if hasattr(input_embeddings, 'weight'):
+                output_embeddings.weight = input_embeddings.weight
+            else:
+                # maybe returned input_embeddings a tensor directly
+                output_embeddings.weight = input_embeddings
+
+            if getattr(output_embeddings, "bias", None) is not None:
+                output_embeddings.bias.data = nn.functional.pad(
+                    output_embeddings.bias.data,
+                    (0, output_embeddings.weight.shape[0] - output_embeddings.bias.shape[0]),
+                    "constant",
+                    0,
+                )
+            print("Tied input and output embeddings using standard assignment.")
+        else:
+            print("tie_word_embeddings is False, not tying weights.")
+
+    # Also, ensure set_output_embeddings is safe, though your implementation looks okay.
+    # The key is to avoid calling it after accelerator.prepare().
+    def set_output_embeddings(self, new_embeddings):
+        # Your current implementation using data.copy_ is good practice,
+        # but the best way is to not call this after prepare().
+        self.lm_head = new_embeddings
+
+    def forward_speech_features(
+            self, 
+            speech_tensors=None, 
+            speech_masks=None, 
+            speech_type="audio", 
+            return_unmask=False
+        ):
+        if speech_tensors is None:
+            # Use config to get vae_dim instead of non-existent self.args
+            vae_dim = self.config.acoustic_tokenizer_config.vae_dim
+            audio_features = torch.zeros(1, 1, vae_dim).to(self.get_input_embeddings().weight)
+            connect_features = self.model.acoustic_connector(audio_features)
+            return audio_features, connect_features
+        else:
+            with torch.no_grad():
+                if speech_type == "audio":
+                    with torch.no_grad():
+                        frames = self.model.acoustic_tokenizer.encode(speech_tensors.unsqueeze(1))[0][0]
+                    audio_tokens = frames.sample(self.model.acoustic_tokenizer.std_dist_type)[0]
+
+                elif speech_type == "vae":
+                    # Use config to get vae_dim instead of non-existent self.args
+                    vae_dim = self.config.acoustic_tokenizer_config.vae_dim
+                    speech_mode = speech_tensors.reshape(speech_tensors.size(0), -1, vae_dim)
+
+                    # gaussian sample from the speech_mode
+                    batch_size = speech_mode.size(0)
+                    value = self.model.acoustic_tokenizer.fix_std / 0.8
+                    std = torch.randn(batch_size, dtype=speech_mode.dtype, device=speech_mode.device) * value
+                    std = std.view(-1, *[1] * (speech_mode.dim() - 1))
+                    audio_tokens = speech_mode + std * torch.randn(speech_mode.shape).to(speech_mode)
+                else:
+                    raise NotImplementedError(f"Speech type {speech_type} not implemented")
+                
+                if torch.isnan(self.model.speech_scaling_factor) or torch.isnan(self.model.speech_bias_factor):
+                    scaling_factor = 1. / audio_tokens[speech_masks].flatten().std()
+                    bias_factor = -audio_tokens[speech_masks].flatten().mean()
+                    
+                    # Only use distributed operations if the process group is initialized
+                    if dist.is_available() and dist.is_initialized():
+                        dist.all_reduce(scaling_factor, op=dist.ReduceOp.SUM)
+                        dist.all_reduce(bias_factor, op=dist.ReduceOp.SUM)
+                        world_size = dist.get_world_size()
+                        self.model.speech_scaling_factor.copy_(scaling_factor / world_size)  
+                        self.model.speech_bias_factor.copy_(bias_factor / world_size)
+                        print(f"Speech scaling factor (distributed): {self.model.speech_scaling_factor}, bias factor: {self.model.speech_bias_factor}", flush=True)
+                    else:
+                        # Single process case
+                        self.model.speech_scaling_factor.copy_(scaling_factor)  
+                        self.model.speech_bias_factor.copy_(bias_factor)
+                        print(f"Speech scaling factor (single process): {self.model.speech_scaling_factor}, bias factor: {self.model.speech_bias_factor}", flush=True)
+                    
+                audio_features = (audio_tokens + self.model.speech_bias_factor) * self.model.speech_scaling_factor
+            
+            connect_features = self.model.acoustic_connector(audio_features)
+            if return_unmask:
+                return audio_features, connect_features
+            return audio_features[speech_masks], connect_features[speech_masks]
+        
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = False,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        # New arguments for speech processing and loss calculation
+        speech_tensors: Optional[torch.FloatTensor] = None,
+        speech_masks: Optional[torch.BoolTensor] = None,
+        speeches_loss_input: Optional[torch.FloatTensor] = None,
+        speech_semantic_tensors: Optional[torch.FloatTensor] = None, 
+        acoustic_input_mask: Optional[torch.BoolTensor] = None,
+        acoustic_loss_mask: Optional[torch.BoolTensor] = None,
+        ddpm_batch_mul: int = 1,
+        **kwargs: Optional[Dict[str, Union[torch.Tensor, str]]],
+        ) -> Union[Tuple, VibeVoiceCausalLMOutputWithPast]:
+        
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        
+        x = self.get_input_embeddings()(input_ids)
+
+        semantic_speech_all_connect_features = self.model.semantic_connector(speech_semantic_tensors)
+        if speeches_loss_input is not None:
+            # only part audio need diffuse
+            speech_all_features, speech_all_connect_features = self.forward_speech_features(
+                    speech_tensors=speech_tensors.type_as(x) if speech_tensors is not None else None,
+                    speech_masks=speech_masks,
+                    speech_type=kwargs.get("speech_type", "audio"),
+                    return_unmask=True
+                )
+            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]
+                    )
+                else:
+                    x[acoustic_input_mask] = speech_all_connect_features[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,
+                    speech_masks=speech_masks,
+                    speech_type=kwargs.get("speech_type", "audio"),
+                )
+            if speech_tensors is not None:
+                x[acoustic_input_mask] = speech_connect_features
+
+        outputs = self.model(
+            input_ids=None,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=x,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=False,
+            return_dict=return_dict,
+            cache_position=cache_position,
+        )
+
+        hidden_states = outputs.last_hidden_state
+        logits = self.lm_head(hidden_states)
+        # logits = logits.float()
+
+        loss = None
+        if labels is not None:
+            # The custom CE loss with masking is calculated in the training script.
+            # We leave the standard loss calculation here as None.
+            pass
+
+        # --- Diffusion Loss Calculation ---
+        diffusion_loss = None
+        # This block is executed only if we are in a context that involves speech.
+        if speech_tensors is not None and acoustic_loss_mask.sum().item() > 0:
+            condition_features = hidden_states[acoustic_loss_mask]
+            
+            speech_len, latent_size = speech_features.shape
+            
+            noise = torch.randn(
+                (speech_len * ddpm_batch_mul, latent_size),
+                device=hidden_states.device,
+                dtype=hidden_states.dtype
+            )
+            
+            timesteps = torch.multinomial(
+                torch.ones(self.config.diffusion_head_config.ddpm_num_steps),
+                speech_len * ddpm_batch_mul,
+                replacement=True,
+            ).to(hidden_states.device)
+
+            speech_features_repeated = speech_features.repeat_interleave(ddpm_batch_mul, dim=0)
+            condition_features_repeated = condition_features.repeat_interleave(ddpm_batch_mul, dim=0)
+
+            noisy_speech_features = self.model.noise_scheduler.add_noise(
+                speech_features_repeated, noise, timesteps
+            )
+            
+            model_output = self.model.prediction_head(
+                noisy_speech_features, 
+                timesteps.type_as(x), 
+                condition_features_repeated
+            )
+
+            prediction_type = self.config.diffusion_head_config.prediction_type
+            if prediction_type == "epsilon":
+                target_for_loss = noise
+            elif prediction_type == "v_prediction":
+                target_for_loss = self.model.noise_scheduler.get_velocity(
+                    speech_features_repeated, noise, timesteps
+                )
+            else:
+                raise NotImplementedError(f"Prediction type {prediction_type} not implemented")
+
+            diffusion_loss = F.mse_loss(model_output.float(), target_for_loss.float(), reduction='sum')
+            if latent_size > 0 and ddpm_batch_mul > 0:
+                diffusion_loss = diffusion_loss / latent_size / ddpm_batch_mul
+            else:
+                diffusion_loss = torch.tensor(0.0, device=diffusion_loss.device)
+        
+        else:
+            # Dummy loss for DDP to work when there are no speech samples in a batch,
+            # but we are in a speech context.
+            diffusion_loss = sum(p.sum() for p in self.model.prediction_head.parameters()) * 0.0
+            diffusion_loss += sum(p.sum() for p in self.model.acoustic_connector.parameters()) * 0.0
+            diffusion_loss += sum(p.sum() for p in self.model.semantic_connector.parameters()) * 0.0
+        # --- End Diffusion Loss Calculation ---
+
+        if not return_dict:
+            output = (logits, speech_len) + outputs.to_tuple()[1:]
+            return (loss, diffusion_loss) + output
+
+        return VibeVoiceCausalLMOutputWithPast(
+            loss=loss,
+            diffusion_loss=diffusion_loss,
+            speech_token_num=speech_len if speech_tensors is not None else 0,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+AutoModel.register(VibeVoiceConfig, VibeVoiceModel)
+AutoModelForCausalLM.register(VibeVoiceConfig, VibeVoiceForConditionalGeneration)
+
+__all__ = [
+    "VibeVoiceModel",
+    "VibeVoicePreTrainedModel",
+    "VibeVoiceForConditionalGeneration",
+    "VibeVoiceCausalLMOutputWithPast",
+    "VibeVoiceGenerationOutput",
+]

vibevoice/modular/modeling_vibevoice_asr.py

@@ -0,0 +1,522 @@
+from typing import List, Optional, Tuple, Union
+import torch
+import torch.nn as nn
+
+from transformers.models.auto import AutoModel, AutoModelForCausalLM
+
+from transformers.modeling_outputs import CausalLMOutput, BaseModelOutputWithPast
+from transformers import modeling_utils
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import logging
+from transformers.generation import GenerationMixin
+
+from .modular_vibevoice_tokenizer import (
+    VibeVoiceTokenizerStreamingCache, 
+    VibeVoiceTokenizerEncoderOutput
+)
+
+from .configuration_vibevoice import VibeVoiceASRConfig
+from .modeling_vibevoice import (
+    VibeVoiceCausalLMOutputWithPast,    
+    SpeechConnector
+)
+
+logger = logging.get_logger(__name__)
+
+if not hasattr(modeling_utils, "ALL_PARALLEL_STYLES") or modeling_utils.ALL_PARALLEL_STYLES is None:
+    modeling_utils.ALL_PARALLEL_STYLES = ["tp", "none", "colwise", "rowwise"]
+
+# @auto_docstring
+class VibeVoiceASRPreTrainedModel(PreTrainedModel):
+    config_class = VibeVoiceASRConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _skip_keys_device_placement = "past_key_values"
+    _supports_cache_class = True
+    _supports_flash_attn = True
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+    _supports_quantized_cache = True
+    _supports_static_cache = True
+    _supports_attention_backend = True
+
+    def _init_weights(self, module):
+
+        # Use the language model's initializer_range if available
+        if hasattr(self.config, 'language_model_config') and hasattr(self.config.language_model_config, 'initializer_range'):
+            std = self.config.language_model_config.initializer_range
+        elif hasattr(self.config, 'decoder_config') and hasattr(self.config.decoder_config, 'initializer_range'):
+            std = self.config.decoder_config.initializer_range
+        else:
+            std = 0.02  # Default value
+            
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.LayerNorm):
+            module.weight.data.fill_(1.0)
+            module.bias.data.zero_()
+
+# @auto_docstring
+class VibeVoiceASRModel(VibeVoiceASRPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        
+        if hasattr(config, 'torch_dtype') and config.torch_dtype is not None:
+            if isinstance(config.torch_dtype, str):
+                dtype = getattr(torch, config.torch_dtype)
+            else:
+                dtype = config.torch_dtype
+        else:
+            dtype = torch.float32
+        
+        # Initialize Qwen2 model for language modeling
+        lm_config = config.decoder_config 
+        self.language_model = AutoModel.from_config(lm_config)
+        
+        # Initialize speech components if needed
+        self.acoustic_tokenizer = AutoModel.from_config(config.acoustic_tokenizer_config).to(dtype)
+        self.semantic_tokenizer = AutoModel.from_config(config.semantic_tokenizer_config).to(dtype)
+
+        self.acoustic_connector = SpeechConnector(config.acoustic_vae_dim, lm_config.hidden_size).to(dtype)
+        self.semantic_connector = SpeechConnector(config.semantic_vae_dim, lm_config.hidden_size).to(dtype)
+    
+    def get_input_embeddings(self):
+        if hasattr(self.language_model, 'embed_tokens'):
+            # If the language model has an embed_tokens attribute, return it
+            return self.language_model.embed_tokens
+        
+        for name, attr in self.language_model.fullmap.items(): # parallel by nnscaler, the name is changed
+            if attr.orig_name == 'embed_tokens.weight':
+                return getattr(self.language_model, name)
+        assert False, 'should not arrive here'
+
+    def set_input_embeddings(self, value):
+        self.language_model.embed_tokens = value
+    
+    def set_speech_tokenizers(self, acoustic_tokenizer=None, semantic_tokenizer=None):
+        """Set the speech tokenizers used for encoding and decoding speech."""
+        self.acoustic_tokenizer = acoustic_tokenizer
+        self.semantic_tokenizer = semantic_tokenizer
+        
+        # Reset the encoder to evaluation mode
+        if self.acoustic_tokenizer is not None:
+            self.acoustic_tokenizer.eval()
+            
+        if self.semantic_tokenizer is not None:
+            self.semantic_tokenizer.eval()
+    
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs,
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+        
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        
+        # Forward through language model
+        outputs = self.language_model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            cache_position=cache_position,
+            **kwargs,
+        )
+        
+        if not return_dict:
+            return outputs
+            
+        return BaseModelOutputWithPast(
+            last_hidden_state=outputs.last_hidden_state,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+class VibeVoiceASRForConditionalGeneration(VibeVoiceASRPreTrainedModel, GenerationMixin):
+    """
+    VibeVoice model for Automatic Speech Recognition (ASR) with language modeling head for conditional generation.
+    This class is designed for inference and generation tasks.
+    """
+    _tied_weights_keys = ["lm_head.weight"]
+    _tp_plan = {"lm_head": "colwise_rep"}
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = VibeVoiceASRModel(config)
+        self.vocab_size = config.decoder_config.vocab_size
+        
+        # Determine the dtype to use
+        if hasattr(config, 'torch_dtype') and config.torch_dtype is not None:
+            if isinstance(config.torch_dtype, str):
+                dtype = getattr(torch, config.torch_dtype)
+            else:
+                dtype = config.torch_dtype
+        else:
+            dtype = torch.float32
+        
+        # Initialize lm_head with the correct dtype
+        self.lm_head = nn.Linear(config.decoder_config.hidden_size, self.vocab_size, bias=False).to(dtype)
+        
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.get_input_embeddings()
+
+    def set_input_embeddings(self, value):
+        self.model.set_input_embeddings(value)
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def set_decoder(self, decoder):
+        self.model.language_model = decoder
+
+    def get_decoder(self):
+        return self.model.language_model
+
+    def tie_weights(self):
+        """Tie the weights between the input embeddings and the output embeddings."""
+        if getattr(self.config.decoder_config, 'tie_word_embeddings', False):
+            output_embeddings = self.get_output_embeddings()
+            input_embeddings = self.get_input_embeddings()
+            if hasattr(input_embeddings, 'weight'):
+                output_embeddings.weight = input_embeddings.weight
+            else:
+                output_embeddings.weight = input_embeddings
+
+    def encode_speech(
+        self,
+        speech_tensors: torch.FloatTensor,
+        speech_masks: Optional[torch.BoolTensor] = None,
+        speech_semantic_tensors: Optional[torch.FloatTensor] = None,
+        streaming_segment_duration: float = 60.0,  # seconds
+    ):
+        """
+        Encode speech input into features that can be used by the language model.
+        This method is called once before generation to process the speech input.
+        
+        For long audio (>600s by default), uses streaming processing to avoid conv overflow (>2^32).
+        Segments are processed independently, then concatenated before final sampling.
+        
+        Args:
+            speech_tensors: Input audio tensor [batch_size, samples]
+            speech_masks: Optional mask for speech features
+            speech_semantic_tensors: Optional pre-computed semantic tokens
+            streaming_segment_duration: Segment duration in seconds for streaming processing (default: 60s)
+        """
+        if hasattr(self.config, 'torch_dtype') and self.config.torch_dtype is not None:
+            if isinstance(self.config.torch_dtype, str):
+                dtype = getattr(torch, self.config.torch_dtype)
+            else:
+                dtype = self.config.torch_dtype
+        else:
+            dtype = torch.float32
+            
+        speech_tensors = speech_tensors.to(dtype)
+        
+        # Ensure proper shape: (batch, samples)
+        if speech_tensors.ndim == 1:
+            speech_tensors = speech_tensors.unsqueeze(0)
+        
+        batch_size, total_samples = speech_tensors.shape
+        sample_rate = 24000  # fix 24kHz sample rate
+        
+        # Calculate segment size in samples
+        segment_samples = int(streaming_segment_duration * sample_rate)
+        
+        # Decide whether to use streaming based on audio length
+        use_streaming = total_samples > segment_samples
+        
+        with torch.no_grad():
+            if not use_streaming:
+                # Short audio: direct processing (original behavior)
+                encoder_output = self.model.acoustic_tokenizer.encode(speech_tensors.unsqueeze(1))
+                audio_tokens = encoder_output.sample(dist_type=self.model.acoustic_tokenizer.std_dist_type)[0]
+                acoustic_features = self.model.acoustic_connector(audio_tokens)
+                
+                # Encode semantic features
+                if speech_semantic_tensors is not None:
+                    semantic_features = self.model.semantic_connector(speech_semantic_tensors)
+                else:
+                    semantic_tokens = self.model.semantic_tokenizer.encode(speech_tensors.unsqueeze(1)).mean
+                    semantic_features = self.model.semantic_connector(semantic_tokens)
+            else:
+                # Long audio: streaming processing
+                # print(f"Using streaming processing for long audio: {total_samples/sample_rate:.1f}s "
+                #       f"(segment size: {streaming_segment_duration}s)")
+                
+                # Initialize caches for both tokenizers
+                acoustic_encoder_cache = VibeVoiceTokenizerStreamingCache()
+                semantic_encoder_cache = VibeVoiceTokenizerStreamingCache()
+                acoustic_mean_segments = []
+                semantic_mean_segments = []
+                sample_indices = torch.arange(batch_size, device=speech_tensors.device)
+                
+                # Helper function from batch_asr_sft_cache.py
+                def _iter_segments(total_length: int, segment_length: int):
+                    """Iterate over audio segments with a given segment length."""
+                    if segment_length <= 0:
+                        raise ValueError("segment_length must be positive")
+                    for start in range(0, total_length, segment_length):
+                        end = min(start + segment_length, total_length)
+                        if end > start:
+                            yield start, end
+                
+                # Process each segment for both acoustic and semantic tokenizers
+                segments = list(_iter_segments(total_samples, segment_samples))
+                num_segments = len(segments)
+                for seg_idx, (start, end) in enumerate(segments):
+                    chunk = speech_tensors[:, start:end].contiguous()
+                    if chunk.numel() == 0:
+                        continue
+                    
+                    # Check if this is the final segment
+                    is_final = (seg_idx == num_segments - 1)
+                    
+                    # Encode chunk for acoustic tokenizer (don't sample yet)
+                    acoustic_encoder_output = self.model.acoustic_tokenizer.encode(
+                        chunk.unsqueeze(1),
+                        cache=acoustic_encoder_cache,
+                        sample_indices=sample_indices,
+                        use_cache=True,
+                        is_final_chunk=is_final,
+                    )
+                    acoustic_mean_segments.append(acoustic_encoder_output.mean)
+                    
+                    # Encode chunk for semantic tokenizer (take mean directly)
+                    semantic_encoder_output = self.model.semantic_tokenizer.encode(
+                        chunk.unsqueeze(1),
+                        cache=semantic_encoder_cache,
+                        sample_indices=sample_indices,
+                        use_cache=True,
+                        is_final_chunk=is_final,
+                    )
+                    semantic_mean_segments.append(semantic_encoder_output.mean)
+                
+                # print(f"Processed {len(acoustic_mean_segments)} segments.")
+                # Concatenate all acoustic means and sample once
+                acoustic_mean_full = torch.cat(acoustic_mean_segments, dim=1).contiguous()
+                acoustic_encoder_output = VibeVoiceTokenizerEncoderOutput(
+                    mean=acoustic_mean_full,
+                    std=self.model.acoustic_tokenizer.fix_std
+                )
+                audio_tokens = acoustic_encoder_output.sample(
+                    dist_type=self.model.acoustic_tokenizer.std_dist_type
+                )[0]
+                acoustic_features = self.model.acoustic_connector(audio_tokens)
+                
+                # Concatenate all semantic means
+                semantic_tokens = torch.cat(semantic_mean_segments, dim=1).contiguous()
+                semantic_features = self.model.semantic_connector(semantic_tokens)
+            
+            # Combine acoustic and semantic features
+            if speech_masks is not None:
+                combined_features = acoustic_features[speech_masks] + semantic_features[speech_masks]
+            else:
+                combined_features = acoustic_features + semantic_features
+                
+        return combined_features
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        # Speech-specific arguments
+        speech_tensors: Optional[torch.FloatTensor] = None,
+        speech_masks: Optional[torch.BoolTensor] = None,
+        speech_semantic_tensors: Optional[torch.FloatTensor] = None,
+        acoustic_input_mask: Optional[torch.BoolTensor] = None,
+        **kwargs,
+    ) -> Union[Tuple, CausalLMOutput]:
+        """
+        Forward pass for the model. Handles both training and generation scenarios.
+        """
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        use_cache = use_cache if use_cache is not None else getattr(self.config, 'use_cache', False)
+
+        # Process inputs
+        if inputs_embeds is None and input_ids is not None:
+            inputs_embeds = self.get_input_embeddings()(input_ids)
+            
+        # If we have speech input and acoustic_input_mask, encode and insert speech features
+        if speech_tensors is not None and acoustic_input_mask is not None:
+            speech_features = self.encode_speech(
+                speech_tensors=speech_tensors,
+                speech_masks=speech_masks,
+                speech_semantic_tensors=speech_semantic_tensors,
+            )
+            # Clone to avoid in-place operation on leaf variable during training
+            inputs_embeds = inputs_embeds.clone()
+            inputs_embeds[acoustic_input_mask] = speech_features
+
+        # Forward through the model
+        outputs = self.model(
+            input_ids=None,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            cache_position=cache_position,
+        )
+
+        hidden_states = outputs[0] if not return_dict else outputs.last_hidden_state
+        logits = self.lm_head(hidden_states)
+
+        loss = None
+        if labels is not None:
+            # Shift so that tokens < n predict n
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            # Flatten the tokens
+            loss_fct = nn.CrossEntropyLoss(ignore_index=-100)
+            shift_logits = shift_logits.view(-1, self.vocab_size)
+            shift_labels = shift_labels.view(-1)
+            # Enable model parallelism
+            shift_labels = shift_labels.to(shift_logits.device)
+            loss = loss_fct(shift_logits, shift_labels)
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+
+        return VibeVoiceCausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids,
+        past_key_values=None,
+        attention_mask=None,
+        inputs_embeds=None,
+        cache_position=None,
+        position_ids=None,
+        use_cache=True,
+        speech_tensors=None,
+        speech_masks=None,
+        speech_semantic_tensors=None,
+        acoustic_input_mask=None,
+        **kwargs,
+    ):
+        """
+        Prepare inputs for generation step. This method is called by generate() 
+        for each token generation step.
+        
+        Following Qwen2-VL's approach: speech inputs are only forwarded on the first pass
+        (when cache_position[0] == 0), and are excluded in subsequent generation steps.
+        """
+        # If we have past key values, we only need to process the new tokens
+        if past_key_values is not None:
+            if isinstance(past_key_values, tuple):
+                past_length = past_key_values[0][0].shape[2]
+            else:
+                past_length = past_key_values.get_seq_length()
+            
+            # Keep only the new tokens
+            if input_ids is not None and input_ids.shape[1] > past_length:
+                input_ids = input_ids[:, past_length:]
+
+        # Prepare position ids
+        if position_ids is None and attention_mask is not None:
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            if past_key_values is not None and input_ids is not None:
+                position_ids = position_ids[:, -input_ids.shape[1]:]
+
+        # Prepare cache position
+        if cache_position is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(
+                past_seen_tokens, 
+                past_seen_tokens + (input_ids.shape[1] if input_ids is not None else inputs_embeds.shape[1]), 
+                device=input_ids.device if input_ids is not None else inputs_embeds.device
+            )
+
+        # Prepare model inputs
+        if inputs_embeds is not None and past_key_values is None:
+            model_inputs = {"inputs_embeds": inputs_embeds}
+        else:
+            model_inputs = {"input_ids": input_ids}
+
+        model_inputs.update(
+            {
+                "position_ids": position_ids,
+                "cache_position": cache_position,
+                "past_key_values": past_key_values,
+                "use_cache": use_cache,
+                "attention_mask": attention_mask,
+            }
+        )
+        
+        # Following Qwen2-VL pattern: only include speech inputs on the first forward pass
+        # (when cache_position[0] == 0), exclude them in subsequent generation steps
+        if cache_position is not None and len(cache_position) > 0 and cache_position[0] == 0:
+            # First forward pass - include speech inputs if provided
+            model_inputs.update({
+                "speech_tensors": speech_tensors,
+                "speech_masks": speech_masks,
+                "speech_semantic_tensors": speech_semantic_tensors,
+                "acoustic_input_mask": acoustic_input_mask,
+            })
+        else:
+            # Subsequent generation steps - exclude speech inputs
+            model_inputs.update({
+                "speech_tensors": None,
+                "speech_masks": None, 
+                "speech_semantic_tensors": None,
+                "acoustic_input_mask": None,
+            })
+        
+        # Include any remaining kwargs that might be needed
+        model_inputs.update(kwargs)
+        
+        return model_inputs
+
+AutoModel.register(VibeVoiceASRConfig, VibeVoiceASRModel)
+AutoModelForCausalLM.register(VibeVoiceASRConfig, VibeVoiceASRForConditionalGeneration)
+
+__all__ = [
+    "VibeVoiceASRPreTrainedModel",
+    "VibeVoiceASRModel",
+    "VibeVoiceASRForConditionalGeneration",
+]

vibevoice/modular/modeling_vibevoice_streaming.py

@@ -0,0 +1,190 @@
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Tuple, Union, Callable
+from tqdm import tqdm
+import copy
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.distributed as dist
+
+from transformers.models.auto import AutoModel, AutoModelForCausalLM
+
+from transformers.activations import ACT2FN
+from transformers.modeling_outputs import CausalLMOutput, BaseModelOutputWithPast, ModelOutput
+from transformers.models.llama.modeling_llama import LlamaRMSNorm
+from transformers import modeling_utils
+from transformers.modeling_utils import PreTrainedModel
+from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
+from transformers.utils import logging
+
+from .modular_vibevoice_diffusion_head import VibeVoiceDiffusionHead
+from vibevoice.schedule.dpm_solver import DPMSolverMultistepScheduler
+
+from .configuration_vibevoice_streaming import VibeVoiceStreamingConfig
+
+
+logger = logging.get_logger(__name__)
+
+if not hasattr(modeling_utils, "ALL_PARALLEL_STYLES") or modeling_utils.ALL_PARALLEL_STYLES is None:
+    modeling_utils.ALL_PARALLEL_STYLES = ["tp", "none", "colwise", "rowwise"]
+
+
+class BinaryClassifier(nn.Module):
+    def __init__(self, hidden_size):
+        super(BinaryClassifier, self).__init__()
+        self.fc1 = nn.Linear(hidden_size, hidden_size)
+        self.fc2 = nn.Linear(hidden_size, 1)
+
+    def forward(self, x):
+        x = torch.relu(self.fc1(x))
+        x = self.fc2(x)
+        return x
+
+
+class SpeechConnector(nn.Module):
+    def __init__(self, input_dim, output_dim):
+        super().__init__()
+        self.fc1 = nn.Linear(input_dim, output_dim)
+        self.norm = LlamaRMSNorm(output_dim, eps=1e-6)
+        self.fc2 = nn.Linear(output_dim, output_dim)
+
+    def forward(self, features, **kwargs):    
+        x = self.fc1(features)
+        x = self.norm(x)
+        x = self.fc2(x)
+        return x
+
+
+# @auto_docstring
+class VibeVoiceStreamingPreTrainedModel(PreTrainedModel):
+    config_class = VibeVoiceStreamingConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _skip_keys_device_placement = "past_key_values"
+    _supports_cache_class = True
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+    _supports_quantized_cache = True
+    _supports_static_cache = True
+    _supports_attention_backend = True
+
+    def _init_weights(self, module):
+        if isinstance(module, VibeVoiceDiffusionHead):
+            module.initialize_weights()
+            return
+
+        # Use the language model's initializer_range if available
+        if hasattr(self.config, 'language_model_config') and hasattr(self.config.language_model_config, 'initializer_range'):
+            std = self.config.language_model_config.initializer_range
+        elif hasattr(self.config, 'decoder_config') and hasattr(self.config.decoder_config, 'initializer_range'):
+            std = self.config.decoder_config.initializer_range
+        else:
+            std = 0.02  # Default value
+            
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.LayerNorm):
+            module.weight.data.fill_(1.0)
+            module.bias.data.zero_()
+
+
+# @auto_docstring
+class VibeVoiceStreamingModel(VibeVoiceStreamingPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        
+        if hasattr(config, 'torch_dtype') and config.torch_dtype is not None:
+            if isinstance(config.torch_dtype, str):
+                dtype = getattr(torch, config.torch_dtype)
+            else:
+                dtype = config.torch_dtype
+        else:
+            dtype = torch.float32
+        
+        # Initialize Qwen2 model for language modeling. 
+        # The lower Transformer layers are only used for encoding text, while the upper Transformer layers are used for encoding text and generating speech. 
+        # To keep the code clean, we constructs two language models.
+        # The final norm layer of the first language_model is set to identity and will not be used in inference.
+        lm_config = copy.deepcopy(config.decoder_config)
+        lm_backbone_num_hidden_layers = getattr(lm_config, 'num_hidden_layers', 24) - config.tts_backbone_num_hidden_layers
+        lm_config.num_hidden_layers = lm_backbone_num_hidden_layers
+        self.language_model = AutoModel.from_config(lm_config)
+        self.language_model.norm = nn.Identity()
+        
+        # We only need the Transformer layers here. Note that embed_tokens in tts_language_model is unused
+        tts_lm_config = copy.deepcopy(lm_config)
+        tts_lm_config.num_hidden_layers = config.tts_backbone_num_hidden_layers
+        self.tts_language_model = AutoModel.from_config(tts_lm_config)
+
+        # Marks the text that needs to be spoken by the TTS model.
+        self.tts_input_types = nn.Embedding(num_embeddings=2, embedding_dim=config.decoder_config.hidden_size)
+        
+        # Initialize speech components if needed
+        self.acoustic_tokenizer = AutoModel.from_config(config.acoustic_tokenizer_config).to(dtype)
+        self.acoustic_connector = SpeechConnector(config.acoustic_vae_dim, lm_config.hidden_size).to(dtype)
+        
+        # Register scaling factors as buffers - use 1D tensors for FSDP compatibility
+        self.register_buffer('speech_scaling_factor', torch.tensor(float('nan')))  
+        self.register_buffer('speech_bias_factor', torch.tensor(float('nan')))
+
+        # Initialize prediction head for speech generation
+        self.prediction_head = AutoModel.from_config(config.diffusion_head_config).to(dtype)
+
+        # Initialize noise scheduler
+        self.noise_scheduler = DPMSolverMultistepScheduler(
+            num_train_timesteps=config.diffusion_head_config.ddpm_num_steps,
+            beta_schedule=config.diffusion_head_config.ddpm_beta_schedule,
+            prediction_type=config.diffusion_head_config.prediction_type
+        )
+    
+    def get_input_embeddings(self):
+        if hasattr(self.language_model, 'embed_tokens'):
+            # If the language model has an embed_tokens attribute, return it
+            return self.language_model.embed_tokens
+        
+        for name, attr in self.language_model.fullmap.items(): # parallel by nnscaler, the name is changed
+            if attr.orig_name == 'embed_tokens.weight':
+                return getattr(self.language_model, name)
+        assert False, 'should not arrive here'
+
+    def set_input_embeddings(self, value):
+        self.language_model.embed_tokens = value
+    
+    def set_speech_tokenizers(self, acoustic_tokenizer=None):
+        """Set the speech tokenizers used for encoding and decoding speech."""
+        self.acoustic_tokenizer = acoustic_tokenizer
+        
+        # Reset the encoder to evaluation mode
+        if self.acoustic_tokenizer is not None:
+            self.acoustic_tokenizer.eval()
+    
+    def forward(self, *args, **kwargs):
+        """
+        Intentionally not implemented.
+
+        This streaming model is split into two explicit submodules:
+          - `language_model`      for plain text processing (lower layers).
+          - `tts_language_model`  for TTS-related upper layers.
+
+        We deliberately avoid a unified `forward` to prevent accidental calls
+        that mix responsibilities.
+
+        To use the model:
+          - Call `self.language_model(...)` for text embeddings / hidden states.
+          - Call `self.tts_language_model(...)` for the TTS portion.
+          - Use the dedicated inference class for combined generation logic.
+        """
+        raise RuntimeError(
+            "VibeVoiceStreamingModel.forward is intentionally disabled. "
+            "Use `model.language_model(...)` or `model.tts_language_model(...)` instead."
+        )
+
+
+AutoModel.register(VibeVoiceStreamingConfig, VibeVoiceStreamingModel)
+
+__all__ = [
+    "VibeVoiceStreamingPreTrainedModel",
+    "VibeVoiceStreamingModel",
+]

vibevoice/modular/modeling_vibevoice_streaming_inference.py

@@ -0,0 +1,906 @@
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple, Union, Callable
+from tqdm import tqdm
+import inspect
+import torch
+import torch.nn as nn
+
+from transformers.models.auto import AutoModel, AutoModelForCausalLM
+from transformers.generation import GenerationMixin, GenerationConfig, LogitsProcessor, LogitsProcessorList, StoppingCriteriaList
+from transformers.modeling_outputs import BaseModelOutputWithPast, ModelOutput
+from transformers import modeling_utils
+from transformers.modeling_utils import PreTrainedModel
+from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
+from transformers.utils import logging
+
+from .modular_vibevoice_tokenizer import VibeVoiceTokenizerStreamingCache
+from .modular_vibevoice_diffusion_head import VibeVoiceDiffusionHead
+from vibevoice.schedule.dpm_solver import DPMSolverMultistepScheduler
+from .configuration_vibevoice_streaming import VibeVoiceStreamingConfig
+from .modular_vibevoice_text_tokenizer import VibeVoiceTextTokenizer, VibeVoiceTextTokenizerFast
+from .modeling_vibevoice_streaming import VibeVoiceStreamingPreTrainedModel, VibeVoiceStreamingModel, BinaryClassifier
+from .streamer import AudioStreamer, AsyncAudioStreamer
+
+logger = logging.get_logger(__name__)
+
+if not hasattr(modeling_utils, "ALL_PARALLEL_STYLES") or modeling_utils.ALL_PARALLEL_STYLES is None:
+    modeling_utils.ALL_PARALLEL_STYLES = ["tp", "none", "colwise", "rowwise"]
+
+TTS_TEXT_WINDOW_SIZE = 5
+TTS_SPEECH_WINDOW_SIZE = 6
+
+
+# ============================================================================
+# Transformers >= 4.57 Compatibility Layer
+# The cache system was refactored in transformers 4.57, requiring these helpers.
+# ============================================================================
+
+class MockCacheLayer:
+    """
+    Mock cache layer for transformers >= 4.57 compatibility.
+    Provides the `layers` interface expected by DynamicCache in newer versions.
+    """
+    
+    def __init__(self, key_cache, value_cache, parent_cache=None, layer_idx=0):
+        self.key_cache = key_cache
+        self.value_cache = value_cache
+        self._parent_cache = parent_cache
+        self._layer_idx = layer_idx
+    
+    def get_mask_sizes(self, cache_position):
+        """Return KV length and offset for mask creation."""
+        seq_length = self.key_cache.shape[2] if self.key_cache is not None else 0
+        query_length = cache_position.shape[0]
+        return seq_length + query_length, 0
+    
+    def update(self, key_states, value_states, cache_kwargs=None):
+        """Update the cache with new key/value states."""
+        if self._parent_cache is None:
+            return self.key_cache, self.value_cache
+        
+        parent = self._parent_cache
+        idx = self._layer_idx
+        
+        # Extend cache lists if needed
+        while len(parent.key_cache) <= idx:
+            parent.key_cache.append(None)
+            parent.value_cache.append(None)
+        
+        # Concatenate or initialize cache
+        if parent.key_cache[idx] is not None:
+            parent.key_cache[idx] = torch.cat([parent.key_cache[idx], key_states], dim=2)
+            parent.value_cache[idx] = torch.cat([parent.value_cache[idx], value_states], dim=2)
+        else:
+            parent.key_cache[idx] = key_states
+            parent.value_cache[idx] = value_states
+        
+        # Update local references
+        self.key_cache = parent.key_cache[idx]
+        self.value_cache = parent.value_cache[idx]
+        return self.key_cache, self.value_cache
+
+
+def _ensure_cache_has_layers(cache):
+    """
+    Ensure the cache has all required attributes for transformers >= 4.57.
+    Creates MockCacheLayer wrappers to provide the expected `layers` interface.
+    """
+    if cache is None:
+        return cache
+    
+    # Add required attributes (skip if read-only)
+    for attr, default in [('layer_class_to_replicate', None), ('offloading', False), ('is_compileable', False)]:
+        if not hasattr(cache, attr):
+            try:
+                setattr(cache, attr, default)
+            except AttributeError:
+                pass
+    
+    # Build layers list from key_cache/value_cache
+    if hasattr(cache, 'key_cache') and hasattr(cache, 'value_cache'):
+        try:
+            cache.layers = [
+                MockCacheLayer(cache.key_cache[i], cache.value_cache[i], parent_cache=cache, layer_idx=i)
+                for i in range(len(cache.key_cache))
+            ]
+        except AttributeError:
+            pass
+    elif not hasattr(cache, 'layers'):
+        try:
+            cache.layers = []
+        except AttributeError:
+            pass
+    
+    return cache
+
+
+def _update_model_kwargs_for_generation(
+    outputs: ModelOutput,
+    model_kwargs: Dict[str, Any],
+    num_new_tokens: int = 1,
+) -> Dict[str, Any]:
+    """
+    Update model_kwargs after adding new tokens (supports multi-token windows).
+    
+    Updates past_key_values, attention_mask, and cache_position for the next forward pass.
+    """
+    model_kwargs["past_key_values"] = _ensure_cache_has_layers(outputs.past_key_values)
+    
+    attention_mask = model_kwargs["attention_mask"]
+    model_kwargs["attention_mask"] = torch.cat(
+        [attention_mask, attention_mask.new_ones((attention_mask.shape[0], num_new_tokens))], dim=-1
+    )
+    
+    cache_pos = model_kwargs["cache_position"]
+    model_kwargs["cache_position"] = torch.arange(
+        cache_pos[-1] + 1, cache_pos[-1] + num_new_tokens + 1, device=cache_pos.device
+    )
+    
+    return model_kwargs
+
+
+@dataclass
+class VibeVoiceCausalLMOutputWithPast(BaseModelOutputWithPast):
+    logits: Optional[torch.FloatTensor] = None
+
+
+@dataclass
+class VibeVoiceGenerationOutput(ModelOutput):
+    """
+    Output type for VibeVoice generation.
+    
+    Args:
+        sequences (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+            The generated sequences. 
+        speech_outputs (`List[torch.FloatTensor]`, *optional*):
+            List of generated speech waveforms or latents for each speech segment.
+    """
+    sequences: torch.LongTensor = None
+    speech_outputs: Optional[List[torch.FloatTensor]] = None
+    reach_max_step_sample: Optional[torch.BoolTensor] = None
+
+
+class VibeVoiceStreamingForConditionalGenerationInference(VibeVoiceStreamingPreTrainedModel, GenerationMixin):
+
+    def __init__(self, config):
+        super().__init__(config)
+        
+        # Initialize the base model
+        self.model = VibeVoiceStreamingModel(config)
+
+        # TTS generation EOS classifier
+        self.tts_eos_classifier = BinaryClassifier(config.decoder_config.hidden_size)
+        
+        # inference configuration
+        self.ddpm_inference_steps = config.diffusion_head_config.ddpm_num_inference_steps
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @property
+    def noise_scheduler(self):
+        return self.model.noise_scheduler
+
+    @property
+    def prediction_head(self):
+        return self.model.prediction_head
+    
+    @property
+    def speech_scaling_factor(self):
+        return self.model.speech_scaling_factor
+
+    @property
+    def speech_bias_factor(self):
+        return self.model.speech_bias_factor
+
+    @property
+    def acoustic_tokenizer(self):
+        return self.model.acoustic_tokenizer
+    
+    @property
+    def acoustic_connector(self):
+        return self.model.acoustic_connector
+        
+    def tie_weights(self):
+        """
+        Tie the weights between the input embeddings and the output embeddings.
+        """
+        # Tie lm_head.weight to language_model.embed_tokens.weight
+        if not getattr(self.config, 'tie_word_embeddings', False):
+            return
+         
+        if hasattr(self, 'lm_head') and hasattr(self.model.language_model, 'embed_tokens'):
+            self.lm_head.weight = self.model.language_model.embed_tokens.weight
+        
+    def get_input_embeddings(self):
+        return self.model.get_input_embeddings()
+    
+    def set_input_embeddings(self, value):
+        self.model.set_input_embeddings(value)
+    
+    def get_output_embeddings(self):
+        """
+        This model does not define an `lm_head` (vocabulary projection).
+        """
+        return None
+    
+    def set_output_embeddings(self, new_embeddings):
+        """
+        No-op because there is no `lm_head`. Provided only to satisfy optional API calls.
+        To enable, first create `self.lm_head` then allow assignment.
+        """
+        raise RuntimeError("Output embeddings (lm_head) are not defined for this model. "
+                           "Create one before calling set_output_embeddings if needed.")
+    
+    def set_speech_tokenizers(self, acoustic_tokenizer=None):
+        """Set the speech tokenizers used for encoding and decoding speech."""
+        self.model.set_speech_tokenizers(acoustic_tokenizer)
+    
+    def set_ddpm_inference_steps(self, num_steps=None):
+        self.ddpm_inference_steps = num_steps or self.config.diffusion_head_config.ddpm_num_inference_steps
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: torch.LongTensor,
+        past_key_values=None,
+        attention_mask=None,
+        inputs_embeds=None,
+        cache_position=None,
+        **kwargs,
+    ):
+        """Prepare model inputs for generation (transformers >= 4.57 compatible)."""
+        model_inputs = {"cache_position": cache_position}
+
+        # Slice inputs when using cache
+        if past_key_values is not None:
+            model_inputs["past_key_values"] = past_key_values
+            if inputs_embeds is not None and input_ids.shape[1] == 0:
+                inputs_embeds = inputs_embeds[:, -cache_position.shape[0]:]
+            elif inputs_embeds is not None or (cache_position is not None and cache_position[-1] >= input_ids.shape[1]):
+                input_ids = input_ids[:, -cache_position.shape[0]:]
+            elif cache_position is not None and input_ids.shape[1] != cache_position.shape[0]:
+                input_ids = input_ids[:, cache_position]
+
+        # Set input_ids or inputs_embeds
+        use_embeds = inputs_embeds is not None and (
+            past_key_values is None or (cache_position is not None and len(cache_position) == inputs_embeds.shape[1])
+        )
+        if use_embeds:
+            model_inputs["input_ids"] = None
+            model_inputs["inputs_embeds"] = inputs_embeds
+        else:
+            model_inputs["input_ids"] = input_ids.clone(memory_format=torch.contiguous_format) if input_ids is not None else None
+            model_inputs["inputs_embeds"] = None
+
+        if attention_mask is not None:
+            model_inputs["attention_mask"] = attention_mask
+
+        # Create position_ids from attention_mask
+        if attention_mask is not None and kwargs.get("position_ids") is None:
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            kwargs["position_ids"] = position_ids
+
+        # Slice position_ids when using cache
+        if kwargs.get("position_ids") is not None:
+            if past_key_values is not None:
+                seq_len = model_inputs["inputs_embeds"].shape[1] if model_inputs.get("inputs_embeds") is not None else model_inputs["input_ids"].shape[1]
+                model_inputs["position_ids"] = kwargs["position_ids"][:, -seq_len:].clone(memory_format=torch.contiguous_format)
+            else:
+                model_inputs["position_ids"] = kwargs.pop("position_ids").clone(memory_format=torch.contiguous_format)
+
+        # Forward remaining kwargs
+        for key, value in kwargs.items():
+            if key not in model_inputs:
+                model_inputs[key] = value
+
+        model_inputs.pop("labels", None)
+        return model_inputs
+
+    def _update_model_kwargs_for_generation(
+        self,
+        outputs,
+        model_kwargs,
+        is_encoder_decoder=False,
+        num_new_tokens=1,
+    ):
+        """Override to ensure cache compatibility with transformers >= 4.57."""
+        model_kwargs = super()._update_model_kwargs_for_generation(
+            outputs, model_kwargs, is_encoder_decoder=is_encoder_decoder, num_new_tokens=num_new_tokens
+        )
+        if "past_key_values" in model_kwargs:
+            model_kwargs["past_key_values"] = _ensure_cache_has_layers(model_kwargs["past_key_values"])
+        return model_kwargs
+
+    def _init_cache_for_generation(self, generation_config, model_kwargs, batch_size, max_cache_length, device):
+        """
+        Initialize cache for generation, handling different transformers versions.
+        For transformers >= 4.57, returns None to let the model create the cache dynamically.
+        """
+        try:
+            from transformers.cache_utils import DynamicCache
+            sig = inspect.signature(DynamicCache.__init__)
+            if 'config' in sig.parameters:
+                # transformers >= 4.57: let model handle cache creation
+                return None
+            else:
+                # Older versions: use parent method
+                prep_sig = inspect.signature(self._prepare_cache_for_generation)
+                if 'device' in prep_sig.parameters:
+                    self._prepare_cache_for_generation(generation_config, model_kwargs, None, batch_size, max_cache_length, device)
+                else:
+                    self._prepare_cache_for_generation(generation_config, model_kwargs, None, batch_size, max_cache_length)
+                return model_kwargs.get("past_key_values")
+        except Exception:
+            return None
+
+    # @can_return_tuple
+    def forward_lm(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs,
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+        """
+        Single pass of the base text LM.
+
+        - Builds embeddings if `inputs_embeds` not provided.
+        - Uses (and returns) `past_key_values` when `use_cache=True`.
+        - No loss / no lm_head / no speech logic.
+
+        Args:
+            input_ids: (B, S) token ids.
+            attention_mask: (B, S) mask.
+            past_key_values: cache from previous steps.
+            cache_position: positions for cached tokens.
+            labels: unsupported (will raise).
+
+        Returns:
+            BaseModelOutputWithPast with `last_hidden_state` and `past_key_values`.
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        
+        # Get embeddings
+        if inputs_embeds is None:
+            inputs_embeds = self.model.get_input_embeddings()(input_ids)
+
+        outputs = self.model.language_model(
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            cache_position=cache_position,
+            **kwargs,
+        )
+
+        hidden_states = outputs[0] if not return_dict else outputs.last_hidden_state
+                
+        if labels is not None:
+            raise NotImplementedError("Loss computation is not implemented in this version.")
+
+        return BaseModelOutputWithPast(
+            past_key_values=outputs.past_key_values,
+            last_hidden_state=hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    # @can_return_tuple
+    def forward_tts_lm(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        lm_last_hidden_state: Optional[torch.FloatTensor] = None,
+        tts_text_masks: Optional[torch.BoolTensor] = None,
+        **kwargs,
+    ) -> Union[Tuple, VibeVoiceCausalLMOutputWithPast]:
+        """
+        Single pass of the TTS LM.
+
+        - Overwrites tail embeddings with `lm_last_hidden_state`.
+        - Adds type embedding via `tts_text_masks` (1=text, 0=speech).
+        - Predicts EOS from last hidden state (binary classifier).
+        - No loss / no full acoustic decoding here.
+
+        Args:
+            input_ids: (B, S) token ids.
+            attention_mask: (B, S) mask.
+            lm_last_hidden_state: (B, K, H) hidden states to splice into the tail.
+            tts_text_masks: (B, 1) mask marking current position as text(1)/speech(0).
+            past_key_values: cache from previous TTS steps.
+            cache_position: positions for cached tokens.
+            labels: unsupported (will raise).
+
+        Returns:
+            VibeVoiceCausalLMOutputWithPast with `logits` (EOS), `last_hidden_state`, `past_key_values`.
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        
+        # Get embeddings
+        if inputs_embeds is None:
+            # Will be replaced with lm_last_hidden_state
+            inputs_embeds = self.model.get_input_embeddings()(input_ids)
+        
+        # Replace the last part of inputs_embeds with lm_last_hidden_state
+        start_idx = inputs_embeds.shape[1] - lm_last_hidden_state.shape[1]
+        inputs_embeds[:, start_idx:, :] = lm_last_hidden_state
+        
+        # Adds type embedding via `tts_text_masks`.
+        inputs_embeds = inputs_embeds + self.model.tts_input_types(tts_text_masks.long())
+
+        outputs = self.model.tts_language_model(
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            cache_position=cache_position,
+            **kwargs,
+        )
+
+        hidden_states = outputs[0] if not return_dict else outputs.last_hidden_state
+        logits = self.tts_eos_classifier(hidden_states[:, -1, :])
+                
+        if labels is not None:
+            raise NotImplementedError("Loss computation is not implemented in this version.")
+
+        return VibeVoiceCausalLMOutputWithPast(
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            last_hidden_state=hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    def forward(self, *args, **kwargs):
+        """
+        Unified forward is intentionally disabled.
+
+        Reasons:
+          1. The inference pipeline is staged: base text LM, then TTS LM, plus streaming & diffusion handled in `generate`.
+          2. A monolithic call would hide required sequencing (prefill, window stepping, speech diffusion sampling).
+
+        Use instead:
+          - self.forward_lm(...)       for a base text LM step (prefill or incremental).
+          - self.forward_tts_lm(...)   for a single TTS LM step (needs LM hidden states).
+          - self.generate(...)         for full streaming (text + speech + diffusion + audio assembly).
+
+        Raises:
+            RuntimeError: Always (by design).
+        """
+        raise RuntimeError(
+            "Unified forward is disabled. Use `forward_lm`, `forward_tts_lm`, or `generate` instead."
+        )
+
+    def _build_generate_config_model_kwargs(self, generation_config, inputs, tokenizer, return_processors=False, **kwargs):
+        if generation_config is None:
+            generation_config = GenerationConfig(
+                bos_token_id=tokenizer.bos_token_id,
+                eos_token_id=tokenizer.eos_token_id,
+                pad_token_id = tokenizer.pad_token_id
+            )
+        else:
+            generation_config = GenerationConfig(
+                **generation_config,
+                bos_token_id=tokenizer.bos_token_id,
+                eos_token_id=tokenizer.eos_token_id,
+                pad_token_id = tokenizer.pad_token_id
+            )
+
+        generation_config, model_kwargs = self._prepare_generation_config(
+            generation_config, 
+            True, 
+            speech_start_id=tokenizer.speech_start_id, 
+            speech_end_id=tokenizer.speech_end_id, 
+            speech_diffusion_id=tokenizer.speech_diffusion_id, 
+            **kwargs
+        )
+        generation_config.speech_start_id = tokenizer.speech_start_id
+        generation_config.speech_end_id = tokenizer.speech_end_id
+        generation_config.speech_diffusion_id = tokenizer.speech_diffusion_id
+
+        inputs_tensor, model_input_name, model_kwargs = self._prepare_model_inputs(inputs, generation_config.bos_token_id, model_kwargs)
+        batch_size = inputs_tensor.shape[0]
+        device = self.device
+        
+        self._prepare_special_tokens(generation_config, True, device=device)
+        generation_config.use_cache = True
+        model_kwargs["use_cache"] = generation_config.use_cache
+        input_ids = inputs_tensor.to(self.device)
+
+        input_ids_length = input_ids.shape[1]
+        has_default_max_length = kwargs.get("max_length") is None and generation_config.max_length is not None
+        has_default_min_length = kwargs.get("min_length") is None and generation_config.min_length is not None
+        generation_config = self._prepare_generated_length(
+            generation_config=generation_config,
+            has_default_max_length=has_default_max_length,
+            has_default_min_length=has_default_min_length,
+            model_input_name=model_input_name,
+            inputs_tensor=inputs_tensor,
+            input_ids_length=input_ids_length,
+        )
+
+        max_cache_length = generation_config.max_length - 1
+        # Handle cache initialization for different transformers versions
+        model_kwargs["past_key_values"] = self._init_cache_for_generation(
+            generation_config, model_kwargs, batch_size, max_cache_length, device
+        )
+        model_kwargs['cache_position'] = torch.arange(input_ids_length, device=device, dtype=torch.long)
+        for k, v in model_kwargs.items():
+            if isinstance(v, torch.Tensor):
+                model_kwargs[k] = v.to(device=device)
+        
+        if return_processors:
+            logits_processor = self._get_logits_processor(
+                generation_config=generation_config,
+                input_ids_seq_length=input_ids_length,
+                encoder_input_ids=inputs_tensor,
+                prefix_allowed_tokens_fn=None,
+                logits_processor=LogitsProcessorList(),
+                device=inputs_tensor.device,
+                model_kwargs=model_kwargs,
+            )
+
+            stopping_criteria = self._get_stopping_criteria(generation_config=generation_config, stopping_criteria=StoppingCriteriaList())
+        
+            return generation_config, model_kwargs, input_ids, logits_processor, stopping_criteria
+        else:
+            return generation_config, model_kwargs, input_ids
+
+    @torch.no_grad()
+    def generate(
+        self,
+        inputs: Optional[torch.Tensor] = None,
+        generation_config: Optional[GenerationConfig] = None,
+        logits_processor: Optional[LogitsProcessorList] = None,
+        stopping_criteria: Optional[StoppingCriteriaList] = None,
+        prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor], List[int]]] = None,
+        synced_gpus: Optional[bool] = None,
+        assistant_model: Optional["PreTrainedModel"] = None,
+        audio_streamer: Optional[Union[AudioStreamer, AsyncAudioStreamer]] = None,
+        negative_prompt_ids: Optional[torch.Tensor] = None,
+        negative_prompt_attention_mask: Optional[torch.Tensor] = None,
+        speech_tensors: Optional[torch.FloatTensor] = None,
+        speech_masks: Optional[torch.BoolTensor] = None,
+        speech_input_mask: Optional[torch.BoolTensor] = None,
+        tts_text_ids: Optional[torch.LongTensor] = None,
+        return_speech: bool = True,
+        cfg_scale: float = 1.0,
+        stop_check_fn: Optional[Callable[[], bool]] = None,
+        **kwargs,
+    ) -> Union[torch.LongTensor, VibeVoiceGenerationOutput]:
+        """
+        Text is fed in small windows (dynamic slicing of `tts_text_ids`), which enables streaming text input: you don’t need the full text upfront. After each text window, a loop samples several speech latents (diffusion). The interleaved text encoding + speech generation enables streaming text input and realtime speech output.
+        The function only supports batch size = 1 currently.
+
+        - Windowed text prefill → incremental LM + TTS LM updates.
+        - Interleave speech token diffusion sampling (`sample_speech_tokens`).
+        - Stops on EOS (binary classifier) or max length / external `stop_check_fn`.
+        - Returns final token `sequences` and (optionally) concatenated speech audio.
+
+        Args (selected):
+            tts_text_ids: Full text tokens to stream in windows.
+            audio_streamer: If provided, emits audio chunks during generation.
+            cfg_scale: Classifier-free guidance scale for speech diffusion.
+            return_speech: If False, skips audio decode concatenation.
+            stop_check_fn: External early-stop hook (returns True to halt).
+
+        Returns:
+            VibeVoiceGenerationOutput with:
+              - sequences: final token ids
+              - speech_outputs: list of concatenated audio tensors (or None)
+              - reach_max_step_sample: flags for samples stopped by max length
+        """
+        # 1. Handle `generation_config` and kwargs that might update it, and validate the `.generate()` call
+        tokenizer = kwargs.pop("tokenizer", None)
+        neg_text_input_id = tokenizer.convert_tokens_to_ids("<|image_pad|>")
+        
+        tts_lm_input_ids = kwargs.pop("tts_lm_input_ids", None)
+        tts_lm_attention_mask = kwargs.pop("tts_lm_attention_mask", None)
+        # all_prefilled_outputs: cached prefilled prompt outputs for lm, tts_lm, neg_lm, neg_tts_lm
+        all_prefilled_outputs = kwargs.pop("all_prefilled_outputs", None)
+        tts_text_ids = tts_text_ids.to(self.device)
+
+        if kwargs.get('max_new_tokens', None) is None:
+            kwargs['max_new_tokens'] = self.config.decoder_config.max_position_embeddings - tts_lm_input_ids.shape[-1]
+
+        generation_config, model_kwargs, input_ids, logits_processor, stopping_criteria = self._build_generate_config_model_kwargs(
+            generation_config, inputs, tokenizer, return_processors=True, **kwargs
+        )
+        
+        negative_kwargs = {
+            'input_ids': torch.full((kwargs['input_ids'].shape[0], 1), neg_text_input_id, dtype=torch.long, device=kwargs['input_ids'].device),
+            'attention_mask':  torch.ones((kwargs['input_ids'].shape[0], 1), dtype=torch.long, device=kwargs['input_ids'].device),
+            'max_new_tokens': kwargs.get('max_new_tokens', 100) 
+        }
+        negative_generation_config, negative_model_kwargs, negative_input_ids = self._build_generate_config_model_kwargs(
+            None, None, tokenizer, return_processors=False, **negative_kwargs
+        )
+
+        tts_lm_kwargs = {
+            'input_ids': tts_lm_input_ids,
+            'attention_mask': tts_lm_attention_mask,
+            'max_new_tokens': kwargs.get('max_new_tokens', 100) 
+        }
+        tts_lm_generation_config, tts_lm_model_kwargs, tts_lm_input_ids = self._build_generate_config_model_kwargs(
+            None, None, tokenizer, return_processors=False, **tts_lm_kwargs
+        )
+
+        tts_lm_negative_kwargs = {
+            'input_ids': torch.full((kwargs['input_ids'].shape[0], 1), neg_text_input_id, dtype=torch.long, device=kwargs['input_ids'].device),
+            'attention_mask':  torch.ones((kwargs['input_ids'].shape[0], 1), dtype=torch.long, device=kwargs['input_ids'].device),
+            'max_new_tokens': kwargs.get('max_new_tokens', 100) 
+        }
+        tts_lm_negative_generation_config, tts_lm_negative_model_kwargs, tts_lm_negative_input_ids = self._build_generate_config_model_kwargs(
+            None, None, tokenizer, return_processors=False, **tts_lm_negative_kwargs
+        )
+
+        acoustic_cache = VibeVoiceTokenizerStreamingCache()
+        batch_size = input_ids.shape[0]
+        assert batch_size == 1, "Currently only supports batch size == 1"
+        device = input_ids.device
+        finished_tags = torch.zeros(batch_size, dtype=torch.bool, device=device)
+        verbose = kwargs.get("verbose", False)
+
+        # Initialize audio chunks storage for each sample
+        audio_chunks = [[] for _ in range(batch_size)]
+        tts_text_window_index = 0
+        reach_max_step_sample = torch.zeros(batch_size, dtype=torch.bool, device=device)
+        first_text_window_size = TTS_TEXT_WINDOW_SIZE if tts_text_ids.shape[1] >= TTS_TEXT_WINDOW_SIZE else tts_text_ids.shape[1]
+
+        outputs = all_prefilled_outputs["lm"]
+        tts_lm_outputs = all_prefilled_outputs["tts_lm"]
+        negative_outputs = all_prefilled_outputs["neg_lm"]
+        tts_lm_negative_outputs = all_prefilled_outputs["neg_tts_lm"]
+
+        model_kwargs = _update_model_kwargs_for_generation(
+            outputs, model_kwargs, num_new_tokens=first_text_window_size,
+        )
+        tts_lm_model_kwargs = _update_model_kwargs_for_generation(
+            tts_lm_outputs, tts_lm_model_kwargs, num_new_tokens=first_text_window_size,
+        )
+        negative_model_kwargs = self._update_model_kwargs_for_generation(
+            negative_outputs, negative_model_kwargs, is_encoder_decoder=False,
+        )
+        tts_lm_negative_model_kwargs = self._update_model_kwargs_for_generation(
+            tts_lm_negative_outputs, tts_lm_negative_model_kwargs, is_encoder_decoder=False,
+        )
+
+        step = tts_lm_input_ids.shape[1]
+        total_generated_speech_tokens = 0
+        total_prefilled_text_tokens = 0
+        if kwargs.get("show_progress_bar", True):
+            progress_bar = tqdm(
+                total=tts_lm_generation_config.max_length,
+                desc=f"Prefilled {step} tokens, current step ({step} / {tts_lm_generation_config.max_length})",
+                initial=step,
+                leave=False
+            )
+        else:
+            progress_bar = None
+
+        while True:
+            # Check for external stop signal
+            if stop_check_fn is not None and stop_check_fn():
+                if verbose:
+                    print(f"Generation stopped externally at step {step + 1}")
+                # End the audio streamer if it exists
+                if audio_streamer is not None:
+                    audio_streamer.end()
+                break
+            
+            # # Check if audio_streamer has been ended (stopped externally)
+            # if audio_streamer is not None and hasattr(audio_streamer, 'finished_flags'):
+            #     if any(audio_streamer.finished_flags):
+            #         if verbose:
+            #             print(f"Audio generation stopped externally at step {step + 1}")
+            #         break
+            
+            if finished_tags.all():
+                if hasattr(progress_bar, 'set_description'):
+                    progress_bar.set_description("Generation complete")
+                break
+
+            cur_input_tts_text_ids = tts_text_ids[:, tts_text_window_index*TTS_TEXT_WINDOW_SIZE:(tts_text_window_index+1)*TTS_TEXT_WINDOW_SIZE]
+            next_text_window_size = tts_text_ids[:, (tts_text_window_index+1)*TTS_TEXT_WINDOW_SIZE:(tts_text_window_index+2)*TTS_TEXT_WINDOW_SIZE].shape[1]
+            tts_text_window_index += 1
+
+            if cur_input_tts_text_ids.shape[1] > 0:
+                input_ids = torch.cat([input_ids, cur_input_tts_text_ids], dim=-1)
+                tts_lm_input_ids = torch.cat([tts_lm_input_ids, cur_input_tts_text_ids], dim=-1)
+
+                if tts_lm_input_ids.shape[1] > tts_lm_generation_config.max_length:
+                    if verbose:
+                        print(f"Reached maximum generation length {generation_config.max_length}, stopped it.")
+                    reached_samples = torch.arange(batch_size, device=device)[~finished_tags]
+                    if reached_samples.numel() > 0:
+                        reach_max_step_sample[reached_samples] = True
+                    break
+                
+                step += cur_input_tts_text_ids.shape[1]
+                total_prefilled_text_tokens += cur_input_tts_text_ids.shape[1]
+                if progress_bar is not None:
+                    progress_bar.update(cur_input_tts_text_ids.shape[1])
+                    progress_bar.set_description(f"Prefilled {total_prefilled_text_tokens} text tokens, generated {total_generated_speech_tokens} speech tokens, current step ({step} / {tts_lm_generation_config.max_length})")
+
+                model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
+                # Forward pass through the model
+                outputs = self.forward_lm(
+                    **model_inputs, return_dict=True, output_attentions=False, output_hidden_states=False,
+                )
+                model_kwargs = _update_model_kwargs_for_generation(
+                    outputs, model_kwargs, num_new_tokens=next_text_window_size,
+                )
+
+                tts_lm_model_inputs = self.prepare_inputs_for_generation(tts_lm_input_ids, **tts_lm_model_kwargs)
+                tts_lm_additional_inputs = {
+                    "tts_text_masks": torch.ones_like(tts_lm_input_ids[:, -1:]),
+                    "lm_last_hidden_state": outputs.last_hidden_state,
+                }
+                # Forward pass through the model
+                tts_lm_outputs = self.forward_tts_lm(
+                    **tts_lm_model_inputs, **tts_lm_additional_inputs, return_dict=True, output_attentions=False, output_hidden_states=False,
+                )
+                tts_lm_model_kwargs = self._update_model_kwargs_for_generation(
+                    tts_lm_outputs, tts_lm_model_kwargs, is_encoder_decoder=False,
+                )
+
+            diffusion_indices = torch.LongTensor([0])
+            for cur_speech_index in range(TTS_SPEECH_WINDOW_SIZE):
+                positive_condition = tts_lm_outputs.last_hidden_state[diffusion_indices, -1, :]
+                negative_condition = tts_lm_negative_outputs.last_hidden_state[diffusion_indices, -1, :]
+                
+                speech_latent = self.sample_speech_tokens(
+                    positive_condition,
+                    negative_condition,
+                    cfg_scale=cfg_scale,
+                ).unsqueeze(1)
+                                
+                # Decode acoustic latent to audio using acoustic streaming cache
+                scaled_latent = speech_latent / self.model.speech_scaling_factor.to(speech_latent.device) - self.model.speech_bias_factor.to(speech_latent.device)
+                audio_chunk = self.model.acoustic_tokenizer.decode(
+                    scaled_latent.to(self.model.acoustic_tokenizer.device),
+                    cache=acoustic_cache,  # Use acoustic-specific cache
+                    sample_indices=diffusion_indices.to(self.model.acoustic_tokenizer.device),
+                    use_cache=True,
+                    debug=False
+                )
+                
+                # Store audio chunks for each sample
+                for i, sample_idx in enumerate(diffusion_indices):
+                    idx = sample_idx.item()
+                    # Only append audio chunk if the sample is not finished
+                    if not finished_tags[idx]:
+                        audio_chunks[idx].append(audio_chunk[i])
+
+                 # Add streaming support here
+                if audio_streamer is not None:
+                    # Stream the audio chunks immediately
+                    audio_streamer.put(audio_chunk, diffusion_indices)
+
+                acoustic_embed = self.model.acoustic_connector(speech_latent)
+                tts_lm_input_ids = torch.cat([tts_lm_input_ids, torch.ones_like(tts_lm_input_ids[:, -1:])], dim=-1)
+
+                if tts_lm_input_ids.shape[1] > tts_lm_generation_config.max_length:
+                    break
+                
+                step += 1
+                total_generated_speech_tokens += 1
+                if progress_bar is not None:
+                    progress_bar.update(1)
+                    progress_bar.set_description(f"Prefilled {total_prefilled_text_tokens} text tokens, generated {total_generated_speech_tokens} speech tokens, current step ({step} / {tts_lm_generation_config.max_length})")
+
+                tts_lm_model_inputs = self.prepare_inputs_for_generation(tts_lm_input_ids, **tts_lm_model_kwargs)
+                tts_lm_additional_inputs = {
+                    "tts_text_masks": torch.zeros_like(tts_lm_input_ids[:, -1:]),
+                    "lm_last_hidden_state": acoustic_embed,
+                }
+                # Forward pass through the model
+                tts_lm_outputs = self.forward_tts_lm(
+                    **tts_lm_model_inputs, **tts_lm_additional_inputs, return_dict=True, output_attentions=False, output_hidden_states=False,
+                )
+                if cur_speech_index == TTS_SPEECH_WINDOW_SIZE - 1 and next_text_window_size > 0:
+                    tts_lm_model_kwargs = _update_model_kwargs_for_generation(
+                        tts_lm_outputs, tts_lm_model_kwargs, num_new_tokens=next_text_window_size,
+                    )
+                else:
+                    tts_lm_model_kwargs = self._update_model_kwargs_for_generation(
+                        tts_lm_outputs, tts_lm_model_kwargs, is_encoder_decoder=False,
+                    )
+
+                tts_lm_negative_input_ids = torch.cat([tts_lm_negative_input_ids, torch.ones_like(tts_lm_input_ids[:, -1:])], dim=-1)
+                tts_lm_negative_model_inputs = self.prepare_inputs_for_generation(tts_lm_negative_input_ids, **tts_lm_negative_model_kwargs)
+                # Forward negative pass through the model
+                tts_lm_negative_additional_inputs = {
+                    "tts_text_masks": torch.zeros_like(tts_lm_negative_input_ids[:, -1:]),
+                    "lm_last_hidden_state": acoustic_embed,
+                }
+                tts_lm_negative_outputs = self.forward_tts_lm(
+                    **tts_lm_negative_model_inputs, **tts_lm_negative_additional_inputs, return_dict=True, output_attentions=False, output_hidden_states=False,
+                )
+                tts_lm_negative_model_kwargs = self._update_model_kwargs_for_generation(
+                    tts_lm_negative_outputs, tts_lm_negative_model_kwargs, is_encoder_decoder=False,
+                )
+
+                tts_eos_logits = torch.sigmoid(self.tts_eos_classifier(tts_lm_outputs.last_hidden_state[diffusion_indices, -1, :]))
+                if tts_eos_logits[0].item() > 0.5:
+                    # If EOS token is predicted, we can stop generation for this sample
+                    finished_tags[diffusion_indices] = True
+                    if audio_streamer is not None:
+                        audio_streamer.end(diffusion_indices)
+
+            if tts_lm_input_ids.shape[1] > tts_lm_generation_config.max_length:
+                if verbose:
+                    print(f"Reached maximum generation length {tts_lm_generation_config.max_length}, stopped it.")
+                reached_samples = torch.arange(batch_size, device=device)[~finished_tags]
+                if reached_samples.numel() > 0:
+                    reach_max_step_sample[reached_samples] = True
+                break
+
+        if audio_streamer is not None:
+            audio_streamer.end()
+
+        # Concatenate audio chunks for each sample
+        final_audio_outputs = []
+        for sample_chunks in audio_chunks:
+            if sample_chunks:
+                # Concatenate all chunks along the time dimension (assumed to be the last dimension)
+                concatenated_audio = torch.cat(sample_chunks, dim=-1)
+                final_audio_outputs.append(concatenated_audio)
+            else:
+                # If no audio was generated for this sample, append None
+                final_audio_outputs.append(None)
+        
+        if reach_max_step_sample is not None and reach_max_step_sample.any():
+            print(f"Reached maximum generation length {tts_lm_generation_config.max_length}, stopped it.")
+
+        return VibeVoiceGenerationOutput(
+            sequences=tts_lm_input_ids,
+            speech_outputs=final_audio_outputs if return_speech else None,
+            reach_max_step_sample=reach_max_step_sample,
+        )
+
+    @torch.no_grad()
+    def sample_speech_tokens(self, condition, neg_condition, cfg_scale=3.0):
+        self.model.noise_scheduler.set_timesteps(self.ddpm_inference_steps)
+        condition = torch.cat([condition, neg_condition], dim=0).to(self.model.prediction_head.device)
+        speech = torch.randn(condition.shape[0], self.config.acoustic_vae_dim).to(condition)
+        for t in self.model.noise_scheduler.timesteps:
+            half = speech[: len(speech) // 2]
+            combined = torch.cat([half, half], dim=0)
+            eps = self.model.prediction_head(combined, t.repeat(combined.shape[0]).to(combined), condition=condition)
+            cond_eps, uncond_eps = torch.split(eps, len(eps) // 2, dim=0)
+            half_eps = uncond_eps + cfg_scale * (cond_eps - uncond_eps)
+            eps = torch.cat([half_eps, half_eps], dim=0)
+            speech = self.model.noise_scheduler.step(eps, t, speech).prev_sample
+        return speech[: len(speech) // 2]
+    
+
+AutoModelForCausalLM.register(VibeVoiceStreamingConfig, VibeVoiceStreamingForConditionalGenerationInference)
+
+__all__ = [
+    "VibeVoiceStreamingForConditionalGenerationInference",
+]

vibevoice/modular/modular_vibevoice_diffusion_head.py

@@ -0,0 +1,287 @@
+import math
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers.models.auto import AutoModel
+from transformers.modeling_utils import PreTrainedModel
+# from transformers.modeling_layers import GradientCheckpointingLayer
+from transformers.activations import ACT2FN
+from transformers.utils import logging
+
+from .configuration_vibevoice import VibeVoiceDiffusionHeadConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6, elementwise_affine=True, memory_efficient=False):
+        super().__init__()
+        self.dim = dim
+        self.eps = eps
+        self.elementwise_affine = elementwise_affine
+        if self.elementwise_affine:
+            self.weight = nn.Parameter(torch.ones(dim))
+        else:
+            self.register_parameter('weight', None)
+
+    def _norm(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        output = self._norm(x.float()).type_as(x)
+        if self.weight is not None:
+            output = output * self.weight
+        return output
+
+    def extra_repr(self) -> str:
+        return f'dim={self.dim}, eps={self.eps}, elementwise_affine={self.elementwise_affine}'
+    
+def modulate(x, shift, scale):
+    """Apply modulation to input tensor."""
+    return x * (1 + scale) + shift
+
+
+class TimestepEmbedder(nn.Module):
+    """
+    Embeds scalar timesteps into vector representations.
+    
+    Args:
+        hidden_size (`int`): Size of the output embedding
+        frequency_embedding_size (`int`, optional): Size of the intermediate frequency embedding
+    """
+    def __init__(self, hidden_size, frequency_embedding_size=256):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(frequency_embedding_size, hidden_size, bias=False),
+            # nn.SiLU(),
+            ACT2FN['silu'],
+            nn.Linear(hidden_size, hidden_size, bias=False),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+
+    @staticmethod
+    def timestep_embedding(t, dim, max_period=10000):
+        """
+        Create sinusoidal timestep embeddings.
+        
+        Args:
+            t (`torch.Tensor`): A 1-D Tensor of N indices, one per batch element.
+                            These may be fractional.
+            dim (`int`): The dimension of the output.
+            max_period (`int`, optional): Controls the minimum frequency of the embeddings.
+            
+        Returns:
+            `torch.Tensor`: An [N, D] Tensor of positional embeddings.
+        """
+        half = dim // 2
+        freqs = torch.exp(
+            -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
+        ).to(t.device)
+        args = t[:, None].float() * freqs[None]
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2:
+            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+        return embedding.to(t.dtype)
+
+    def forward(self, t):
+        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+        t_emb = self.mlp(t_freq)
+        return t_emb
+
+
+class FeedForwardNetwork(nn.Module):
+    """
+    Standard feed-forward network with SwiGLU activation.
+    
+    Args:
+        embed_dim (`int`): Input dimension
+        ffn_dim (`int`): Hidden dimension
+    """
+    def __init__(
+        self,
+        embed_dim,
+        ffn_dim,
+    ):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.gate_proj = nn.Linear(self.embed_dim, ffn_dim, bias=False)
+        self.up_proj = nn.Linear(self.embed_dim, ffn_dim, bias=False)
+        self.down_proj = nn.Linear(ffn_dim, self.embed_dim, bias=False)
+        self.act_fn = ACT2FN['silu']  # Using SiLU as the activation function
+
+    def forward(self, x):
+        gate = self.gate_proj(x)
+        up = self.up_proj(x)
+        
+        # SwiGLU activation
+        # gate = F.silu(gate)
+        gate = self.act_fn(gate)
+        return self.down_proj(gate * up)
+
+    
+class HeadLayer(nn.Module):
+    """
+    A layer in the diffusion head.
+    
+    Args:
+        embed_dim (`int`): Input dimension
+        ffn_dim (`int`): Hidden dimension
+        cond_dim (`int`): Condition embedding dimension
+        norm_eps (`float`, optional): Epsilon for normalization
+    """
+    def __init__(
+        self,
+        embed_dim,
+        ffn_dim,
+        cond_dim,
+        norm_eps=1e-5,
+    ):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.cond_dim = cond_dim
+        self.ffn_dim = ffn_dim
+        self.ffn = FeedForwardNetwork(
+            self.embed_dim,
+            self.ffn_dim,
+        )
+        self.norm = RMSNorm(self.embed_dim, eps=norm_eps)
+        self.adaLN_modulation = nn.Sequential(
+            # nn.SiLU(),
+            ACT2FN['silu'],
+            nn.Linear(cond_dim, 3 * self.embed_dim, bias=False)
+        )
+
+    def forward(self, x, c):
+        shift_ffn, scale_ffn, gate_ffn = self.adaLN_modulation(c).chunk(3, dim=-1)
+        x = x + gate_ffn * self.ffn(modulate(self.norm(x), shift_ffn, scale_ffn))
+        return x
+
+
+class FinalLayer(nn.Module):
+    """
+    Final layer in the diffusion head.
+    
+    Args:
+        hidden_size (`int`): Input dimension
+        output_size (`int`): Output dimension
+        cond_size (`int`): Condition embedding dimension
+        norm_eps (`float`, optional): Epsilon for normalization
+    """
+    def __init__(self, hidden_size, output_size, cond_size, norm_eps=1e-5):
+        super().__init__()
+        self.norm_final = RMSNorm(hidden_size, eps=norm_eps, elementwise_affine=False)
+        self.linear = nn.Linear(hidden_size, output_size, bias=False)
+        self.adaLN_modulation = nn.Sequential(
+            # nn.SiLU(),
+            ACT2FN['silu'],
+            nn.Linear(cond_size, 2 * hidden_size, bias=False)
+        )
+
+    def forward(self, x, c):
+        shift, scale = self.adaLN_modulation(c).chunk(2, dim=-1)
+        x = modulate(self.norm_final(x), shift, scale)
+        x = self.linear(x)
+        return x
+
+
+class VibeVoiceDiffusionHead(PreTrainedModel):
+    """
+    Diffusion head model for vibevoice.
+    
+    Args:
+        config (`VibeVoiceDiffusionHeadConfig`): Model configuration
+        latent_size (`int`, optional): Size of the latent space. If not provided, uses `config.latent_size`.
+    """
+    config_class = VibeVoiceDiffusionHeadConfig
+    supports_gradient_checkpointing = True
+    _supports_flash_attn_2 = True  
+    _supports_sdpa = True  
+    
+    def __init__(
+        self,
+        config,
+    ):
+        super().__init__(config)
+        self.config = config
+        self.cond_dim = config.hidden_size
+        latent_size = config.latent_size
+        
+        self.noisy_images_proj = nn.Linear(latent_size, config.hidden_size, bias=False)
+        self.cond_proj = nn.Linear(config.hidden_size, self.cond_dim, bias=False)
+        self.t_embedder = TimestepEmbedder(self.cond_dim)
+        
+        ffn_dim = int(config.hidden_size * config.head_ffn_ratio)
+        
+        # Create the intermediate layers
+        self.layers = nn.ModuleList([
+            HeadLayer(
+                embed_dim=config.hidden_size,
+                ffn_dim=ffn_dim,
+                cond_dim=self.cond_dim,
+                norm_eps=config.rms_norm_eps
+            )
+            for _ in range(config.head_layers)
+        ])
+        
+        # Final layer for output
+        self.final_layer = FinalLayer(
+            hidden_size=config.hidden_size, 
+            output_size=latent_size,
+            cond_size=self.cond_dim,
+            norm_eps=config.rms_norm_eps
+        )
+        
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        """Initialize the weights of the model."""
+        # Initialize timestep embedder
+        nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
+        nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
+
+        # Zero-out adaLN modulation layers
+        for layer in self.layers:
+            nn.init.constant_(layer.adaLN_modulation[-1].weight, 0)
+
+        # Zero-out output layers
+        nn.init.constant_(self.final_layer.adaLN_modulation[-1].weight, 0)
+        nn.init.constant_(self.final_layer.linear.weight, 0)
+
+    def forward(
+        self,
+        noisy_images,
+        timesteps,
+        condition,
+    ):
+        """
+        Forward pass of the prediction head.
+        
+        Args:
+            noisy_images (`torch.Tensor`): Noisy images/latents to denoise
+            timesteps (`torch.Tensor`): Timesteps for diffusion
+            condition (`torch.Tensor`): Conditioning information
+            
+        Returns:
+            `torch.Tensor`: The predicted noise/velocity
+        """
+        x = self.noisy_images_proj(noisy_images)
+        t = self.t_embedder(timesteps)
+        condition = self.cond_proj(condition)
+        c = condition + t
+        
+        for layer in self.layers:
+            x = layer(x, c)
+            
+        x = self.final_layer(x, c)
+        return x
+
+
+AutoModel.register(VibeVoiceDiffusionHeadConfig, VibeVoiceDiffusionHead)
+
+__all__ = [
+    "VibeVoiceDiffusionHead",
+]

vibevoice/modular/modular_vibevoice_text_tokenizer.py

@@ -0,0 +1,313 @@
+"""Tokenization classes for vibevoice."""
+
+from typing import List, Optional, Union
+
+from transformers.utils import logging
+from transformers.models.qwen2.tokenization_qwen2 import Qwen2Tokenizer
+from transformers.models.qwen2.tokenization_qwen2_fast import Qwen2TokenizerFast
+
+logger = logging.get_logger(__name__)
+
+
+class VibeVoiceTextTokenizer(Qwen2Tokenizer):
+    """
+    Construct a VibeVoice tokenizer. Based on the Qwen2 tokenizer with additional special tokens for speech.
+    
+    Args:
+        vocab_file (`str`):
+            Path to the vocabulary file.
+        merges_file (`str`):
+            Path to the merges file.
+        errors (`str`, *optional*, defaults to `"replace"`):
+            Paradigm to follow when decoding bytes to UTF-8.
+        unk_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
+            The unknown token.
+        bos_token (`str`, *optional*):
+            The beginning of sequence token. Not used for vibevoice.
+        eos_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
+            The end of sequence token.
+        pad_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
+            The token used for padding.
+        add_special_tokens (`bool`, *optional*, defaults to `True`):
+            Whether or not to add special tokens when encoding.
+    """
+
+    model_input_names = ["input_ids", "attention_mask"]
+
+    def __init__(
+        self,
+        vocab_file,
+        merges_file,
+        errors="replace",
+        unk_token="<|endoftext|>",
+        bos_token=None,
+        eos_token="<|endoftext|>",
+        pad_token="<|endoftext|>",
+        add_prefix_space=False,
+        add_special_tokens=True,
+        **kwargs,
+    ):
+        super().__init__(
+            vocab_file=vocab_file,
+            merges_file=merges_file,
+            errors=errors,
+            unk_token=unk_token,
+            bos_token=bos_token,
+            eos_token=eos_token,
+            pad_token=pad_token,
+            add_prefix_space=add_prefix_space,
+            add_special_tokens=add_special_tokens,
+            **kwargs,
+        )
+        
+        # Add VibeVoice-specific special tokens
+        self._add_vibevoice_special_tokens()
+        
+    def _add_vibevoice_special_tokens(self):
+        """Add VibeVoice-specific special tokens."""
+        special_tokens = {
+            "additional_special_tokens": [
+                "<|vision_start|>",  # Speech start (reusing vision tokens)
+                "<|vision_end|>",  # Speech end
+                "<|vision_pad|>",  # Speech diffusion pad
+            ]
+        }
+        num_added = self.add_special_tokens(special_tokens)
+        
+        # Cache special token IDs
+        self._speech_start_id = self.convert_tokens_to_ids("<|vision_start|>")
+        self._speech_end_id = self.convert_tokens_to_ids("<|vision_end|>")
+        self._speech_diffusion_id = self.convert_tokens_to_ids("<|vision_pad|>")
+        
+        self._eos_id = self.convert_tokens_to_ids('<|endoftext|>')
+
+        return num_added
+    
+    @property
+    def eos_id(self) -> int:
+        """ID of the end of sequence token."""
+        return self._eos_id
+    
+    @property
+    def speech_start_id(self) -> int:
+        """ID of the speech start token."""
+        return self._speech_start_id
+    
+    @property
+    def speech_end_id(self) -> int:
+        """ID of the speech end token."""
+        return self._speech_end_id
+    
+    @property
+    def speech_diffusion_id(self) -> int:
+        """ID of the speech diffusion token."""
+        return self._speech_diffusion_id
+    
+    @property
+    def pad_id(self) -> int:
+        """ID used for padding (returns -100 for loss masking)."""
+        return -100
+
+
+class VibeVoiceTextTokenizerFast(Qwen2TokenizerFast):
+    """
+    Construct a "fast" VibeVoice tokenizer (backed by HuggingFace's *tokenizers* library).
+    Based on the Qwen2 tokenizer with additional special tokens for speech.
+    
+    Args:
+        vocab_file (`str`, *optional*):
+            Path to the vocabulary file.
+        merges_file (`str`, *optional*):
+            Path to the merges file.
+        tokenizer_file (`str`, *optional*):
+            Path to [tokenizers](https://github.com/huggingface/tokenizers) file.
+        unk_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
+            The unknown token.
+        bos_token (`str`, *optional*):
+            The beginning of sequence token. Not used for vibevoice.
+        eos_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
+            The end of sequence token.
+        pad_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
+            The token used for padding.
+    """
+
+    model_input_names = ["input_ids", "attention_mask"]
+
+    def __init__(
+        self,
+        vocab_file=None,
+        merges_file=None,
+        tokenizer_file=None,
+        unk_token="<|endoftext|>",
+        bos_token=None,
+        eos_token="<|endoftext|>",
+        pad_token="<|endoftext|>",
+        add_prefix_space=False,
+        **kwargs,
+    ):
+        super().__init__(
+            vocab_file=vocab_file,
+            merges_file=merges_file,
+            tokenizer_file=tokenizer_file,
+            unk_token=unk_token,
+            bos_token=bos_token,
+            eos_token=eos_token,
+            pad_token=pad_token,
+            add_prefix_space=add_prefix_space,
+            **kwargs,
+        )
+        
+        # Add VibeVoice-specific special tokens
+        self._add_vibevoice_special_tokens()
+        
+    def _add_vibevoice_special_tokens(self):
+        """Add VibeVoice-specific special tokens."""
+        special_tokens = {
+            "additional_special_tokens": [
+                "<|vision_start|>",  # Speech start (reusing vision tokens)
+                "<|vision_end|>",  # Speech end
+                "<|vision_pad|>",  # Speech diffusion pad
+            ]
+        }
+        num_added = self.add_special_tokens(special_tokens)
+        
+        # Cache special token IDs
+        self._speech_start_id = self.convert_tokens_to_ids("<|vision_start|>")
+        self._speech_end_id = self.convert_tokens_to_ids("<|vision_end|>")
+        self._speech_diffusion_id = self.convert_tokens_to_ids("<|vision_pad|>")
+
+        # self._eos_id = self.convert_tokens_to_ids('<|endoftext|>')
+        self._eos_id = self.eos_token_id # qwen2 / qwen3
+        self._pad_id = self.convert_tokens_to_ids('<|image_pad|>')
+        
+        return num_added
+    
+    @property
+    def eos_id(self) -> int:
+        """ID of the end of sequence token."""
+        return self._eos_id
+    
+    @property
+    def speech_start_id(self) -> int:
+        """ID of the speech start token."""
+        return self._speech_start_id
+    
+    @property
+    def speech_end_id(self) -> int:
+        """ID of the speech end token."""
+        return self._speech_end_id
+    
+    @property
+    def speech_diffusion_id(self) -> int:
+        """ID of the speech diffusion token."""
+        return self._speech_diffusion_id
+    
+    @property
+    def pad_id(self) -> int:
+        """ID used for padding (returns -100 for loss masking)."""
+        return self._pad_id
+
+class VibeVoiceASRTextTokenizerFast(Qwen2TokenizerFast):
+    """
+    Construct a "fast" VibeVoice tokenizer (backed by HuggingFace's *tokenizers* library).
+    Based on the Qwen2 tokenizer with additional special tokens for speech.
+    
+    Args:
+        vocab_file (`str`, *optional*):
+            Path to the vocabulary file.
+        merges_file (`str`, *optional*):
+            Path to the merges file.
+        tokenizer_file (`str`, *optional*):
+            Path to [tokenizers](https://github.com/huggingface/tokenizers) file.
+        unk_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
+            The unknown token.
+        bos_token (`str`, *optional*):
+            The beginning of sequence token. Not used for vibevoice.
+        eos_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
+            The end of sequence token.
+        pad_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
+            The token used for padding.
+    """
+
+    model_input_names = ["input_ids", "attention_mask"]
+
+    def __init__(
+        self,
+        vocab_file=None,
+        merges_file=None,
+        tokenizer_file=None,
+        unk_token="<|endoftext|>",
+        bos_token=None,
+        eos_token="<|endoftext|>",
+        pad_token="<|endoftext|>",
+        add_prefix_space=False,
+        **kwargs,
+    ):
+        super().__init__(
+            vocab_file=vocab_file,
+            merges_file=merges_file,
+            tokenizer_file=tokenizer_file,
+            unk_token=unk_token,
+            bos_token=bos_token,
+            eos_token=eos_token,
+            pad_token=pad_token,
+            add_prefix_space=add_prefix_space,
+            **kwargs,
+        )
+        
+        # Add VibeVoice-specific special tokens
+        self._add_vibevoice_special_tokens()
+        
+        # https://github.com/QwenLM/Qwen2.5-VL/blob/d2240f11656bfe404b9ba56db4e51cd09f522ff1/qwen-vl-finetune/qwenvl/data/data_qwen_packed.py#L57C5-L57C222
+        self.chat_template = "{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}"
+        
+    def _add_vibevoice_special_tokens(self):
+        """Add VibeVoice-specific special tokens."""
+        special_tokens = {
+            "additional_special_tokens": [
+                "<|object_ref_start|>",  # Speech start (reusing vision tokens)
+                "<|object_ref_end|>",  # Speech end
+                "<|box_start|>",  # Speech diffusion pad
+            ]
+        }
+        num_added = self.add_special_tokens(special_tokens)
+        
+        # Cache special token IDs
+        self._speech_start_id = self.convert_tokens_to_ids("<|object_ref_start|>")
+        self._speech_end_id = self.convert_tokens_to_ids("<|object_ref_end|>")
+        self._speech_pad_id = self.convert_tokens_to_ids("<|box_start|>")
+
+        self._eos_id = self.eos_token_id # qwen2 / qwen3
+        self._pad_id = self.convert_tokens_to_ids('<|image_pad|>')
+        
+        return num_added    
+    
+    @property
+    def eos_id(self) -> int:
+        """ID of the end of sequence token."""
+        return self._eos_id
+    
+    @property
+    def speech_start_id(self) -> int:
+        """ID of the speech start token."""
+        return self._speech_start_id
+    
+    @property
+    def speech_end_id(self) -> int:
+        """ID of the speech end token."""
+        return self._speech_end_id
+    
+    @property
+    def speech_pad_id(self) -> int:
+        """ID of the speech diffusion token."""
+        return self._speech_pad_id
+    
+    @property
+    def pad_id(self) -> int:
+        return self._pad_id
+    
+__all__ = [
+    "VibeVoiceTextTokenizer", 
+    "VibeVoiceTextTokenizerFast", 
+    "VibeVoiceASRTextTokenizerFast",
+]

vibevoice/modular/modular_vibevoice_tokenizer.py

@@ -0,0 +1,1207 @@
+import math
+import typing as tp
+from functools import partial
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional, Tuple, Union
+import copy
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers.models.auto import AutoModel
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+from transformers.modeling_utils import PreTrainedModel
+from transformers.activations import ACT2FN
+
+from .configuration_vibevoice import VibeVoiceAcousticTokenizerConfig, VibeVoiceSemanticTokenizerConfig
+
+logger = logging.get_logger(__name__)
+
+import os
+# Try to import APEX FusedRMSNorm
+try:
+    from apex.normalization.fused_layer_norm import fused_rms_norm_affine
+    APEX_AVAILABLE = True
+    # logger.info("APEX FusedRMSNorm is available and will be used for optimization")
+    if int(os.getenv("OPTIMIZE_FOR_SPEED", "0")) == 0:
+        APEX_AVAILABLE = False
+        # logger.warning("APEX FusedRMSNorm is disabled by environment variable OPTIMIZE_FOR_SPEED=0")
+except ImportError:
+    APEX_AVAILABLE = False
+    # logger.warning("APEX FusedRMSNorm not available, using native implementation")
+
+# Normalization modules
+class ConvLayerNorm(nn.LayerNorm):
+    """
+    Convolution-friendly LayerNorm that moves channels to last dimensions
+    before running the normalization and moves them back to original position right after.
+    """
+    def __init__(self, normalized_shape: tp.Union[int, tp.List[int], torch.Size], **kwargs):
+        super().__init__(normalized_shape, **kwargs)
+
+    def forward(self, x):
+        x = x.transpose(1, 2)  # b ... t -> b t ...
+        x = nn.functional.layer_norm(x.float(), self.normalized_shape, self.weight.float(), self.bias.float(), self.eps).type_as(x) 
+        x = x.transpose(1, 2)  # b t ... -> b ... t
+        return x
+    
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-5, elementwise_affine=True, weight_shape=None):
+        super().__init__()
+        self.dim = dim
+        self.eps = eps
+        self.elementwise_affine = elementwise_affine
+        if self.elementwise_affine:
+            weight_shape = (dim,) if weight_shape is None else weight_shape
+            self.weight = nn.Parameter(torch.ones(weight_shape))
+        else:
+            self.register_parameter('weight', None)
+
+    def _norm(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        output = self._norm(x.float()).type_as(x)
+        if self.weight is not None:
+            output = output * self.weight
+        return output
+
+    def extra_repr(self) -> str:
+        return f'dim={self.dim}, eps={self.eps}, elementwise_affine={self.elementwise_affine}'
+
+class ConvRMSNorm(RMSNorm):
+    def __init__(self, dim: int, eps: float = 1e-5, elementwise_affine=True, weight_shape=None):
+        super().__init__(dim, eps, elementwise_affine, weight_shape)
+
+    def forward(self, x):
+        x = x.transpose(1, 2)  # b ... t -> b t ...
+        if (not APEX_AVAILABLE) or (not self.elementwise_affine):
+            # Fallback to native implementation
+            output = self._norm(x.float()).type_as(x)
+            if self.weight is not None:
+                output = output * self.weight
+        else:
+            output = fused_rms_norm_affine(x, self.weight, self.weight.shape, self.eps)
+        output = output.transpose(1, 2)  # b t ... -> b ... t
+        return output
+
+# Convolutional layers and utilities
+CONV_NORMALIZATIONS = frozenset(['none', 'weight_norm', 'spectral_norm',
+                                'time_layer_norm', 'layer_norm', 'time_group_norm'])
+
+
+def apply_parametrization_norm(module: nn.Module, norm: str = 'none') -> nn.Module:
+    assert norm in CONV_NORMALIZATIONS
+    if norm == 'weight_norm':
+        return nn.utils.weight_norm(module)
+    elif norm == 'spectral_norm':
+        return nn.utils.spectral_norm(module)
+    else:
+        # We already check was in CONV_NORMALIZATION, so any other choice
+        # doesn't need reparametrization.
+        return module
+
+
+def get_norm_module(module: nn.Module, causal: bool = False, norm: str = 'none', **norm_kwargs) -> nn.Module:
+    """Return the proper normalization module. If causal is True, this will ensure the returned
+    module is causal, or return an error if the normalization doesn't support causal evaluation.
+    """
+    assert norm in CONV_NORMALIZATIONS
+    if norm == 'layer_norm':
+        assert isinstance(module, nn.modules.conv._ConvNd)
+        return ConvLayerNorm(module.out_channels, **norm_kwargs)
+    elif norm == 'time_group_norm':
+        if causal:
+            raise ValueError("GroupNorm doesn't support causal evaluation.")
+        assert isinstance(module, nn.modules.conv._ConvNd)
+        return nn.GroupNorm(1, module.out_channels, **norm_kwargs)
+    else:
+        return nn.Identity()
+
+
+def get_extra_padding_for_conv1d(x: torch.Tensor, kernel_size: int, stride: int,
+                                padding_total: int = 0) -> int:
+    """Calculate extra padding needed for convolution to have the same output length"""
+    length = x.shape[-1]
+    n_frames = (length - kernel_size + padding_total) / stride + 1
+    ideal_length = (math.ceil(n_frames) - 1) * stride + (kernel_size - padding_total)
+    return ideal_length - length
+
+
+def pad1d(x: torch.Tensor, paddings: tp.Tuple[int, int], mode: str = 'zero', value: float = 0.):
+    """Pad 1D input with handling for small inputs in reflect mode"""
+    length = x.shape[-1]
+    padding_left, padding_right = paddings
+    assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right)
+    if mode == 'reflect':
+        max_pad = max(padding_left, padding_right)
+        extra_pad = 0
+        if length <= max_pad:
+            extra_pad = max_pad - length + 1
+            x = F.pad(x, (0, extra_pad))
+        padded = F.pad(x, paddings, mode, value)
+        end = padded.shape[-1] - extra_pad
+        return padded[..., :end]
+    else:
+        return F.pad(x, paddings, mode, value)
+
+
+def unpad1d(x: torch.Tensor, paddings: tp.Tuple[int, int]):
+    """Remove padding from x, handling properly zero padding. Only for 1d!"""
+    padding_left, padding_right = paddings
+    assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right)
+    assert (padding_left + padding_right) <= x.shape[-1]
+    end = x.shape[-1] - padding_right
+    return x[..., padding_left: end]
+
+
+class NormConv1d(nn.Module):
+    """Wrapper around Conv1d and normalization applied to this conv"""
+    def __init__(self, *args, causal: bool = False, norm: str = 'none',
+                norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs):
+        super().__init__()
+        self.conv = apply_parametrization_norm(nn.Conv1d(*args, **kwargs), norm)
+        self.norm = get_norm_module(self.conv, causal, norm, **norm_kwargs)
+        self.norm_type = norm
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.norm(x)
+        return x
+
+
+class NormConvTranspose1d(nn.Module):
+    """Wrapper around ConvTranspose1d and normalization applied to this conv"""
+    def __init__(self, *args, causal: bool = False, norm: str = 'none',
+                norm_kwargs: tp.Dict[str, tp.Any] = {}, **kwargs):
+        super().__init__()
+        self.convtr = apply_parametrization_norm(nn.ConvTranspose1d(*args, **kwargs), norm)
+        self.norm = get_norm_module(self.convtr, causal, norm, **norm_kwargs)
+        self.norm_type = norm
+
+    def forward(self, x):
+        x = self.convtr(x)
+        x = self.norm(x)
+        return x
+
+
+class VibeVoiceTokenizerStreamingCache:
+    """Cache for streaming convolution, similar to KV cache in attention"""
+    def __init__(self):
+        self.cache = {}  # Dict mapping (layer_id, sample_idx) to state tensor
+        
+    def get(self, layer_id: str, sample_indices: torch.Tensor) -> Optional[torch.Tensor]:
+        """Get cached states for given layer and sample indices"""
+        states = []
+        max_length = 0
+        
+        # First pass: collect states and find max length
+        for idx in sample_indices.tolist():
+            key = (layer_id, idx)
+            if key not in self.cache:
+                return None  # If any sample is missing, return None
+            state = self.cache[key]
+            states.append(state)
+            max_length = max(max_length, state.shape[-1])
+        
+        # Second pass: pad states to max length if needed
+        if len(states) > 0 and states[0].dim() >= 2:
+            padded_states = []
+            for state in states:
+                if state.shape[-1] < max_length:
+                    # Pad on the time dimension (last dimension)
+                    pad_size = max_length - state.shape[-1]
+                    # Pad with zeros on the LEFT to align the most recent samples
+                    padded_state = F.pad(state, (pad_size, 0), mode='constant', value=0)
+                    padded_states.append(padded_state)
+                else:
+                    padded_states.append(state)
+            return torch.stack(padded_states, dim=0)
+        else:
+            return torch.stack(states, dim=0)
+    
+    def set(self, layer_id: str, sample_indices: torch.Tensor, states: torch.Tensor):
+        """Set cached states for given layer and sample indices"""
+        for i, idx in enumerate(sample_indices.tolist()):
+            key = (layer_id, idx)
+            self.cache[key] = states[i].detach()
+
+    def set_to_zero(self, sample_indices: torch.Tensor):
+        """Set all cached states to zero for given sample indices"""
+        for key in list(self.cache.keys()):
+            layer_id, sample_idx = key
+            if sample_idx in sample_indices.tolist():
+                # Create zero tensor with same shape and dtype as cached tensor
+                cached_tensor = self.cache[key]
+                self.cache[key] = torch.zeros_like(cached_tensor)
+                
+    def clear(self, layer_id: Optional[str] = None, sample_indices: Optional[torch.Tensor] = None):
+        """Clear cache for specific layer/samples or everything"""
+        if layer_id is None and sample_indices is None:
+            self.cache.clear()
+        elif layer_id is not None and sample_indices is None:
+            # Clear all samples for a specific layer
+            keys_to_remove = [k for k in self.cache.keys() if k[0] == layer_id]
+            for k in keys_to_remove:
+                del self.cache[k]
+        elif layer_id is not None and sample_indices is not None:
+            # Clear specific samples for a specific layer
+            for idx in sample_indices.tolist():
+                key = (layer_id, idx)
+                self.cache.pop(key, None)
+
+class SConv1d(nn.Module):
+    """Conv1d with built-in handling of asymmetric or causal padding and normalization."""
+    def __init__(self, in_channels: int, out_channels: int,
+                kernel_size: int, stride: int = 1, dilation: int = 1,
+                groups: int = 1, bias: bool = True, causal: bool = False,
+                norm: str = 'none', norm_kwargs: tp.Dict[str, tp.Any] = {},
+                pad_mode: str = 'reflect'):
+        super().__init__()
+        self.conv = NormConv1d(in_channels, out_channels, kernel_size, stride,
+                            dilation=dilation, groups=groups, bias=bias, causal=causal,
+                            norm=norm, norm_kwargs=norm_kwargs)
+        self.causal = causal
+        self.pad_mode = pad_mode
+        
+        # Store configuration
+        self.kernel_size = kernel_size
+        self.dilation = dilation
+        self.stride = stride
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        
+        # For causal convolution, we need to maintain kernel_size - 1 samples as context
+        # need to check use which context_size is more suitable
+        # self.context_size = (kernel_size - 1) * dilation
+        self.context_size = (kernel_size - 1) * dilation - (stride - 1)
+        
+        # For non-streaming mode, calculate padding
+        self.padding_total = (kernel_size - 1) * dilation - (stride - 1)
+        
+        # Create a unique layer ID for cache management
+        self._layer_id = None
+                  
+    @property
+    def layer_id(self):
+        if self._layer_id is None:
+            self._layer_id = f"sconv1d_{id(self)}"
+        return self._layer_id
+        
+    def forward(self, x: torch.Tensor, 
+                cache: Optional[VibeVoiceTokenizerStreamingCache] = None,
+                sample_indices: Optional[torch.Tensor] = None,
+                use_cache: bool = False,
+                debug: bool = False,
+                is_final_chunk: bool = False) -> torch.Tensor:
+        """
+        Forward pass with optional streaming support via cache.
+        
+        Args:
+            x: Input tensor [batch_size, channels, time]
+            cache: VibeVoiceTokenizerStreamingCache object for maintaining states
+            sample_indices: Indices identifying each sample for cache management
+            use_cache: Whether to use cached states for streaming
+            debug: Whether to print debug information
+            is_final_chunk: Whether this is the final chunk (adds extra padding for alignment)
+            
+        Returns:
+            Output tensor
+        """
+        B, C, T = x.shape
+        
+        # Non-streaming mode
+        if not use_cache or cache is None:
+            return self._forward_non_streaming(x, debug=debug)
+        
+        # Streaming mode
+        assert self.causal, "Streaming mode is only supported for causal convolutions"
+        assert sample_indices is not None, "sample_indices must be provided for streaming mode"
+        assert len(sample_indices) == B, "sample_indices must match batch size"
+        
+        return self._forward_streaming(x, cache, sample_indices, debug, is_final_chunk)
+    
+    def _forward_streaming(self, x: torch.Tensor, 
+                          cache: VibeVoiceTokenizerStreamingCache,
+                          sample_indices: torch.Tensor,
+                          debug: bool = False,
+                          is_final_chunk: bool = False) -> torch.Tensor:
+        """Streaming forward pass with cache operations kept separate from compiled code"""
+        B, C, T = x.shape
+        
+        # Cache operations (not compiled)
+        cached_states = cache.get(self.layer_id, sample_indices)
+        
+        if cached_states is None:
+            # First chunk - initialize with zeros for context
+            if self.context_size > 0:
+                cached_states = torch.zeros(B, C, self.context_size, device=x.device, dtype=x.dtype)
+                if debug:
+                    print(f"[DEBUG] Initialized cache with shape: {cached_states.shape}, context_size={self.context_size}")
+            else:
+                cached_states = torch.zeros(B, C, 0, device=x.device, dtype=x.dtype)
+                if debug:
+                    print(f"[DEBUG] No context needed (kernel_size=stride)")
+        
+        # Concatenate cached states with input
+        if cached_states.shape[2] > 0:
+            input_with_context = torch.cat([cached_states, x], dim=2)
+        else:
+            input_with_context = x
+        
+        # For final chunk, add extra padding to ensure ceil behavior (same as non-streaming)
+        if is_final_chunk:
+            extra_padding = get_extra_padding_for_conv1d(
+                input_with_context, self.kernel_size, self.stride, self.padding_total
+            )
+            if extra_padding > 0:
+                input_with_context = pad1d(input_with_context, (0, extra_padding), mode=self.pad_mode)
+                if debug:
+                    print(f"[DEBUG] Final chunk: added extra_padding={extra_padding}")
+            
+        if debug:
+            print(f"[DEBUG] Input shape: {x.shape}, Cache shape: {cached_states.shape}, Combined: {input_with_context.shape}")
+        
+        # Apply convolution directly - no extra padding in streaming mode
+        # The conv layer will handle its own padding internally
+        output = self.conv(input_with_context)
+
+        if debug:
+            print(f"[DEBUG] Output shape: {output.shape}")
+        
+        # Update cache for next chunk
+        if self.context_size > 0:
+            # Calculate how many samples to keep
+            total_input_length = input_with_context.shape[2]
+            
+            # Keep the last context_size samples
+            if total_input_length >= self.context_size:
+                new_cache_start = total_input_length - self.context_size
+                new_cache = input_with_context[:, :, new_cache_start:]
+            else:
+                # If we have less than context_size samples, keep everything
+                new_cache = input_with_context
+                
+            if debug:
+                print(f"[DEBUG] New cache shape: {new_cache.shape}")
+                
+            cache.set(self.layer_id, sample_indices, new_cache)
+        
+        return output
+    
+    def _forward_non_streaming(self, x: torch.Tensor, debug: bool = False) -> torch.Tensor:
+        """Standard forward pass without streaming"""
+        B, C, T = x.shape
+        kernel_size = self.kernel_size
+        stride = self.stride
+        dilation = self.dilation
+        padding_total = self.padding_total
+        
+        # Compute extra padding for stride alignment
+        extra_padding = get_extra_padding_for_conv1d(x, kernel_size, stride, padding_total)
+        
+        if debug:
+            print(f"[DEBUG NON-STREAMING] Input shape: {x.shape}, padding_total={padding_total}, extra_padding={extra_padding}")
+        
+        if self.causal:
+            # Left padding for causal
+            if self.pad_mode == 'constant':
+                x = pad1d(x, (padding_total, extra_padding), mode=self.pad_mode, value=0)
+            else:
+                x = pad1d(x, (padding_total, extra_padding), mode=self.pad_mode)
+        else:
+            # Symmetric padding for non-causal
+            padding_right = padding_total // 2
+            padding_left = padding_total - padding_right
+            x = pad1d(x, (padding_left, padding_right + extra_padding), mode=self.pad_mode)
+        
+        if debug:
+            print(f"[DEBUG NON-STREAMING] After padding: {x.shape}")
+            
+        output = self.conv(x)
+        
+        if debug:
+            print(f"[DEBUG NON-STREAMING] Output shape: {output.shape}")
+        
+        return output
+
+
+class SConvTranspose1d(nn.Module):
+    """ConvTranspose1d with built-in handling of asymmetric or causal padding and normalization."""
+    def __init__(self, in_channels: int, out_channels: int,
+                kernel_size: int, stride: int = 1, causal: bool = False,
+                norm: str = 'none', trim_right_ratio: float = 1.,
+                norm_kwargs: tp.Dict[str, tp.Any] = {}, bias: bool = True):
+        super().__init__()
+        self.convtr = NormConvTranspose1d(in_channels, out_channels, kernel_size, stride,
+                                        causal=causal, norm=norm, norm_kwargs=norm_kwargs, bias=bias)
+        self.causal = causal
+        self.trim_right_ratio = trim_right_ratio
+        assert self.causal or self.trim_right_ratio == 1., \
+            "`trim_right_ratio` != 1.0 only makes sense for causal convolutions"
+        assert self.trim_right_ratio >= 0. and self.trim_right_ratio <= 1.
+
+        # Store configuration
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        
+        # For transposed convolution, padding calculation is different
+        self.padding_total = kernel_size - stride
+        
+        # For streaming, we need to keep track of input history
+        # Transposed conv needs to see multiple input samples to produce correct output
+        self.context_size = kernel_size - 1
+        
+        # Create a unique layer ID for cache management
+        self._layer_id = None
+
+    @property
+    def layer_id(self):
+        if self._layer_id is None:
+            self._layer_id = f"sconvtr1d_{id(self)}"
+        return self._layer_id
+    
+    def forward(self, x: torch.Tensor,
+                cache: Optional[VibeVoiceTokenizerStreamingCache] = None,
+                sample_indices: Optional[torch.Tensor] = None,
+                use_cache: bool = False,
+                debug: bool = False) -> torch.Tensor:
+        """
+        Forward pass with optional streaming support via cache.
+        """
+        B, C, T = x.shape
+        
+        # Non-streaming mode
+        if not use_cache or cache is None:
+            return self._forward_non_streaming(x, debug=debug)
+        
+        # Streaming mode
+        assert sample_indices is not None, "sample_indices must be provided for streaming mode"
+        assert len(sample_indices) == B, "sample_indices must match batch size"
+        
+        return self._forward_streaming(x, cache, sample_indices, debug)
+    
+    def _forward_streaming(self, x: torch.Tensor,
+                          cache: VibeVoiceTokenizerStreamingCache,
+                          sample_indices: torch.Tensor,
+                          debug: bool = False) -> torch.Tensor:
+        """Streaming forward pass with cache operations kept separate from compiled code"""
+        B, C, T = x.shape
+        
+        # Cache operations (not compiled)
+        cached_input = cache.get(self.layer_id, sample_indices)
+        
+        if cached_input is None:
+            # First chunk - no history yet
+            cached_input = torch.zeros(B, C, 0, device=x.device, dtype=x.dtype)
+            if debug:
+                print(f"[DEBUG] Initialized empty cache for transposed conv")
+        
+        # Concatenate cached input with new input
+        full_input = torch.cat([cached_input, x], dim=2)
+        
+        if debug:
+            print(f"[DEBUG] Input shape: {x.shape}, Cache shape: {cached_input.shape}, Combined: {full_input.shape}")
+        
+        # First chunk or debug mode - use uncompiled version
+        full_output = self.convtr(full_input)
+        
+        if debug:
+            print(f"[DEBUG] Full transposed conv output shape: {full_output.shape}")
+        
+        # Calculate padding to remove
+        if self.causal:
+            padding_right = math.ceil(self.padding_total * self.trim_right_ratio)
+            padding_left = self.padding_total - padding_right
+        else:
+            padding_right = self.padding_total // 2
+            padding_left = self.padding_total - padding_right
+        
+        # Remove padding
+        if padding_left + padding_right > 0:
+            full_output = unpad1d(full_output, (padding_left, padding_right))
+        
+        if debug:
+            print(f"[DEBUG] After unpadding: {full_output.shape}")
+        
+        # Determine which part of the output corresponds to the new input
+        if cached_input.shape[2] == 0:
+            # First chunk - return all output
+            output = full_output
+        else:
+            # Subsequent chunks - return only the new output
+            expected_new_output = T * self.stride
+            
+            # Take the last expected_new_output samples
+            if full_output.shape[2] >= expected_new_output:
+                output = full_output[:, :, -expected_new_output:]
+            else:
+                output = full_output
+        
+        if debug:
+            print(f"[DEBUG] Final streaming output shape: {output.shape}")
+        
+        # Update cache
+        if full_input.shape[2] > self.context_size:
+            new_cache = full_input[:, :, -self.context_size:]
+        else:
+            new_cache = full_input
+        
+        if debug:
+            print(f"[DEBUG] New cache shape: {new_cache.shape}")
+            
+        cache.set(self.layer_id, sample_indices, new_cache)
+        
+        return output
+    
+    def _forward_non_streaming(self, x: torch.Tensor, debug: bool = False) -> torch.Tensor:
+        """Standard forward pass without streaming"""
+        if debug:
+            print(f"[DEBUG NON-STREAMING] Input shape: {x.shape}")
+        
+        # Apply transposed convolution
+        y = self.convtr(x)
+        
+        if debug:
+            print(f"[DEBUG NON-STREAMING] After transposed conv: {y.shape}")
+        
+        # Calculate and remove padding
+        if self.causal:
+            padding_right = math.ceil(self.padding_total * self.trim_right_ratio)
+            padding_left = self.padding_total - padding_right
+        else:
+            padding_right = self.padding_total // 2
+            padding_left = self.padding_total - padding_right
+        
+        if padding_left + padding_right > 0:
+            y = unpad1d(y, (padding_left, padding_right))
+        
+        if debug:
+            print(f"[DEBUG NON-STREAMING] Final output shape: {y.shape}")
+            
+        return y
+    
+# FFN 
+class FFN(nn.Module):
+    def __init__(
+        self,
+        embed_dim,
+        ffn_dim,
+        bias=False,
+    ):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.linear1 = nn.Linear(self.embed_dim, ffn_dim, bias=bias) 
+        self.gelu = ACT2FN["gelu"]
+        self.linear2 = nn.Linear(ffn_dim, self.embed_dim, bias=bias)
+
+    def forward(self, x):
+        x = self.linear1(x)
+        x = self.gelu(x)
+        x = self.linear2(x)
+        return x
+
+
+class Convlayer(nn.Module):
+    def __init__(
+            self, 
+            in_channels, 
+            out_channels, 
+            kernel_size, 
+            stride=1, 
+            dilation=1, 
+            groups=1, 
+            bias=True, 
+            pad_mode='zeros', 
+            norm='weight_norm', 
+            causal=True, 
+        ):
+        super().__init__()
+        self.conv = SConv1d(in_channels, out_channels, kernel_size, stride=stride, dilation=dilation, 
+                           groups=groups, bias=bias, pad_mode=pad_mode, norm=norm, causal=causal)
+
+    def forward(self, x):
+        return self.conv(x)
+
+class Block1D(nn.Module):
+    def __init__(self, dim, kernel_size=7, drop_path=0., mixer_layer='conv',  
+                layer_scale_init_value=1e-6, **kwargs):
+        super().__init__()
+        
+        if kwargs.get('layernorm', 'LN') == 'LN':
+            self.norm = ConvLayerNorm(dim, eps=kwargs.get('eps', 1e-6))
+            self.ffn_norm = ConvLayerNorm(dim, eps=kwargs.get('eps', 1e-6))               
+        elif kwargs.get('layernorm', 'RMSNorm') == 'RMSNorm':
+            self.norm = ConvRMSNorm(dim, eps=kwargs.get('eps', 1e-6))
+            self.ffn_norm = ConvRMSNorm(dim, eps=kwargs.get('eps', 1e-6))
+
+        if mixer_layer == 'conv':
+            self.mixer = Convlayer(dim, dim, groups=kwargs.get('groups', 1),
+                                kernel_size=kernel_size, 
+                                pad_mode=kwargs.get('pad_mode', 'reflect'), 
+                                norm=kwargs.get('norm', 'none'), 
+                                causal=kwargs.get('causal', True), 
+                                bias=kwargs.get('bias', True),
+                                )
+        elif mixer_layer == 'depthwise_conv':
+            self.mixer = Convlayer(dim, dim, groups=dim,
+                                kernel_size=kernel_size, 
+                                pad_mode=kwargs.get('pad_mode', 'reflect'), 
+                                norm=kwargs.get('norm', 'none'), 
+                                causal=kwargs.get('causal', True), 
+                                bias=kwargs.get('bias', True),
+                                )
+        else:
+            raise ValueError(f"Unsupported mixer layer: {mixer_layer}")
+        
+        self.ffn = FFN(
+            dim, 
+            kwargs.get('ffn_expansion', 4) * dim, 
+            bias=kwargs.get('bias', False),
+        )
+        self.drop_path = nn.Identity() if drop_path <= 0. else nn.modules.DropPath(drop_path)
+
+        if layer_scale_init_value > 0:
+            self.gamma = nn.Parameter(layer_scale_init_value * torch.ones((dim)), requires_grad=True)
+            self.ffn_gamma = nn.Parameter(layer_scale_init_value * torch.ones((dim)), requires_grad=True)
+        else:
+            self.gamma = None
+            self.ffn_gamma = None
+
+    def forward(self, x):
+        # mixer
+        residual = x
+        x = self.norm(x)
+        x = self.mixer(x)
+        if self.gamma is not None:
+            x = x * self.gamma.unsqueeze(-1)
+        x = residual + self.drop_path(x)
+
+        # ffn
+        residual = x
+        x = self.ffn_norm(x)
+        x = x.permute(0, 2, 1)
+        x = self.ffn(x)
+        x = x.permute(0, 2, 1)
+        if self.ffn_gamma is not None:
+            x = x * self.ffn_gamma.unsqueeze(-1)
+        x = residual + self.drop_path(x)
+
+        return x
+
+
+class TokenizerEncoder(nn.Module):
+    """
+    Encoder component for the VibeVoice tokenizer that converts audio to latent representations.
+    
+    Args:
+        config: Configuration object with model parameters
+    """
+    def __init__(self, config):
+        super().__init__()
+        
+        # Extract parameters from config
+        self.channels = config.channels
+        self.dimension = config.dimension
+        self.n_filters = config.n_filters
+        self.ratios = list(reversed(config.ratios))
+        self.depths = config.depths
+        self.n_residual_layers = getattr(config, "n_residual_layers", 1)
+        self.hop_length = np.prod(self.ratios)
+        self.causal = config.causal
+        
+        # Additional config parameters with defaults
+        kernel_size = getattr(config, "kernel_size", 7)
+        last_kernel_size = getattr(config, "last_kernel_size", 7)
+        norm = getattr(config, "norm", "none")
+        norm_params = getattr(config, "norm_params", {})
+        pad_mode = getattr(config, "pad_mode", "reflect")
+        bias = getattr(config, "bias", True)
+        layernorm = getattr(config, "layernorm", "LN")
+        layernorm_eps = getattr(config, "layernorm_eps", 1e-6)
+        layernorm_elementwise_affine = getattr(config, "layernorm_elementwise_affine", True)
+        drop_path_rate = getattr(config, "drop_path_rate", 0.0)
+        mixer_layer = getattr(config, "mixer_layer", "conv")
+        layer_scale_init_value = getattr(config, "layer_scale_init_value", 0)
+        disable_last_norm = getattr(config, "disable_last_norm", False)
+        
+        # determine the norm type based on layernorm
+        if layernorm == 'LN':
+            norm_type = ConvLayerNorm
+        elif layernorm == 'RMSNorm':
+            norm_type = partial(ConvRMSNorm, elementwise_affine=layernorm_elementwise_affine)
+        else:
+            raise ValueError(f"Unsupported norm type: {layernorm}")
+        
+        # stem and intermediate downsampling conv layers
+        stem = nn.Sequential(
+                SConv1d(self.channels, self.n_filters, kernel_size, norm=norm, norm_kwargs=norm_params, causal=self.causal, pad_mode=pad_mode, bias=bias),
+            )
+        
+        self.downsample_layers = nn.ModuleList()
+        self.downsample_layers.append(stem)
+        for i in range(len(self.ratios)):
+            in_ch = self.n_filters * (2 ** i)
+            out_ch = self.n_filters * (2 ** (i + 1))
+            downsample_layer = nn.Sequential(
+                SConv1d(in_ch, out_ch, kernel_size=self.ratios[i] * 2, stride=self.ratios[i], causal=self.causal, pad_mode=pad_mode, norm=norm, bias=bias)
+            )
+            self.downsample_layers.append(downsample_layer)
+
+        # configure the transformer blocks
+        layer_type = partial(
+            Block1D,
+            mixer_layer=mixer_layer,
+            layernorm=layernorm,
+            eps=layernorm_eps,
+            causal=self.causal,
+            pad_mode=pad_mode,
+            norm=norm,
+            bias=bias,
+            layer_scale_init_value=layer_scale_init_value,
+        )
+        
+        self.stages = nn.ModuleList()
+        dp_rates = [x.item() for x in torch.linspace(0, drop_path_rate, sum(self.depths))] 
+        cur = 0
+
+        for i in range(len(self.depths)):
+            in_ch = self.n_filters * (2 ** i)
+            stage = nn.Sequential(
+                *[layer_type(dim=in_ch, drop_path=dp_rates[cur + j]) for j in range(self.depths[i])]
+            )
+            self.stages.append(stage)
+            cur += self.depths[i]
+        
+        if not disable_last_norm:
+            self.norm = norm_type(in_ch, eps=layernorm_eps)
+        else:
+            self.norm = nn.Identity()
+        self.head = SConv1d(in_ch, self.dimension, kernel_size=last_kernel_size, causal=self.causal, pad_mode=pad_mode, norm=norm, bias=bias)
+
+    def forward_features(self, x, cache=None, sample_indices=None, use_cache=False, debug=False, is_final_chunk=False):
+        for i in range(len(self.depths)):
+            # Apply downsampling
+            for layer in self.downsample_layers[i]:
+                if isinstance(layer, SConv1d):
+                    x = layer(x, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug, is_final_chunk=is_final_chunk)
+                else:
+                    x = layer(x)
+            
+            # Apply stage (Block1D contains Convlayer which contains SConv1d)
+            for block in self.stages[i]:
+                if hasattr(block, 'mixer') and hasattr(block.mixer, 'conv') and isinstance(block.mixer.conv, SConv1d):
+                    # Block1D forward with cache support
+                    residual = x
+                    x = block.norm(x)
+                    x = block.mixer.conv(x, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug, is_final_chunk=is_final_chunk)
+                    if block.gamma is not None:
+                        x = x * block.gamma.unsqueeze(-1)
+                    x = residual + x
+                    
+                    # FFN part
+                    residual = x
+                    x = block.ffn_norm(x)
+                    x = x.permute(0, 2, 1)
+                    x = block.ffn(x)
+                    x = x.permute(0, 2, 1)
+                    if block.ffn_gamma is not None:
+                        x = x * block.ffn_gamma.unsqueeze(-1)
+                    x = residual + x
+                else:
+                    x = block(x)
+
+        return self.norm(x)
+
+    def forward(self, x, cache=None, sample_indices=None, use_cache=False, debug=False, is_final_chunk=False):
+        x = self.forward_features(x, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug, is_final_chunk=is_final_chunk)
+        x = self.head(x, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug, is_final_chunk=is_final_chunk)
+        return x
+
+
+class TokenizerDecoder(nn.Module):
+    """
+    Decoder component for the VibeVoice tokenizer that converts latent representations back to audio.
+    
+    Args:
+        config: Configuration object with model parameters
+    """
+    def __init__(self, config):
+        super().__init__()
+        
+        # Extract parameters from config
+        self.dimension = config.dimension
+        self.channels = config.channels
+        self.n_filters = config.n_filters
+        self.ratios = config.ratios
+        
+        # IMPORTANT CHANGE: Don't reverse depths again since they're already reversed in VibeVoiceAcousticTokenizerModel
+        self.depths = config.depths  # Changed from list(reversed(config.depths))
+        
+        self.n_residual_layers = getattr(config, "n_residual_layers", 1)
+        self.hop_length = np.prod(self.ratios)
+        self.causal = config.causal
+        
+        # Additional config parameters with defaults
+        kernel_size = getattr(config, "kernel_size", 7)
+        last_kernel_size = getattr(config, "last_kernel_size", 7)
+        norm = getattr(config, "norm", "none")
+        norm_params = getattr(config, "norm_params", {})
+        pad_mode = getattr(config, "pad_mode", "reflect")
+        bias = getattr(config, "bias", True)
+        layernorm = getattr(config, "layernorm", "LN")
+        layernorm_eps = getattr(config, "layernorm_eps", 1e-6)
+        trim_right_ratio = getattr(config, "trim_right_ratio", 1.0)
+        layernorm_elementwise_affine = getattr(config, "layernorm_elementwise_affine", True)
+        drop_path_rate = getattr(config, "drop_path_rate", 0.0)
+        mixer_layer = getattr(config, "mixer_layer", "conv")
+        layer_scale_init_value = getattr(config, "layer_scale_init_value", 0)
+        disable_last_norm = getattr(config, "disable_last_norm", False)
+
+        # determine the norm type based on layernorm
+        if layernorm == 'LN':
+            norm_type = ConvLayerNorm
+        elif layernorm == 'RMSNorm':
+            norm_type = partial(ConvRMSNorm, elementwise_affine=layernorm_elementwise_affine)
+        else:
+            raise ValueError(f"Unsupported norm type: {layernorm}")
+        
+        # stem and upsampling layers
+        stem = nn.Sequential(
+                SConv1d(self.dimension, self.n_filters * 2 ** (len(self.depths) - 1), kernel_size, norm=norm, 
+                        norm_kwargs=norm_params, causal=self.causal, pad_mode=pad_mode, bias=bias),
+            )
+        
+        self.upsample_layers = nn.ModuleList()
+        self.upsample_layers.append(stem)
+        for i in range(len(self.ratios)):
+            in_ch = self.n_filters * (2 ** (len(self.depths) - 1 - i))
+            out_ch = self.n_filters * (2 ** (len(self.depths) - 1 - i - 1))
+            upsample_layer = nn.Sequential(
+                SConvTranspose1d(in_ch, out_ch,
+                                kernel_size=self.ratios[i] * 2, stride=self.ratios[i],
+                                norm=norm, norm_kwargs=norm_params, bias=bias,
+                                causal=self.causal, trim_right_ratio=trim_right_ratio),
+            )
+            self.upsample_layers.append(upsample_layer)
+
+        # configure transformer blocks
+        layer_type = partial(
+            Block1D,
+            mixer_layer=mixer_layer,
+            layernorm=layernorm,
+            eps=layernorm_eps,
+            causal=self.causal,
+            pad_mode=pad_mode,
+            norm=norm,
+            bias=bias,
+            layer_scale_init_value=layer_scale_init_value,
+        )
+
+        self.stages = nn.ModuleList()
+        dp_rates = [x.item() for x in torch.linspace(0, drop_path_rate, sum(self.depths))] 
+        cur = 0
+        
+        # Create stages in the same order as the original model
+        for i in range(len(self.depths)):
+            in_ch = self.n_filters * (2 ** (len(self.depths) - 1 - i))
+            stage = nn.Sequential(
+                *[layer_type(dim=in_ch, drop_path=dp_rates[cur + j]) for j in range(self.depths[i])]
+            )
+            self.stages.append(stage)
+            cur += self.depths[i]
+
+        if not disable_last_norm:
+            self.norm = norm_type(in_ch, eps=layernorm_eps)
+        else:
+            self.norm = nn.Identity()
+        self.head = SConv1d(in_ch, self.channels, kernel_size=last_kernel_size, causal=self.causal, pad_mode=pad_mode, norm=norm, bias=bias)
+
+    def forward_features(self, x, cache=None, sample_indices=None, use_cache=False, debug=False):
+        for i in range(len(self.depths)):
+            # Apply upsampling
+            for layer in self.upsample_layers[i]:
+                if isinstance(layer, (SConv1d, SConvTranspose1d)):
+                    x = layer(x, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug)
+                else:
+                    x = layer(x)
+            
+            # Apply stage (Block1D contains Convlayer which contains SConv1d)
+            for block in self.stages[i]:
+                if hasattr(block, 'mixer') and hasattr(block.mixer, 'conv') and isinstance(block.mixer.conv, SConv1d):
+                    # Block1D forward with cache support
+                    residual = x
+                    x = block.norm(x)
+                    x = block.mixer.conv(x, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug)
+                    if block.gamma is not None:
+                        x = x * block.gamma.unsqueeze(-1)
+                    x = residual + x
+                    
+                    # FFN part
+                    residual = x
+                    x = block.ffn_norm(x)
+                    x = x.permute(0, 2, 1)
+                    x = block.ffn(x)
+                    x = x.permute(0, 2, 1)
+                    if block.ffn_gamma is not None:
+                        x = x * block.ffn_gamma.unsqueeze(-1)
+                    x = residual + x
+                else:
+                    x = block(x)
+
+        return self.norm(x)
+    
+    def forward(self, x, cache=None, sample_indices=None, use_cache=False, debug=False):
+        x = self.forward_features(x, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug)
+        x = self.head(x, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug)
+        return x
+    
+
+@dataclass
+class VibeVoiceTokenizerEncoderOutput:
+    """
+    Output of VibeVoice tokenizer encoder, representing a Gaussian distribution with fixed variance.
+    
+    Args:
+        mean (`torch.FloatTensor`): The mean parameters of the distribution.
+        std (`float` or `torch.FloatTensor`): Fixed standard deviation value.
+    """
+    mean: torch.Tensor
+    std: Optional[Union[float, torch.Tensor]] = None
+    
+    def sample(self, dist_type='fix'):
+        """
+        Sample from the distribution.
+        
+        Args:
+            dist_type (`str`): Sampling method, either 'fix' or 'gaussian'.
+                
+        Returns:
+            `torch.FloatTensor`: Sampled values.
+            `torch.FloatTensor` (optional): Standard deviation used (only when dist_type='gaussian').
+        """
+        if dist_type == 'fix':
+            x = self.mean + self.std * torch.randn_like(self.mean)
+            return x, self.std
+        elif dist_type == 'gaussian':
+            batch_size = self.mean.size(0)
+            value = self.std / 0.8
+            std = torch.randn(batch_size, device=self.mean.device, dtype=self.mean.dtype) * value
+
+            while std.dim() < self.mean.dim():
+                std = std.unsqueeze(-1)
+
+            x = self.mean + std * torch.randn_like(self.mean)
+            return x, std
+        else:
+            return self.mean, self.std
+
+    def kl(self):
+        """Compute KL divergence between this distribution and a standard normal."""
+        target = torch.zeros_like(self.mean)
+        return F.mse_loss(self.mean, target, reduction='none')
+
+    def mode(self):
+        """Return the distribution mode (which is the mean for Gaussian)."""
+        return self.mean
+    
+class VibeVoiceAcousticTokenizerModel(PreTrainedModel):
+    """VibeVoice speech tokenizer model combining encoder and decoder for acoustic tokens"""
+    
+    config_class = VibeVoiceAcousticTokenizerConfig
+    base_model_prefix = "vibevoice_acoustic_tokenizer"
+    _supports_flash_attn_2 = True  
+    _supports_sdpa = True  
+    _no_split_modules = ["TokenizerEncoder", "TokenizerDecoder"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        
+        self.register_buffer('fix_std', torch.tensor(config.fix_std), persistent=False)
+        self.std_dist_type = getattr(config, "std_dist_type", "fix")
+        
+        # Parse encoder depths
+        if isinstance(config.encoder_depths, str):
+            encoder_depths = [int(d) for d in config.encoder_depths.split('-')]
+        else:
+            encoder_depths = config.encoder_depths
+            
+        # Parse decoder depths if provided
+        if config.decoder_depths is not None and isinstance(config.decoder_depths, str):
+            decoder_depths = [int(d) for d in config.decoder_depths.split('-')]
+        else:
+            # Default: use reversed encoder depths if decoder_depths is None
+            decoder_depths = list(reversed(encoder_depths))
+        
+        # Create encoder config
+        encoder_config = copy.deepcopy(config)
+        encoder_config.dimension = config.vae_dim
+        encoder_config.n_filters = config.encoder_n_filters
+        encoder_config.ratios = config.encoder_ratios
+        encoder_config.depths = encoder_depths
+        encoder_config.norm = config.conv_norm
+        encoder_config.pad_mode = config.pad_mode
+        encoder_config.bias = config.conv_bias
+        encoder_config.layernorm_eps = config.layernorm_eps
+        encoder_config.layernorm_elementwise_affine = config.layernorm_elementwise_affine
+        encoder_config.mixer_layer = config.mixer_layer
+        encoder_config.layer_scale_init_value = config.layer_scale_init_value
+        encoder_config.disable_last_norm = config.disable_last_norm
+        
+        # Create decoder config
+        decoder_config = copy.deepcopy(config)
+        decoder_config.dimension = config.vae_dim
+        decoder_config.n_filters = config.decoder_n_filters
+        decoder_config.ratios = config.decoder_ratios
+        decoder_config.depths = decoder_depths
+        decoder_config.norm = config.conv_norm
+        decoder_config.pad_mode = config.pad_mode
+        decoder_config.bias = config.conv_bias
+        decoder_config.layernorm_eps = config.layernorm_eps
+        decoder_config.layernorm_elementwise_affine = config.layernorm_elementwise_affine
+        decoder_config.mixer_layer = config.mixer_layer
+        decoder_config.layer_scale_init_value = config.layer_scale_init_value
+        decoder_config.disable_last_norm = config.disable_last_norm
+        
+        # Initialize encoder and decoder
+        self.encoder = TokenizerEncoder(encoder_config)
+        self.decoder = TokenizerDecoder(decoder_config)
+        
+        # Initialize weights
+        self.apply(self._init_weights)
+    
+    def _init_weights(self, module):
+        """Initialize weights for the model"""
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, std=self.config.weight_init_value)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.LayerNorm):
+            nn.init.ones_(module.weight)
+            nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Conv1d):
+            nn.init.normal_(module.weight, std=self.config.weight_init_value)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+    
+    @torch.no_grad()
+    def encode(self, audio, cache=None, sample_indices=None, use_cache=False, debug=False, is_final_chunk=False):
+        """Convert audio to latent representations"""
+        latents = self.encoder(audio, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug, is_final_chunk=is_final_chunk)
+        return VibeVoiceTokenizerEncoderOutput(mean=latents.permute(0, 2, 1), std=self.fix_std)
+    
+    @torch.no_grad()
+    def sampling(self, encoder_output, dist_type=None):
+        """Sample from the encoder output distribution"""
+        dist_type = dist_type or self.std_dist_type
+    
+        if dist_type == 'fix':
+            return encoder_output.sample(dist_type='fix')
+        elif dist_type == 'gaussian':
+            return encoder_output.sample(dist_type='gaussian')
+        else:
+            raise ValueError(f"Unsupported dist_type: {dist_type}, expected 'fix' or 'gaussian'")
+    
+    @torch.no_grad()
+    def decode(self, latents, cache=None, sample_indices=None, use_cache=False, debug=False):
+        """Convert latent representations back to audio"""
+        if latents.shape[1] == self.config.vae_dim:
+            pass
+        else:
+            latents = latents.permute(0, 2, 1)
+
+        audio = self.decoder(latents, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug)
+        return audio
+
+    def forward(self, audio, cache=None, sample_indices=None, use_cache=False, debug=False):
+        """Full forward pass: encode audio to latents, then decode back to audio"""
+        encoder_output = self.encode(audio, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug)
+        sampled_latents, _ = self.sampling(encoder_output)
+        reconstructed = self.decode(sampled_latents, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug)
+        return reconstructed, sampled_latents
+
+
+class VibeVoiceSemanticTokenizerModel(PreTrainedModel):
+    """VibeVoice speech tokenizer model with only encoder for semantic tokens"""
+    
+    config_class = VibeVoiceSemanticTokenizerConfig
+    base_model_prefix = "vibevoice_semantic_tokenizer"
+    _supports_flash_attn_2 = True  
+    _supports_sdpa = True  
+    _no_split_modules = ["TokenizerEncoder"]
+    
+    def __init__(self, config):
+        super().__init__(config)
+        
+        # Parse encoder depths
+        if isinstance(config.encoder_depths, str):
+            encoder_depths = [int(d) for d in config.encoder_depths.split('-')]
+        else:
+            encoder_depths = config.encoder_depths
+        
+        # Create encoder config
+        encoder_config = copy.deepcopy(config)
+        encoder_config.dimension = config.vae_dim
+        encoder_config.n_filters = config.encoder_n_filters
+        encoder_config.ratios = config.encoder_ratios
+        encoder_config.depths = encoder_depths
+        encoder_config.norm = config.conv_norm
+        encoder_config.pad_mode = config.pad_mode
+        encoder_config.bias = config.conv_bias
+        encoder_config.layernorm_eps = config.layernorm_eps
+        encoder_config.layernorm_elementwise_affine = config.layernorm_elementwise_affine
+        encoder_config.mixer_layer = config.mixer_layer
+        encoder_config.layer_scale_init_value = config.layer_scale_init_value
+        encoder_config.disable_last_norm = config.disable_last_norm
+        
+        # Initialize encoder and decoder
+        self.encoder = TokenizerEncoder(encoder_config)
+        
+        # Initialize weights
+        self.apply(self._init_weights)
+    
+    def _init_weights(self, module):
+        """Initialize weights for the model"""
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, std=self.config.weight_init_value)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.LayerNorm):
+            nn.init.ones_(module.weight)
+            nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Conv1d):
+            nn.init.normal_(module.weight, std=self.config.weight_init_value)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+    
+    @torch.no_grad()
+    def encode(self, audio, cache=None, sample_indices=None, use_cache=False, debug=False, is_final_chunk=False):
+        """Convert audio to latent representations"""
+        latents = self.encoder(audio, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug, is_final_chunk=is_final_chunk)
+        return VibeVoiceTokenizerEncoderOutput(mean=latents.permute(0, 2, 1))
+    
+    @torch.no_grad()
+    def sampling(self, encoder_output, dist_type=None):
+        """Sample from the encoder output distribution"""
+        return encoder_output.sample(dist_type='none')
+
+    def forward(self, audio, cache=None, sample_indices=None, use_cache=False, debug=False):
+        """Full forward pass: encode audio to latents, then decode back to audio"""
+        encoder_output = self.encode(audio, cache=cache, sample_indices=sample_indices, use_cache=use_cache, debug=debug)
+        sampled_latents, _ = self.sampling(encoder_output, dist_type='none')
+        return None, sampled_latents
+
+AutoModel.register(VibeVoiceAcousticTokenizerConfig, VibeVoiceAcousticTokenizerModel)
+AutoModel.register(VibeVoiceSemanticTokenizerConfig, VibeVoiceSemanticTokenizerModel)
+
+__all__ = [
+    "VibeVoiceTokenizerStreamingCache",
+    "VibeVoiceAcousticTokenizerModel",
+    "VibeVoiceSemanticTokenizerModel",
+]

vibevoice/modular/streamer.py

@@ -0,0 +1,264 @@
+from __future__ import annotations
+
+import torch
+
+import asyncio
+from queue import Empty, Queue
+from typing import TYPE_CHECKING, Optional
+
+
+from transformers.generation import BaseStreamer
+
+
+class AudioStreamer(BaseStreamer):
+    """
+    Audio streamer that stores audio chunks in queues for each sample in the batch.
+    This allows streaming audio generation for multiple samples simultaneously.
+    
+    Parameters:
+        batch_size (`int`):
+            The batch size for generation
+        stop_signal (`any`, *optional*):
+            The signal to put in the queue when generation ends. Defaults to None.
+        timeout (`float`, *optional*):
+            The timeout for the audio queue. If `None`, the queue will block indefinitely.
+    """
+    
+    def __init__(
+        self, 
+        batch_size: int,
+        stop_signal: Optional[any] = None,
+        timeout: Optional[float] = None,
+    ):
+        self.batch_size = batch_size
+        self.stop_signal = stop_signal
+        self.timeout = timeout
+        
+        # Create a queue for each sample in the batch
+        self.audio_queues = [Queue() for _ in range(batch_size)]
+        self.finished_flags = [False for _ in range(batch_size)]
+        self.sample_indices_map = {}  # Maps from sample index to queue index
+        
+    def put(self, audio_chunks: torch.Tensor, sample_indices: torch.Tensor):
+        """
+        Receives audio chunks and puts them in the appropriate queues.
+        
+        Args:
+            audio_chunks: Tensor of shape (num_samples, ...) containing audio chunks
+            sample_indices: Tensor indicating which samples these chunks belong to
+        """
+        for i, sample_idx in enumerate(sample_indices):
+            idx = sample_idx.item()
+            if idx < self.batch_size and not self.finished_flags[idx]:
+                # Convert to numpy or keep as tensor based on preference
+                audio_chunk = audio_chunks[i].detach().cpu()
+                self.audio_queues[idx].put(audio_chunk, timeout=self.timeout)
+    
+    def end(self, sample_indices: Optional[torch.Tensor] = None):
+        """
+        Signals the end of generation for specified samples or all samples.
+        
+        Args:
+            sample_indices: Optional tensor of sample indices to end. If None, ends all.
+        """
+        if sample_indices is None:
+            # End all samples
+            for idx in range(self.batch_size):
+                if not self.finished_flags[idx]:
+                    self.audio_queues[idx].put(self.stop_signal, timeout=self.timeout)
+                    self.finished_flags[idx] = True
+        else:
+            # End specific samples
+            for sample_idx in sample_indices:
+                idx = sample_idx.item() if torch.is_tensor(sample_idx) else sample_idx
+                if idx < self.batch_size and not self.finished_flags[idx]:
+                    self.audio_queues[idx].put(self.stop_signal, timeout=self.timeout)
+                    self.finished_flags[idx] = True
+    
+    def __iter__(self):
+        """Returns an iterator over the batch of audio streams."""
+        return AudioBatchIterator(self)
+    
+    def get_stream(self, sample_idx: int):
+        """Get the audio stream for a specific sample."""
+        if sample_idx >= self.batch_size:
+            raise ValueError(f"Sample index {sample_idx} exceeds batch size {self.batch_size}")
+        return AudioSampleIterator(self, sample_idx)
+
+
+class AudioSampleIterator:
+    """Iterator for a single audio stream from the batch."""
+    
+    def __init__(self, streamer: AudioStreamer, sample_idx: int):
+        self.streamer = streamer
+        self.sample_idx = sample_idx
+        
+    def __iter__(self):
+        return self
+    
+    def __next__(self):
+        value = self.streamer.audio_queues[self.sample_idx].get(timeout=self.streamer.timeout)
+        if value == self.streamer.stop_signal:
+            raise StopIteration()
+        return value
+
+
+class AudioBatchIterator:
+    """Iterator that yields audio chunks for all samples in the batch."""
+    
+    def __init__(self, streamer: AudioStreamer):
+        self.streamer = streamer
+        self.active_samples = set(range(streamer.batch_size))
+        
+    def __iter__(self):
+        return self
+    
+    def __next__(self):
+        if not self.active_samples:
+            raise StopIteration()
+            
+        batch_chunks = {}
+        samples_to_remove = set()
+        
+        # Try to get chunks from all active samples
+        for idx in self.active_samples:
+            try:
+                value = self.streamer.audio_queues[idx].get(block=False)
+                if value == self.streamer.stop_signal:
+                    samples_to_remove.add(idx)
+                else:
+                    batch_chunks[idx] = value
+            except Empty:
+                # Queue is empty for this sample, skip it this iteration
+                pass
+        
+        # Remove finished samples
+        self.active_samples -= samples_to_remove
+        
+        if batch_chunks:
+            return batch_chunks
+        elif self.active_samples:
+            # If no chunks were ready but we still have active samples, 
+            # wait a bit and try again
+            import time
+            time.sleep(0.01)
+            return self.__next__()
+        else:
+            raise StopIteration()
+
+
+class AsyncAudioStreamer(AudioStreamer):
+    """
+    Async version of AudioStreamer for use in async contexts.
+    """
+    
+    def __init__(
+        self, 
+        batch_size: int,
+        stop_signal: Optional[any] = None,
+        timeout: Optional[float] = None,
+    ):
+        super().__init__(batch_size, stop_signal, timeout)
+        # Replace regular queues with async queues
+        self.audio_queues = [asyncio.Queue() for _ in range(batch_size)]
+        self.loop = asyncio.get_running_loop()
+        
+    def put(self, audio_chunks: torch.Tensor, sample_indices: torch.Tensor):
+        """Put audio chunks in the appropriate async queues."""
+        for i, sample_idx in enumerate(sample_indices):
+            idx = sample_idx.item()
+            if idx < self.batch_size and not self.finished_flags[idx]:
+                audio_chunk = audio_chunks[i].detach().cpu()
+                self.loop.call_soon_threadsafe(
+                    self.audio_queues[idx].put_nowait, audio_chunk
+                )
+    
+    def end(self, sample_indices: Optional[torch.Tensor] = None):
+        """Signal the end of generation for specified samples."""
+        if sample_indices is None:
+            indices_to_end = range(self.batch_size)
+        else:
+            indices_to_end = [s.item() if torch.is_tensor(s) else s for s in sample_indices]
+            
+        for idx in indices_to_end:
+            if idx < self.batch_size and not self.finished_flags[idx]:
+                self.loop.call_soon_threadsafe(
+                    self.audio_queues[idx].put_nowait, self.stop_signal
+                )
+                self.finished_flags[idx] = True
+    
+    async def get_stream(self, sample_idx: int):
+        """Get async iterator for a specific sample's audio stream."""
+        if sample_idx >= self.batch_size:
+            raise ValueError(f"Sample index {sample_idx} exceeds batch size {self.batch_size}")
+            
+        while True:
+            value = await self.audio_queues[sample_idx].get()
+            if value == self.stop_signal:
+                break
+            yield value
+    
+    def __aiter__(self):
+        """Returns an async iterator over all audio streams."""
+        return AsyncAudioBatchIterator(self)
+
+
+class AsyncAudioBatchIterator:
+    """Async iterator for batch audio streaming."""
+    
+    def __init__(self, streamer: AsyncAudioStreamer):
+        self.streamer = streamer
+        self.active_samples = set(range(streamer.batch_size))
+        
+    def __aiter__(self):
+        return self
+        
+    async def __anext__(self):
+        if not self.active_samples:
+            raise StopAsyncIteration()
+            
+        batch_chunks = {}
+        samples_to_remove = set()
+        
+        # Create tasks for all active samples
+        tasks = {
+            idx: asyncio.create_task(self._get_chunk(idx)) 
+            for idx in self.active_samples
+        }
+        
+        # Wait for at least one chunk to be ready
+        done, pending = await asyncio.wait(
+            tasks.values(), 
+            return_when=asyncio.FIRST_COMPLETED,
+            timeout=self.streamer.timeout
+        )
+        
+        # Cancel pending tasks
+        for task in pending:
+            task.cancel()
+            
+        # Process completed tasks
+        for idx, task in tasks.items():
+            if task in done:
+                try:
+                    value = await task
+                    if value == self.streamer.stop_signal:
+                        samples_to_remove.add(idx)
+                    else:
+                        batch_chunks[idx] = value
+                except asyncio.CancelledError:
+                    pass
+                    
+        self.active_samples -= samples_to_remove
+        
+        if batch_chunks:
+            return batch_chunks
+        elif self.active_samples:
+            # Try again if we still have active samples
+            return await self.__anext__()
+        else:
+            raise StopAsyncIteration()
+    
+    async def _get_chunk(self, idx):
+        """Helper to get a chunk from a specific queue."""
+        return await self.streamer.audio_queues[idx].get()

vibevoice/processor/__init__.py

@@ -0,0 +1,11 @@
+# vibevoice/processor/__init__.py
+from .vibevoice_processor import VibeVoiceProcessor
+from .vibevoice_streaming_processor import VibeVoiceStreamingProcessor
+from .vibevoice_tokenizer_processor import VibeVoiceTokenizerProcessor, AudioNormalizer
+
+__all__ = [
+    "VibeVoiceProcessor",
+    "VibeVoiceStreamingProcessor",
+    "VibeVoiceTokenizerProcessor",
+    "AudioNormalizer",
+]

vibevoice/processor/audio_utils.py

@@ -0,0 +1,217 @@
+import os
+import threading
+
+import numpy as np
+from subprocess import run
+from typing import List, Optional, Union, Dict, Any
+
+COMMON_AUDIO_EXTS = [
+    '.mp3', '.MP3', '.Mp3',  # All case variations of mp3
+    '.m4a', 
+    '.mp4', '.MP4',
+    '.wav', '.WAV',
+    '.m4v',
+    '.aac',
+    '.ogg',
+    '.mov', '.MOV',
+    '.opus',
+    '.m4b',
+    '.flac',
+    '.wma', '.WMA',
+    '.rm', '.3gp', '.mpeg', '.flv', '.webm', '.mp2', '.aif', '.aiff', '.oga', '.ogv', '.mpga', '.m3u8', '.amr'
+]
+
+def load_audio_use_ffmpeg(file: str, resample: bool = False, target_sr: int = 24000):
+    """
+    Open an audio file and read as mono waveform, optionally resampling.
+    Returns both the audio data and the original sample rate.
+
+    Parameters
+    ----------
+    file: str
+        The audio file to open
+    resample: bool
+        Whether to resample the audio
+    target_sr: int
+        The target sample rate if resampling is requested
+
+    Returns
+    -------
+    A tuple containing:
+    - A NumPy array with the audio waveform in float32 dtype
+    - The original sample rate of the audio file
+    """
+    if not resample:
+        # First, get the original sample rate
+        cmd_probe = [
+            "ffprobe",
+            "-v", "quiet",
+            "-show_entries", "stream=sample_rate",
+            "-of", "default=noprint_wrappers=1:nokey=1",
+            file
+        ]
+        
+        original_sr = int(run(cmd_probe, capture_output=True, check=True).stdout.decode().strip())
+    else:
+        original_sr = None
+
+    # Now load the audio
+    sr_to_use = target_sr if resample else original_sr
+    
+    cmd = [
+        "ffmpeg",
+        "-loglevel", "error",
+        "-nostdin",
+        "-threads", "0",
+        "-i", file,
+        "-f", "s16le",
+        "-ac", "1",
+        "-acodec", "pcm_s16le",
+        "-ar", str(sr_to_use),
+        "-",
+    ]
+    
+    out = _run_ffmpeg(cmd).stdout
+    audio_data = np.frombuffer(out, np.int16).flatten().astype(np.float32) / 32768.0
+    
+    return audio_data, sr_to_use
+
+
+def _get_ffmpeg_max_concurrency() -> int:
+    """Get the maximum FFmpeg concurrency from environment variable."""
+    v = os.getenv("VIBEVOICE_FFMPEG_MAX_CONCURRENCY", "")
+    try:
+        n = int(v) if v.strip() else 0
+    except Exception:
+        n = 0
+    # 0/negative means no explicit limit.
+    return n
+
+
+_FFMPEG_MAX_CONCURRENCY = _get_ffmpeg_max_concurrency()
+_FFMPEG_SEM = threading.Semaphore(_FFMPEG_MAX_CONCURRENCY) if _FFMPEG_MAX_CONCURRENCY > 0 else None
+
+
+def _run_ffmpeg(cmd: list, *, stdin_bytes: bytes = None):
+    """Run ffmpeg with optional global concurrency limiting.
+
+    This is important for vLLM multi-request concurrency: spawning too many
+    ffmpeg processes can saturate CPU/IO and cause request failures/timeouts.
+    """
+    if _FFMPEG_SEM is None:
+        return run(cmd, capture_output=True, check=True, input=stdin_bytes)
+    with _FFMPEG_SEM:
+        return run(cmd, capture_output=True, check=True, input=stdin_bytes)
+
+
+def load_audio_bytes_use_ffmpeg(data: bytes, *, resample: bool = False, target_sr: int = 24000):
+    """Decode audio bytes via ffmpeg stdin pipe.
+
+    Compared to writing bytes to a temp file, this avoids filesystem IO and
+    reduces contention under high request concurrency.
+    
+    Parameters
+    ----------
+    data: bytes
+        The audio data bytes
+    resample: bool
+        Whether to resample the audio (must be True)
+    target_sr: int
+        The target sample rate if resampling is requested
+
+    Returns
+    -------
+    A tuple containing:
+    - A NumPy array with the audio waveform in float32 dtype
+    - The sample rate
+    """
+    if not resample:
+        # For stdin bytes, we don't have a cheap/robust way to probe original sr.
+        # Keep behavior explicit.
+        raise ValueError("load_audio_bytes_use_ffmpeg requires resample=True")
+
+    cmd = [
+        "ffmpeg",
+        "-loglevel", "error",
+        "-threads", "0",
+        "-i", "pipe:0",
+        "-f", "s16le",
+        "-ac", "1",
+        "-acodec", "pcm_s16le",
+        "-ar", str(target_sr),
+        "-",
+    ]
+    out = _run_ffmpeg(cmd, stdin_bytes=data).stdout
+    audio_data = np.frombuffer(out, np.int16).flatten().astype(np.float32) / 32768.0
+    return audio_data, target_sr
+
+
+class AudioNormalizer:
+    """
+    Audio normalization class for VibeVoice tokenizer.
+    
+    This class provides audio normalization to ensure consistent input levels
+    for the VibeVoice tokenizer while maintaining audio quality.
+    """
+    
+    def __init__(self, target_dB_FS: float = -25, eps: float = 1e-6):
+        """
+        Initialize the audio normalizer.
+        
+        Args:
+            target_dB_FS (float): Target dB FS level for the audio. Default: -25
+            eps (float): Small value to avoid division by zero. Default: 1e-6
+        """
+        self.target_dB_FS = target_dB_FS
+        self.eps = eps
+    
+    def tailor_dB_FS(self, audio: np.ndarray) -> tuple:
+        """
+        Adjust the audio to the target dB FS level.
+        
+        Args:
+            audio (np.ndarray): Input audio signal
+            
+        Returns:
+            tuple: (normalized_audio, rms, scalar)
+        """
+        rms = np.sqrt(np.mean(audio**2))
+        scalar = 10 ** (self.target_dB_FS / 20) / (rms + self.eps)
+        normalized_audio = audio * scalar
+        return normalized_audio, rms, scalar
+    
+    def avoid_clipping(self, audio: np.ndarray, scalar: Optional[float] = None) -> tuple:
+        """
+        Avoid clipping by scaling down if necessary.
+        
+        Args:
+            audio (np.ndarray): Input audio signal
+            scalar (float, optional): Explicit scaling factor
+            
+        Returns:
+            tuple: (normalized_audio, scalar)
+        """
+        if scalar is None:
+            max_val = np.max(np.abs(audio))
+            if max_val > 1.0:
+                scalar = max_val + self.eps
+            else:
+                scalar = 1.0
+        
+        return audio / scalar, scalar
+    
+    def __call__(self, audio: np.ndarray) -> np.ndarray:
+        """
+        Normalize the audio by adjusting to target dB FS and avoiding clipping.
+        
+        Args:
+            audio (np.ndarray): Input audio signal
+            
+        Returns:
+            np.ndarray: Normalized audio signal
+        """
+        # First adjust to target dB FS
+        audio, _, _ = self.tailor_dB_FS(audio)
+        # Then avoid clipping
+        audio, _ = self.avoid_clipping(audio)
+        return audio

vibevoice/processor/vibevoice_asr_processor.py

@@ -0,0 +1,572 @@
+"""
+Processor class for VibeVoice ASR models.
+"""
+
+import os
+import json
+import math
+import warnings
+from typing import List, Optional, Union, Dict, Any, Tuple
+
+import numpy as np
+import torch
+
+from transformers.tokenization_utils_base import BatchEncoding
+from transformers.utils import TensorType, logging
+from .vibevoice_tokenizer_processor import VibeVoiceTokenizerProcessor, AudioNormalizer
+
+try:
+    from .audio_utils import load_audio_use_ffmpeg
+    HAS_FFMPEG_UTILS = True
+except ImportError:
+    HAS_FFMPEG_UTILS = False
+    warnings.warn("audio_utils not available, will fall back to soundfile for audio loading")
+
+logger = logging.get_logger(__name__)
+
+SYSTEM_PROMPT = "You are a helpful assistant that transcribes audio input into text output in JSON format."
+
+
+class VibeVoiceASRProcessor: 
+    """
+    Processor for VibeVoice ASR (Automatic Speech Recognition) models.
+    
+    This processor handles audio preprocessing and tokenization for ASR tasks,
+    following the exact format used in training with proper chat templates.
+    
+    Args:
+        tokenizer: The text tokenizer for processing text
+        audio_processor: The audio processor for processing speech
+        speech_tok_compress_ratio (int): Compression ratio for speech tokenization. Default: 3200 (product of encoder ratios [8,5,5,4,2,2])
+        target_sample_rate (int): Target sample rate for audio
+        normalize_audio (bool): Whether to normalize audio input
+    """
+    
+    def __init__(
+        self,
+        tokenizer=None,
+        audio_processor=None,
+        speech_tok_compress_ratio=3200,
+        target_sample_rate=24000,
+        normalize_audio=True,
+        **kwargs
+    ):
+        self.tokenizer = tokenizer
+        self.audio_processor = audio_processor or VibeVoiceTokenizerProcessor(
+            sampling_rate=target_sample_rate,
+            normalize_audio=normalize_audio
+        )
+        self.speech_tok_compress_ratio = speech_tok_compress_ratio
+        self.target_sample_rate = target_sample_rate
+        self.normalize_audio = normalize_audio
+        
+        if normalize_audio:
+            self.audio_normalizer = AudioNormalizer()
+        else:
+            self.audio_normalizer = None
+        
+        # Cache special token IDs
+        self._cache_special_tokens()
+    
+    def _cache_special_tokens(self):
+        """Cache special token IDs for efficiency."""
+        # Add safety checks for special tokens
+        if hasattr(self.tokenizer, 'speech_start_id'):
+            self.speech_start_id = self.tokenizer.speech_start_id
+        else:
+            self.speech_start_id = self.tokenizer.convert_tokens_to_ids("<|speech_start|>")
+            
+        if hasattr(self.tokenizer, 'speech_end_id'):
+            self.speech_end_id = self.tokenizer.speech_end_id
+        else:
+            self.speech_end_id = self.tokenizer.convert_tokens_to_ids("<|speech_end|>")
+            
+        if hasattr(self.tokenizer, 'speech_pad_id'):
+            self.speech_pad_id = self.tokenizer.speech_pad_id
+        else:
+            self.speech_pad_id = self.tokenizer.convert_tokens_to_ids("<|speech_pad|>")
+            
+        if hasattr(self.tokenizer, 'pad_id'):
+            self.pad_id = self.tokenizer.pad_id
+        elif hasattr(self.tokenizer, 'pad_token_id'):
+            self.pad_id = self.tokenizer.pad_token_id
+        else:
+            self.pad_id = self.tokenizer.convert_tokens_to_ids("<|endoftext|>")
+        
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+        """
+        Load processor from a pretrained model path.
+        
+        Args:
+            pretrained_model_name_or_path: Path to the pretrained model
+            **kwargs: Additional keyword arguments
+            
+        Returns:
+            VibeVoiceASRProcessor: The loaded processor
+        """
+        import json
+        from transformers.utils import cached_file
+        from vibevoice.modular.modular_vibevoice_text_tokenizer import VibeVoiceASRTextTokenizerFast
+        
+        # Try to load configuration
+        config_path = os.path.join(pretrained_model_name_or_path, "preprocessor_config.json")
+        config = {}
+        
+        if os.path.exists(config_path):
+            with open(config_path, 'r') as f:
+                config = json.load(f)
+        else:
+            try:
+                config_file = cached_file(
+                    pretrained_model_name_or_path,
+                    "preprocessor_config.json",
+                    **kwargs
+                )
+                with open(config_file, 'r') as f:
+                    config = json.load(f)
+            except Exception as e:
+                logger.warning(f"Could not load preprocessor_config.json: {e}")
+                logger.warning("Using default configuration")
+        
+        # Extract parameters
+        speech_tok_compress_ratio = config.get("speech_tok_compress_ratio", 3200)
+        target_sample_rate = config.get("target_sample_rate", 24000)
+        normalize_audio = config.get("normalize_audio", True)
+        
+        # Load tokenizer
+        language_model_pretrained_name = config.get("language_model_pretrained_name", None) or kwargs.pop("language_model_pretrained_name", "Qwen/Qwen2.5-1.5B")
+        logger.info(f"Loading tokenizer from {language_model_pretrained_name}")
+        
+        if 'qwen' in language_model_pretrained_name.lower():
+            tokenizer = VibeVoiceASRTextTokenizerFast.from_pretrained(
+                language_model_pretrained_name,
+                **kwargs
+            )
+        else:
+            raise ValueError(f"Unsupported tokenizer type for {language_model_pretrained_name}")
+        
+        # Load audio processor
+        audio_processor = VibeVoiceTokenizerProcessor(
+            sampling_rate=target_sample_rate,
+            normalize_audio=normalize_audio,
+            target_dB_FS=config.get("target_dB_FS", -25),
+            eps=config.get("eps", 1e-6),
+        )
+        
+        return cls(
+            tokenizer=tokenizer,
+            audio_processor=audio_processor,
+            speech_tok_compress_ratio=speech_tok_compress_ratio,
+            target_sample_rate=target_sample_rate,
+            normalize_audio=normalize_audio,
+        )
+    
+    def save_pretrained(self, save_directory: Union[str, os.PathLike], **kwargs):
+        """
+        Save processor configuration to a directory.
+        
+        Args:
+            save_directory: Directory to save the configuration
+            **kwargs: Additional keyword arguments
+        """
+        import json
+        
+        os.makedirs(save_directory, exist_ok=True)
+        
+        # Save processor configuration
+        processor_config = {
+            "processor_class": "VibeVoiceASRProcessor",
+            "speech_tok_compress_ratio": self.speech_tok_compress_ratio,
+            "target_sample_rate": self.target_sample_rate,
+            "normalize_audio": self.normalize_audio,
+            "target_dB_FS": -25,
+            "eps": 1e-6,
+        }
+        
+        config_path = os.path.join(save_directory, "preprocessor_config.json")
+        with open(config_path, 'w') as f:
+            json.dump(processor_config, f, indent=2)
+        
+        logger.info(f"Processor configuration saved in {config_path}")
+    
+    def __call__(
+        self,
+        audio: Optional[Union[str, np.ndarray, torch.Tensor, List[Union[str, np.ndarray, torch.Tensor]]]] = None,
+        sampling_rate: Optional[int] = None,
+        return_tensors: Optional[Union[str, TensorType]] = None,
+        padding: bool = True,
+        max_length: Optional[int] = None,
+        truncation: bool = False,
+        add_generation_prompt: bool = True,
+        use_streaming: bool = True,
+        context_info: Optional[str] = None,
+        **kwargs
+    ) -> BatchEncoding:
+        """
+        Process audio input for ASR model.
+        
+        Args:
+            audio: Audio input(s). Can be:
+                - str: Path to audio file
+                - np.ndarray: Audio array
+                - torch.Tensor: Audio tensor
+                - List of the above for batch processing
+            sampling_rate: Sampling rate of input audio
+            return_tensors: Output format ('pt' for PyTorch, 'np' for NumPy)
+            padding: Whether to pad batch inputs
+            max_length: Maximum sequence length
+            truncation: Whether to truncate long sequences
+            add_generation_prompt: Whether to add generation prompt for inference
+            use_streaming: Whether to use streaming mode (True by default, auto False if <60s)
+            context_info: Optional context information (e.g., hotwords, metadata) to help transcription
+            
+        Returns:
+            BatchEncoding with:
+                - input_ids: Token IDs for the model
+                - attention_mask: Attention mask
+                - acoustic_input_mask: Mask indicating speech token positions
+                - speech_tensors: Processed speech features
+                - speech_masks: Valid speech masks
+                - vae_tok_seqlens: Length of each speech segment in tokens
+        """
+        if audio is None:
+            raise ValueError("Audio input is required for ASR processing")
+        
+        # Handle single vs batch input
+        if isinstance(audio, list):
+            is_batched = True
+            audio_list = audio
+        else:
+            is_batched = False
+            audio_list = [audio]
+        
+        # Process each audio input
+        all_encodings = []
+        for audio_input in audio_list:
+            encoding = self._process_single_audio(
+                audio_input,
+                sampling_rate=sampling_rate,
+                add_generation_prompt=add_generation_prompt,
+                use_streaming=use_streaming,
+                context_info=context_info,
+            )
+            all_encodings.append(encoding)
+        
+        # Combine into batch
+        batch_encoding = self._batch_encode(
+            all_encodings,
+            padding=padding,
+            max_length=max_length,
+            truncation=truncation,
+            return_tensors=return_tensors,
+        )
+        
+        return batch_encoding
+    
+    def _process_single_audio(
+        self,
+        audio: Union[str, np.ndarray, torch.Tensor],
+        sampling_rate: Optional[int] = None,
+        add_generation_prompt: bool = True,
+        use_streaming: bool = True,
+        context_info: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """
+        Process a single audio input.
+        
+        Args:
+            audio: Single audio input
+            sampling_rate: Audio sampling rate
+            add_generation_prompt: Whether to add generation prompt
+            context_info: Optional context information (e.g., hotwords, metadata) to help transcription
+            
+        Returns:
+            Dictionary with processed tokens and audio features
+        """
+        # Process audio through audio processor
+        if isinstance(audio, str):
+            # Load from file using ffmpeg for better format support
+            if HAS_FFMPEG_UTILS:
+                try:
+                    audio_array, file_sr = load_audio_use_ffmpeg(audio, resample=False)
+                except Exception as e:
+                    # Fall back to soundfile if ffmpeg fails
+                    warnings.warn(f"ffmpeg loading failed, falling back to soundfile: {e}")
+                    import soundfile as sf
+                    audio_array, file_sr = sf.read(audio)
+                    if audio_array.ndim > 1:
+                        audio_array = audio_array.mean(axis=1)  # Convert to mono
+            else:
+                import soundfile as sf
+                audio_array, file_sr = sf.read(audio)
+                if audio_array.ndim > 1:
+                    audio_array = audio_array.mean(axis=1)  # Convert to mono
+            
+            # Resample if needed
+            if file_sr != self.target_sample_rate:
+                import librosa
+                audio_array = librosa.resample(
+                    audio_array, 
+                    orig_sr=file_sr, 
+                    target_sr=self.target_sample_rate
+                )
+        elif isinstance(audio, torch.Tensor):
+            audio_array = audio.cpu().numpy()
+            if audio_array.ndim > 1:
+                audio_array = audio_array.squeeze()
+        else:
+            audio_array = np.array(audio, dtype=np.float32)
+            if audio_array.ndim > 1:
+                audio_array = audio_array.squeeze()
+        
+        # Ensure float32
+        audio_array = audio_array.astype(np.float32)
+        
+        # Normalize if needed
+        if self.normalize_audio and self.audio_normalizer:
+            audio_array = self.audio_normalizer(audio_array)
+        
+        # Calculate audio duration
+        audio_duration = len(audio_array) / self.target_sample_rate
+        
+        # Auto-disable streaming for short audio (<60s)
+        if use_streaming and audio_duration < 60.0:
+            use_streaming = False
+        
+        # Calculate token length based on streaming mode
+        # Non-streaming: uses ceil (encoder adds extra_padding for stride alignment)
+        # Streaming: uses floor (segments processed independently, no global alignment)
+        # if use_streaming:
+        #     vae_tok_len = len(audio_array) // self.speech_tok_compress_ratio
+        # else:
+        vae_tok_len = math.ceil(len(audio_array) / self.speech_tok_compress_ratio)
+        
+        # Build token sequence following training format
+        # 1. System prompt - use apply_chat_template then encode like in training
+        system_prompt_text = self.tokenizer.apply_chat_template(
+            [{"role": "system", "content": SYSTEM_PROMPT}],
+            tokenize=False
+        )
+        system_tokens = self.tokenizer.encode(system_prompt_text)
+        
+        # 2. User input with speech tokens
+        # Build speech placeholder string
+        sp_start_token = self.tokenizer.convert_ids_to_tokens(self.speech_start_id)
+        sp_pad_token = self.tokenizer.convert_ids_to_tokens(self.speech_pad_id)
+        sp_end_token = self.tokenizer.convert_ids_to_tokens(self.speech_end_id)
+        
+        # User suffix with audio duration info
+        show_keys = ['Start time', 'End time', 'Speaker ID', 'Content']
+        if context_info and context_info.strip():
+            user_suffix = f"This is a {audio_duration:.2f} seconds audio, with extra info: {context_info.strip()}\n\nPlease transcribe it with these keys: " + ", ".join(show_keys)
+        else:
+            user_suffix = f"This is a {audio_duration:.2f} seconds audio, please transcribe it with these keys: " + ", ".join(show_keys)
+        
+        user_input_string = ''.join(
+            [sp_start_token] + [sp_pad_token] * vae_tok_len + [sp_end_token]
+        ) + '\n' + user_suffix
+        
+        user_tokens = self.tokenizer.apply_chat_template(
+            [{"role": "user", "content": user_input_string}],
+            tokenize=True
+        )
+        
+        # Combine tokens
+        full_tokens = system_tokens + user_tokens
+        
+        # Create acoustic input mask
+        acoustic_input_mask = [1 if token == self.speech_pad_id else 0 for token in full_tokens]
+        
+        return {
+            "input_ids": full_tokens,
+            "acoustic_input_mask": acoustic_input_mask,
+            "speech": audio_array,
+            "vae_tok_len": vae_tok_len,
+        }
+    
+    def _batch_encode(
+        self,
+        encodings: List[Dict[str, Any]],
+        padding: bool = True,
+        max_length: Optional[int] = None,
+        truncation: bool = False,
+        return_tensors: Optional[str] = None,
+    ) -> BatchEncoding:
+        """
+        Combine multiple encodings into a batch.
+        
+        Args:
+            encodings: List of encoded samples
+            padding: Whether to pad sequences
+            max_length: Maximum sequence length
+            truncation: Whether to truncate
+            return_tensors: Output format
+            
+        Returns:
+            BatchEncoding with batched data
+        """
+        # Extract components
+        input_ids_list = [enc["input_ids"] for enc in encodings]
+        acoustic_masks_list = [enc["acoustic_input_mask"] for enc in encodings]
+        speech_list = [enc["speech"] for enc in encodings]
+        vae_tok_lens = [enc["vae_tok_len"] for enc in encodings]
+        
+        # Determine max length for padding
+        if padding:
+            if max_length is not None:
+                target_length = max_length
+            else:
+                target_length = max(len(ids) for ids in input_ids_list)
+            
+            # Pad sequences
+            padded_input_ids = []
+            padded_acoustic_masks = []
+            attention_masks = []
+            
+            for input_ids, acoustic_mask in zip(input_ids_list, acoustic_masks_list):
+                # Truncate if needed
+                if truncation and len(input_ids) > target_length:
+                    input_ids = input_ids[:target_length]
+                    acoustic_mask = acoustic_mask[:target_length]
+                
+                # Pad sequences to left (for autoregressive generation)
+                padding_length = target_length - len(input_ids)
+                padded_ids = [self.pad_id] * padding_length + input_ids
+                padded_acoustic = [0] * padding_length + acoustic_mask
+                attention_mask = [0] * padding_length + [1] * len(input_ids)
+                
+                padded_input_ids.append(padded_ids)
+                padded_acoustic_masks.append(padded_acoustic)
+                attention_masks.append(attention_mask)
+            
+            input_ids_list = padded_input_ids
+            acoustic_masks_list = padded_acoustic_masks
+        else:
+            attention_masks = [[1] * len(ids) for ids in input_ids_list]
+        
+        # Process speech tensors - raw audio is 1D, so we keep it as is
+        max_speech_length = max(len(s) for s in speech_list)
+        padded_speeches = np.zeros((len(speech_list), max_speech_length), dtype=np.float32)
+        speech_masks = np.zeros((len(speech_list), max(vae_tok_lens)), dtype=bool)
+        
+        for i, (speech, vae_len) in enumerate(zip(speech_list, vae_tok_lens)):
+            padded_speeches[i, :len(speech)] = speech
+            speech_masks[i, :vae_len] = True
+        
+        # Create batch encoding
+        batch_encoding = BatchEncoding()
+        
+        if return_tensors == "pt":
+            batch_encoding["input_ids"] = torch.tensor(input_ids_list, dtype=torch.long)
+            batch_encoding["attention_mask"] = torch.tensor(attention_masks, dtype=torch.long)
+            batch_encoding["acoustic_input_mask"] = torch.tensor(acoustic_masks_list, dtype=torch.bool)
+            batch_encoding["speech_tensors"] = torch.tensor(padded_speeches, dtype=torch.float32)
+            batch_encoding["speech_masks"] = torch.tensor(speech_masks, dtype=torch.bool)
+            # Note: vae_tok_seqlens and speech_type are not included as they are not model inputs
+        else:
+            batch_encoding["input_ids"] = input_ids_list if len(input_ids_list) > 1 else input_ids_list[0]
+            batch_encoding["attention_mask"] = attention_masks if len(attention_masks) > 1 else attention_masks[0]
+            batch_encoding["acoustic_input_mask"] = acoustic_masks_list if len(acoustic_masks_list) > 1 else acoustic_masks_list[0]
+            batch_encoding["speech_tensors"] = padded_speeches if len(padded_speeches) > 1 else padded_speeches[0]
+            batch_encoding["speech_masks"] = speech_masks if len(speech_masks) > 1 else speech_masks[0]
+        
+        return batch_encoding
+    
+    def batch_decode(self, *args, **kwargs):
+        """
+        Decode batch of token IDs to text.
+        Forwards to tokenizer's batch_decode method.
+        """
+        return self.tokenizer.batch_decode(*args, **kwargs)
+    
+    def decode(self, *args, **kwargs):
+        """
+        Decode token IDs to text.
+        Forwards to tokenizer's decode method.
+        """
+        return self.tokenizer.decode(*args, **kwargs)
+    
+    def post_process_transcription(self, text: str) -> List[Dict[str, Any]]:
+        """
+        Post-process the generated transcription text to extract structured data.
+        
+        Args:
+            text: Generated text from the model
+            
+        Returns:
+            List of dictionaries with transcription segments
+        """
+        try:
+            # Try to parse as JSON
+            if "```json" in text:
+                # Extract JSON from markdown code block
+                json_start = text.find("```json") + 7
+                json_end = text.find("```", json_start)
+                json_str = text[json_start:json_end].strip()
+            else:
+                # Try to find JSON array or object
+                json_start = text.find("[")
+                if json_start == -1:
+                    json_start = text.find("{")
+                if json_start != -1:
+                    # Find matching closing bracket
+                    bracket_count = 0
+                    json_end = json_start
+                    for i in range(json_start, len(text)):
+                        if text[i] in "[{":
+                            bracket_count += 1
+                        elif text[i] in "]}":
+                            bracket_count -= 1
+                            if bracket_count == 0:
+                                json_end = i + 1
+                                break
+                    json_str = text[json_start:json_end]
+                else:
+                    json_str = text
+            
+            # Parse JSON
+            result = json.loads(json_str)
+            
+            # Ensure it's a list
+            if isinstance(result, dict):
+                result = [result]
+            
+            # Validate and clean up the result
+            cleaned_result = []
+            for item in result:
+                if isinstance(item, dict):
+                    cleaned_item = {}
+                    # Map keys to expected format
+                    key_mapping = {
+                        "Start time": "start_time",
+                        "Start": "start_time",
+                        "End time": "end_time",
+                        "End": "end_time",
+                        "Speaker ID": "speaker_id",
+                        "Speaker": "speaker_id",
+                        "Content": "text",
+                    }
+                    for key, mapped_key in key_mapping.items():
+                        if key in item:
+                            cleaned_item[mapped_key] = item[key]
+                    
+                    if cleaned_item:
+                        cleaned_result.append(cleaned_item)
+            
+            return cleaned_result
+            
+        except json.JSONDecodeError as e:
+            logger.warning(f"Failed to parse JSON from transcription: {e}")
+            logger.debug(f"Raw text: {text}")
+            return []
+        except Exception as e:
+            logger.warning(f"Error post-processing transcription: {e}")
+            return []
+    
+    @property
+    def model_input_names(self):
+        """Return the list of inputs accepted by the model."""
+        return ["input_ids", "attention_mask", "acoustic_input_mask", "speech_tensors", "speech_masks"]
+
+__all__ = ["VibeVoiceASRProcessor"]

vibevoice/processor/vibevoice_processor.py

@@ -0,0 +1,692 @@
+import math
+import warnings
+from typing import List, Optional, Union, Dict, Any, Tuple
+import os
+import re
+
+import numpy as np
+import torch
+
+from transformers.tokenization_utils_base import BatchEncoding, PaddingStrategy, PreTokenizedInput, TextInput, TruncationStrategy
+from transformers.utils import TensorType, logging
+from .vibevoice_tokenizer_processor import AudioNormalizer
+
+logger = logging.get_logger(__name__)
+
+
+class VibeVoiceProcessor:
+    r"""
+    Constructs a VibeVoice processor which wraps a VibeVoice tokenizer and audio processor into a single processor.
+
+    [`VibeVoiceProcessor`] offers all the functionalities of [`VibeVoiceTokenizer`] and [`VibeVoiceTokenizerProcessor`]. 
+    See the [`~VibeVoiceProcessor.__call__`] and [`~VibeVoiceProcessor.decode`] for more information.
+
+    Args:
+        tokenizer (`VibeVoiceTextTokenizer` or `VibeVoiceTextTokenizerFast`):
+            The tokenizer for text processing.
+        audio_processor (`VibeVoiceTokenizerProcessor`):
+            The audio processor for speech processing.
+        speech_tok_compress_ratio (`int`, *optional*, defaults to 3200):
+            The compression ratio for speech tokenization.
+        db_normalize (`bool`, *optional*, defaults to True):
+            Whether to apply decibel normalization to audio inputs.
+    """
+
+    def __init__(self, tokenizer=None, audio_processor=None, speech_tok_compress_ratio=3200, db_normalize=True, **kwargs):
+        self.tokenizer = tokenizer
+        self.audio_processor = audio_processor
+        self.speech_tok_compress_ratio = speech_tok_compress_ratio
+        self.db_normalize = db_normalize
+        self.audio_normalizer = AudioNormalizer() if db_normalize else None
+        self.system_prompt = " Transform the text provided by various speakers into speech output, utilizing the distinct voice of each respective speaker.\n"
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+        """
+        Instantiate a VibeVoiceProcessor from a pretrained VibeVoice processor.
+
+        Args:
+            pretrained_model_name_or_path (`str` or `os.PathLike`):
+                This can be either:
+                - a string, the *model id* of a pretrained model
+                - a path to a *directory* containing processor config
+
+        Returns:
+            [`VibeVoiceProcessor`]: The processor object instantiated from pretrained model.
+        """
+        import os
+        import json
+        from transformers.utils import cached_file
+        from .vibevoice_tokenizer_processor import VibeVoiceTokenizerProcessor
+        from vibevoice.modular.modular_vibevoice_text_tokenizer import (
+            VibeVoiceTextTokenizer, 
+            VibeVoiceTextTokenizerFast
+        )
+        
+        # Try to load from local path first, then from HF hub
+        config_path = os.path.join(pretrained_model_name_or_path, "preprocessor_config.json")
+        config = None
+        
+        if os.path.exists(config_path):
+            # Local path exists
+            with open(config_path, 'r') as f:
+                config = json.load(f)
+        else:
+            # Try to load from HF hub
+            try:
+                config_file = cached_file(
+                    pretrained_model_name_or_path,
+                    "preprocessor_config.json",
+                    **kwargs
+                )
+                with open(config_file, 'r') as f:
+                    config = json.load(f)
+            except Exception as e:
+                logger.warning(f"Could not load preprocessor_config.json from {pretrained_model_name_or_path}: {e}")
+                logger.warning("Using default configuration")
+                config = {
+                    "speech_tok_compress_ratio": 3200,
+                    "db_normalize": True,
+                }
+        
+        # Extract main processor parameters
+        speech_tok_compress_ratio = config.get("speech_tok_compress_ratio", 3200)
+        db_normalize = config.get("db_normalize", True)
+        
+        # Load tokenizer - try from model path first, then fall back to Qwen        
+        language_model_pretrained_name = config.get("language_model_pretrained_name", None) or kwargs.pop("language_model_pretrained_name", "Qwen/Qwen2.5-1.5B")
+        logger.info(f"Loading tokenizer from {language_model_pretrained_name}")
+        if 'qwen' in language_model_pretrained_name.lower():
+            tokenizer = VibeVoiceTextTokenizerFast.from_pretrained(
+                language_model_pretrained_name,
+                **kwargs
+            )
+        else:
+            raise ValueError(f"Unsupported tokenizer type for {language_model_pretrained_name}. Supported types: Qwen, Llama, Gemma.")
+        
+        # Load audio processor
+        if "audio_processor" in config:
+            # Create audio processor from config
+            audio_config = config["audio_processor"]
+            audio_processor = VibeVoiceTokenizerProcessor(
+                sampling_rate=audio_config.get("sampling_rate", 24000),
+                normalize_audio=audio_config.get("normalize_audio", True),
+                target_dB_FS=audio_config.get("target_dB_FS", -25),
+                eps=audio_config.get("eps", 1e-6),
+            )
+        else:
+            # Create default audio processor
+            audio_processor = VibeVoiceTokenizerProcessor()
+        
+        # Create and return the processor
+        return cls(
+            tokenizer=tokenizer,
+            audio_processor=audio_processor,
+            speech_tok_compress_ratio=speech_tok_compress_ratio,
+            db_normalize=db_normalize,
+        )
+    
+    def save_pretrained(self, save_directory: Union[str, os.PathLike], **kwargs):
+        """
+        Save a processor to a directory, so that it can be re-loaded using the
+        [`~VibeVoiceProcessor.from_pretrained`] class method.
+
+        Args:
+            save_directory (`str` or `os.PathLike`):
+                Directory where the processor will be saved.
+        """
+        import os
+        import json
+        
+        os.makedirs(save_directory, exist_ok=True)
+        
+        # Save processor configuration
+        processor_config = {
+            "processor_class": "VibeVoiceProcessor",
+            "speech_tok_compress_ratio": self.speech_tok_compress_ratio,
+            "db_normalize": self.db_normalize,
+            "audio_processor": {
+                "feature_extractor_type": "VibeVoiceTokenizerProcessor",
+                "sampling_rate": getattr(self.audio_processor, 'sampling_rate', 24000),
+                "normalize_audio": getattr(self.audio_processor, 'normalize_audio', True),
+                "target_dB_FS": getattr(self.audio_processor, 'target_dB_FS', -25),
+                "eps": getattr(self.audio_processor, 'eps', 1e-6),
+            }
+        }
+        
+        config_path = os.path.join(save_directory, "preprocessor_config.json")
+        with open(config_path, 'w') as f:
+            json.dump(processor_config, f, indent=2)
+        
+        logger.info(f"Processor configuration saved in {config_path}")
+    
+    def __call__(
+        self,
+        text: Optional[Union[str, List[str], TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]]] = None,
+        voice_samples: Optional[Union[List[Union[str, np.ndarray]], List[List[Union[str, np.ndarray]]]]] = None,
+        padding: Union[bool, str, PaddingStrategy] = True,
+        truncation: Union[bool, str, TruncationStrategy] = False,
+        max_length: Optional[int] = None,
+        return_tensors: Optional[Union[str, TensorType]] = None,
+        return_attention_mask: bool = True,
+        **kwargs,
+    ) -> BatchEncoding:
+        """
+        Main method to process one or more podcast scripts with optional voice samples.
+
+        Args:
+            text (`str`, `List[str]`):
+                The input text(s) to process. Can be:
+                - A single script string
+                - A list of script strings for batch processing
+                - A path to a .json or .txt file
+                - A list of paths
+            voice_samples (`List[Union[str, np.ndarray]]`, `List[List[Union[str, np.ndarray]]]`, *optional*):
+                Voice samples for each script. Can be:
+                - A list of samples for a single script
+                - A list of lists for batch processing
+            padding (`bool`, `str` or `PaddingStrategy`, defaults to `True`):
+                Whether to pad sequences to the same length
+            truncation (`bool`, `str` or `TruncationStrategy`, defaults to `False`):
+                Whether to truncate sequences
+            max_length (`int`, *optional*):
+                Maximum length of the returned sequences
+            return_tensors (`str` or `TensorType`, *optional*):
+                If set, will return tensors of a particular framework
+            return_attention_mask (`bool`, defaults to `True`):
+                Whether to return the attention mask
+
+        Returns:
+            `BatchEncoding`: A BatchEncoding with the following fields:
+                - **input_ids** -- List of token id sequences or tensor
+                - **attention_mask** -- List of attention masks or tensor
+                - **speech_tensors** -- Padded speech inputs (if voice_samples provided)
+                - **speech_masks** -- Speech masks (if voice_samples provided)
+                - **speech_input_mask** -- Boolean masks indicating speech token positions
+        """
+        # Handle single vs batch input
+        if isinstance(text, str) or (isinstance(text, list) and len(text) > 0 and not isinstance(text[0], str)):
+            # Single input
+            texts = [text]
+            is_batched = False
+        else:
+            # Batch input
+            texts = text
+            is_batched = True
+            
+        # Handle voice samples
+        if voice_samples is not None:
+            if not is_batched or (isinstance(voice_samples[0], (str, np.ndarray))):
+                # Single set of voice samples
+                voice_samples_list = [voice_samples]
+            else:
+                # Batch of voice samples
+                voice_samples_list = voice_samples
+        else:
+            voice_samples_list = [None] * len(texts)
+        
+        # Process each input
+        all_encodings = []
+        for text_input, voice_input in zip(texts, voice_samples_list):
+            encoding = self._process_single(text_input, voice_input)
+            all_encodings.append(encoding)
+            
+        # Combine batch
+        batch_encoding = self._batch_encode(
+            all_encodings,
+            padding=padding,
+            truncation=truncation,
+            max_length=max_length,
+            return_tensors=return_tensors,
+            return_attention_mask=return_attention_mask,
+        )
+        
+        return batch_encoding
+    
+    def _process_single(
+        self,
+        text: Union[str, TextInput],
+        voice_samples: Optional[List[Union[str, np.ndarray]]] = None,
+    ) -> Dict[str, Any]:
+        """Process a single podcast script."""
+        # Determine if text is a file path or direct script
+        script = None
+        if isinstance(text, str):
+            # Check if it's a file path
+            if text.endswith('.json') and os.path.exists(text):
+                script = self._convert_json_to_script(text)
+            elif text.endswith('.txt') and os.path.exists(text):
+                script = self._convert_text_to_script(text)
+            else:
+                # Assume it's the script content directly
+                script = text
+        
+        if script is None:
+            raise ValueError(f"Could not process input text: {text}")
+        
+        # Parse the script
+        parsed_lines = self._parse_script(script)
+        all_speakers = list(set(speaker_id for speaker_id, _ in parsed_lines))
+        
+        # Create system prompt
+        # system_tokens = self.tokenizer.encode(self.system_prompt, add_special_tokens=False)
+        system_tokens = self.tokenizer.encode(self.system_prompt)
+        
+        # Process voice samples if provided
+        if voice_samples:
+            voice_tokens, voice_speech_inputs, voice_speech_masks = self._create_voice_prompt(voice_samples[:len(all_speakers)])
+        else:
+            voice_tokens, voice_speech_inputs, voice_speech_masks = [], [], []
+        
+        # Build full token sequence
+        full_tokens = system_tokens + voice_tokens
+        speech_input_mask = [False] * len(system_tokens) + voice_speech_masks
+        
+        # Add text input section
+        full_tokens += self.tokenizer.encode(' Text input:\n', add_special_tokens=False)
+        speech_input_mask += [False] * len(self.tokenizer.encode(' Text input:\n', add_special_tokens=False))
+        
+        for speaker_id, speaker_text in parsed_lines:
+            speaker_text_tokens = self.tokenizer.encode(f" Speaker {speaker_id}:{speaker_text}\n", add_special_tokens=False)
+            full_tokens += speaker_text_tokens
+            speech_input_mask += [False] * len(speaker_text_tokens)
+        
+        # Add speech output section
+        full_tokens += self.tokenizer.encode(' Speech output:\n', add_special_tokens=False) + [self.tokenizer.speech_start_id]
+        speech_input_mask += [False] * (len(self.tokenizer.encode(' Speech output:\n', add_special_tokens=False)) + 1)
+        
+        return {
+            "input_ids": full_tokens,
+            "speech_inputs": voice_speech_inputs if voice_speech_inputs else None,
+            "speech_input_mask": speech_input_mask,
+            "parsed_script": parsed_lines,
+            "all_speakers": all_speakers,
+        }
+    
+    def _batch_encode(
+        self,
+        encodings: List[Dict[str, Any]],
+        padding: Union[bool, str, PaddingStrategy] = True,
+        truncation: Union[bool, str, TruncationStrategy] = False,
+        max_length: Optional[int] = None,
+        return_tensors: Optional[Union[str, TensorType]] = None,
+        return_attention_mask: bool = True,
+    ) -> BatchEncoding:
+        """Combine multiple encodings into a batch with padding."""
+        # Extract input_ids and create attention_mask
+        input_ids_list = [enc["input_ids"] for enc in encodings]
+        speech_input_masks_list = [enc["speech_input_mask"] for enc in encodings]
+        
+        # Determine padding strategy
+        if isinstance(padding, bool):
+            padding_strategy = PaddingStrategy.LONGEST if padding else PaddingStrategy.DO_NOT_PAD
+        elif isinstance(padding, str):
+            padding_strategy = PaddingStrategy(padding)
+        else:
+            padding_strategy = padding
+            
+        # Apply padding to input_ids
+        if padding_strategy != PaddingStrategy.DO_NOT_PAD:
+            if padding_strategy == PaddingStrategy.LONGEST:
+                max_len = max(len(ids) for ids in input_ids_list)
+            elif padding_strategy == PaddingStrategy.MAX_LENGTH and max_length is not None:
+                max_len = max_length
+            else:
+                max_len = max(len(ids) for ids in input_ids_list)
+                
+            # Pad sequences
+            padded_input_ids = []
+            attention_masks = []
+            padded_speech_input_masks = []
+            
+            for input_ids, speech_mask in zip(input_ids_list, speech_input_masks_list):
+                # Truncate if needed
+                if truncation and len(input_ids) > max_len:
+                    input_ids = input_ids[:max_len]
+                    speech_mask = speech_mask[:max_len]
+                    
+                # Pad
+                padding_length = max_len - len(input_ids)
+                # padded_ids = [self.tokenizer.pad_token_id] * padding_length + input_ids
+                padded_ids = [self.tokenizer.pad_id] * padding_length + input_ids
+                attention_mask = [0] * padding_length + [1] * len(input_ids)
+                padded_speech_mask = [False] * padding_length + speech_mask
+                
+                padded_input_ids.append(padded_ids)
+                attention_masks.append(attention_mask)
+                padded_speech_input_masks.append(padded_speech_mask)
+                
+            input_ids_list = padded_input_ids
+            speech_input_masks_list = padded_speech_input_masks
+        else:
+            # No padding, just create attention masks
+            attention_masks = [[1] * len(ids) for ids in input_ids_list] if return_attention_mask else None
+            
+        # Process speech inputs
+        all_speech_inputs = []
+        has_speech = False
+        for enc in encodings:
+            if enc["speech_inputs"] is not None:
+                all_speech_inputs.extend(enc["speech_inputs"])
+                has_speech = True
+                
+        # Prepare batch encoding
+        batch_encoding = BatchEncoding()
+        
+        # Handle tensor conversion
+        if return_tensors is not None:
+            batch_encoding["input_ids"] = torch.tensor(input_ids_list, dtype=torch.long)
+            if return_attention_mask and attention_masks is not None:
+                batch_encoding["attention_mask"] = torch.tensor(attention_masks, dtype=torch.long)
+            batch_encoding["speech_input_mask"] = torch.tensor(speech_input_masks_list, dtype=torch.bool)
+        else:
+            batch_encoding["input_ids"] = input_ids_list
+            if return_attention_mask and attention_masks is not None:
+                batch_encoding["attention_mask"] = attention_masks
+            batch_encoding["speech_input_mask"] = speech_input_masks_list
+            
+        # Process speech tensors if present
+        if has_speech:
+            speech_dict = self.prepare_speech_inputs(
+                all_speech_inputs,
+                return_tensors=return_tensors,
+            )
+            batch_encoding["speech_tensors"] = speech_dict["padded_speeches"]
+            batch_encoding["speech_masks"] = speech_dict["speech_masks"]
+        else:
+            batch_encoding["speech_tensors"] = None
+            batch_encoding["speech_masks"] = None
+            
+        # Add metadata
+        batch_encoding["parsed_scripts"] = [enc["parsed_script"] for enc in encodings]
+        batch_encoding["all_speakers_list"] = [enc["all_speakers"] for enc in encodings]
+        
+        return batch_encoding
+
+    def _create_voice_prompt(
+        self, 
+        speaker_samples: List[Union[str, np.ndarray]]
+    ) -> Tuple[List[int], List[np.ndarray], List[bool]]:
+        """
+        Create voice prompt tokens and process audio samples.
+        
+        Returns:
+            tuple: (voice_tokens, voice_speech_inputs, voice_speech_masks)
+        """
+        vae_token_id = self.tokenizer.speech_diffusion_id
+        
+        voice_full_tokens = self.tokenizer.encode(' Voice input:\n', add_special_tokens=False)
+        voice_speech_inputs = []
+        voice_speech_masks = [False] * len(voice_full_tokens)
+        
+        for speaker_id, speaker_audio in enumerate(speaker_samples):
+            prefix_tokens = self.tokenizer.encode(f" Speaker {speaker_id}:", add_special_tokens=False)
+            
+            # Process audio
+            if isinstance(speaker_audio, str):
+                # Load audio from file
+                wav = self.audio_processor._load_audio_from_path(speaker_audio)
+            else:
+                wav = np.array(speaker_audio, dtype=np.float32)
+            
+            # Apply normalization if needed
+            if self.db_normalize and self.audio_normalizer:
+                wav = self.audio_normalizer(wav)
+            
+            # Calculate token length based on compression ratio
+            # if speaker_audio.endswith('.pt') or speaker_audio.endswith('.npy'):
+            #     vae_tok_len = wav.shape[0]
+            # else:
+            vae_tok_len = math.ceil(wav.shape[0] / self.speech_tok_compress_ratio)
+            
+            # Build tokens and masks
+            speaker_tokens = (prefix_tokens + 
+                            [self.tokenizer.speech_start_id] + 
+                            [vae_token_id] * vae_tok_len + 
+                            [self.tokenizer.speech_end_id] + 
+                            self.tokenizer.encode('\n', add_special_tokens=False))
+            
+            vae_input_mask = ([False] * len(prefix_tokens) + 
+                            [False] + 
+                            [True] * vae_tok_len + 
+                            [False] + 
+                            [False])
+            
+            voice_full_tokens.extend(speaker_tokens)
+            voice_speech_masks.extend(vae_input_mask)
+            voice_speech_inputs.append(wav)
+            
+        return voice_full_tokens, voice_speech_inputs, voice_speech_masks
+
+    def prepare_speech_inputs(
+        self,
+        speech_inputs: List[np.ndarray],
+        return_tensors: Optional[Union[str, TensorType]] = None,
+        device: Optional[Union[str, torch.device]] = None,
+        dtype: Optional[torch.dtype] = None,
+    ) -> Dict[str, Any]:
+        """
+        Prepare speech inputs for model consumption.
+        
+        Args:
+            speech_inputs: List of speech arrays
+            return_tensors: Output tensor type
+            device: Device to place tensors on
+            dtype: Data type for tensors
+            
+        Returns:
+            Dictionary with padded_speeches and speech_masks
+        """
+        if not speech_inputs:
+            return {"padded_speeches": None, "speech_masks": None}
+        
+        # Calculate sequence lengths
+        vae_tok_seqlens = [math.ceil(s.shape[0] / self.speech_tok_compress_ratio) for s in speech_inputs]
+        # vae_tok_seqlens = [math.ceil(s.shape[0] / self.speech_tok_compress_ratio) if s.ndim == 1 else s.shape[0] for s in speech_inputs]
+        max_speech_length = max(s.shape[0] for s in speech_inputs)
+        
+        # Pad speeches
+        if speech_inputs[0].ndim == 1:
+            padded_speeches = np.full((len(speech_inputs), max_speech_length), fill_value=0, dtype=np.float32)
+        else:
+            padded_speeches = np.full((len(speech_inputs), max_speech_length, speech_inputs[0].shape[-1]), fill_value=0, dtype=np.float32)
+        speech_masks = np.zeros((len(speech_inputs), max(vae_tok_seqlens)), dtype=np.bool_)
+        
+        for i, (speech, vae_tok_length) in enumerate(zip(speech_inputs, vae_tok_seqlens)):
+            padded_speeches[i, :len(speech)] = speech
+            speech_masks[i, :vae_tok_length] = True
+        
+        result = {
+            "padded_speeches": padded_speeches,
+            "speech_masks": speech_masks,
+        }
+        
+        # Convert to tensors if requested
+        if return_tensors == "pt":
+            result["padded_speeches"] = torch.tensor(padded_speeches, device=device, dtype=dtype or torch.float32)
+            result["speech_masks"] = torch.tensor(speech_masks, device=device, dtype=torch.bool)
+            
+        return result
+        
+    def _convert_json_to_script(self, json_file: str) -> str:
+        """
+        Convert JSON format to script format.
+        Expected JSON format:
+        [
+            {"speaker": "1", "text": "Hello everyone..."},
+            {"speaker": "2", "text": "Great to be here..."}
+        ]
+        """
+        import json
+        
+        with open(json_file, 'r', encoding='utf-8') as f:
+            data = json.load(f)
+        
+        if not isinstance(data, list):
+            raise ValueError("JSON file must contain a list of speaker entries")
+        
+        script_lines = []
+        for item in data:
+            if not isinstance(item, dict):
+                logger.warning(f"Skipping non-dict entry: {item}")
+                continue
+                
+            speaker = item.get('speaker')
+            text = item.get('text')
+            
+            if speaker is None or text is None:
+                logger.warning(f"Skipping entry missing speaker or text: {item}")
+                continue
+            
+            # Ensure speaker ID is valid
+            try:
+                speaker_id = int(speaker)
+            except (ValueError, TypeError):
+                logger.warning(f"Invalid speaker ID: {speaker}, skipping entry")
+                continue
+            
+            # Clean up text
+            text = text.strip()
+            if text:
+                script_lines.append(f"Speaker {speaker_id}: {text}")
+        
+        if not script_lines:
+            raise ValueError("No valid entries found in JSON file")
+            
+        return "\n".join(script_lines)
+
+    def _convert_text_to_script(self, text_file: str) -> str:
+        """
+        Convert text file to script format.
+        Handles multiple formats:
+        1. Already formatted as "Speaker X: text"
+        2. Plain text (assigns to Speaker 1)
+        
+        Handles edge cases like multiple colons in a line.
+        """
+        with open(text_file, 'r', encoding='utf-8') as f:
+            lines = f.readlines()
+        
+        script_lines = []
+        current_speaker = 1
+        
+        for line in lines:
+            line = line.strip()
+            if not line:
+                continue
+            
+            # Try to parse as "Speaker X: text" format
+            # Use regex to be more robust
+            speaker_match = re.match(r'^Speaker\s+(\d+)\s*:\s*(.*)$', line, re.IGNORECASE)
+            
+            if speaker_match:
+                speaker_id = int(speaker_match.group(1))
+                text = speaker_match.group(2).strip()
+                if text:
+                    script_lines.append(f"Speaker {speaker_id}: {text}")
+            else:
+                # Treat as plain text - assign to current speaker
+                script_lines.append(f"Speaker {current_speaker}: {line}")
+        
+        if not script_lines:
+            raise ValueError("No valid content found in text file")
+            
+        return "\n".join(script_lines)
+
+    def _parse_script(self, script: str) -> List[Tuple[int, str]]:
+        """Parse script into list of (speaker_id, text) tuples."""
+        lines = script.strip().split("\n")
+        parsed_lines = []
+        speaker_ids = []
+                
+        # First pass: parse all lines and collect speaker IDs
+        for line in lines:
+            if not line.strip():
+                continue
+                
+            # Use regex to handle edge cases like multiple colons
+            match = re.match(r'^Speaker\s+(\d+)\s*:\s*(.*)$', line.strip(), re.IGNORECASE)
+            
+            if match:
+                speaker_id = int(match.group(1))
+                text = ' ' + match.group(2).strip()
+                parsed_lines.append((speaker_id, text))
+                speaker_ids.append(speaker_id)
+            else:
+                logger.warning(f"Could not parse line: '{line}'")
+        
+        if not parsed_lines:
+            raise ValueError("No valid speaker lines found in script")
+        
+        # Check if we need to normalize speaker IDs (only if all are > 0)
+        min_speaker_id = min(speaker_ids)
+        if min_speaker_id > 0:
+            # Normalize to start from 0
+            normalized_lines = []
+            for speaker_id, text in parsed_lines:
+                normalized_lines.append((speaker_id - 1, text))
+            return normalized_lines
+        else:
+            # Keep original IDs
+            return parsed_lines
+
+    def _merge_inputs(self, text_inputs: BatchEncoding, audio_inputs: Dict) -> BatchEncoding:
+        """Merge text and audio inputs into a single BatchEncoding."""
+        # Start with text inputs
+        merged = BatchEncoding(text_inputs)
+        
+        # Add audio-specific fields
+        if "audio" in audio_inputs:
+            merged["speech_inputs"] = audio_inputs["audio"]
+        if "streaming" in audio_inputs:
+            merged["streaming"] = audio_inputs["streaming"]
+            
+        return merged
+
+    def batch_decode(self, *args, **kwargs):
+        """
+        This method forwards all its arguments to VibeVoiceTextTokenizer's [`~PreTrainedTokenizer.batch_decode`].
+        Please refer to the docstring of this method for more information.
+        """
+        return self.tokenizer.batch_decode(*args, **kwargs)
+
+    def decode(self, *args, **kwargs):
+        """
+        This method forwards all its arguments to VibeVoiceTextTokenizer's [`~PreTrainedTokenizer.decode`].
+        Please refer to the docstring of this method for more information.
+        """
+        return self.tokenizer.decode(*args, **kwargs)
+
+    @property
+    def model_input_names(self):
+        """
+        Return the list of inputs accepted by the model.
+        """
+        tokenizer_input_names = self.tokenizer.model_input_names
+        audio_processor_input_names = self.audio_processor.model_input_names
+        return list(dict.fromkeys(tokenizer_input_names + audio_processor_input_names + ["speech_inputs", "speech_input_mask"]))
+
+    def save_audio(self, 
+        audio: Union[torch.Tensor, np.ndarray, List[Union[torch.Tensor, np.ndarray]]],
+        output_path: str = "output.wav",
+        sampling_rate: Optional[int] = None,
+        normalize: bool = False,
+        batch_prefix: str = "audio_",
+    ) -> str:
+        """
+        Save audio data to a file.
+        Args:
+            audio (Union[torch.Tensor, np.ndarray, List[Union[torch.Tensor, np.ndarray]]]):
+                The audio data to save. Can be a single tensor/array or a list of them.
+            output_path (str, optional): Path to save the audio file. Defaults to "output.wav".
+            sampling_rate (int, optional): Sampling rate for the audio. If None, uses the processor's default.
+            normalize (bool, optional): Whether to normalize the audio before saving. Defaults to False.
+            batch_prefix (str, optional): Prefix for batch audio files. Defaults to "audio_".
+        Returns:
+            str: The path to the saved audio file.
+        """
+        return self.audio_processor.save_audio(audio, output_path=output_path, sampling_rate=sampling_rate, normalize=normalize, batch_prefix=batch_prefix)
+    
+__all__ = [
+    "VibeVoiceProcessor",
+]

vibevoice/processor/vibevoice_streaming_processor.py

@@ -0,0 +1,409 @@
+import math
+import warnings
+from typing import List, Optional, Union, Dict, Any, Tuple
+import os
+import re
+
+import numpy as np
+import torch
+
+from transformers.tokenization_utils_base import BatchEncoding, PaddingStrategy, PreTokenizedInput, TextInput, TruncationStrategy
+from transformers.utils import TensorType, logging
+from .vibevoice_tokenizer_processor import AudioNormalizer
+
+logger = logging.get_logger(__name__)
+
+
+class VibeVoiceStreamingProcessor:
+    r"""
+    Constructs a VibeVoice Streaming processor which wraps a VibeVoice tokenizer and audio processor into a single processor.
+
+    Args:
+        tokenizer (`VibeVoiceTextTokenizer` or `VibeVoiceTextTokenizerFast`):
+            The tokenizer for text processing.
+        audio_processor (`VibeVoiceTokenizerProcessor`):
+            The audio processor for speech processing.
+        speech_tok_compress_ratio (`int`, *optional*, defaults to 3200):
+            The compression ratio for speech tokenization.
+        db_normalize (`bool`, *optional*, defaults to True):
+            Whether to apply decibel normalization to audio inputs.
+    """
+
+    def __init__(self, tokenizer=None, audio_processor=None, speech_tok_compress_ratio=3200, db_normalize=True, **kwargs):
+        self.tokenizer = tokenizer
+        self.audio_processor = audio_processor
+        self.speech_tok_compress_ratio = speech_tok_compress_ratio
+        self.db_normalize = db_normalize
+        self.audio_normalizer = AudioNormalizer() if db_normalize else None
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+        """
+        Instantiate a VibeVoiceStreamingProcessor from a pretrained VibeVoice Streaming processor.
+
+        Args:
+            pretrained_model_name_or_path (`str` or `os.PathLike`):
+                This can be either:
+                - a string, the *model id* of a pretrained model
+                - a path to a *directory* containing processor config
+
+        Returns:
+            [`VibeVoiceStreamingProcessor`]: The processor object instantiated from pretrained model.
+        """
+        import os
+        import json
+        from transformers.utils import cached_file
+        from .vibevoice_tokenizer_processor import VibeVoiceTokenizerProcessor
+        from vibevoice.modular.modular_vibevoice_text_tokenizer import (
+            VibeVoiceTextTokenizer, 
+            VibeVoiceTextTokenizerFast
+        )
+        
+        # Try to load from local path first, then from HF hub
+        config_path = os.path.join(pretrained_model_name_or_path, "preprocessor_config.json")
+        config = None
+        
+        if os.path.exists(config_path):
+            # Local path exists
+            with open(config_path, 'r') as f:
+                config = json.load(f)
+        else:
+            # Try to load from HF hub
+            try:
+                config_file = cached_file(
+                    pretrained_model_name_or_path,
+                    "preprocessor_config.json",
+                    **kwargs
+                )
+                with open(config_file, 'r') as f:
+                    config = json.load(f)
+            except Exception as e:
+                logger.warning(f"Could not load preprocessor_config.json from {pretrained_model_name_or_path}: {e}")
+                logger.warning("Using default configuration")
+                config = {
+                    "speech_tok_compress_ratio": 3200,
+                    "db_normalize": True,
+                }
+        
+        # Extract main processor parameters
+        speech_tok_compress_ratio = config.get("speech_tok_compress_ratio", 3200)
+        db_normalize = config.get("db_normalize", True)
+        
+        # Load tokenizer - try from model path first, then fall back to Qwen        
+        language_model_pretrained_name = config.get("language_model_pretrained_name", None) or kwargs.pop("language_model_pretrained_name", "Qwen/Qwen2.5-1.5B")
+        logger.info(f"Loading tokenizer from {language_model_pretrained_name}")
+        if 'qwen' in language_model_pretrained_name.lower():
+            tokenizer = VibeVoiceTextTokenizerFast.from_pretrained(
+                language_model_pretrained_name,
+                **kwargs
+            )
+        else:
+            raise ValueError(f"Unsupported tokenizer type for {language_model_pretrained_name}. Supported types: Qwen, Llama, Gemma.")
+        
+        # Load audio processor
+        if "audio_processor" in config:
+            # Create audio processor from config
+            audio_config = config["audio_processor"]
+            audio_processor = VibeVoiceTokenizerProcessor(
+                sampling_rate=audio_config.get("sampling_rate", 24000),
+                normalize_audio=audio_config.get("normalize_audio", True),
+                target_dB_FS=audio_config.get("target_dB_FS", -25),
+                eps=audio_config.get("eps", 1e-6),
+            )
+        else:
+            # Create default audio processor
+            audio_processor = VibeVoiceTokenizerProcessor()
+        
+        # Create and return the processor
+        return cls(
+            tokenizer=tokenizer,
+            audio_processor=audio_processor,
+            speech_tok_compress_ratio=speech_tok_compress_ratio,
+            db_normalize=db_normalize,
+        )
+    
+    def save_pretrained(self, save_directory: Union[str, os.PathLike], **kwargs):
+        """
+        Save a processor to a directory, so that it can be re-loaded using the
+        [`~VibeVoiceStreamingProcessor.from_pretrained`] class method.
+
+        Args:
+            save_directory (`str` or `os.PathLike`):
+                Directory where the processor will be saved.
+        """
+        import os
+        import json
+        
+        os.makedirs(save_directory, exist_ok=True)
+        
+        # Save processor configuration
+        processor_config = {
+            "processor_class": "VibeVoiceStreamingProcessor",
+            "speech_tok_compress_ratio": self.speech_tok_compress_ratio,
+            "db_normalize": self.db_normalize,
+            "audio_processor": {
+                "feature_extractor_type": "VibeVoiceTokenizerProcessor",
+                "sampling_rate": getattr(self.audio_processor, 'sampling_rate', 24000),
+                "normalize_audio": getattr(self.audio_processor, 'normalize_audio', True),
+                "target_dB_FS": getattr(self.audio_processor, 'target_dB_FS', -25),
+                "eps": getattr(self.audio_processor, 'eps', 1e-6),
+            }
+        }
+        
+        config_path = os.path.join(save_directory, "preprocessor_config.json")
+        with open(config_path, 'w') as f:
+            json.dump(processor_config, f, indent=2)
+        
+        logger.info(f"Processor configuration saved in {config_path}")
+    
+    def __call__(self) -> BatchEncoding:
+        """
+        Note:
+            This method is intentionally not implemented in the streaming processor.
+            Use `process_input_with_cached_prompt` for streaming use cases.
+        """
+        raise NotImplementedError(
+            "VibeVoiceStreamingProcessor.__call__ is not implemented. "
+            "Use process_input_with_cached_prompt for streaming inputs."
+        )
+
+    def process_input_with_cached_prompt(
+        self,
+        text: Optional[str] = None,
+        cached_prompt: Optional[Dict[str, Any]] = None,
+        padding: Union[bool, str, PaddingStrategy] = True,
+        truncation: Union[bool, str, TruncationStrategy] = False,
+        max_length: Optional[int] = None,
+        return_tensors: Optional[Union[str, TensorType]] = None,
+        return_attention_mask: bool = True,
+        **kwargs,
+    ) -> BatchEncoding:
+        """
+        Main method to process one text script based on cached prompt. The function currently only supports single examples.
+
+        Args:
+            text (`str`):
+                The input text to process.
+            cached_prompt (`Dict[str, Any]`, *optional*):
+                The cached prompt to use for processing. It contains the kv cache of the voice prompt.
+            padding (`bool`, `str` or `PaddingStrategy`, defaults to `True`):
+                Whether to pad sequences to the same length
+            truncation (`bool`, `str` or `TruncationStrategy`, defaults to `False`):
+                Whether to truncate sequences
+            max_length (`int`, *optional*):
+                Maximum length of the returned sequences
+            return_tensors (`str` or `TensorType`, *optional*):
+                If set, will return tensors of a particular framework
+            return_attention_mask (`bool`, defaults to `True`):
+                Whether to return the attention mask
+
+        Returns:
+            `BatchEncoding`: A BatchEncoding with the following fields:
+                - **input_ids** -- List of token id sequences or tensor
+                - **attention_mask** -- List of attention masks or tensor
+                - **tts_lm_input_ids** -- List of token id sequences or tensor used for TTS LM
+                - **tts_lm_attention_mask** -- List of attention masks or tensor used for TTS LM
+                - **tts_text_ids** -- List of token id sequences or tensor for TTS text input
+                - **speech_tensors** -- Padded speech inputs (if voice_samples provided)
+                - **speech_masks** -- Speech masks (if voice_samples provided)
+                - **speech_input_mask** -- Boolean masks indicating speech token positions
+        """
+        # Only support single example
+        texts = [text]
+        cached_prompts = [cached_prompt]
+        is_batched = False
+        
+        # Process each input
+        all_encodings = []
+        for text_input, cached_prompt_input in zip(texts, cached_prompts):
+            script_tokens = self.tokenizer.encode(text_input.strip() + "\n", add_special_tokens=False)
+            input_id_length = cached_prompt_input['lm']['last_hidden_state'].size(1)
+            tts_lm_input_id_length = cached_prompt_input['tts_lm']['last_hidden_state'].size(1)
+
+            # pseudo input ids and masks
+            input_ids = [self.tokenizer.pad_id] * input_id_length
+            tts_lm_input_ids = [self.tokenizer.pad_id] * tts_lm_input_id_length
+            speech_input_mask = [False] * tts_lm_input_id_length
+
+            encoding = {
+                        "input_ids": input_ids,
+                        "tts_lm_input_ids": tts_lm_input_ids,
+                        "tts_text_ids": script_tokens,
+                        "speech_inputs": None,
+                        "speech_input_mask": speech_input_mask,
+                    }
+            all_encodings.append(encoding)
+            
+        # Combine batch
+        batch_encoding = self._batch_encode(
+            all_encodings,
+            padding=padding,
+            truncation=truncation,
+            max_length=max_length,
+            return_tensors=return_tensors,
+            return_attention_mask=return_attention_mask,
+        )
+        
+        return batch_encoding
+    
+    def _batch_encode(
+        self,
+        encodings: List[Dict[str, Any]],
+        padding: Union[bool, str, PaddingStrategy] = True,
+        truncation: Union[bool, str, TruncationStrategy] = False,
+        max_length: Optional[int] = None,
+        return_tensors: Optional[Union[str, TensorType]] = None,
+        return_attention_mask: bool = True,
+    ) -> BatchEncoding:
+        """Combine multiple encodings into a batch with padding."""
+        # Extract input_ids and create attention_mask
+        input_ids_list = [enc["input_ids"] for enc in encodings]
+        tts_lm_input_ids_list = [enc["tts_lm_input_ids"] for enc in encodings]
+        tts_text_ids_list = [enc["tts_text_ids"] for enc in encodings]
+        speech_input_masks_list = [enc["speech_input_mask"] for enc in encodings]
+        
+        attention_masks = [[1] * len(ids) for ids in input_ids_list] if return_attention_mask else None
+        tts_lm_attention_masks = [[1] * len(ids) for ids in tts_lm_input_ids_list] if return_attention_mask else None
+            
+        # Process speech inputs
+        all_speech_inputs = []
+        has_speech = False
+        for enc in encodings:
+            if enc["speech_inputs"] is not None:
+                all_speech_inputs.extend(enc["speech_inputs"])
+                has_speech = True
+                
+        # Prepare batch encoding
+        batch_encoding = BatchEncoding()
+        
+        # Handle tensor conversion
+        if return_tensors is not None:
+            batch_encoding["input_ids"] = torch.tensor(input_ids_list, dtype=torch.long)
+            batch_encoding["tts_lm_input_ids"] = torch.tensor(tts_lm_input_ids_list, dtype=torch.long)
+            batch_encoding["tts_text_ids"] = torch.tensor(tts_text_ids_list, dtype=torch.long)
+
+            if return_attention_mask and attention_masks is not None:
+                batch_encoding["attention_mask"] = torch.tensor(attention_masks, dtype=torch.long)
+                batch_encoding["tts_lm_attention_mask"] = torch.tensor(tts_lm_attention_masks, dtype=torch.long)
+            
+            batch_encoding["speech_input_mask"] = torch.tensor(speech_input_masks_list, dtype=torch.bool)
+        else:
+            batch_encoding["input_ids"] = input_ids_list
+            batch_encoding["tts_lm_input_ids"] = tts_lm_input_ids_list
+            batch_encoding["tts_text_ids"] = tts_text_ids_list
+            if return_attention_mask and attention_masks is not None:
+                batch_encoding["attention_mask"] = attention_masks
+                batch_encoding["tts_lm_attention_mask"] = tts_lm_attention_masks
+            batch_encoding["speech_input_mask"] = speech_input_masks_list
+            
+        # Process speech tensors if present
+        if has_speech:
+            speech_dict = self.prepare_speech_inputs(
+                all_speech_inputs,
+                return_tensors=return_tensors,
+            )
+            batch_encoding["speech_tensors"] = speech_dict["padded_speeches"]
+            batch_encoding["speech_masks"] = speech_dict["speech_masks"]
+        else:
+            batch_encoding["speech_tensors"] = None
+            batch_encoding["speech_masks"] = None
+            
+        return batch_encoding
+
+    def prepare_speech_inputs(
+        self,
+        speech_inputs: List[np.ndarray],
+        return_tensors: Optional[Union[str, TensorType]] = None,
+        device: Optional[Union[str, torch.device]] = None,
+        dtype: Optional[torch.dtype] = None,
+    ) -> Dict[str, Any]:
+        """
+        Prepare speech inputs for model consumption.
+        
+        Args:
+            speech_inputs: List of speech arrays
+            return_tensors: Output tensor type
+            device: Device to place tensors on
+            dtype: Data type for tensors
+            
+        Returns:
+            Dictionary with padded_speeches and speech_masks
+        """
+        if not speech_inputs:
+            return {"padded_speeches": None, "speech_masks": None}
+        
+        # Calculate sequence lengths
+        vae_tok_seqlens = [math.ceil(s.shape[0] / self.speech_tok_compress_ratio) for s in speech_inputs]
+        # vae_tok_seqlens = [math.ceil(s.shape[0] / self.speech_tok_compress_ratio) if s.ndim == 1 else s.shape[0] for s in speech_inputs]
+        max_speech_length = max(s.shape[0] for s in speech_inputs)
+        
+        # Pad speeches
+        if speech_inputs[0].ndim == 1:
+            padded_speeches = np.full((len(speech_inputs), max_speech_length), fill_value=0, dtype=np.float32)
+        else:
+            padded_speeches = np.full((len(speech_inputs), max_speech_length, speech_inputs[0].shape[-1]), fill_value=0, dtype=np.float32)
+        speech_masks = np.zeros((len(speech_inputs), max(vae_tok_seqlens)), dtype=np.bool_)
+        
+        for i, (speech, vae_tok_length) in enumerate(zip(speech_inputs, vae_tok_seqlens)):
+            padded_speeches[i, :len(speech)] = speech
+            speech_masks[i, :vae_tok_length] = True
+        
+        result = {
+            "padded_speeches": padded_speeches,
+            "speech_masks": speech_masks,
+        }
+        
+        # Convert to tensors if requested
+        if return_tensors == "pt":
+            result["padded_speeches"] = torch.tensor(padded_speeches, device=device, dtype=dtype or torch.float32)
+            result["speech_masks"] = torch.tensor(speech_masks, device=device, dtype=torch.bool)
+            
+        return result
+
+    def batch_decode(self, *args, **kwargs):
+        """
+        This method forwards all its arguments to VibeVoiceTextTokenizer's [`~PreTrainedTokenizer.batch_decode`].
+        Please refer to the docstring of this method for more information.
+        """
+        return self.tokenizer.batch_decode(*args, **kwargs)
+
+    def decode(self, *args, **kwargs):
+        """
+        This method forwards all its arguments to VibeVoiceTextTokenizer's [`~PreTrainedTokenizer.decode`].
+        Please refer to the docstring of this method for more information.
+        """
+        return self.tokenizer.decode(*args, **kwargs)
+
+    @property
+    def model_input_names(self):
+        """
+        Return the list of inputs accepted by the model.
+        """
+        tokenizer_input_names = self.tokenizer.model_input_names
+        audio_processor_input_names = self.audio_processor.model_input_names
+        return list(dict.fromkeys(tokenizer_input_names + audio_processor_input_names + ["speech_inputs", "speech_input_mask"]))
+
+    def save_audio(self, 
+        audio: Union[torch.Tensor, np.ndarray, List[Union[torch.Tensor, np.ndarray]]],
+        output_path: str = "output.wav",
+        sampling_rate: Optional[int] = None,
+        normalize: bool = False,
+        batch_prefix: str = "audio_",
+    ) -> str:
+        """
+        Save audio data to a file.
+        Args:
+            audio (Union[torch.Tensor, np.ndarray, List[Union[torch.Tensor, np.ndarray]]]):
+                The audio data to save. Can be a single tensor/array or a list of them.
+            output_path (str, optional): Path to save the audio file. Defaults to "output.wav".
+            sampling_rate (int, optional): Sampling rate for the audio. If None, uses the processor's default.
+            normalize (bool, optional): Whether to normalize the audio before saving. Defaults to False.
+            batch_prefix (str, optional): Prefix for batch audio files. Defaults to "audio_".
+        Returns:
+            str: The path to the saved audio file.
+        """
+        return self.audio_processor.save_audio(audio, output_path=output_path, sampling_rate=sampling_rate, normalize=normalize, batch_prefix=batch_prefix)
+    
+__all__ = [
+    "VibeVoiceStreamingProcessor",
+]

vibevoice/processor/vibevoice_tokenizer_processor.py

@@ -0,0 +1,413 @@
+"""
+Processor class for VibeVoice models.
+"""
+
+import os
+import json
+import warnings
+from typing import List, Optional, Union, Dict, Any
+
+import numpy as np
+import torch
+
+from transformers.feature_extraction_utils import FeatureExtractionMixin
+from transformers.utils import logging
+
+from .audio_utils import AudioNormalizer
+
+logger = logging.get_logger(__name__)
+
+# Change from ProcessorMixin to FeatureExtractionMixin which is designed for single components
+class VibeVoiceTokenizerProcessor(FeatureExtractionMixin):
+    """
+    Processor for VibeVoice acoustic tokenizer models.
+    
+    This processor handles audio preprocessing for VibeVoice models, including:
+    - Audio format conversion (stereo to mono)
+    - Optional audio normalization
+    - Streaming support for infinite-length audio
+    
+    Args:
+        sampling_rate (int, optional): Expected sampling rate. Defaults to 24000.
+        normalize_audio (bool, optional): Whether to normalize audio. Defaults to True.
+        target_dB_FS (float, optional): Target dB FS for normalization. Defaults to -25.
+        eps (float, optional): Small value for numerical stability. Defaults to 1e-6.
+    """
+    model_input_names = ["input_features"]
+    
+    def __init__(
+        self,
+        sampling_rate: int = 24000,
+        normalize_audio: bool = True,
+        target_dB_FS: float = -25,
+        eps: float = 1e-6,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        
+        self.sampling_rate = sampling_rate
+        self.normalize_audio = normalize_audio
+        
+        # Initialize audio normalizer if needed
+        if self.normalize_audio:
+            self.normalizer = AudioNormalizer(target_dB_FS=target_dB_FS, eps=eps)
+        else:
+            self.normalizer = None
+        
+        # Save config
+        self.feature_extractor_dict = {
+            "sampling_rate": sampling_rate,
+            "normalize_audio": normalize_audio,
+            "target_dB_FS": target_dB_FS,
+            "eps": eps,
+        }
+    
+    def _ensure_mono(self, audio: np.ndarray) -> np.ndarray:
+        """
+        Convert stereo audio to mono if needed.
+        
+        Args:
+            audio (np.ndarray): Input audio array
+            
+        Returns:
+            np.ndarray: Mono audio array
+        """
+        if len(audio.shape) == 1:
+            return audio
+        elif len(audio.shape) == 2:
+            if audio.shape[0] == 2:  # (2, time)
+                return np.mean(audio, axis=0)
+            elif audio.shape[1] == 2:  # (time, 2)
+                return np.mean(audio, axis=1)
+            else:
+                # If one dimension is 1, squeeze it
+                if audio.shape[0] == 1:
+                    return audio.squeeze(0)
+                elif audio.shape[1] == 1:
+                    return audio.squeeze(1)
+                else:
+                    raise ValueError(f"Unexpected audio shape: {audio.shape}")
+        else:
+            raise ValueError(f"Audio should be 1D or 2D, got shape: {audio.shape}")
+    
+    def _process_single_audio(self, audio: Union[np.ndarray, List[float]]) -> np.ndarray:
+        """
+        Process a single audio array.
+        
+        Args:
+            audio: Single audio input
+            
+        Returns:
+            np.ndarray: Processed audio
+        """
+        # Convert to numpy array
+        if not isinstance(audio, np.ndarray):
+            audio = np.array(audio, dtype=np.float32)
+        else:
+            audio = audio.astype(np.float32)
+        
+        # Ensure mono
+        audio = self._ensure_mono(audio)
+        
+        # Normalize if requested
+        if self.normalize_audio and self.normalizer is not None:
+            audio = self.normalizer(audio)
+        
+        return audio
+    
+    def __call__(
+        self,
+        audio: Union[str, np.ndarray, List[float], List[np.ndarray], List[List[float]], List[str]] = None,
+        sampling_rate: Optional[int] = None,
+        return_tensors: Optional[str] = None,
+        **kwargs,
+    ):
+        """
+        Process audio for VibeVoice models.
+        
+        Args:
+            audio: Audio input(s) to process. Can be:
+                - str: Path to audio file
+                - np.ndarray: Audio array
+                - List[float]: Audio as list of floats
+                - List[np.ndarray]: Batch of audio arrays
+                - List[str]: Batch of audio file paths
+            sampling_rate (int, optional): Sampling rate of the input audio
+            return_tensors (str, optional): Return format ('pt' for PyTorch, 'np' for NumPy)
+            
+        Returns:
+            dict: Processed audio inputs with keys:
+                - input_features: Audio tensor(s) ready for the model
+        """
+        if audio is None:
+            raise ValueError("Audio input is required")
+        
+        # Validate sampling rate
+        if sampling_rate is not None and sampling_rate != self.sampling_rate:
+            logger.warning(
+                f"Input sampling rate ({sampling_rate}) differs from expected "
+                f"sampling rate ({self.sampling_rate}). Please resample your audio."
+            )
+        
+        # Handle different input types
+        if isinstance(audio, str):
+            # Single audio file path
+            audio = self._load_audio_from_path(audio)
+            is_batched = False
+        elif isinstance(audio, list):
+            if len(audio) == 0:
+                raise ValueError("Empty audio list provided")
+            
+            # Check if it's a list of file paths
+            if all(isinstance(item, str) for item in audio):
+                # Batch of audio file paths
+                audio = [self._load_audio_from_path(path) for path in audio]
+                is_batched = True
+            else:
+                # Check if it's batched audio arrays
+                is_batched = isinstance(audio[0], (np.ndarray, list))
+        else:
+            # Single audio array or list
+            is_batched = False
+        
+        # Process audio
+        if is_batched:
+            processed_audio = [self._process_single_audio(a) for a in audio]
+        else:
+            processed_audio = [self._process_single_audio(audio)]
+        
+        # Convert to tensors if requested
+        if return_tensors == "pt":
+            if len(processed_audio) == 1:
+                # Create a proper batch dimension (B, T)
+                input_features = torch.from_numpy(processed_audio[0]).unsqueeze(0).unsqueeze(1)
+            else:
+                # For batched input with different lengths, create a batch properly
+                input_features = torch.stack([torch.from_numpy(a) for a in processed_audio]).unsqueeze(1)
+        elif return_tensors == "np":
+            if len(processed_audio) == 1:
+                input_features = processed_audio[0][np.newaxis, np.newaxis, :]
+            else:
+                input_features = np.stack(processed_audio)[:, np.newaxis, :]
+        else:
+            input_features = processed_audio[0] if len(processed_audio) == 1 else processed_audio
+        
+        outputs = {
+            "audio": input_features,  # Use "audio" instead of "input_features"
+        }
+        
+        return outputs
+
+    def _load_audio_from_path(self, audio_path: str) -> np.ndarray:
+        """
+        Load audio from file path.
+        
+        Args:
+            audio_path (str): Path to audio file
+            
+        Returns:
+            np.ndarray: Loaded audio array
+        """
+        # Get file extension to determine loading method
+        file_ext = os.path.splitext(audio_path)[1].lower()
+        
+        if file_ext in ['.wav', '.mp3', '.flac', '.m4a', '.ogg']:
+            # Audio file - use librosa
+            import librosa
+            audio_array, sr = librosa.load(
+                audio_path, 
+                sr=self.sampling_rate, 
+                mono=True
+            )
+            return audio_array
+        elif file_ext == '.pt':
+            # PyTorch tensor file
+            audio_tensor = torch.load(audio_path, map_location='cpu', weights_only=True).squeeze()
+            if isinstance(audio_tensor, torch.Tensor):
+                audio_array = audio_tensor.numpy()
+            else:
+                audio_array = np.array(audio_tensor)
+            return audio_array.astype(np.float32)
+        elif file_ext == '.npy':
+            # NumPy file
+            audio_array = np.load(audio_path)
+            return audio_array.astype(np.float32)
+        else:
+            raise ValueError(
+                f"Unsupported file format: {file_ext}. "
+                f"Supported formats: .wav, .mp3, .flac, .m4a, .ogg, .pt, .npy, .npz"
+            )
+    
+    def preprocess_audio(
+        self, 
+        audio_path_or_array: Union[str, np.ndarray],
+        normalize: Optional[bool] = None,
+    ) -> np.ndarray:
+        """
+        Convenience method to preprocess audio from file path or array.
+        This method is kept for backward compatibility but __call__ is recommended.
+        
+        Args:
+            audio_path_or_array: Path to audio file or numpy array
+            normalize: Whether to normalize (overrides default setting)
+            
+        Returns:
+            np.ndarray: Preprocessed audio array
+        """
+        if isinstance(audio_path_or_array, str):
+            audio_array = self._load_audio_from_path(audio_path_or_array)
+        else:
+            audio_array = np.array(audio_path_or_array, dtype=np.float32)
+        
+        # Override normalization setting if specified
+        original_normalize = self.normalize_audio
+        if normalize is not None:
+            self.normalize_audio = normalize
+        
+        try:
+            processed = self._process_single_audio(audio_array)
+        finally:
+            # Restore original setting
+            self.normalize_audio = original_normalize
+        
+        return processed
+    
+    # Override to_dict method for configuration saving
+    def to_dict(self) -> Dict[str, Any]:
+        """
+        Convert the object to a dict containing all attributes needed for serialization.
+        """
+        return self.feature_extractor_dict
+
+    def save_audio(
+        self,
+        audio: Union[torch.Tensor, np.ndarray, List[Union[torch.Tensor, np.ndarray]]],
+        output_path: str = "output.wav",
+        sampling_rate: Optional[int] = None,
+        normalize: bool = False,
+        batch_prefix: str = "audio_",
+    ):
+        """
+        Save audio data to WAV file(s).
+        
+        Args:
+            audio: Audio data to save. Can be:
+                - torch.Tensor: PyTorch tensor with shape (B, C, T) or (B, T) or (T)
+                - np.ndarray: NumPy array with shape (B, C, T) or (B, T) or (T)
+                - List of tensors or arrays
+            output_path: Path where to save the audio. If saving multiple files,
+                this is treated as a directory and individual files will be saved inside.
+            sampling_rate: Sampling rate for the saved audio. Defaults to the processor's rate.
+            normalize: Whether to normalize audio before saving.
+            batch_prefix: Prefix for batch files when saving multiple audios.
+                
+        Returns:
+            List[str]: Paths to the saved audio files.
+        """
+        if sampling_rate is None:
+            sampling_rate = self.sampling_rate
+        
+        try:
+            import soundfile as sf
+        except ImportError:
+            raise ImportError(
+                "soundfile is required to save audio files. "
+                "Install it with: pip install soundfile"
+            )
+        
+        # Ensure audio is in the right format
+        if isinstance(audio, torch.Tensor):
+            # Convert PyTorch tensor to numpy
+            audio_np = audio.float().detach().cpu().numpy()
+        elif isinstance(audio, np.ndarray):
+            audio_np = audio
+        elif isinstance(audio, list):
+            # Handle list of tensors or arrays
+            if all(isinstance(a, torch.Tensor) for a in audio):
+                audio_np = [a.float().detach().cpu().numpy() for a in audio]
+            else:
+                audio_np = audio
+        else:
+            raise ValueError(f"Unsupported audio type: {type(audio)}")
+        
+        saved_paths = []
+        
+        # Handle based on shape or type
+        if isinstance(audio_np, list):
+            # Multiple separate audios to save
+            output_dir = output_path
+            
+            # Ensure output directory exists
+            os.makedirs(output_dir, exist_ok=True)
+            
+            # Save each audio
+            for i, audio_item in enumerate(audio_np):
+                audio_item = self._prepare_audio_for_save(audio_item, normalize)
+                file_path = os.path.join(output_dir, f"{batch_prefix}{i}.wav")
+                sf.write(file_path, audio_item, sampling_rate)
+                saved_paths.append(file_path)
+                
+        else:
+            # Handle different dimensions
+            if len(audio_np.shape) >= 3:  # (B, C, T) or similar
+                # Get batch size
+                batch_size = audio_np.shape[0]
+                
+                if batch_size > 1:
+                    # Multiple audios in a batch
+                    output_dir = output_path
+                    
+                    # Ensure output directory exists
+                    os.makedirs(output_dir, exist_ok=True)
+                    
+                    # Save each audio in the batch
+                    for i in range(batch_size):
+                        # Extract single audio and remove channel dim if present
+                        single_audio = audio_np[i]
+                        if len(single_audio.shape) > 1:
+                            if single_audio.shape[0] == 1:  # (1, T)
+                                single_audio = single_audio.squeeze(0)
+                        
+                        single_audio = self._prepare_audio_for_save(single_audio, normalize)
+                        file_path = os.path.join(output_dir, f"{batch_prefix}{i}.wav")
+                        sf.write(file_path, single_audio, sampling_rate)
+                        saved_paths.append(file_path)
+                else:
+                    # Single audio with batch and channel dims
+                    audio_item = audio_np.squeeze()  # Remove batch and channel dimensions
+                    audio_item = self._prepare_audio_for_save(audio_item, normalize)
+                    sf.write(output_path, audio_item, sampling_rate)
+                    saved_paths.append(output_path)
+            else:
+                # Single audio without batch dimension
+                audio_item = self._prepare_audio_for_save(audio_np, normalize)
+                sf.write(output_path, audio_item, sampling_rate)
+                saved_paths.append(output_path)
+        
+        return saved_paths
+
+    def _prepare_audio_for_save(self, audio: np.ndarray, normalize: bool) -> np.ndarray:
+        """
+        Prepare audio for saving by ensuring it's the right shape and optionally normalizing.
+        
+        Args:
+            audio: Audio data as numpy array
+            normalize: Whether to normalize audio
+            
+        Returns:
+            np.ndarray: Processed audio ready for saving
+        """
+        # Ensure right dimensionality
+        if len(audio.shape) > 1 and audio.shape[0] == 1:  # (1, T)
+            audio = audio.squeeze(0)
+        
+        # Normalize if requested
+        if normalize:
+            max_val = np.abs(audio).max()
+            if max_val > 0:
+                audio = audio / max_val
+        
+        return audio
+
+
+__all__ = ["VibeVoiceTokenizerProcessor", "AudioNormalizer"]

vibevoice/schedule/dpm_solver.py

@@ -0,0 +1,1065 @@
+# Copyright 2024 TSAIL Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# DISCLAIMER: This file is strongly influenced by https://github.com/LuChengTHU/dpm-solver
+
+import math
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.utils import deprecate
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.schedulers.scheduling_utils import KarrasDiffusionSchedulers, SchedulerMixin, SchedulerOutput
+
+def betas_for_alpha_bar(
+    num_diffusion_timesteps,
+    max_beta=0.999,
+    alpha_transform_type="cosine",
+):
+    """
+    Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of
+    (1-beta) over time from t = [0,1].
+
+    Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up
+    to that part of the diffusion process.
+
+
+    Args:
+        num_diffusion_timesteps (`int`): the number of betas to produce.
+        max_beta (`float`): the maximum beta to use; use values lower than 1 to
+                     prevent singularities.
+        alpha_transform_type (`str`, *optional*, default to `cosine`): the type of noise schedule for alpha_bar.
+                     Choose from `cosine` or `exp`
+
+    Returns:
+        betas (`np.ndarray`): the betas used by the scheduler to step the model outputs
+    """
+    if alpha_transform_type == "cosine":
+
+        def alpha_bar_fn(t):
+            return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2
+            # return math.cos(t * math.pi / 2 * 0.95) ** 2
+
+    elif alpha_transform_type == "exp":
+
+        def alpha_bar_fn(t):
+            return math.exp(t * -12.0)
+
+    elif alpha_transform_type == "cauchy":
+        # µ + γ tan (π (0.5 - x))  γ = 1, µ = 3
+        # alpha^2 = 1-1/(exp(λ)+1)
+        def alpha_bar_fn(t, gamma=1, mu=3):
+            snr = mu + gamma * math.tan(math.pi * (0.5 - t) * 0.9)
+            return 1 - 1 / (math.exp(snr) + 1.1)
+
+    elif alpha_transform_type == "laplace":
+        # µ − bsgn(0.5 − t) log(1 − 2|t − 0.5|) µ = 0, b = 1
+        def alpha_bar_fn(t, mu=0, b=1):
+            snr = mu - b * math.copysign(1, 0.5 - t) * math.log(1 - 2 * abs(t - 0.5) * 0.98)
+            return 1 - 1 / (math.exp(snr) + 1.02)
+
+    else:
+        raise ValueError(f"Unsupported alpha_transform_type: {alpha_transform_type}")
+
+    betas = []
+    for i in range(num_diffusion_timesteps):
+        t1 = i / num_diffusion_timesteps
+        t2 = (i + 1) / num_diffusion_timesteps
+        betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta))
+    return torch.tensor(betas, dtype=torch.float32)
+
+
+# Copied from diffusers.schedulers.scheduling_ddim.rescale_zero_terminal_snr
+def rescale_zero_terminal_snr(betas):
+    """
+    Rescales betas to have zero terminal SNR Based on https://arxiv.org/pdf/2305.08891.pdf (Algorithm 1)
+
+
+    Args:
+        betas (`torch.Tensor`):
+            the betas that the scheduler is being initialized with.
+
+    Returns:
+        `torch.Tensor`: rescaled betas with zero terminal SNR
+    """
+    # Convert betas to alphas_bar_sqrt
+    alphas = 1.0 - betas
+    alphas_cumprod = torch.cumprod(alphas, dim=0)
+    alphas_bar_sqrt = alphas_cumprod.sqrt()
+
+    # Store old values.
+    alphas_bar_sqrt_0 = alphas_bar_sqrt[0].clone()
+    alphas_bar_sqrt_T = alphas_bar_sqrt[-1].clone()
+
+    # Shift so the last timestep is zero.
+    alphas_bar_sqrt -= alphas_bar_sqrt_T
+
+    # Scale so the first timestep is back to the old value.
+    alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T)
+
+    # Convert alphas_bar_sqrt to betas
+    alphas_bar = alphas_bar_sqrt**2  # Revert sqrt
+    alphas = alphas_bar[1:] / alphas_bar[:-1]  # Revert cumprod
+    alphas = torch.cat([alphas_bar[0:1], alphas])
+    betas = 1 - alphas
+
+    return betas
+
+class DPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin):
+    """
+    `DPMSolverMultistepScheduler` is a fast dedicated high-order solver for diffusion ODEs.
+
+    This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
+    methods the library implements for all schedulers such as loading and saving.
+
+    Args:
+        num_train_timesteps (`int`, defaults to 1000):
+            The number of diffusion steps to train the model.
+        beta_start (`float`, defaults to 0.0001):
+            The starting `beta` value of inference.
+        beta_end (`float`, defaults to 0.02):
+            The final `beta` value.
+        beta_schedule (`str`, defaults to `"linear"`):
+            The beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from
+            `linear`, `scaled_linear`, or `squaredcos_cap_v2`.
+        trained_betas (`np.ndarray`, *optional*):
+            Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`.
+        solver_order (`int`, defaults to 2):
+            The DPMSolver order which can be `1` or `2` or `3`. It is recommended to use `solver_order=2` for guided
+            sampling, and `solver_order=3` for unconditional sampling.
+        prediction_type (`str`, defaults to `epsilon`, *optional*):
+            Prediction type of the scheduler function; can be `epsilon` (predicts the noise of the diffusion process),
+            `sample` (directly predicts the noisy sample`) or `v_prediction` (see section 2.4 of [Imagen
+            Video](https://imagen.research.google/video/paper.pdf) paper).
+        thresholding (`bool`, defaults to `False`):
+            Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such
+            as Stable Diffusion.
+        dynamic_thresholding_ratio (`float`, defaults to 0.995):
+            The ratio for the dynamic thresholding method. Valid only when `thresholding=True`.
+        sample_max_value (`float`, defaults to 1.0):
+            The threshold value for dynamic thresholding. Valid only when `thresholding=True` and
+            `algorithm_type="dpmsolver++"`.
+        algorithm_type (`str`, defaults to `dpmsolver++`):
+            Algorithm type for the solver; can be `dpmsolver`, `dpmsolver++`, `sde-dpmsolver` or `sde-dpmsolver++`. The
+            `dpmsolver` type implements the algorithms in the [DPMSolver](https://huggingface.co/papers/2206.00927)
+            paper, and the `dpmsolver++` type implements the algorithms in the
+            [DPMSolver++](https://huggingface.co/papers/2211.01095) paper. It is recommended to use `dpmsolver++` or
+            `sde-dpmsolver++` with `solver_order=2` for guided sampling like in Stable Diffusion.
+        solver_type (`str`, defaults to `midpoint`):
+            Solver type for the second-order solver; can be `midpoint` or `heun`. The solver type slightly affects the
+            sample quality, especially for a small number of steps. It is recommended to use `midpoint` solvers.
+        lower_order_final (`bool`, defaults to `True`):
+            Whether to use lower-order solvers in the final steps. Only valid for < 15 inference steps. This can
+            stabilize the sampling of DPMSolver for steps < 15, especially for steps <= 10.
+        euler_at_final (`bool`, defaults to `False`):
+            Whether to use Euler's method in the final step. It is a trade-off between numerical stability and detail
+            richness. This can stabilize the sampling of the SDE variant of DPMSolver for small number of inference
+            steps, but sometimes may result in blurring.
+        use_karras_sigmas (`bool`, *optional*, defaults to `False`):
+            Whether to use Karras sigmas for step sizes in the noise schedule during the sampling process. If `True`,
+            the sigmas are determined according to a sequence of noise levels {σi}.
+        use_lu_lambdas (`bool`, *optional*, defaults to `False`):
+            Whether to use the uniform-logSNR for step sizes proposed by Lu's DPM-Solver in the noise schedule during
+            the sampling process. If `True`, the sigmas and time steps are determined according to a sequence of
+            `lambda(t)`.
+        final_sigmas_type (`str`, defaults to `"zero"`):
+            The final `sigma` value for the noise schedule during the sampling process. If `"sigma_min"`, the final
+            sigma is the same as the last sigma in the training schedule. If `zero`, the final sigma is set to 0.
+        lambda_min_clipped (`float`, defaults to `-inf`):
+            Clipping threshold for the minimum value of `lambda(t)` for numerical stability. This is critical for the
+            cosine (`squaredcos_cap_v2`) noise schedule.
+        variance_type (`str`, *optional*):
+            Set to "learned" or "learned_range" for diffusion models that predict variance. If set, the model's output
+            contains the predicted Gaussian variance.
+        timestep_spacing (`str`, defaults to `"linspace"`):
+            The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and
+            Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.
+        steps_offset (`int`, defaults to 0):
+            An offset added to the inference steps, as required by some model families.
+        rescale_betas_zero_snr (`bool`, defaults to `False`):
+            Whether to rescale the betas to have zero terminal SNR. This enables the model to generate very bright and
+            dark samples instead of limiting it to samples with medium brightness. Loosely related to
+            [`--offset_noise`](https://github.com/huggingface/diffusers/blob/74fd735eb073eb1d774b1ab4154a0876eb82f055/examples/dreambooth/train_dreambooth.py#L506).
+    """
+
+    _compatibles = [e.name for e in KarrasDiffusionSchedulers]
+    order = 1
+
+    @register_to_config
+    def __init__(
+        self,
+        num_train_timesteps: int = 1000,
+        beta_start: float = 0.0001,
+        beta_end: float = 0.02,
+        beta_schedule: str = "linear",
+        trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
+        solver_order: int = 2,
+        prediction_type: str = "epsilon",
+        thresholding: bool = False,
+        dynamic_thresholding_ratio: float = 0.995,
+        sample_max_value: float = 1.0,
+        algorithm_type: str = "dpmsolver++",
+        solver_type: str = "midpoint",
+        lower_order_final: bool = True,
+        euler_at_final: bool = False,
+        use_karras_sigmas: Optional[bool] = False,
+        use_lu_lambdas: Optional[bool] = False,
+        final_sigmas_type: Optional[str] = "zero",  # "zero", "sigma_min"
+        lambda_min_clipped: float = -float("inf"),
+        variance_type: Optional[str] = None,
+        timestep_spacing: str = "linspace",
+        steps_offset: int = 0,
+        rescale_betas_zero_snr: bool = False,
+    ):
+        if algorithm_type in ["dpmsolver", "sde-dpmsolver"]:
+            deprecation_message = f"algorithm_type {algorithm_type} is deprecated and will be removed in a future version. Choose from `dpmsolver++` or `sde-dpmsolver++` instead"
+            deprecate("algorithm_types dpmsolver and sde-dpmsolver", "1.0.0", deprecation_message)
+
+        if trained_betas is not None:
+            self.betas = torch.tensor(trained_betas, dtype=torch.float32)
+        elif beta_schedule == "linear":
+            self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32)
+        elif beta_schedule == "scaled_linear":
+            # this schedule is very specific to the latent diffusion model.
+            self.betas = torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32) ** 2
+        elif beta_schedule == "squaredcos_cap_v2" or beta_schedule == "cosine":
+            # Glide cosine schedule
+            self.betas = betas_for_alpha_bar(num_train_timesteps, alpha_transform_type="cosine")
+        elif beta_schedule == "cauchy":
+            self.betas = betas_for_alpha_bar(num_train_timesteps, alpha_transform_type="cauchy")
+        elif beta_schedule == "laplace":
+            self.betas = betas_for_alpha_bar(num_train_timesteps, alpha_transform_type="laplace")
+        else:
+            raise NotImplementedError(f"{beta_schedule} is not implemented for {self.__class__}")
+
+        if rescale_betas_zero_snr:
+            self.betas = rescale_zero_terminal_snr(self.betas)
+
+        self.alphas = 1.0 - self.betas
+        self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
+
+        if rescale_betas_zero_snr:
+            # Close to 0 without being 0 so first sigma is not inf
+            # FP16 smallest positive subnormal works well here
+            self.alphas_cumprod[-1] = 2**-24
+
+        # Currently we only support VP-type noise schedule
+        self.alpha_t = torch.sqrt(self.alphas_cumprod)
+        self.sigma_t = torch.sqrt(1 - self.alphas_cumprod)
+        self.lambda_t = torch.log(self.alpha_t) - torch.log(self.sigma_t)
+        self.sigmas = ((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5
+
+        # standard deviation of the initial noise distribution
+        self.init_noise_sigma = 1.0
+
+        # settings for DPM-Solver
+        if algorithm_type not in ["dpmsolver", "dpmsolver++", "sde-dpmsolver", "sde-dpmsolver++"]:
+            if algorithm_type == "deis":
+                self.register_to_config(algorithm_type="dpmsolver++")
+            else:
+                raise NotImplementedError(f"{algorithm_type} is not implemented for {self.__class__}")
+
+        if solver_type not in ["midpoint", "heun"]:
+            if solver_type in ["logrho", "bh1", "bh2"]:
+                self.register_to_config(solver_type="midpoint")
+            else:
+                raise NotImplementedError(f"{solver_type} is not implemented for {self.__class__}")
+
+        if algorithm_type not in ["dpmsolver++", "sde-dpmsolver++"] and final_sigmas_type == "zero":
+            raise ValueError(
+                f"`final_sigmas_type` {final_sigmas_type} is not supported for `algorithm_type` {algorithm_type}. Please choose `sigma_min` instead."
+            )
+
+        # settable values
+        self.num_inference_steps = None
+        timesteps = np.linspace(0, num_train_timesteps - 1, num_train_timesteps, dtype=np.float32)[::-1].copy()
+        self.timesteps = torch.from_numpy(timesteps)
+        self.model_outputs = [None] * solver_order
+        self.lower_order_nums = 0
+        self._step_index = None
+        self._begin_index = None
+        self.sigmas = self.sigmas.to("cpu")  # to avoid too much CPU/GPU communication
+
+    @property
+    def step_index(self):
+        """
+        The index counter for current timestep. It will increase 1 after each scheduler step.
+        """
+        return self._step_index
+
+    @property
+    def begin_index(self):
+        """
+        The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
+        """
+        return self._begin_index
+
+    def set_begin_index(self, begin_index: int = 0):
+        """
+        Sets the begin index for the scheduler. This function should be run from pipeline before the inference.
+
+        Args:
+            begin_index (`int`):
+                The begin index for the scheduler.
+        """
+        self._begin_index = begin_index
+
+    def set_timesteps(
+        self,
+        num_inference_steps: int = None,
+        device: Union[str, torch.device] = None,
+        timesteps: Optional[List[int]] = None,
+    ):
+        """
+        Sets the discrete timesteps used for the diffusion chain (to be run before inference).
+
+        Args:
+            num_inference_steps (`int`):
+                The number of diffusion steps used when generating samples with a pre-trained model.
+            device (`str` or `torch.device`, *optional*):
+                The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+            timesteps (`List[int]`, *optional*):
+                Custom timesteps used to support arbitrary timesteps schedule. If `None`, timesteps will be generated
+                based on the `timestep_spacing` attribute. If `timesteps` is passed, `num_inference_steps` and `sigmas`
+                must be `None`, and `timestep_spacing` attribute will be ignored.
+        """
+        if num_inference_steps is None and timesteps is None:
+            raise ValueError("Must pass exactly one of `num_inference_steps` or `timesteps`.")
+        if num_inference_steps is not None and timesteps is not None:
+            raise ValueError("Can only pass one of `num_inference_steps` or `custom_timesteps`.")
+        if timesteps is not None and self.config.use_karras_sigmas:
+            raise ValueError("Cannot use `timesteps` with `config.use_karras_sigmas = True`")
+        if timesteps is not None and self.config.use_lu_lambdas:
+            raise ValueError("Cannot use `timesteps` with `config.use_lu_lambdas = True`")
+
+        if timesteps is not None:
+            timesteps = np.array(timesteps).astype(np.int64)
+        else:
+            # Clipping the minimum of all lambda(t) for numerical stability.
+            # This is critical for cosine (squaredcos_cap_v2) noise schedule.
+            clipped_idx = torch.searchsorted(torch.flip(self.lambda_t, [0]), self.config.lambda_min_clipped)
+            last_timestep = ((self.config.num_train_timesteps - clipped_idx).numpy()).item()
+
+            # "linspace", "leading", "trailing" corresponds to annotation of Table 2. of https://arxiv.org/abs/2305.08891
+            if self.config.timestep_spacing == "linspace":
+                timesteps = (
+                    np.linspace(0, last_timestep - 1, num_inference_steps + 1)
+                    .round()[::-1][:-1]
+                    .copy()
+                    .astype(np.int64)
+                )
+            elif self.config.timestep_spacing == "leading":
+                step_ratio = last_timestep // (num_inference_steps + 1)
+                # creates integer timesteps by multiplying by ratio
+                # casting to int to avoid issues when num_inference_step is power of 3
+                timesteps = (
+                    (np.arange(0, num_inference_steps + 1) * step_ratio).round()[::-1][:-1].copy().astype(np.int64)
+                )
+                timesteps += self.config.steps_offset
+            elif self.config.timestep_spacing == "trailing":
+                step_ratio = self.config.num_train_timesteps / num_inference_steps
+                # creates integer timesteps by multiplying by ratio
+                # casting to int to avoid issues when num_inference_step is power of 3
+                timesteps = np.arange(last_timestep, 0, -step_ratio).round().copy().astype(np.int64)
+                timesteps -= 1
+            else:
+                raise ValueError(
+                    f"{self.config.timestep_spacing} is not supported. Please make sure to choose one of 'linspace', 'leading' or 'trailing'."
+                )
+
+        sigmas = np.array(((1 - self.alphas_cumprod) / self.alphas_cumprod) ** 0.5)
+        log_sigmas = np.log(sigmas)
+
+        if self.config.use_karras_sigmas:
+            sigmas = np.flip(sigmas).copy()
+            sigmas = self._convert_to_karras(in_sigmas=sigmas, num_inference_steps=num_inference_steps)
+            timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]).round()
+        elif self.config.use_lu_lambdas:
+            lambdas = np.flip(log_sigmas.copy())
+            lambdas = self._convert_to_lu(in_lambdas=lambdas, num_inference_steps=num_inference_steps)
+            sigmas = np.exp(lambdas)
+            timesteps = np.array([self._sigma_to_t(sigma, log_sigmas) for sigma in sigmas]).round()
+        else:
+            sigmas = np.interp(timesteps, np.arange(0, len(sigmas)), sigmas)
+
+        if self.config.final_sigmas_type == "sigma_min":
+            sigma_last = ((1 - self.alphas_cumprod[0]) / self.alphas_cumprod[0]) ** 0.5
+        elif self.config.final_sigmas_type == "zero":
+            sigma_last = 0
+        else:
+            raise ValueError(
+                f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}"
+            )
+
+        sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32)
+
+        self.sigmas = torch.from_numpy(sigmas)
+        self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=torch.int64)
+
+        self.num_inference_steps = len(timesteps)
+
+        self.model_outputs = [
+            None,
+        ] * self.config.solver_order
+        self.lower_order_nums = 0
+
+        # add an index counter for schedulers that allow duplicated timesteps
+        self._step_index = None
+        self._begin_index = None
+        self.sigmas = self.sigmas.to("cpu")  # to avoid too much CPU/GPU communication
+
+    # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample
+    def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor:
+        """
+        "Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the
+        prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by
+        s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing
+        pixels from saturation at each step. We find that dynamic thresholding results in significantly better
+        photorealism as well as better image-text alignment, especially when using very large guidance weights."
+
+        https://arxiv.org/abs/2205.11487
+        """
+        dtype = sample.dtype
+        batch_size, channels, *remaining_dims = sample.shape
+
+        if dtype not in (torch.float32, torch.float64):
+            sample = sample.float()  # upcast for quantile calculation, and clamp not implemented for cpu half
+
+        # Flatten sample for doing quantile calculation along each image
+        sample = sample.reshape(batch_size, channels * np.prod(remaining_dims))
+
+        abs_sample = sample.abs()  # "a certain percentile absolute pixel value"
+
+        s = torch.quantile(abs_sample, self.config.dynamic_thresholding_ratio, dim=1)
+        s = torch.clamp(
+            s, min=1, max=self.config.sample_max_value
+        )  # When clamped to min=1, equivalent to standard clipping to [-1, 1]
+        s = s.unsqueeze(1)  # (batch_size, 1) because clamp will broadcast along dim=0
+        sample = torch.clamp(sample, -s, s) / s  # "we threshold xt0 to the range [-s, s] and then divide by s"
+
+        sample = sample.reshape(batch_size, channels, *remaining_dims)
+        sample = sample.to(dtype)
+
+        return sample
+
+    # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._sigma_to_t
+    def _sigma_to_t(self, sigma, log_sigmas):
+        # get log sigma
+        log_sigma = np.log(np.maximum(sigma, 1e-10))
+
+        # get distribution
+        dists = log_sigma - log_sigmas[:, np.newaxis]
+
+        # get sigmas range
+        low_idx = np.cumsum((dists >= 0), axis=0).argmax(axis=0).clip(max=log_sigmas.shape[0] - 2)
+        high_idx = low_idx + 1
+
+        low = log_sigmas[low_idx]
+        high = log_sigmas[high_idx]
+
+        # interpolate sigmas
+        w = (low - log_sigma) / (low - high)
+        w = np.clip(w, 0, 1)
+
+        # transform interpolation to time range
+        t = (1 - w) * low_idx + w * high_idx
+        t = t.reshape(sigma.shape)
+        return t
+
+    def _sigma_to_alpha_sigma_t(self, sigma):
+        alpha_t = 1 / ((sigma**2 + 1) ** 0.5)
+        sigma_t = sigma * alpha_t
+
+        return alpha_t, sigma_t
+
+    # Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._convert_to_karras
+    def _convert_to_karras(self, in_sigmas: torch.Tensor, num_inference_steps) -> torch.Tensor:
+        """Constructs the noise schedule of Karras et al. (2022)."""
+
+        # Hack to make sure that other schedulers which copy this function don't break
+        # TODO: Add this logic to the other schedulers
+        if hasattr(self.config, "sigma_min"):
+            sigma_min = self.config.sigma_min
+        else:
+            sigma_min = None
+
+        if hasattr(self.config, "sigma_max"):
+            sigma_max = self.config.sigma_max
+        else:
+            sigma_max = None
+
+        sigma_min = sigma_min if sigma_min is not None else in_sigmas[-1].item()
+        sigma_max = sigma_max if sigma_max is not None else in_sigmas[0].item()
+
+        rho = 7.0  # 7.0 is the value used in the paper
+        ramp = np.linspace(0, 1, num_inference_steps)
+        min_inv_rho = sigma_min ** (1 / rho)
+        max_inv_rho = sigma_max ** (1 / rho)
+        sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
+        return sigmas
+
+    def _convert_to_lu(self, in_lambdas: torch.Tensor, num_inference_steps) -> torch.Tensor:
+        """Constructs the noise schedule of Lu et al. (2022)."""
+
+        lambda_min: float = in_lambdas[-1].item()
+        lambda_max: float = in_lambdas[0].item()
+
+        rho = 1.0  # 1.0 is the value used in the paper
+        ramp = np.linspace(0, 1, num_inference_steps)
+        min_inv_rho = lambda_min ** (1 / rho)
+        max_inv_rho = lambda_max ** (1 / rho)
+        lambdas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** rho
+        return lambdas
+
+    def convert_model_output(
+        self,
+        model_output: torch.Tensor,
+        *args,
+        sample: torch.Tensor = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        """
+        Convert the model output to the corresponding type the DPMSolver/DPMSolver++ algorithm needs. DPM-Solver is
+        designed to discretize an integral of the noise prediction model, and DPM-Solver++ is designed to discretize an
+        integral of the data prediction model.
+
+        <Tip>
+
+        The algorithm and model type are decoupled. You can use either DPMSolver or DPMSolver++ for both noise
+        prediction and data prediction models.
+
+        </Tip>
+
+        Args:
+            model_output (`torch.Tensor`):
+                The direct output from the learned diffusion model.
+            sample (`torch.Tensor`):
+                A current instance of a sample created by the diffusion process.
+
+        Returns:
+            `torch.Tensor`:
+                The converted model output.
+        """
+        timestep = args[0] if len(args) > 0 else kwargs.pop("timestep", None)
+        if sample is None:
+            if len(args) > 1:
+                sample = args[1]
+            else:
+                raise ValueError("missing `sample` as a required keyword argument")
+        if timestep is not None:
+            deprecate(
+                "timesteps",
+                "1.0.0",
+                "Passing `timesteps` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+            )
+
+        # DPM-Solver++ needs to solve an integral of the data prediction model.
+        if self.config.algorithm_type in ["dpmsolver++", "sde-dpmsolver++"]:
+            if self.config.prediction_type == "epsilon":
+                # DPM-Solver and DPM-Solver++ only need the "mean" output.
+                if self.config.variance_type in ["learned", "learned_range"]:
+                    model_output = model_output[:, :3]
+                sigma = self.sigmas[self.step_index]
+                alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)
+                x0_pred = (sample - sigma_t * model_output) / alpha_t
+            elif self.config.prediction_type == "sample":
+                x0_pred = model_output
+            elif self.config.prediction_type == "v_prediction":
+                sigma = self.sigmas[self.step_index]
+                alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)
+                x0_pred = alpha_t * sample - sigma_t * model_output
+            else:
+                raise ValueError(
+                    f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, or"
+                    " `v_prediction` for the DPMSolverMultistepScheduler."
+                )
+
+            if self.config.thresholding:
+                x0_pred = self._threshold_sample(x0_pred)
+
+            return x0_pred
+
+        # DPM-Solver needs to solve an integral of the noise prediction model.
+        elif self.config.algorithm_type in ["dpmsolver", "sde-dpmsolver"]:
+            if self.config.prediction_type == "epsilon":
+                # DPM-Solver and DPM-Solver++ only need the "mean" output.
+                if self.config.variance_type in ["learned", "learned_range"]:
+                    epsilon = model_output[:, :3]
+                else:
+                    epsilon = model_output
+            elif self.config.prediction_type == "sample":
+                sigma = self.sigmas[self.step_index]
+                alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)
+                epsilon = (sample - alpha_t * model_output) / sigma_t
+            elif self.config.prediction_type == "v_prediction":
+                sigma = self.sigmas[self.step_index]
+                alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)
+                epsilon = alpha_t * model_output + sigma_t * sample
+            else:
+                raise ValueError(
+                    f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, or"
+                    " `v_prediction` for the DPMSolverMultistepScheduler."
+                )
+
+            if self.config.thresholding:
+                sigma = self.sigmas[self.step_index]
+                alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)
+                x0_pred = (sample - sigma_t * epsilon) / alpha_t
+                x0_pred = self._threshold_sample(x0_pred)
+                epsilon = (sample - alpha_t * x0_pred) / sigma_t
+
+            return epsilon
+
+    def dpm_solver_first_order_update(
+        self,
+        model_output: torch.Tensor,
+        *args,
+        sample: torch.Tensor = None,
+        noise: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        """
+        One step for the first-order DPMSolver (equivalent to DDIM).
+
+        Args:
+            model_output (`torch.Tensor`):
+                The direct output from the learned diffusion model.
+            sample (`torch.Tensor`):
+                A current instance of a sample created by the diffusion process.
+
+        Returns:
+            `torch.Tensor`:
+                The sample tensor at the previous timestep.
+        """
+        timestep = args[0] if len(args) > 0 else kwargs.pop("timestep", None)
+        prev_timestep = args[1] if len(args) > 1 else kwargs.pop("prev_timestep", None)
+        if sample is None:
+            if len(args) > 2:
+                sample = args[2]
+            else:
+                raise ValueError(" missing `sample` as a required keyword argument")
+        if timestep is not None:
+            deprecate(
+                "timesteps",
+                "1.0.0",
+                "Passing `timesteps` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+            )
+
+        if prev_timestep is not None:
+            deprecate(
+                "prev_timestep",
+                "1.0.0",
+                "Passing `prev_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+            )
+
+        sigma_t, sigma_s = self.sigmas[self.step_index + 1], self.sigmas[self.step_index]
+        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
+        alpha_s, sigma_s = self._sigma_to_alpha_sigma_t(sigma_s)
+        lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
+        lambda_s = torch.log(alpha_s) - torch.log(sigma_s)
+
+        h = lambda_t - lambda_s
+        if self.config.algorithm_type == "dpmsolver++":
+            x_t = (sigma_t / sigma_s) * sample - (alpha_t * (torch.exp(-h) - 1.0)) * model_output
+        elif self.config.algorithm_type == "dpmsolver":
+            x_t = (alpha_t / alpha_s) * sample - (sigma_t * (torch.exp(h) - 1.0)) * model_output
+        elif self.config.algorithm_type == "sde-dpmsolver++":
+            assert noise is not None
+            x_t = (
+                (sigma_t / sigma_s * torch.exp(-h)) * sample
+                + (alpha_t * (1 - torch.exp(-2.0 * h))) * model_output
+                + sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) * noise
+            )
+        elif self.config.algorithm_type == "sde-dpmsolver":
+            assert noise is not None
+            x_t = (
+                (alpha_t / alpha_s) * sample
+                - 2.0 * (sigma_t * (torch.exp(h) - 1.0)) * model_output
+                + sigma_t * torch.sqrt(torch.exp(2 * h) - 1.0) * noise
+            )
+        return x_t
+
+    def multistep_dpm_solver_second_order_update(
+        self,
+        model_output_list: List[torch.Tensor],
+        *args,
+        sample: torch.Tensor = None,
+        noise: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        """
+        One step for the second-order multistep DPMSolver.
+
+        Args:
+            model_output_list (`List[torch.Tensor]`):
+                The direct outputs from learned diffusion model at current and latter timesteps.
+            sample (`torch.Tensor`):
+                A current instance of a sample created by the diffusion process.
+
+        Returns:
+            `torch.Tensor`:
+                The sample tensor at the previous timestep.
+        """
+        timestep_list = args[0] if len(args) > 0 else kwargs.pop("timestep_list", None)
+        prev_timestep = args[1] if len(args) > 1 else kwargs.pop("prev_timestep", None)
+        if sample is None:
+            if len(args) > 2:
+                sample = args[2]
+            else:
+                raise ValueError(" missing `sample` as a required keyword argument")
+        if timestep_list is not None:
+            deprecate(
+                "timestep_list",
+                "1.0.0",
+                "Passing `timestep_list` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+            )
+
+        if prev_timestep is not None:
+            deprecate(
+                "prev_timestep",
+                "1.0.0",
+                "Passing `prev_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+            )
+
+        sigma_t, sigma_s0, sigma_s1 = (
+            self.sigmas[self.step_index + 1],
+            self.sigmas[self.step_index],
+            self.sigmas[self.step_index - 1],
+        )
+
+        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
+        alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)
+        alpha_s1, sigma_s1 = self._sigma_to_alpha_sigma_t(sigma_s1)
+
+        lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
+        lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0)
+        lambda_s1 = torch.log(alpha_s1) - torch.log(sigma_s1)
+
+        m0, m1 = model_output_list[-1], model_output_list[-2]
+
+        h, h_0 = lambda_t - lambda_s0, lambda_s0 - lambda_s1
+        r0 = h_0 / h
+        D0, D1 = m0, (1.0 / r0) * (m0 - m1)
+        if self.config.algorithm_type == "dpmsolver++":
+            # See https://arxiv.org/abs/2211.01095 for detailed derivations
+            if self.config.solver_type == "midpoint":
+                x_t = (
+                    (sigma_t / sigma_s0) * sample
+                    - (alpha_t * (torch.exp(-h) - 1.0)) * D0
+                    - 0.5 * (alpha_t * (torch.exp(-h) - 1.0)) * D1
+                )
+            elif self.config.solver_type == "heun":
+                x_t = (
+                    (sigma_t / sigma_s0) * sample
+                    - (alpha_t * (torch.exp(-h) - 1.0)) * D0
+                    + (alpha_t * ((torch.exp(-h) - 1.0) / h + 1.0)) * D1
+                )
+        elif self.config.algorithm_type == "dpmsolver":
+            # See https://arxiv.org/abs/2206.00927 for detailed derivations
+            if self.config.solver_type == "midpoint":
+                x_t = (
+                    (alpha_t / alpha_s0) * sample
+                    - (sigma_t * (torch.exp(h) - 1.0)) * D0
+                    - 0.5 * (sigma_t * (torch.exp(h) - 1.0)) * D1
+                )
+            elif self.config.solver_type == "heun":
+                x_t = (
+                    (alpha_t / alpha_s0) * sample
+                    - (sigma_t * (torch.exp(h) - 1.0)) * D0
+                    - (sigma_t * ((torch.exp(h) - 1.0) / h - 1.0)) * D1
+                )
+        elif self.config.algorithm_type == "sde-dpmsolver++":
+            assert noise is not None
+            if self.config.solver_type == "midpoint":
+                x_t = (
+                    (sigma_t / sigma_s0 * torch.exp(-h)) * sample
+                    + (alpha_t * (1 - torch.exp(-2.0 * h))) * D0
+                    + 0.5 * (alpha_t * (1 - torch.exp(-2.0 * h))) * D1
+                    + sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) * noise
+                )
+            elif self.config.solver_type == "heun":
+                x_t = (
+                    (sigma_t / sigma_s0 * torch.exp(-h)) * sample
+                    + (alpha_t * (1 - torch.exp(-2.0 * h))) * D0
+                    + (alpha_t * ((1.0 - torch.exp(-2.0 * h)) / (-2.0 * h) + 1.0)) * D1
+                    + sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) * noise
+                )
+        elif self.config.algorithm_type == "sde-dpmsolver":
+            assert noise is not None
+            if self.config.solver_type == "midpoint":
+                x_t = (
+                    (alpha_t / alpha_s0) * sample
+                    - 2.0 * (sigma_t * (torch.exp(h) - 1.0)) * D0
+                    - (sigma_t * (torch.exp(h) - 1.0)) * D1
+                    + sigma_t * torch.sqrt(torch.exp(2 * h) - 1.0) * noise
+                )
+            elif self.config.solver_type == "heun":
+                x_t = (
+                    (alpha_t / alpha_s0) * sample
+                    - 2.0 * (sigma_t * (torch.exp(h) - 1.0)) * D0
+                    - 2.0 * (sigma_t * ((torch.exp(h) - 1.0) / h - 1.0)) * D1
+                    + sigma_t * torch.sqrt(torch.exp(2 * h) - 1.0) * noise
+                )
+        return x_t
+
+    def multistep_dpm_solver_third_order_update(
+        self,
+        model_output_list: List[torch.Tensor],
+        *args,
+        sample: torch.Tensor = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        """
+        One step for the third-order multistep DPMSolver.
+
+        Args:
+            model_output_list (`List[torch.Tensor]`):
+                The direct outputs from learned diffusion model at current and latter timesteps.
+            sample (`torch.Tensor`):
+                A current instance of a sample created by diffusion process.
+
+        Returns:
+            `torch.Tensor`:
+                The sample tensor at the previous timestep.
+        """
+
+        timestep_list = args[0] if len(args) > 0 else kwargs.pop("timestep_list", None)
+        prev_timestep = args[1] if len(args) > 1 else kwargs.pop("prev_timestep", None)
+        if sample is None:
+            if len(args) > 2:
+                sample = args[2]
+            else:
+                raise ValueError(" missing`sample` as a required keyword argument")
+        if timestep_list is not None:
+            deprecate(
+                "timestep_list",
+                "1.0.0",
+                "Passing `timestep_list` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+            )
+
+        if prev_timestep is not None:
+            deprecate(
+                "prev_timestep",
+                "1.0.0",
+                "Passing `prev_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+            )
+
+        sigma_t, sigma_s0, sigma_s1, sigma_s2 = (
+            self.sigmas[self.step_index + 1],
+            self.sigmas[self.step_index],
+            self.sigmas[self.step_index - 1],
+            self.sigmas[self.step_index - 2],
+        )
+
+        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
+        alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)
+        alpha_s1, sigma_s1 = self._sigma_to_alpha_sigma_t(sigma_s1)
+        alpha_s2, sigma_s2 = self._sigma_to_alpha_sigma_t(sigma_s2)
+
+        lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
+        lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0)
+        lambda_s1 = torch.log(alpha_s1) - torch.log(sigma_s1)
+        lambda_s2 = torch.log(alpha_s2) - torch.log(sigma_s2)
+
+        m0, m1, m2 = model_output_list[-1], model_output_list[-2], model_output_list[-3]
+
+        h, h_0, h_1 = lambda_t - lambda_s0, lambda_s0 - lambda_s1, lambda_s1 - lambda_s2
+        r0, r1 = h_0 / h, h_1 / h
+        D0 = m0
+        D1_0, D1_1 = (1.0 / r0) * (m0 - m1), (1.0 / r1) * (m1 - m2)
+        D1 = D1_0 + (r0 / (r0 + r1)) * (D1_0 - D1_1)
+        D2 = (1.0 / (r0 + r1)) * (D1_0 - D1_1)
+        if self.config.algorithm_type == "dpmsolver++":
+            # See https://arxiv.org/abs/2206.00927 for detailed derivations
+            x_t = (
+                (sigma_t / sigma_s0) * sample
+                - (alpha_t * (torch.exp(-h) - 1.0)) * D0
+                + (alpha_t * ((torch.exp(-h) - 1.0) / h + 1.0)) * D1
+                - (alpha_t * ((torch.exp(-h) - 1.0 + h) / h**2 - 0.5)) * D2
+            )
+        elif self.config.algorithm_type == "dpmsolver":
+            # See https://arxiv.org/abs/2206.00927 for detailed derivations
+            x_t = (
+                (alpha_t / alpha_s0) * sample
+                - (sigma_t * (torch.exp(h) - 1.0)) * D0
+                - (sigma_t * ((torch.exp(h) - 1.0) / h - 1.0)) * D1
+                - (sigma_t * ((torch.exp(h) - 1.0 - h) / h**2 - 0.5)) * D2
+            )
+        return x_t
+
+    def index_for_timestep(self, timestep, schedule_timesteps=None):
+        if schedule_timesteps is None:
+            schedule_timesteps = self.timesteps
+
+        index_candidates = (schedule_timesteps == timestep).nonzero()
+
+        if len(index_candidates) == 0:
+            step_index = len(self.timesteps) - 1
+        # The sigma index that is taken for the **very** first `step`
+        # is always the second index (or the last index if there is only 1)
+        # This way we can ensure we don't accidentally skip a sigma in
+        # case we start in the middle of the denoising schedule (e.g. for image-to-image)
+        elif len(index_candidates) > 1:
+            step_index = index_candidates[1].item()
+        else:
+            step_index = index_candidates[0].item()
+
+        return step_index
+
+    def _init_step_index(self, timestep):
+        """
+        Initialize the step_index counter for the scheduler.
+        """
+
+        if self.begin_index is None:
+            if isinstance(timestep, torch.Tensor):
+                timestep = timestep.to(self.timesteps.device)
+            self._step_index = self.index_for_timestep(timestep)
+        else:
+            self._step_index = self._begin_index
+
+    def step(
+        self,
+        model_output: torch.Tensor,
+        timestep: int,
+        sample: torch.Tensor,
+        generator=None,
+        variance_noise: Optional[torch.Tensor] = None,
+        return_dict: bool = True,
+    ) -> Union[SchedulerOutput, Tuple]:
+        """
+        Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with
+        the multistep DPMSolver.
+
+        Args:
+            model_output (`torch.Tensor`):
+                The direct output from learned diffusion model.
+            timestep (`int`):
+                The current discrete timestep in the diffusion chain.
+            sample (`torch.Tensor`):
+                A current instance of a sample created by the diffusion process.
+            generator (`torch.Generator`, *optional*):
+                A random number generator.
+            variance_noise (`torch.Tensor`):
+                Alternative to generating noise with `generator` by directly providing the noise for the variance
+                itself. Useful for methods such as [`LEdits++`].
+            return_dict (`bool`):
+                Whether or not to return a [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`.
+
+        Returns:
+            [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`:
+                If return_dict is `True`, [`~schedulers.scheduling_utils.SchedulerOutput`] is returned; otherwise, a
+                tuple is returned where the first element is the sample tensor.
+
+        """
+        if self.num_inference_steps is None:
+            raise ValueError(
+                "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler"
+            )
+
+        if self.step_index is None:
+            self._init_step_index(timestep)
+
+        # Improve numerical stability for small number of steps
+        lower_order_final = (self.step_index == len(self.timesteps) - 1) and (
+            self.config.euler_at_final
+            or (self.config.lower_order_final and len(self.timesteps) < 15)
+            or self.config.final_sigmas_type == "zero"
+        )
+        lower_order_second = (
+            (self.step_index == len(self.timesteps) - 2) and self.config.lower_order_final and len(self.timesteps) < 15
+        )
+
+        model_output = self.convert_model_output(model_output, sample=sample)
+        for i in range(self.config.solver_order - 1):
+            self.model_outputs[i] = self.model_outputs[i + 1]
+        self.model_outputs[-1] = model_output
+
+        # Upcast to avoid precision issues when computing prev_sample
+        sample = sample.to(torch.float32)
+        if self.config.algorithm_type in ["sde-dpmsolver", "sde-dpmsolver++"] and variance_noise is None:
+            noise = randn_tensor(
+                model_output.shape, generator=generator, device=model_output.device, dtype=torch.float32
+            )
+        elif self.config.algorithm_type in ["sde-dpmsolver", "sde-dpmsolver++"]:
+            noise = variance_noise.to(device=model_output.device, dtype=torch.float32)
+        else:
+            noise = None
+
+        if self.config.solver_order == 1 or self.lower_order_nums < 1 or lower_order_final:
+            prev_sample = self.dpm_solver_first_order_update(model_output, sample=sample, noise=noise)
+        elif self.config.solver_order == 2 or self.lower_order_nums < 2 or lower_order_second:
+            prev_sample = self.multistep_dpm_solver_second_order_update(self.model_outputs, sample=sample, noise=noise)
+        else:
+            prev_sample = self.multistep_dpm_solver_third_order_update(self.model_outputs, sample=sample)
+
+        if self.lower_order_nums < self.config.solver_order:
+            self.lower_order_nums += 1
+
+        # Cast sample back to expected dtype
+        prev_sample = prev_sample.to(model_output.dtype)
+
+        # upon completion increase step index by one
+        self._step_index += 1
+
+        if not return_dict:
+            return (prev_sample,)
+
+        return SchedulerOutput(prev_sample=prev_sample)
+
+    def add_noise(
+        self,
+        original_samples: torch.Tensor,
+        noise: torch.Tensor,
+        timesteps: torch.IntTensor,
+    ) -> torch.Tensor:
+        # Make sure sigmas and timesteps have the same device and dtype as original_samples
+        # alpha_t = self.alpha_t.to(device=original_samples.device, dtype=original_samples.dtype)
+        # sigma_t = self.sigma_t.to(device=original_samples.device, dtype=original_samples.dtype)
+        alpha_t = self.alpha_t.to(original_samples.device).to(original_samples.dtype)
+        sigma_t = self.sigma_t.to(original_samples.device).to(original_samples.dtype)
+        timesteps = timesteps.to(original_samples.device)
+        alpha_t = alpha_t[timesteps].flatten()
+        while len(alpha_t.shape) < len(original_samples.shape):
+            alpha_t = alpha_t.unsqueeze(-1)
+
+        sigma_t = sigma_t[timesteps].flatten()
+        while len(sigma_t.shape) < len(original_samples.shape):
+            sigma_t = sigma_t.unsqueeze(-1)
+        noisy_samples = alpha_t * original_samples + sigma_t * noise
+        return noisy_samples
+    
+    def get_velocity(self, original_samples: torch.Tensor, noise: torch.Tensor, timesteps: torch.IntTensor) -> torch.Tensor:
+        # alpha_t = self.alpha_t.to(device=original_samples.device, dtype=original_samples.dtype)
+        # sigma_t = self.sigma_t.to(device=original_samples.device, dtype=original_samples.dtype)
+        alpha_t = self.alpha_t.to(original_samples.device).to(original_samples.dtype)
+        sigma_t = self.sigma_t.to(original_samples.device).to(original_samples.dtype)
+
+        timesteps = timesteps.to(original_samples.device)
+        alpha_t = alpha_t[timesteps].flatten()
+        while len(alpha_t.shape) < len(original_samples.shape):
+            alpha_t = alpha_t.unsqueeze(-1)
+
+        sigma_t = sigma_t[timesteps].flatten()
+        while len(sigma_t.shape) < len(original_samples.shape):
+            sigma_t = sigma_t.unsqueeze(-1)
+
+        velocity = alpha_t * noise - sigma_t * original_samples
+        return velocity
+
+    def __len__(self):
+        return self.config.num_train_timesteps

vibevoice/schedule/timestep_sampler.py

@@ -0,0 +1,19 @@
+import math
+import torch
+
+
+class UniformSampler:
+    def __init__(self, timesteps = 1000):
+        self.timesteps = timesteps
+    def sample(self, batch_size, device):
+        return torch.randint(0, self.timesteps, (batch_size,), device=device)
+    
+class LogitNormalSampler:
+    def __init__(self, timesteps = 1000, m = 0, s = 1):
+        self.timesteps = timesteps
+        timesteps = torch.linspace(0, 1, timesteps)
+        logit = torch.log(timesteps / (1 - timesteps))
+        self.prob = torch.exp(-0.5 * (logit - m) ** 2 / s ** 2) / (s * math.sqrt(2 * math.pi))
+    def sample(self, batch_size, device):
+        return torch.multinomial(self.prob, batch_size, replacement=True).to(device)
+    

why1.wav

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