mirror of
https://github.com/Nighthawk42/Qwen3-TTS-streaming.git
synced 2026-08-30 08:52:27 +00:00
Merge branch 'main' into fix/repetition-penalty-and-finetune-fixes
This commit is contained in:
@@ -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 |
|
||||
|
||||
@@ -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)
|
||||
@@ -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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user