v1 delay-pattern (k=8): replace flat k=3 with MusicGen-style delay

README.md

@@ -14,52 +14,62 @@ tags:
 - neural-codec
 pipeline_tag: text-to-speech
 datasets:
-- openslr/librispeech_asr
+- mythicinfinity/libritts_r
+- keithito/lj_speech
 ---
 
 # Wren-TTS-360M (v1)
 
 **Wren** is a series of small (<3B) multimodal speech LLMs covering TTS, ASR, and
-speech-language modelling. This is the first public checkpoint: **Wren-TTS-360M-v1**,
-a text-to-speech model that generates [Kyutai Mimi](https://huggingface.co/kyutai/mimi)
-neural-codec tokens from text and decodes them to 24 kHz waveform with the Mimi decoder.
-The autoregressive backbone is [HuggingFaceTB/SmolLM2-360M](https://huggingface.co/HuggingFaceTB/SmolLM2-360M).
+speech-language modelling. **Wren-TTS-360M-v1** generates
+[Kyutai Mimi](https://huggingface.co/kyutai/mimi) neural-codec tokens from text
+using a [HuggingFaceTB/SmolLM2-360M](https://huggingface.co/HuggingFaceTB/SmolLM2-360M)
+backbone, then decodes to 24 kHz waveform with the Mimi decoder.
 
-An early research checkpoint — useful for experimentation, not production.
+An open research checkpoint — useful for experimentation, not production.
 
 ## Architecture
 
 ```
-text ──► SmolLM2-360M ──► k audio-code heads ──► Mimi decoder ──► 24 kHz waveform
+text ──► SmolLM2-360M ──► k=8 parallel Mimi heads ──► Mimi decoder ──► 24 kHz
 ```
 
-- **Backbone:** SmolLM2-360M (causal LM, embeddings shared with text input)
+- **Backbone:** SmolLM2-360M (causal LM; text + audio share the same backbone)
 - **Audio tokenizer:** Mimi (`kyutai/mimi`), 12.5 fps, 2048-entry codebooks
-- **Codebooks used:** 3 (of Mimi's 8 extractable)
-- **Interleaved layout:** `[ text | <audio_sep> | cb0_f0 cb1_f0 cb2_f0 | cb0_f1 ... ]`
-- **Per-codebook heads:** `Linear(hidden, 2048)`. `cb0` has one extra output (index 2048)
-  used as the `AUDIO_EOS` stop token.
-- **Optional speaker conditioning:** prepend `<audio_start> ref_codes <audio_end>` before the
-  text prompt for zero-shot voice cloning from a short reference clip.
+- **Codebooks used:** all 8 Mimi codebooks
+- **Layout:** **MusicGen-style delay pattern** — at each step, k summed codebook
+  embeddings go in, k parallel heads predict k tokens out. Codebook q at frame f
+  lives at step `s = f + q`, so same-frame RVQ conditioning is preserved via the delay.
+- **Per-codebook input tables:** `Embedding(2049, hidden)` — extra row = `AUDIO_PAD` at
+  sequence edges.
+- **Per-codebook output heads:** `Linear(hidden, 2048)` for cb1..cb7.
+  cb0 gets `Linear(hidden, 2049)` with the extra class = `AUDIO_EOS` (stop token).
+- **Speaker conditioning (required):** prepend `<|reference_start|> ref_codes <|reference_end|>`
+  to the prompt; `ref_codes` is the Mimi encoding of a short reference clip. The model was
+  trained multispeaker-only and expects a reference at inference — without one, output quality
+  is poor.
 
 ## Training data
 
-- [LibriSpeech](https://www.openslr.org/12) `train-clean-100` + `train-clean-360` — ~460 h, multi-speaker, English audiobooks
-- Trained for 5 epochs with effective batch size 32
+- [LibriTTS-R](https://huggingface.co/datasets/mythicinfinity/libritts_r)
+  `train-clean-{100,360}` + `train-other-500` — ~960 h multi-speaker English
+- [LJSpeech](https://keithito.com/LJ-Speech-Dataset/) — ~24 h single speaker
 
-Text is lowercased before tokenization. This v1 checkpoint was trained on LibriSpeech
-only; LJSpeech was not included in this run.
+Text casing and punctuation are preserved. Pass text naturally — do not pre-lowercase.
 
 ## Usage
 
 ```bash
-pip install torch torchaudio transformers
+pip install torch torchaudio transformers datasets
 ```
 
-### Text-to-speech
+> **A reference audio clip is required.** The model was trained multispeaker-only; without
+> `ref_codes` it produces poor output. Any 3–12 s English speech clip works as the voice reference.
 
 ```python
 import torch
+import numpy as np
+from datasets import load_dataset
 from transformers import AutoModel, AutoProcessor
 
 model_id = "shangeth/Wren-TTS-360M-v1"
@@ -68,74 +78,75 @@ device   = "cuda" if torch.cuda.is_available() else "cpu"
 processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
 model     = AutoModel.from_pretrained(model_id, trust_remote_code=True).to(device).eval()
 
+# Grab one LibriSpeech test-clean clip (~3.5s, not in training) as the reference voice.
+# Swap this block for `torchaudio.load("your_reference.wav")` to use your own clip.
+sample  = next(iter(load_dataset("openslr/librispeech_asr", "clean", split="test", streaming=True)))
+ref_wav = torch.from_numpy(np.asarray(sample["audio"]["array"], dtype=np.float32)).unsqueeze(0)
+ref_sr  = sample["audio"]["sampling_rate"]
+ref_codes = model.encode_audio(ref_wav, ref_sr)[:, :150]   # cap at ~12s; encode_audio resamples to 24 kHz
+
+# Tokenize the target text and generate speech in the reference voice
 inputs = processor("Hello world, how are you today?")
 inputs = {k: v.to(device) for k, v in inputs.items()}
 
 waveform = model.generate(
     **inputs,
+    ref_codes=ref_codes,
     max_audio_frames=200,
     min_audio_frames=2,
-    temperature=0.8,
-    top_k=50,
-    top_p=0.9,
+    temperature=0.8, top_k=50, top_p=0.9,
     output_audio=True,
 )
 processor.save_audio(waveform, "out.wav")
 ```
 
-### Zero-shot voice cloning
-
-```python
-import torchaudio
+## Sampling tips
 
-ref_wav, sr = torchaudio.load("reference.wav")
-ref_codes   = model.encode_audio(ref_wav, sr)[:, :150]   # cap at ~12 s of reference
+Defaults: `temperature=0.8`, `top_k=50`, `top_p=0.9`, `max_audio_frames=200` (~16 s).
+If you hear the model generating extra speech past the intended text (hallucination):
 
-waveform = model.generate(
-    **inputs,
-    ref_codes=ref_codes,
-    output_audio=True,
-    max_audio_frames=200,
-    min_audio_frames=2,
-    temperature=0.8, top_k=50, top_p=0.9,
-)
-processor.save_audio(waveform, "cloned.wav")
-```
+- Raise `eos_bias` — e.g. 2.0–6.0 — to make the model more eager to stop
+- Lower `temperature` (0.6) and `top_p` (0.8)
+- Set `max_audio_frames` ≈ `12 * len(text_in_chars)`
+- Set `min_audio_frames=1` for very short prompts
 
-## Sampling tips
+## Why delay pattern
 
-Defaults are `temperature=0.8`, `top_k=50`, `top_p=0.9`, `max_audio_frames=200` (~16 s).
-If you hear the model generating extra words past the prompt (hallucination), try:
+Mimi uses **residual vector quantization (RVQ)**: cb0 is semantic, cb1..cb7 encode
+successive residuals. cb_q is only meaningful given cb0..cb_{q-1}, so same-frame
+conditioning matters.
 
-- Lower `temperature` (e.g. 0.6) and `top_p` (e.g. 0.8)
-- Raise `eos_bias` (e.g. 2.0–6.0) to make the model more eager to stop
-- Lower `max_audio_frames` to roughly `12 * len(text)`
-- Set `min_audio_frames=1` for very short prompts
+A flat interleaved layout (`cb0_f0, cb1_f0, ..., cb0_f1, ...`) preserves that
+conditioning best but balloons sequence length by `k×` and forces `k` autoregressive
+LLM calls per frame. The delay pattern keeps RVQ conditioning (cb_q at frame f is
+predicted from a hidden state that has already attended over cb0..cb_{q-1} of the
+same frame) while cutting sequence length to `T + k - 1` and LLM calls to **one per
+step** — enabling all 8 Mimi codebooks without blowing up context.
 
 ## Limitations & known issues
 
-- **Hallucinated continuations:** the model sometimes generates plausible speech beyond
-  the input text. This is a class-imbalance artifact of the `AUDIO_EOS` supervision
-  (one EOS target per ~100–300 audio frames) and is being addressed in training.
+- **Hallucinated continuations**: occasionally generates plausible speech past the
+  input text. Mitigate with `eos_bias` at inference.
 - **English only.**
-- **Limited expressiveness:** read-style prosody inherited from LibriSpeech audiobook data.
-- **Small backbone (360M)** + modest training data — quality is below production TTS systems.
-- **Val loss began to overfit** around epoch 5; v2 will add early stopping, LJSpeech,
-  and EOS-class re-weighting.
+- **Audiobook-style prosody** inherited from LibriTTS-R; not as expressive as modern
+  conversational TTS.
+- **Small backbone (360M)** — quality is below frontier TTS systems.
+- cb0 begins to overfit earlier than cb3–cb7; the released checkpoint is the
+  best-epoch point (by overall val loss) from the full training run.
 
 ## The Wren series
 
-Wren is a family of compact (<3B parameter) multimodal speech LLMs — small enough to run
-on a single consumer GPU, designed for open research on unified speech understanding and
-synthesis. Planned siblings:
+Wren is a family of compact (<3B parameter) multimodal speech LLMs — small enough
+to run on a single consumer GPU, designed for open research on unified speech
+understanding and synthesis. Planned siblings:
 
 - **Wren-TTS** — text → speech (this release)
 - **Wren-ASR** — speech → text
 - **Wren-LM** — speech-language modelling / dialog
 - **Wren-Omni** — unified ASR + TTS + LM in one checkpoint
 
-All Wren models share the same design principles: small backbone LLM + neural audio codec,
-open weights, simple PyTorch checkpoints, reproducible training recipes.
+All Wren models share the same design principles: small backbone LLM + neural
+audio codec, open weights, simple PyTorch checkpoints, reproducible training recipes.
 
 ## Repository contents
 
@@ -143,7 +154,7 @@ open weights, simple PyTorch checkpoints, reproducible training recipes.
 |---|---|
 | `model.safetensors` | Model weights |
 | `config.json` | `WrenConfig` (with `auto_map` for `trust_remote_code`) |
-| `tokenizer.json` + friends | SmolLM2 tokenizer with Wren's 4 special tokens added |
+| `tokenizer.json` + friends | SmolLM2 tokenizer with Wren's 3 special tokens added |
 | `processor_config.json` | `WrenProcessor` auto_map |
 | `configuration_wren.py` | `WrenConfig(PretrainedConfig)` |
 | `modeling_wren.py` | `WrenForTTS(PreTrainedModel)` — loads Mimi codec lazily on first generate |
@@ -160,11 +171,13 @@ open weights, simple PyTorch checkpoints, reproducible training recipes.
   url    = {https://github.com/shangeth/wren}
 }
 
-@inproceedings{panayotov2015librispeech,
-  title     = {Librispeech: an ASR corpus based on public domain audio books},
-  author    = {Panayotov, Vassil and Chen, Guoguo and Povey, Daniel and Khudanpur, Sanjeev},
-  booktitle = {ICASSP},
-  year      = {2015}
+@inproceedings{koizumi2023libritts,
+  title     = {LibriTTS-R: A Restored Multi-Speaker Text-to-Speech Corpus},
+  author    = {Koizumi, Yuma and Zen, Heiga and Karita, Shigeki and Ding, Yifan
+               and Yatabe, Kohei and Morioka, Nobuyuki and Bacchiani, Michiel and
+               Zhang, Yu and Han, Wei and Bapna, Ankur},
+  booktitle = {Interspeech},
+  year      = {2023}
 }
 ```
 

added_tokens.json

@@ -1,6 +1,5 @@
 {
-  "<|audio_end|>": 49155,
-  "<|audio_eos|>": 49153,
-  "<|audio_sep|>": 49152,
-  "<|audio_start|>": 49154
+  "<|audio_start|>": 49152,
+  "<|reference_end|>": 49154,
+  "<|reference_start|>": 49153
 }

config.json

@@ -2,20 +2,20 @@
   "architectures": [
     "WrenForTTS"
   ],
-  "audio_end_id": 49155,
-  "audio_eos_token_id": 49153,
-  "audio_sep_id": 49152,
-  "audio_start_id": 49154,
+  "audio_start_id": 49152,
   "auto_map": {
     "AutoConfig": "configuration_wren.WrenConfig",
     "AutoModel": "modeling_wren.WrenForTTS"
   },
   "codebook_size": 2048,
   "dtype": "bfloat16",
-  "k_codebooks": 3,
+  "k_codebooks": 8,
   "llm_name": "HuggingFaceTB/SmolLM2-360M",
   "mimi_model_name": "kyutai/mimi",
   "model_type": "wren",
+  "pattern": "delay",
+  "reference_end_id": 49154,
+  "reference_start_id": 49153,
   "sampling_rate": 24000,
   "transformers_version": "4.57.6",
   "vocab_size": 49160

configuration_wren.py

@@ -9,13 +9,13 @@ class WrenConfig(PretrainedConfig):
         self,
         llm_name:           str = "HuggingFaceTB/SmolLM2-360M",
         mimi_model_name:    str = "kyutai/mimi",
-        k_codebooks:        int = 3,
+        k_codebooks:        int = 8,
         codebook_size:      int = 2048,
         vocab_size:         int = 49160,
-        audio_sep_id:       int = None,
-        audio_eos_token_id: int = None,
-        audio_start_id:     int = None,
-        audio_end_id:       int = None,
+        # Special-token IDs (tokenized positions in the text vocab)
+        audio_start_id:     int = None,   # <|audio_start|> — text→target-audio boundary
+        reference_start_id: int = None,   # <|reference_start|> — opens speaker-reference block
+        reference_end_id:   int = None,   # <|reference_end|> — closes speaker-reference block
         sampling_rate:      int = 24000,
         **kwargs,
     ):
@@ -24,9 +24,8 @@ class WrenConfig(PretrainedConfig):
         self.k_codebooks        = k_codebooks
         self.codebook_size      = codebook_size
         self.vocab_size         = vocab_size
-        self.audio_sep_id       = audio_sep_id
-        self.audio_eos_token_id = audio_eos_token_id
         self.audio_start_id     = audio_start_id
-        self.audio_end_id       = audio_end_id
+        self.reference_start_id = reference_start_id
+        self.reference_end_id   = reference_end_id
         self.sampling_rate      = sampling_rate
         super().__init__(**kwargs)

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:f8b90adfd5eb6aafeb1bf509833f4560020b7d201f2cb1043af318250f256d6e
-size 770881760
+oid sha256:df687c893d29e93e454908d8a158cb16991e487dbdec86ecf94cccbe4ef37913
+size 849556624

modeling_wren.py

@@ -3,8 +3,19 @@ Wren-TTS model — a transformers-compatible wrapper over SmolLM2 + Mimi codeboo
 
 Designed for use with `AutoModel.from_pretrained(..., trust_remote_code=True)`.
 Self-contained: no imports from a `src/` folder.
+
+Sequence layout (MusicGen-style delay pattern):
+
+  [ text tokens... | <audio_start> | tgt_delay_steps ]
+  + optional [ <reference_start> | ref_delay_steps | <reference_end> ] prefix
+
+At each audio step s the model sees the sum of k per-codebook input embeddings
+(scaled by 1/sqrt(k)) and predicts k tokens via k parallel heads. Codebook q at
+frame f lives at step s = f + q. AUDIO_EOS (=codebook_size) is cb0's stop class;
+cb0 emits it at step T (one past the last real frame).
 """
 
+import math
 from typing import Optional
 
 import torch
@@ -18,15 +29,6 @@ except ImportError:
 
 
 class WrenForTTS(PreTrainedModel):
-    """
-    SmolLM2 backbone + k separate audio-code embedding tables + k audio heads.
-
-    Sequence layout fed to the LLM (via inputs_embeds):
-      [ text tokens... | <audio_sep> | cb0_f0 | cb1_f0 | ... | cbk_f0 | cb0_f1 | ... ]
-
-    cb0's head has an extra output class `AUDIO_EOS = codebook_size` used as stop token.
-    """
-
     config_class      = WrenConfig
     base_model_prefix = "wren"
 
@@ -34,18 +36,21 @@ class WrenForTTS(PreTrainedModel):
         super().__init__(config)
         self.k          = config.k_codebooks
         self.AUDIO_EOS  = config.codebook_size
+        self.AUDIO_PAD  = config.codebook_size
 
         # Build backbone from its config only. Pretrained weights for the backbone
         # are included in our own state_dict, so no need to re-download here.
-        # Set vocab_size directly on the sub-config to avoid a subsequent
-        # resize_token_embeddings call (which breaks under meta-tensor init).
         llm_cfg            = AutoConfig.from_pretrained(config.llm_name)
         llm_cfg.vocab_size = config.vocab_size
         self.llm           = AutoModelForCausalLM.from_config(llm_cfg)
 
         hidden = self.llm.config.hidden_size
+
+        # Input tables: codebook_size + 1 (extra row at index codebook_size = AUDIO_PAD).
+        # Output heads: cb0 has codebook_size + 1 (extra class = AUDIO_EOS); cb1..cb_{k-1}
+        # have codebook_size (no PAD/EOS output — PAD positions have -100 labels in training).
         self.audio_embeds = nn.ModuleList([
-            nn.Embedding(config.codebook_size, hidden)
+            nn.Embedding(config.codebook_size + 1, hidden)
             for _ in range(self.k)
         ])
         self.audio_heads = nn.ModuleList([
@@ -53,6 +58,7 @@ class WrenForTTS(PreTrainedModel):
             for i in range(self.k)
         ])
 
+        self.embed_scale = 1.0 / math.sqrt(self.k)
         self._mimi = None  # lazy-loaded on first use
 
     # --- Mimi codec (lazy-loaded, decoder + encoder used for audio I/O) ---
@@ -94,6 +100,13 @@ class WrenForTTS(PreTrainedModel):
 
     # --- Generation ---
 
+    def _audio_embed_step(self, codes_step: torch.LongTensor) -> torch.Tensor:
+        """Summed per-codebook embedding for one audio step. codes_step: [1, k]. Returns [1, 1, H]."""
+        acc = self.audio_embeds[0](codes_step[:, 0:1])
+        for q in range(1, self.k):
+            acc = acc + self.audio_embeds[q](codes_step[:, q:q + 1])
+        return acc * self.embed_scale
+
     @torch.no_grad()
     def generate(
         self,
@@ -112,11 +125,11 @@ class WrenForTTS(PreTrainedModel):
         Generate Mimi codes (or waveform) from a tokenized prompt.
 
         Args:
-            input_ids:        [1, L] — text tokens ending with <|audio_sep|>, as produced by WrenProcessor.
+            input_ids:        [1, L] — text tokens ending with <|audio_start|>, as produced by WrenProcessor.
             ref_codes:        optional [k, T_ref] reference codes for voice cloning.
             max_audio_frames: hard cap on output length.
-            min_audio_frames: suppress EOS for this many frames (prevents trivially-short outputs).
-            eos_bias:         additive bias on AUDIO_EOS logit; raise (e.g. 2–6) to reduce hallucinated continuations.
+            min_audio_frames: suppress EOS for this many frames.
+            eos_bias:         additive bias on cb0's AUDIO_EOS logit; raise (2–6) to reduce tail hallucination.
             output_audio:     if True, return [1, T] waveform; else return [k, n_frames] codes.
         """
         device = next(self.parameters()).device
@@ -128,67 +141,91 @@ class WrenForTTS(PreTrainedModel):
 
         prompt_embeds_list = []
 
-        # Optional reference-audio block (voice cloning)
+        # Optional reference block: <reference_start> ref_delayed <reference_end>
         if ref_codes is not None:
-            if self.config.audio_start_id is None or self.config.audio_end_id is None:
-                raise ValueError("audio_start_id/audio_end_id missing from config; cannot use ref_codes")
+            if self.config.reference_start_id is None or self.config.reference_end_id is None:
+                raise ValueError("reference_start_id/reference_end_id missing from config; cannot use ref_codes")
             ref_codes = ref_codes.to(device)
-            start_t = torch.tensor([[self.config.audio_start_id]], dtype=torch.long, device=device)
+
+            start_t = torch.tensor([[self.config.reference_start_id]], dtype=torch.long, device=device)
             prompt_embeds_list.append(embed_tokens(start_t.clamp(0, text_vocab - 1)))
 
-            ref_tokens = ref_codes.T.reshape(-1)  # interleaved [T_ref * k]
-            for pos, code in enumerate(ref_tokens):
-                cb_idx = pos % self.k
-                idx    = code.clamp(0, self.config.codebook_size - 1).unsqueeze(0)
-                prompt_embeds_list.append(self.audio_embeds[cb_idx](idx).unsqueeze(0))
+            T_ref = ref_codes.shape[1]
+            L_ref = T_ref + self.k - 1
+            ref_delayed = torch.full((self.k, L_ref), self.AUDIO_PAD, dtype=torch.long, device=device)
+            for q in range(self.k):
+                ref_delayed[q, q:q + T_ref] = ref_codes[q]
+
+            for s in range(L_ref):
+                codes_step = ref_delayed[:, s:s + 1].T.contiguous()  # [1, k]
+                prompt_embeds_list.append(self._audio_embed_step(codes_step))
 
-            end_t = torch.tensor([[self.config.audio_end_id]], dtype=torch.long, device=device)
+            end_t = torch.tensor([[self.config.reference_end_id]], dtype=torch.long, device=device)
             prompt_embeds_list.append(embed_tokens(end_t.clamp(0, text_vocab - 1)))
 
-        # Text prompt — already terminated by <|audio_sep|> via the processor
+        # Text prompt — already terminated by <|audio_start|> via the processor
         ids = input_ids.to(device)
         if ids.dim() == 1:
             ids = ids.unsqueeze(0)
         prompt_embeds_list.append(embed_tokens(ids.clamp(0, text_vocab - 1)))
 
         prompt_embeds = torch.cat(prompt_embeds_list, dim=1).to(llm_dtype)
-        out      = self.llm.model(inputs_embeds=prompt_embeds, use_cache=True)
-        hidden   = out.last_hidden_state
-        past_kv  = out.past_key_values
-
-        all_codes, audio_pos, n_frames = [], 0, 0
-        while n_frames < max_audio_frames:
-            cb_idx = audio_pos % self.k
-            h      = hidden[:, -1:, :]
-            logits = self.audio_heads[cb_idx](h.float()).squeeze(1)
-
-            if cb_idx == 0:
-                logits[:, self.AUDIO_EOS] = logits[:, self.AUDIO_EOS] + eos_bias
-                if n_frames < min_audio_frames:
-                    logits[:, self.AUDIO_EOS] = float("-inf")
-
-            next_code = _sample(logits, temperature, top_k, top_p)
-
-            if cb_idx == 0 and next_code.item() == self.AUDIO_EOS:
+        out     = self.llm.model(inputs_embeds=prompt_embeds, use_cache=True)
+        hidden  = out.last_hidden_state
+        past_kv = out.past_key_values
+
+        # Delay-pattern autoregressive loop
+        outputs: list = [[] for _ in range(self.k)]
+        eos_step: Optional[int] = None
+        max_steps = max_audio_frames + self.k - 1
+
+        for step in range(max_steps):
+            h = hidden[:, -1:, :]
+            logits_per_cb = [self.audio_heads[q](h.float()).squeeze(1) for q in range(self.k)]
+
+            logits_per_cb[0][:, self.AUDIO_EOS] = logits_per_cb[0][:, self.AUDIO_EOS] + eos_bias
+            if step < min_audio_frames:
+                logits_per_cb[0][:, self.AUDIO_EOS] = float("-inf")
+
+            next_codes = torch.empty(self.k, dtype=torch.long, device=device)
+            for q in range(self.k):
+                if step < q:
+                    next_codes[q] = self.AUDIO_PAD
+                    continue
+                if q == 0 and eos_step is not None:
+                    next_codes[q] = self.AUDIO_PAD
+                    continue
+                if q > 0 and eos_step is not None and (step - q) >= eos_step:
+                    next_codes[q] = self.AUDIO_PAD
+                    continue
+                sampled = _sample(logits_per_cb[q], temperature, top_k, top_p)
+                if q == 0 and sampled.item() == self.AUDIO_EOS:
+                    eos_step = step
+                    next_codes[q] = self.AUDIO_PAD
+                else:
+                    next_codes[q] = sampled
+
+            for q in range(self.k):
+                outputs[q].append(next_codes[q].item())
+
+            if eos_step is not None and step >= eos_step + self.k - 1:
                 break
 
-            all_codes.append(next_code.item())
-            audio_pos += 1
-            if audio_pos % self.k == 0:
-                n_frames += 1
-
-            embed_idx = next_code.clamp(0, self.config.codebook_size - 1)
-            emb = self.audio_embeds[cb_idx](embed_idx.unsqueeze(0)).to(hidden.dtype)
-            out     = self.llm.model(inputs_embeds=emb, past_key_values=past_kv, use_cache=True)
+            next_embed = self._audio_embed_step(next_codes.unsqueeze(0)).to(hidden.dtype)
+            out     = self.llm.model(inputs_embeds=next_embed, past_key_values=past_kv, use_cache=True)
             hidden  = out.last_hidden_state
             past_kv = out.past_key_values
 
-        complete = (len(all_codes) // self.k) * self.k
-        if complete == 0:
+        # Un-delay into [k, T]
+        T = eos_step if eos_step is not None else max_audio_frames
+        max_available = min(len(outputs[q]) - q for q in range(self.k))
+        T = min(T, max_available)
+        if T <= 0:
             codes = torch.zeros(self.k, 0, dtype=torch.long)
         else:
-            codes = torch.tensor(all_codes[:complete], dtype=torch.long)
-            codes = codes.reshape(complete // self.k, self.k).T  # [k, n_frames]
+            codes = torch.empty(self.k, T, dtype=torch.long)
+            for q in range(self.k):
+                codes[q] = torch.tensor(outputs[q][q:q + T], dtype=torch.long)
 
         if output_audio:
             return self.decode_audio(codes)

processing_wren.py

@@ -1,8 +1,10 @@
 """
 Wren processor: text tokenization + audio saving.
 
-Text is lowercased to match training distribution. The `<|audio_sep|>` separator is
-always appended so `model.generate(**processor(text))` "just works".
+Text casing is preserved as-is. Pass text naturally ("Hello, World!") — the model
+is trained on mixed-case data (LJSpeech mixed-case, LibriTTS with punctuation).
+The `<|audio_start|>` separator is always appended so `model.generate(**processor(text))`
+"just works".
 """
 
 from typing import List, Union
@@ -17,9 +19,9 @@ class WrenProcessor(ProcessorMixin):
 
     def __init__(self, tokenizer):
         super().__init__(tokenizer=tokenizer)
-        self.audio_sep_id   = tokenizer.convert_tokens_to_ids("<|audio_sep|>")
-        self.audio_start_id = tokenizer.convert_tokens_to_ids("<|audio_start|>")
-        self.audio_end_id   = tokenizer.convert_tokens_to_ids("<|audio_end|>")
+        self.audio_start_id     = tokenizer.convert_tokens_to_ids("<|audio_start|>")
+        self.reference_start_id = tokenizer.convert_tokens_to_ids("<|reference_start|>")
+        self.reference_end_id   = tokenizer.convert_tokens_to_ids("<|reference_end|>")
 
     def __call__(
         self,
@@ -27,9 +29,6 @@ class WrenProcessor(ProcessorMixin):
         return_tensors: str = "pt",
         **kwargs,
     ):
-        # Training data was lowercased; match that distribution.
-        text = text.lower() if isinstance(text, str) else [t.lower() for t in text]
-
         enc = self.tokenizer(
             text,
             add_special_tokens = False,
@@ -40,10 +39,10 @@ class WrenProcessor(ProcessorMixin):
         if ids.dim() == 1:
             ids = ids.unsqueeze(0)
 
-        # Append <|audio_sep|> as the final prompt token
+        # Append <|audio_start|> as the final prompt token
         sep = torch.full(
             (ids.shape[0], 1),
-            self.audio_sep_id,
+            self.audio_start_id,
             dtype=ids.dtype,
             device=ids.device,
         )

special_tokens_map.json

@@ -1,28 +1,21 @@
 {
   "additional_special_tokens": [
     {
-      "content": "<|audio_sep|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false
-    },
-    {
-      "content": "<|audio_eos|>",
+      "content": "<|audio_start|>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
       "single_word": false
     },
     {
-      "content": "<|audio_start|>",
+      "content": "<|reference_start|>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
       "single_word": false
     },
     {
-      "content": "<|audio_end|>",
+      "content": "<|reference_end|>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,

tokenizer.json

@@ -158,7 +158,7 @@
     },
     {
       "id": 49152,
-      "content": "<|audio_sep|>",
+      "content": "<|audio_start|>",
       "single_word": false,
       "lstrip": false,
       "rstrip": false,
@@ -167,7 +167,7 @@
     },
     {
       "id": 49153,
-      "content": "<|audio_eos|>",
+      "content": "<|reference_start|>",
       "single_word": false,
       "lstrip": false,
       "rstrip": false,
@@ -176,16 +176,7 @@
     },
     {
       "id": 49154,
-      "content": "<|audio_start|>",
-      "single_word": false,
-      "lstrip": false,
-      "rstrip": false,
-      "normalized": false,
-      "special": true
-    },
-    {
-      "id": 49155,
-      "content": "<|audio_end|>",
+      "content": "<|reference_end|>",
       "single_word": false,
       "lstrip": false,
       "rstrip": false,

tokenizer_config.json

@@ -138,7 +138,7 @@
       "special": true
     },
     "49152": {
-      "content": "<|audio_sep|>",
+      "content": "<|audio_start|>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -146,7 +146,7 @@
       "special": true
     },
     "49153": {
-      "content": "<|audio_eos|>",
+      "content": "<|reference_start|>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -154,15 +154,7 @@
       "special": true
     },
     "49154": {
-      "content": "<|audio_start|>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "49155": {
-      "content": "<|audio_end|>",
+      "content": "<|reference_end|>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -171,10 +163,9 @@
     }
   },
   "additional_special_tokens": [
-    "<|audio_sep|>",
-    "<|audio_eos|>",
     "<|audio_start|>",
-    "<|audio_end|>"
+    "<|reference_start|>",
+    "<|reference_end|>"
   ],
   "bos_token": "<|endoftext|>",
   "clean_up_tokenization_spaces": false,