Nighthawk/Qwen3-TTS-streaming
1
0
Fork 0
mirror of https://github.com/Nighthawk42/Qwen3-TTS-streaming.git synced 2026-08-30 09:02:26 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor: batch GPU decode calls in batch_stream_generate_voice_clone

Browse Source 
Restructure the batch streaming loop and flush into a 3-phase approach:
1) collect decode windows for all active items, 2) perform a single
batched GPU decode call, 3) per-item post-processing (crossfade, trim).

Add decode_streaming_batch() to Qwen3TTSTokenizer to support B>1 batch
decoding. Remove examples/test_batch_streaming.py test script.
This commit is contained in:
Kedara Studios
2026-02-09 01:57:30 +01:00
parent dc7a8da2b1
commit a48510a58d
3 changed files with 136 additions and 152 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/test_batch_streaming.py
-132
View File
@@ -1,132 +0,0 @@
"""Test batch streaming voice clone generation.
Generates audio for multiple texts (potentially with different voices) in a
single batched pass through the transformer. All items advance in lockstep.
"""
import time
import numpy as np
import soundfile as sf
import torch
from qwen_tts import Qwen3TTSModel
def log_time(start, operation):
elapsed = time.time() - start
print(f"[{elapsed:.2f}s] {operation}")
return time.time()
total_start = time.time()
# Load model
start = time.time()
model = Qwen3TTSModel.from_pretrained(
"Qwen/Qwen3-TTS-12Hz-1.7B-Base",
device_map="cuda:0",
dtype=torch.bfloat16,
attn_implementation="flash_attention_2",
)
start = log_time(start, "Load Base model")
# Create a voice clone prompt (using the same voice for all items in this test)
ref_audio_path = "kuklina-1.wav" # Replace with your reference audio
ref_text = (
"Это брат Кэти, моей одноклассницы. А что у тебя с рукой? И почему ты голая? "
"У него ведь куча наград по боевым искусствам."
)
voice_prompt = model.create_voice_clone_prompt(
ref_audio=ref_audio_path,
ref_text=ref_text,
)
start = log_time(start, "Create voice clone prompt")
# Batch items: different texts, same voice (broadcast)
texts = [
"Hello! This is the first batch item with a short sentence.",
"And this is the second batch item. It has a bit more text to synthesize.",
"Third item here. Testing batch streaming with multiple items at once!",
]
languages = ["English", "English", "English"]
# ============== Batch Streaming ==============
print(f"\n--- Batch streaming ({len(texts)} items) ---")
start = time.time()
# Accumulate per-item chunks
item_chunks: list[list[np.ndarray]] = [[] for _ in range(len(texts))]
chunk_count = 0
sr = 24000
for chunks_list, chunk_sr in model.batch_stream_generate_voice_clone(
text=texts,
language=languages,
voice_clone_prompt=voice_prompt,
emit_every_frames=8,
decode_window_frames=80,
overlap_samples=512,
max_frames=8000,
first_chunk_emit_every=5,
first_chunk_decode_window=48,
first_chunk_frames=48,
):
sr = chunk_sr
chunk_count += 1
sizes = []
for b, chunk in enumerate(chunks_list):
if chunk.size > 0:
item_chunks[b].append(chunk)
sizes.append(f"item{b}={chunk.size}")
else:
sizes.append(f"item{b}=empty")
if chunk_count <= 5 or chunk_count % 10 == 0:
print(f" Chunk {chunk_count}: {', '.join(sizes)}")
start = log_time(start, f"Batch streaming done ({chunk_count} chunks)")
# Save per-item outputs
for i, chunks in enumerate(item_chunks):
if chunks:
combined = np.concatenate(chunks)
filename = f"batch_item_{i}.wav"
sf.write(filename, combined, sr)
duration_ms = len(combined) / sr * 1000
print(f" Saved {filename}: {duration_ms:.0f}ms, {len(combined)} samples")
else:
print(f" Item {i}: no audio generated")
# ============== Compare: Sequential single-item streaming ==============
print(f"\n--- Sequential single-item streaming ({len(texts)} items) ---")
start = time.time()
for i, text in enumerate(texts):
item_single_chunks = []
for chunk, chunk_sr in model.stream_generate_voice_clone(
text=text,
language=languages[i],
voice_clone_prompt=voice_prompt,
emit_every_frames=8,
decode_window_frames=80,
overlap_samples=512,
max_frames=8000,
first_chunk_emit_every=5,
first_chunk_decode_window=48,
first_chunk_frames=48,
):
if chunk.size > 0:
item_single_chunks.append(chunk)
if item_single_chunks:
combined = np.concatenate(item_single_chunks)
filename = f"single_item_{i}.wav"
sf.write(filename, combined, chunk_sr)
duration_ms = len(combined) / chunk_sr * 1000
print(f" Saved {filename}: {duration_ms:.0f}ms")
start = log_time(start, "Sequential streaming done")
total_elapsed = time.time() - total_start
print(f"\nTotal time: {total_elapsed:.2f}s")
qwen_tts/core/models/modeling_qwen3_tts.py
+97 -19
View File
@@ -3069,6 +3069,7 @@ class Qwen3TTSForConditionalGeneration(Qwen3TTSPreTrainedModel, GenerationMixin)
total_frames_emitted: list[int] = [0] * B
generated_token_ids: list[list[int]] = [[] for _ in range(B)]
finished: list[bool] = [False] * B
sr = 24000 # default sample rate, updated on first decode
# Shared frame counter (items advance in lockstep)
frames_since_emit = 0
@@ -3162,31 +3163,57 @@ class Qwen3TTSForConditionalGeneration(Qwen3TTSPreTrainedModel, GenerationMixin)
blend_samples = overlap_samples
chunks_list: list[np.ndarray] = []
# Phase 1: Collect windows for all active items
active_indices: list[int] = []
active_windows: list[torch.Tensor] = []
for b in range(B):
if finished[b] or len(codes_buffers[b]) == 0:
chunks_list.append(np.array([], dtype=np.float32))
continue
# Decode window for this item
start = max(0, len(codes_buffers[b]) - current_decode_window)
window_codes = torch.stack(codes_buffers[b][start:], dim=0)
# Add per-item ref_code context
window, _ = _add_ref_code_context(
window_codes, ref_code_contexts[b], ref_code_frames_list[b], current_decode_window
)
active_indices.append(b)
active_windows.append(window)
# Decode (each item independently through compiled decoder)
if current_use_optimized and hasattr(self.speech_tokenizer, 'decode_streaming'):
wavs, sr = self.speech_tokenizer.decode_streaming(
window.to(self.talker.device),
# Phase 2: Batched decode (single GPU call for all active items)
active_wavs: dict[int, np.ndarray] = {}
if active_windows:
# Pad all windows to same length and stack into batch
max_t = max(w.shape[0] for w in active_windows)
padded_windows = []
for w in active_windows:
if w.shape[0] < max_t:
pad = torch.zeros(max_t - w.shape[0], w.shape[1], dtype=w.dtype, device=w.device)
padded_windows.append(torch.cat([pad, w], dim=0))
else:
padded_windows.append(w)
batch_codes = torch.stack(padded_windows, dim=0).to(self.talker.device) # [B_active, max_t, Q]
if current_use_optimized and hasattr(self.speech_tokenizer, 'decode_streaming_batch'):
batch_wavs, sr = self.speech_tokenizer.decode_streaming_batch(
batch_codes,
use_optimized=True,
pad_to_size=decode_window_frames,
)
else:
wavs, sr = self.speech_tokenizer.decode([{"audio_codes": window.to(self.talker.device)}])
batch_wavs = []
for i in range(batch_codes.shape[0]):
wavs_i, sr = self.speech_tokenizer.decode([{"audio_codes": batch_codes[i]}])
batch_wavs.append(wavs_i[0])
wav = wavs[0].astype(np.float32)
for idx, b in enumerate(active_indices):
active_wavs[b] = batch_wavs[idx].astype(np.float32)
# Phase 3: Per-item post-processing (crossfade, fade-in, trim)
for b in range(B):
if b not in active_wavs:
chunks_list.append(np.array([], dtype=np.float32))
continue
wav = active_wavs[b]
chunk = wav[-step_samples:] if step_samples > 0 else wav
# Crossfade with previous chunk tail
@@ -3214,14 +3241,19 @@ class Qwen3TTSForConditionalGeneration(Qwen3TTSPreTrainedModel, GenerationMixin)
yield chunks_list, sr
# Flush: decode remaining per-item frames
# Flush: decode remaining per-item frames (batched)
flush_chunks: list[np.ndarray] = []
flush_sr = 24000 # default
# Phase 1: Collect flush windows and per-item metadata
flush_active_indices: list[int] = []
flush_active_windows: list[torch.Tensor] = []
flush_skip_frames: list[int] = []
flush_window_lengths: list[int] = []
for b in range(B):
remaining_frames = len(codes_buffers[b]) - total_frames_emitted[b]
if remaining_frames <= 0:
flush_chunks.append(np.array([], dtype=np.float32))
continue
context_frames = min(total_frames_emitted[b], decode_window_frames - remaining_frames)
@@ -3232,13 +3264,59 @@ class Qwen3TTSForConditionalGeneration(Qwen3TTSPreTrainedModel, GenerationMixin)
window_codes, ref_code_contexts[b], ref_code_frames_list[b], decode_window_frames
)
wavs, flush_sr = self.speech_tokenizer.decode([{"audio_codes": window.to(self.talker.device)}])
wav = wavs[0].astype(np.float32)
flush_active_indices.append(b)
flush_active_windows.append(window)
flush_skip_frames.append(flush_ref_prefix_frames + context_frames)
flush_window_lengths.append(window.shape[0])
skip_frames = flush_ref_prefix_frames + context_frames
if skip_frames > 0:
samples_per_frame = len(wav) / window.shape[0]
skip_samples = int(skip_frames * samples_per_frame)
# Phase 2: Batched decode
flush_active_wavs: dict[int, np.ndarray] = {}
if flush_active_windows:
max_t = max(w.shape[0] for w in flush_active_windows)
padded_windows = []
for w in flush_active_windows:
if w.shape[0] < max_t:
pad = torch.zeros(max_t - w.shape[0], w.shape[1], dtype=w.dtype, device=w.device)
padded_windows.append(torch.cat([pad, w], dim=0))
else:
padded_windows.append(w)
batch_codes = torch.stack(padded_windows, dim=0).to(self.talker.device)
if hasattr(self.speech_tokenizer, 'decode_streaming_batch'):
batch_wavs, flush_sr = self.speech_tokenizer.decode_streaming_batch(
batch_codes,
use_optimized=True,
pad_to_size=decode_window_frames,
)
else:
batch_wavs = []
for i in range(batch_codes.shape[0]):
wavs_i, flush_sr = self.speech_tokenizer.decode([{"audio_codes": batch_codes[i]}])
batch_wavs.append(wavs_i[0])
for idx, b in enumerate(flush_active_indices):
flush_active_wavs[b] = batch_wavs[idx].astype(np.float32)
# Phase 3: Per-item post-processing (skip, crossfade, fade-out)
for b in range(B):
if b not in flush_active_wavs:
flush_chunks.append(np.array([], dtype=np.float32))
continue
wav = flush_active_wavs[b]
idx_in_active = flush_active_indices.index(b)
skip = flush_skip_frames[idx_in_active]
win_len = flush_window_lengths[idx_in_active]
# Account for batch left-padding: items padded from win_len to max_t
# have extra decoded samples at the front that need skipping
max_t_flush = max(w.shape[0] for w in flush_active_windows) if flush_active_windows else win_len
batch_pad_frames = max_t_flush - win_len
total_skip = skip + batch_pad_frames
if total_skip > 0:
samples_per_frame = len(wav) / max_t_flush
skip_samples = int(total_skip * samples_per_frame)
wav = wav[skip_samples:]
blend_samples = overlap_samples
qwen_tts/inference/qwen3_tts_tokenizer.py
+38
View File
@@ -484,3 +484,41 @@ class Qwen3TTSTokenizer:
# Convert to numpy and return
wav = wav_tensor[0].to(torch.float32).detach().cpu().numpy()
return [wav], int(self.model.get_output_sample_rate())
def decode_streaming_batch(
self,
audio_codes: torch.Tensor,
use_optimized: bool = True,
pad_to_size: Optional[int] = None,
) -> Tuple[List[np.ndarray], int]:
"""
Batched streaming decode for multiple items.
Same as decode_streaming() but accepts B>1 batch dimension.
For B>1, the manual CUDA graph path is skipped and torch.compile
with reduce-overhead handles CUDA graphs internally.
Args:
audio_codes: [B, T, num_quantizers] tensor
use_optimized: Whether to use optimized path
pad_to_size: Pad to fixed frame count for consistent compilation
Returns:
Tuple[List[np.ndarray], int]: (list of B waveforms, sample_rate)
"""
model_type = self.model.get_model_type()
if model_type != "qwen3_tts_tokenizer_12hz":
return self.decode({"audio_codes": audio_codes})
assert audio_codes.dim() == 3, f"Expected [B, T, Q], got {audio_codes.shape}"
wav_tensor = self.model.decode_streaming(
audio_codes,
use_optimized=use_optimized,
pad_to_size=pad_to_size,
)
# wav_tensor is [B, samples]
wavs = [wav_tensor[b].to(torch.float32).detach().cpu().numpy()
for b in range(wav_tensor.shape[0])]
return wavs, int(self.model.get_output_sample_rate())