Merge branch 'main' into fix/repetition-penalty-and-finetune-fixes

This commit is contained in:
rekuenkdr
2026-02-07 10:42:24 +01:00
committed by GitHub
3 changed files with 173 additions and 108 deletions
+100 -7
View File
@@ -8,7 +8,6 @@ From [dffdeeq/Qwen3-TTS-streaming](https://github.com/dffdeeq/Qwen3-TTS-streamin
- `stream_generate_voice_clone()` - streaming with voice cloning
- `stream_generate_pcm()` - real-time PCM audio streaming
- `torch.compile` + CUDA graphs optimization
- Crossfade overlap for seamless chunk transitions
Added in this fork:
- **Two-phase streaming** - faster first-chunk latency
@@ -25,7 +24,6 @@ Standard streaming with Qwen's TTS library waits for `emit_every_frames` (e.g.,
│ PHASE 1 (First N frames) │ PHASE 2 (Rest of audio) │
│ - emit_every = 5 (fast) │ - emit_every = 12 (stable) │
│ - decode_window = 48 │ - decode_window = 80 │
│ - optimized = OFF │ - optimized = ON │
│ → FAST first chunk │ → QUALITY for rest │
└─────────────────────────────────────────────────────────────────┘
```
@@ -46,22 +44,116 @@ User hears audio **362ms earlier** vs baseline, **174ms earlier** vs only optimi
- vs Optimized: **1.87x faster** (389ms → 208ms, saves 181ms)
- vs Optimized_2: **1.84x faster** (382ms → 208ms, saves 174ms)
## Audio Quality Fixes
Streaming TTS can produce clicks, pops, and artifacts at chunk boundaries. This fork implements several fixes:
### Crossfade Blending
Chunks are blended using a Hann window crossfade to eliminate boundary discontinuities:
```python
# ~21ms at 24kHz, matches RMS check window
# Lower values may cause clicks, set to 0 to disable
DEFAULT_BLEND_SAMPLES = 512
# Hann crossfade
fade_out = 0.5 * (1 + np.cos(np.pi * t))
fade_in = 0.5 * (1 - np.cos(np.pi * t))
blended = prev_tail * fade_out + curr_head * fade_in
```
### Overlap Trimming
Each chunk is processed in this order to prevent audio duplication (echo artifacts):
1. Crossfade current chunk's HEAD with previous chunk's saved TAIL
2. Apply fade-in (first chunk only)
3. Save FULL processed chunk for next iteration's crossfade
4. Trim END of chunk before emission (this region will be replaced by next chunk's crossfade)
5. Yield trimmed chunk
### First/Last Chunk Fades
- **First chunk**: Hann fade-in prevents pop at audio start
- **Final chunk**: Hann fade-out prevents pop at audio end
## Optimization API
### enable_streaming_optimizations()
Call after loading the model to enable torch.compile and CUDA graphs:
```python
from qwen_tts import Qwen3TTSModel
model = Qwen3TTSModel.from_pretrained(
"Qwen/Qwen3-TTS-12Hz-1.7B-Base",
torch_dtype=torch.bfloat16,
device_map="cuda",
)
# Enable optimizations (recommended)
model.enable_streaming_optimizations(
decode_window_frames=80, # Must match streaming parameter
use_compile=True, # torch.compile the decoder
compile_mode="reduce-overhead", # Includes CUDA graphs automatically
)
```
**Parameters:**
| Parameter | Default | Description |
|-----------|---------|-------------|
| `decode_window_frames` | 80 | Window size (must match streaming call) |
| `use_compile` | True | Apply torch.compile to decoder |
| `use_cuda_graphs` | True | Capture CUDA graphs for fixed window |
| `compile_mode` | "reduce-overhead" | torch.compile mode |
| `use_fast_codebook` | False | Use fast codebook generation (experimental) |
| `compile_codebook_predictor` | True | Apply torch.compile to codebook predictor |
## Usage
```python
import torch
import sounddevice as sd
from qwen_tts import Qwen3TTSModel
# Load model
model = Qwen3TTSModel.from_pretrained(
"Qwen/Qwen3-TTS-12Hz-Base",
torch_dtype=torch.bfloat16,
device_map="cuda",
)
# Enable optimizations (recommended for streaming)
model.enable_streaming_optimizations(
decode_window_frames=80,
use_compile=True,
compile_mode="reduce-overhead",
)
# Create voice clone prompt from reference audio
prompt = model.create_voice_clone_prompt(
ref_audio="reference.wav",
ref_text="Transcript of the reference audio.",
)
# Stream audio with two-phase settings
for chunk, sr in model.stream_generate_voice_clone(
text="Hello!",
text="Hello, this is a streaming TTS demo!",
language="en",
voice_clone_prompt=prompt,
# Phase 2 settings
# Phase 2 settings (stable)
emit_every_frames=12,
decode_window_frames=80,
# Phase 1 settings (two-phase)
# Phase 1 settings (fast first chunk)
first_chunk_emit_every=5,
first_chunk_decode_window=48,
first_chunk_frames=48,
):
play_audio(chunk, sr)
sd.play(chunk, sr)
sd.wait()
```
## Parameters
@@ -70,7 +162,8 @@ for chunk, sr in model.stream_generate_voice_clone(
|-----------|---------|-------------|
| `emit_every_frames` | 8 | Emit audio every N frames |
| `decode_window_frames` | 80 | Decoder context window |
| `overlap_samples` | 512 | Crossfade overlap between chunks |
| `overlap_samples` | 512 | Crossfade overlap between chunks (0 to disable) |
| `max_frames` | 10000 | Maximum codec frames to generate |
| `first_chunk_emit_every` | 0 | Phase 1 emit interval (0 = disabled) |
| `first_chunk_decode_window` | 48 | Phase 1 decode window |
| `first_chunk_frames` | 48 | Switch to phase 2 after N frames |
-87
View File
@@ -1,87 +0,0 @@
# Qwen3-TTS Streaming
Real-time streaming audio generation for [Qwen3-TTS](https://github.com/QwenLM/Qwen3-TTS).
## Features
From [dffdeeq/Qwen3-TTS-streaming](https://github.com/dffdeeq/Qwen3-TTS-streaming):
- `stream_generate_voice_clone()` - streaming with voice cloning
- `stream_generate_pcm()` - real-time PCM audio streaming
- `torch.compile` + CUDA graphs optimization
- Crossfade overlap for seamless chunk transitions
Added in this fork:
- **Two-phase streaming** - faster first-chunk latency
## Two-Phase Streaming
Standard streaming with Qwen's TTS library waits for `emit_every_frames` (e.g., 12) before emitting the first audio. Two-phase uses aggressive settings for the first chunk to improve latency, then switches to stable settings.
```
┌─────────────────────────────────────────────────────────────────┐
│ PHASE 1 (First N frames) │ PHASE 2 (Rest of audio) │
│ - emit_every = 5 (fast) │ - emit_every = 12 (stable) │
│ - decode_window = 48 │ - decode_window = 80 │
│ - optimized = OFF │ - optimized = ON │
│ → FAST first chunk │ → QUALITY for rest │
└─────────────────────────────────────────────────────────────────┘
```
### Benchmarks
| Test | Method | emit | 1st Chunk | 1st Spdup | Total | Tot Spdup | RTF |
|------|--------|------|-----------|-----------|-------|-----------|-----|
| 2 | Baseline (no opt) | 12 | 570ms | 1.00x | 3.16s | 1.00x | 0.56 |
| 3 | Optimized | 12 | 389ms | 1.47x | 2.37s | 1.34x | 0.37 |
| 4 | Optimized_2 (stable) | 12 | 382ms | 1.49x | 2.27s | 1.39x | 0.36 |
| 5 | **Two-phase (5→12)** | 5→12 | **208ms** | **2.75x** | 2.58s | 1.23x | 0.39 |
User hears audio **362ms earlier** vs baseline, **174ms earlier** vs only optimized.
**First-chunk latency improvement:**
- vs Baseline: **2.75x faster** (570ms → 208ms, saves 362ms)
- vs Optimized: **1.87x faster** (389ms → 208ms, saves 181ms)
- vs Optimized_2: **1.84x faster** (382ms → 208ms, saves 174ms)
## Usage
```python
for chunk, sr in model.stream_generate_voice_clone(
text="Hello!",
language="en",
voice_clone_prompt=prompt,
# Phase 2 settings
emit_every_frames=12,
decode_window_frames=80,
# Phase 1 settings (two-phase)
first_chunk_emit_every=5,
first_chunk_decode_window=48,
first_chunk_frames=48,
):
play_audio(chunk, sr)
```
## Parameters
| Parameter | Default | Description |
|-----------|---------|-------------|
| `emit_every_frames` | 8 | Emit audio every N frames |
| `decode_window_frames` | 80 | Decoder context window |
| `overlap_samples` | 512 | Crossfade overlap between chunks |
| `first_chunk_emit_every` | 0 | Phase 1 emit interval (0 = disabled) |
| `first_chunk_decode_window` | 48 | Phase 1 decode window |
| `first_chunk_frames` | 48 | Switch to phase 2 after N frames |
## Installation
```bash
sudo apt install sox
pip install torch torchaudio flash-attn
pip install -e .
```
---
Based on:
- [QwenLM/Qwen3-TTS](https://github.com/QwenLM/Qwen3-TTS)
- [dffdeeq/Qwen3-TTS-streaming](https://github.com/dffdeeq/Qwen3-TTS-streaming)
+73 -14
View File
@@ -93,12 +93,20 @@ def _sample_next_token(
def _crossfade(prev_tail: np.ndarray, new_head: np.ndarray) -> np.ndarray:
"""Crossfade between end of previous chunk and start of new chunk."""
"""Crossfade between end of previous chunk and start of new chunk using Hann window."""
n = min(len(prev_tail), len(new_head))
if n <= 0:
return new_head
w = np.linspace(0.0, 1.0, n, dtype=np.float32)
return prev_tail[:n] * (1.0 - w) + new_head[:n] * w
t = np.arange(n, dtype=np.float32) / max(n - 1, 1)
fade_in = 0.5 * (1 - np.cos(np.pi * t))
fade_out = 1 - fade_in
return prev_tail[:n] * fade_out + new_head[:n] * fade_in
# Default blend samples for boundary blending
# ~21ms at 24kHz, matches RMS check window for better coverage
# Lower values may cause clicks, set to 0 to disable
DEFAULT_BLEND_SAMPLES = 512
def _add_ref_code_context(
@@ -2519,6 +2527,17 @@ class Qwen3TTSForConditionalGeneration(Qwen3TTSPreTrainedModel, GenerationMixin)
repetition_penalty: float = 1.05,
**kwargs,
):
# Multiple EOS tokens that can terminate generation
eos_ids = {
self.config.talker_config.codec_eos_token_id, # Primary codec EOS
2150, # Codec EOS (model-specific)
2157, # Secondary codec token
151670, # TTS special token
self.config.tts_eos_token_id, # 151673
self.config.im_end_token_id, # 151645
151643, # <|endoftext|>
}
talker_kwargs = {
"max_new_tokens": max_new_tokens,
"min_new_tokens": 2,
@@ -2526,7 +2545,7 @@ class Qwen3TTSForConditionalGeneration(Qwen3TTSPreTrainedModel, GenerationMixin)
"top_k": top_k,
"top_p": top_p,
"temperature": temperature,
"subtalker_dosample": subtalker_dosample,
"subtalker_dosample": subtalker_dosample,
"subtalker_top_k": subtalker_top_k,
"subtalker_top_p": subtalker_top_p,
"subtalker_temperature": subtalker_temperature,
@@ -2537,7 +2556,7 @@ class Qwen3TTSForConditionalGeneration(Qwen3TTSPreTrainedModel, GenerationMixin)
"suppress_tokens": [
i
for i in range(self.config.talker_config.vocab_size - 1024, self.config.talker_config.vocab_size)
if i not in (self.config.talker_config.codec_eos_token_id,)
if i not in eos_ids
],
"output_hidden_states": kwargs.get("output_hidden_states", True),
"return_dict_in_generate": kwargs.get("return_dict_in_generate", True)
@@ -2568,7 +2587,9 @@ class Qwen3TTSForConditionalGeneration(Qwen3TTSPreTrainedModel, GenerationMixin)
talker_hidden_states = torch.cat([hid[0][-1][:, -1:] for hid in talker_result.hidden_states], dim=1)[:, :-1]
first_codebook = talker_codes[:, :, 0]
is_stop_token = (first_codebook == self.config.talker_config.codec_eos_token_id)
# Check against all EOS tokens
eos_ids_tensor = torch.tensor(list(eos_ids), device=first_codebook.device, dtype=first_codebook.dtype)
is_stop_token = torch.isin(first_codebook, eos_ids_tensor)
stop_indices = torch.argmax(is_stop_token.int(), dim=1)
has_stop_token = is_stop_token.any(dim=1)
effective_lengths = torch.where(has_stop_token, stop_indices, talker_codes.shape[1])
@@ -2652,13 +2673,23 @@ class Qwen3TTSForConditionalGeneration(Qwen3TTSPreTrainedModel, GenerationMixin)
non_streaming_mode=non_streaming_mode,
)
eos_id = self.config.talker_config.codec_eos_token_id
# Multiple EOS tokens that can terminate generation
# Some models may emit different EOS tokens depending on context
eos_ids = {
self.config.talker_config.codec_eos_token_id, # Primary codec EOS
2150, # Codec EOS (model-specific)
2157, # Secondary codec token
151670, # TTS special token
self.config.tts_eos_token_id, # 151673
self.config.im_end_token_id, # 151645
151643, # <|endoftext|>
}
# Build suppress_tokens list (same as in generate())
vocab_size = self.config.talker_config.vocab_size
suppress_tokens = [
i for i in range(vocab_size - 1024, vocab_size)
if i != eos_id
if i not in eos_ids
]
# Mark step begin for CUDA graphs (required for torch.compile with reduce-overhead)
@@ -2745,9 +2776,9 @@ class Qwen3TTSForConditionalGeneration(Qwen3TTSPreTrainedModel, GenerationMixin)
# Get codec_ids from hidden_states tuple: (layer_outputs, codec_ids)
codec_ids = step_out.hidden_states[1] # [B, num_code_groups]
# Check for EOS in first codebook ON GPU (avoids CPU sync bottleneck)
# Check for EOS in first codebook
# EOS token is out of range for speech tokenizer, so we must not include it
if codec_ids[0, 0] == eos_id:
if codec_ids[0, 0].item() in eos_ids:
break
# Keep on GPU to avoid CPU<->GPU transfers during decode
@@ -2819,13 +2850,32 @@ class Qwen3TTSForConditionalGeneration(Qwen3TTSPreTrainedModel, GenerationMixin)
chunk = wav[-step_samples:] if step_samples > 0 else wav
# Crossfade with previous chunk tail for smooth transition
if decoded_tail is not None and overlap_samples > 0:
ov = min(overlap_samples, len(decoded_tail), len(chunk))
# Always blend boundaries to prevent clicks from sliding window re-decode artifacts
blend_samples = overlap_samples
if decoded_tail is not None:
ov = min(blend_samples, len(decoded_tail), len(chunk))
if ov > 0:
head = _crossfade(decoded_tail[-ov:], chunk[:ov])
chunk = np.concatenate([head, chunk[ov:]], axis=0)
# Apply Hann fade-in to very first chunk to avoid pop at audio start
# Always apply fade-in on first chunk to prevent pop
blend_samples = overlap_samples
if decoded_tail is None:
fade_len = min(blend_samples, len(chunk))
if fade_len > 0:
t = np.arange(fade_len, dtype=np.float32) / max(fade_len - 1, 1)
fade_in = 0.5 * (1 - np.cos(np.pi * t))
chunk[:fade_len] *= fade_in
# Save FULL chunk for next crossfade reference
decoded_tail = chunk.copy()
# Trim END of chunk - this region will be replaced by next chunk's crossfade
# Don't trim if chunk would become too small
if len(chunk) > blend_samples * 2:
chunk = chunk[:-blend_samples]
total_frames_emitted = len(codes_buffer) # Mark these frames as emitted
yield chunk, sr
@@ -2853,12 +2903,21 @@ class Qwen3TTSForConditionalGeneration(Qwen3TTSPreTrainedModel, GenerationMixin)
wav = wav[skip_samples:]
# Crossfade with previous tail
if decoded_tail is not None and overlap_samples > 0 and len(wav) > 0:
ov = min(overlap_samples, len(decoded_tail), len(wav))
# Always blend flush boundary
blend_samples = overlap_samples
if decoded_tail is not None and len(wav) > 0:
ov = min(blend_samples, len(decoded_tail), len(wav))
if ov > 0:
head = _crossfade(decoded_tail[-ov:], wav[:ov])
wav = np.concatenate([head, wav[ov:]], axis=0)
# Apply fade-out at very end of audio to avoid pop on completion
if len(wav) > blend_samples:
fade_len = min(blend_samples, len(wav))
t = np.arange(fade_len, dtype=np.float32) / max(fade_len - 1, 1)
fade_out = 0.5 * (1 + np.cos(np.pi * t)) # Hann fade-out
wav[-fade_len:] *= fade_out
# Debug removed for performance: flush done
yield wav, sr