diff --git a/examples/test_batch_streaming.py b/examples/test_batch_streaming.py deleted file mode 100644 index be17c8c..0000000 --- a/examples/test_batch_streaming.py +++ /dev/null @@ -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") diff --git a/qwen_tts/core/models/modeling_qwen3_tts.py b/qwen_tts/core/models/modeling_qwen3_tts.py index 5524cd3..8c86fb0 100644 --- a/qwen_tts/core/models/modeling_qwen3_tts.py +++ b/qwen_tts/core/models/modeling_qwen3_tts.py @@ -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 diff --git a/qwen_tts/inference/qwen3_tts_tokenizer.py b/qwen_tts/inference/qwen3_tts_tokenizer.py index 7800bf5..93cbbe6 100644 --- a/qwen_tts/inference/qwen3_tts_tokenizer.py +++ b/qwen_tts/inference/qwen3_tts_tokenizer.py @@ -483,4 +483,42 @@ 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()) \ No newline at end of file + 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()) \ No newline at end of file