mirror of
https://github.com/Nighthawk42/Qwen3-TTS-streaming.git
synced 2026-08-30 08:52:27 +00:00
,,,
This commit is contained in:
dffdeeq
@@ -0,0 +1,211 @@
|
||||
"""
|
||||
Profile Talker forward pass to identify bottlenecks.
|
||||
"""
|
||||
|
||||
import time
|
||||
import torch
|
||||
from qwen_tts import Qwen3TTSModel
|
||||
|
||||
torch.set_float32_matmul_precision('high')
|
||||
|
||||
|
||||
def profile_generate():
|
||||
print("Loading model...")
|
||||
model = Qwen3TTSModel.from_pretrained(
|
||||
"Qwen/Qwen3-TTS-12Hz-1.7B-Base",
|
||||
device_map="cuda:0",
|
||||
dtype=torch.bfloat16,
|
||||
attn_implementation="flash_attention_2",
|
||||
)
|
||||
|
||||
# Get internal model for profiling
|
||||
tts_model = model.model
|
||||
talker = tts_model.talker
|
||||
code_predictor = talker.code_predictor
|
||||
|
||||
print(f"\nModel structure:")
|
||||
print(f" Talker: {talker.__class__.__name__}")
|
||||
print(f" CodePredictor: {code_predictor.__class__.__name__}")
|
||||
print(f" CodePredictor.model: {code_predictor.model.__class__.__name__}")
|
||||
print(f" Num codebook groups: {talker.config.num_code_groups}")
|
||||
|
||||
# Check attention implementation
|
||||
print(f"\nAttention implementation:")
|
||||
print(f" Talker: {talker.config._attn_implementation}")
|
||||
print(f" CodePredictor: {code_predictor.config._attn_implementation}")
|
||||
|
||||
# Create test inputs
|
||||
ref_audio_path = "../neurona-10sec.wav"
|
||||
ref_text = "Тестовый текст для профилирования."
|
||||
|
||||
voice_clone_prompt = model.create_voice_clone_prompt(
|
||||
ref_audio=ref_audio_path,
|
||||
ref_text=ref_text,
|
||||
)
|
||||
|
||||
test_text = "Привет, это тест профилирования генерации речи."
|
||||
|
||||
# Warmup
|
||||
print("\nWarmup run...")
|
||||
for chunk, sr in model.stream_generate_voice_clone(
|
||||
text="Раз два три.",
|
||||
language="Russian",
|
||||
voice_clone_prompt=voice_clone_prompt,
|
||||
emit_every_frames=4,
|
||||
):
|
||||
pass
|
||||
|
||||
# Profile with torch profiler
|
||||
print("\nProfiling with torch.profiler...")
|
||||
|
||||
with torch.profiler.profile(
|
||||
activities=[
|
||||
torch.profiler.ProfilerActivity.CPU,
|
||||
torch.profiler.ProfilerActivity.CUDA,
|
||||
],
|
||||
record_shapes=True,
|
||||
profile_memory=False,
|
||||
with_stack=False,
|
||||
) as prof:
|
||||
chunk_count = 0
|
||||
for chunk, sr in model.stream_generate_voice_clone(
|
||||
text=test_text,
|
||||
language="Russian",
|
||||
voice_clone_prompt=voice_clone_prompt,
|
||||
emit_every_frames=4,
|
||||
):
|
||||
chunk_count += 1
|
||||
if chunk_count >= 5: # Profile first 5 chunks only
|
||||
break
|
||||
|
||||
# Print profiler results
|
||||
print("\n" + "=" * 80)
|
||||
print("TOP 20 CUDA OPERATIONS BY TIME:")
|
||||
print("=" * 80)
|
||||
print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=20))
|
||||
|
||||
print("\n" + "=" * 80)
|
||||
print("TOP 20 CPU OPERATIONS BY TIME:")
|
||||
print("=" * 80)
|
||||
print(prof.key_averages().table(sort_by="cpu_time_total", row_limit=20))
|
||||
|
||||
# Manual timing of code_predictor.generate vs code_predictor.model.forward
|
||||
print("\n" + "=" * 80)
|
||||
print("MANUAL TIMING: code_predictor.generate() breakdown")
|
||||
print("=" * 80)
|
||||
|
||||
# Create dummy inputs for code_predictor
|
||||
batch_size = 1
|
||||
hidden_size = talker.config.hidden_size
|
||||
device = talker.device
|
||||
dtype = next(talker.parameters()).dtype
|
||||
|
||||
# Simulate past_hidden and last_id_hidden
|
||||
past_hidden = torch.randn(batch_size, 1, hidden_size, device=device, dtype=dtype)
|
||||
last_id_hidden = torch.randn(batch_size, 1, hidden_size, device=device, dtype=dtype)
|
||||
inputs_embeds = torch.cat((past_hidden, last_id_hidden), dim=1)
|
||||
|
||||
# Time code_predictor.generate()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
times_generate = []
|
||||
for _ in range(10):
|
||||
start = time.perf_counter()
|
||||
with torch.no_grad():
|
||||
result = code_predictor.generate(
|
||||
inputs_embeds=inputs_embeds,
|
||||
max_new_tokens=talker.config.num_code_groups - 1,
|
||||
do_sample=True,
|
||||
top_k=50,
|
||||
temperature=1.0,
|
||||
output_hidden_states=True,
|
||||
return_dict_in_generate=True,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
times_generate.append(time.perf_counter() - start)
|
||||
|
||||
print(f"code_predictor.generate() (10 runs):")
|
||||
print(f" Mean: {sum(times_generate)/len(times_generate)*1000:.1f}ms")
|
||||
print(f" Min: {min(times_generate)*1000:.1f}ms")
|
||||
print(f" Max: {max(times_generate)*1000:.1f}ms")
|
||||
|
||||
# Time individual forward calls
|
||||
projected = code_predictor.small_to_mtp_projection(inputs_embeds)
|
||||
|
||||
times_forward = []
|
||||
for _ in range(10):
|
||||
start = time.perf_counter()
|
||||
with torch.no_grad():
|
||||
out = code_predictor.model(
|
||||
inputs_embeds=projected,
|
||||
use_cache=True,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
times_forward.append(time.perf_counter() - start)
|
||||
|
||||
print(f"\ncode_predictor.model.forward() single call (10 runs):")
|
||||
print(f" Mean: {sum(times_forward)/len(times_forward)*1000:.1f}ms")
|
||||
print(f" Min: {min(times_forward)*1000:.1f}ms")
|
||||
print(f" Max: {max(times_forward)*1000:.1f}ms")
|
||||
|
||||
# Time 7 sequential forward calls (what generate does internally)
|
||||
times_7_forwards = []
|
||||
for _ in range(10):
|
||||
start = time.perf_counter()
|
||||
with torch.no_grad():
|
||||
# Prefill
|
||||
out = code_predictor.model(inputs_embeds=projected, use_cache=True)
|
||||
past_kv = out.past_key_values
|
||||
|
||||
# 6 more forward calls
|
||||
dummy_embed = torch.randn(1, 1, code_predictor.config.hidden_size,
|
||||
device=device, dtype=dtype)
|
||||
for _ in range(6):
|
||||
out = code_predictor.model(
|
||||
inputs_embeds=dummy_embed,
|
||||
past_key_values=past_kv,
|
||||
use_cache=True,
|
||||
)
|
||||
past_kv = out.past_key_values
|
||||
torch.cuda.synchronize()
|
||||
times_7_forwards.append(time.perf_counter() - start)
|
||||
|
||||
print(f"\n7x code_predictor.model.forward() sequential (10 runs):")
|
||||
print(f" Mean: {sum(times_7_forwards)/len(times_7_forwards)*1000:.1f}ms")
|
||||
print(f" Min: {min(times_7_forwards)*1000:.1f}ms")
|
||||
print(f" Max: {max(times_7_forwards)*1000:.1f}ms")
|
||||
|
||||
overhead = (sum(times_generate)/len(times_generate) - sum(times_7_forwards)/len(times_7_forwards)) * 1000
|
||||
print(f"\nHF generate() overhead: ~{overhead:.1f}ms per call")
|
||||
|
||||
# Time main talker forward
|
||||
print("\n" + "=" * 80)
|
||||
print("MANUAL TIMING: Talker.forward() breakdown")
|
||||
print("=" * 80)
|
||||
|
||||
# Need to prepare proper inputs for talker
|
||||
# This is complex, so let's just time the streaming loop
|
||||
|
||||
times_per_step = []
|
||||
step_count = 0
|
||||
|
||||
import time as _time
|
||||
|
||||
# Patch to measure step time
|
||||
original_forward = talker.forward.__wrapped__ if hasattr(talker.forward, '__wrapped__') else talker.forward
|
||||
|
||||
print("\nMeasuring actual streaming step times...")
|
||||
|
||||
for chunk, sr in model.stream_generate_voice_clone(
|
||||
text=test_text,
|
||||
language="Russian",
|
||||
voice_clone_prompt=voice_clone_prompt,
|
||||
emit_every_frames=4,
|
||||
):
|
||||
pass # Just run to completion
|
||||
|
||||
print("\nDone profiling!")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
profile_generate()
|
||||
@@ -0,0 +1,262 @@
|
||||
"""
|
||||
Test streaming TTS with torch.compile and CUDA graphs optimizations.
|
||||
|
||||
This script compares:
|
||||
1. Standard (non-streaming) generation
|
||||
2. Streaming without optimizations
|
||||
3. Streaming with torch.compile + CUDA graphs
|
||||
|
||||
Usage:
|
||||
cd Qwen3-TTS
|
||||
python examples/test_streaming_optimized.py
|
||||
"""
|
||||
|
||||
import time
|
||||
import numpy as np
|
||||
import torch
|
||||
import soundfile as sf
|
||||
from qwen_tts import Qwen3TTSModel
|
||||
|
||||
# Enable TensorFloat32 for better performance on Ampere+ GPUs
|
||||
torch.set_float32_matmul_precision('high')
|
||||
|
||||
|
||||
def log_time(start, operation):
|
||||
elapsed = time.time() - start
|
||||
print(f"[{elapsed:.2f}s] {operation}")
|
||||
return time.time()
|
||||
|
||||
|
||||
def run_streaming_test(
|
||||
model,
|
||||
text: str,
|
||||
language: str,
|
||||
voice_clone_prompt,
|
||||
emit_every_frames: int = 8,
|
||||
decode_window_frames: int = 80,
|
||||
label: str = "streaming",
|
||||
):
|
||||
"""Run streaming generation and return timing stats."""
|
||||
start = time.time()
|
||||
chunks = []
|
||||
chunk_sizes = []
|
||||
first_chunk_time = None
|
||||
chunk_count = 0
|
||||
sample_rate = 24000
|
||||
|
||||
for chunk, chunk_sr in model.stream_generate_voice_clone(
|
||||
text=text,
|
||||
language=language,
|
||||
voice_clone_prompt=voice_clone_prompt,
|
||||
emit_every_frames=emit_every_frames,
|
||||
decode_window_frames=decode_window_frames,
|
||||
overlap_samples=512,
|
||||
):
|
||||
chunk_count += 1
|
||||
chunks.append(chunk)
|
||||
chunk_sizes.append(len(chunk))
|
||||
sample_rate = chunk_sr
|
||||
if first_chunk_time is None:
|
||||
first_chunk_time = time.time() - start
|
||||
|
||||
total_time = time.time() - start
|
||||
final_audio = np.concatenate(chunks) if chunks else np.array([])
|
||||
|
||||
# Calculate audio duration and chunk stats
|
||||
audio_duration = len(final_audio) / sample_rate if sample_rate > 0 else 0
|
||||
avg_chunk_samples = np.mean(chunk_sizes) if chunk_sizes else 0
|
||||
avg_chunk_duration = avg_chunk_samples / sample_rate if sample_rate > 0 else 0
|
||||
|
||||
return {
|
||||
"label": label,
|
||||
"first_chunk_time": first_chunk_time,
|
||||
"total_time": total_time,
|
||||
"chunk_count": chunk_count,
|
||||
"audio": final_audio,
|
||||
"sample_rate": sample_rate,
|
||||
"audio_duration": audio_duration,
|
||||
"avg_chunk_samples": avg_chunk_samples,
|
||||
"avg_chunk_duration": avg_chunk_duration,
|
||||
}
|
||||
|
||||
|
||||
def main():
|
||||
total_start = time.time()
|
||||
|
||||
# Streaming parameters - KEEP THESE CONSISTENT!
|
||||
EMIT_EVERY = 4 # Reduced from 8 for lower latency
|
||||
DECODE_WINDOW = 80
|
||||
|
||||
print("=" * 60)
|
||||
print("Loading model...")
|
||||
print("=" * 60)
|
||||
|
||||
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",
|
||||
)
|
||||
log_time(start, "Model loaded")
|
||||
|
||||
# Reference audio setup
|
||||
ref_audio_path = "../neurona-10sec.wav"
|
||||
ref_text = (
|
||||
"Обоссышься точно, я короче твои цветы продала, цветы с подоконника, рюкзак, сменку, "
|
||||
"пару парт, ща еще окна и сторожа еще смотри приедут. А У тебя кстати родители. "
|
||||
"Перед тобой Тони Старк, только после пту. Стив Джобс, только с контузией. Илон Макс, не Маск. "
|
||||
)
|
||||
|
||||
start = time.time()
|
||||
voice_clone_prompt = model.create_voice_clone_prompt(
|
||||
ref_audio=ref_audio_path,
|
||||
ref_text=ref_text,
|
||||
)
|
||||
log_time(start, "Voice clone prompt created")
|
||||
|
||||
# Test text
|
||||
test_text = "Привет всем! Я того всё ебала, что за новый голос тут на обзоре у вилсакома? А? Так он мне понравился. Ганс оф буллщит."
|
||||
|
||||
results = []
|
||||
|
||||
# ============== Test 1: Standard generation ==============
|
||||
print("\n" + "=" * 60)
|
||||
print("Test 1: Standard (non-streaming) generation")
|
||||
print("=" * 60)
|
||||
|
||||
start = time.time()
|
||||
wavs, sr = model.generate_voice_clone(
|
||||
text=test_text,
|
||||
language="Russian",
|
||||
voice_clone_prompt=voice_clone_prompt,
|
||||
)
|
||||
standard_time = time.time() - start
|
||||
standard_audio_duration = len(wavs[0]) / sr
|
||||
standard_rtf = standard_time / standard_audio_duration
|
||||
print(f"[{standard_time:.2f}s] Standard generation complete")
|
||||
print(f"Audio duration: {standard_audio_duration:.2f}s, RTF: {standard_rtf:.2f}")
|
||||
sf.write("output_standard.wav", wavs[0], sr)
|
||||
results.append({
|
||||
"label": "standard",
|
||||
"total_time": standard_time,
|
||||
"audio_duration": standard_audio_duration,
|
||||
})
|
||||
|
||||
# ============== Test 2: Streaming WITHOUT optimizations ==============
|
||||
print("\n" + "=" * 60)
|
||||
print("Test 2: Streaming WITHOUT optimizations")
|
||||
print("=" * 60)
|
||||
|
||||
result = run_streaming_test(
|
||||
model, test_text, "Russian", voice_clone_prompt,
|
||||
emit_every_frames=EMIT_EVERY,
|
||||
decode_window_frames=DECODE_WINDOW,
|
||||
label="streaming_baseline",
|
||||
)
|
||||
results.append(result)
|
||||
sf.write("output_streaming_baseline.wav", result["audio"], result["sample_rate"])
|
||||
rtf = result['total_time'] / result['audio_duration'] if result['audio_duration'] > 0 else 0
|
||||
print(f"First chunk: {result['first_chunk_time']:.2f}s, Total: {result['total_time']:.2f}s, Chunks: {result['chunk_count']}")
|
||||
print(f"Audio duration: {result['audio_duration']:.2f}s, Chunk duration: {result['avg_chunk_duration']*1000:.0f}ms, RTF: {rtf:.2f}")
|
||||
|
||||
# ============== Test 3: Streaming WITH optimizations ==============
|
||||
print("\n" + "=" * 60)
|
||||
print("Test 3: Streaming WITH decoder torch.compile")
|
||||
print("=" * 60)
|
||||
|
||||
# Enable optimizations - this is the key step!
|
||||
# - Decoder torch.compile with reduce-overhead mode (includes CUDA graphs)
|
||||
print("\nEnabling streaming optimizations...")
|
||||
model.enable_streaming_optimizations(
|
||||
decode_window_frames=DECODE_WINDOW,
|
||||
use_compile=True,
|
||||
use_cuda_graphs=False, # Not needed with reduce-overhead mode
|
||||
compile_mode="reduce-overhead",
|
||||
)
|
||||
|
||||
# Warmup run (first run after compile is slower due to compilation)
|
||||
print("\nWarmup run (first run after compile)...")
|
||||
warmup_result = run_streaming_test(
|
||||
model, "Тест один два три четыре пять.", "Russian", voice_clone_prompt,
|
||||
emit_every_frames=EMIT_EVERY,
|
||||
decode_window_frames=DECODE_WINDOW,
|
||||
label="warmup",
|
||||
)
|
||||
warmup_rtf = warmup_result['total_time'] / warmup_result['audio_duration'] if warmup_result['audio_duration'] > 0 else 0
|
||||
print(f"Warmup: First chunk: {warmup_result['first_chunk_time']:.2f}s, Total: {warmup_result['total_time']:.2f}s, Audio: {warmup_result['audio_duration']:.2f}s, RTF: {warmup_rtf:.2f}")
|
||||
|
||||
# Actual test run
|
||||
print("\nOptimized test run...")
|
||||
result = run_streaming_test(
|
||||
model, test_text, "Russian", voice_clone_prompt,
|
||||
emit_every_frames=EMIT_EVERY,
|
||||
decode_window_frames=DECODE_WINDOW,
|
||||
label="streaming_optimized",
|
||||
)
|
||||
results.append(result)
|
||||
sf.write("output_streaming_optimized.wav", result["audio"], result["sample_rate"])
|
||||
opt_rtf = result['total_time'] / result['audio_duration'] if result['audio_duration'] > 0 else 0
|
||||
print(f"First chunk: {result['first_chunk_time']:.2f}s, Total: {result['total_time']:.2f}s, Chunks: {result['chunk_count']}")
|
||||
print(f"Audio duration: {result['audio_duration']:.2f}s, Chunk duration: {result['avg_chunk_duration']*1000:.0f}ms, RTF: {opt_rtf:.2f}")
|
||||
|
||||
# Second optimized run to show stable performance
|
||||
print("\nSecond optimized run...")
|
||||
result2 = run_streaming_test(
|
||||
model, test_text, "Russian", voice_clone_prompt,
|
||||
emit_every_frames=EMIT_EVERY,
|
||||
decode_window_frames=DECODE_WINDOW,
|
||||
label="streaming_optimized_2",
|
||||
)
|
||||
results.append(result2)
|
||||
opt2_rtf = result2['total_time'] / result2['audio_duration'] if result2['audio_duration'] > 0 else 0
|
||||
print(f"First chunk: {result2['first_chunk_time']:.2f}s, Total: {result2['total_time']:.2f}s, Audio: {result2['audio_duration']:.2f}s, RTF: {opt2_rtf:.2f}")
|
||||
|
||||
# ============== Summary ==============
|
||||
print("\n" + "=" * 80)
|
||||
print("SUMMARY")
|
||||
print("=" * 80)
|
||||
|
||||
baseline_total = results[1]["total_time"]
|
||||
baseline_first = results[1]["first_chunk_time"]
|
||||
|
||||
print(f"\n{'Method':<25} {'1st Chunk':>10} {'Total':>8} {'Audio':>8} {'RTF':>6} {'Chunks':>7} {'Speedup':>8}")
|
||||
print("-" * 80)
|
||||
|
||||
# Standard generation
|
||||
std = results[0]
|
||||
std_rtf = std['total_time'] / std['audio_duration'] if std.get('audio_duration', 0) > 0 else 0
|
||||
print(f"{'Standard (no streaming)':<25} {'N/A':>10} {std['total_time']:>7.2f}s {std.get('audio_duration', 0):>7.2f}s {std_rtf:>6.2f} {'N/A':>7} {'N/A':>8}")
|
||||
|
||||
for r in results[1:]:
|
||||
first = r.get("first_chunk_time", 0)
|
||||
total = r["total_time"]
|
||||
audio_dur = r.get("audio_duration", 0)
|
||||
rtf = total / audio_dur if audio_dur > 0 else 0
|
||||
chunks = r.get("chunk_count", 0)
|
||||
speedup_total = baseline_total / total if total > 0 else 0
|
||||
print(f"{r['label']:<25} {first:>9.2f}s {total:>7.2f}s {audio_dur:>7.2f}s {rtf:>6.2f} {chunks:>7} {speedup_total:>7.2f}x")
|
||||
|
||||
# Chunk duration info
|
||||
if results[1].get("avg_chunk_duration", 0) > 0:
|
||||
print(f"\nChunk duration: ~{results[1]['avg_chunk_duration']*1000:.0f}ms ({results[1]['avg_chunk_samples']:.0f} samples @ {results[1]['sample_rate']}Hz)")
|
||||
|
||||
print(f"\n[{time.time() - total_start:.2f}s] TOTAL SCRIPT TIME")
|
||||
|
||||
# Tips
|
||||
print("\n" + "=" * 60)
|
||||
print("TIPS FOR BEST PERFORMANCE")
|
||||
print("=" * 60)
|
||||
print("""
|
||||
1. Call enable_streaming_optimizations() ONCE after model loading
|
||||
2. Use compile_mode="reduce-overhead" (default) - it includes CUDA graphs automatically
|
||||
3. First run after compile is slow (compilation), subsequent runs are fast
|
||||
4. For lowest latency: use smaller emit_every_frames (e.g., 4)
|
||||
5. For best quality: use larger decode_window_frames (e.g., 80-100)
|
||||
6. You can also try compile_mode="max-autotune" for potentially better performance
|
||||
(but longer initial compilation time)
|
||||
""")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Binary file not shown.
@@ -1368,6 +1368,122 @@ class Qwen3TTSTalkerCodePredictorModelForConditionalGeneration(Qwen3TTSPreTraine
|
||||
model_kwargs["generation_steps"] = outputs.generation_steps
|
||||
return model_kwargs
|
||||
|
||||
def generate_fast(
|
||||
self,
|
||||
inputs_embeds: torch.Tensor,
|
||||
num_codebooks: int,
|
||||
do_sample: bool = True,
|
||||
temperature: float = 1.0,
|
||||
top_k: int = 50,
|
||||
top_p: float = 1.0,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Fast generation that bypasses HuggingFace's generate() overhead.
|
||||
|
||||
This is ~2-3x faster than using generate() because:
|
||||
1. No GenerationMixin overhead (config creation, stopping criteria, etc.)
|
||||
2. Direct forward calls with minimal wrapper logic
|
||||
3. Simple KV-cache management
|
||||
|
||||
Args:
|
||||
inputs_embeds: Initial embeddings [B, 2, hidden_size] (past_hidden + first_token_embed)
|
||||
num_codebooks: Number of codebook tokens to generate (typically 7)
|
||||
do_sample: Whether to sample or use greedy decoding
|
||||
temperature: Sampling temperature
|
||||
top_k: Top-k filtering
|
||||
top_p: Top-p (nucleus) filtering
|
||||
|
||||
Returns:
|
||||
Generated token IDs [B, num_codebooks]
|
||||
"""
|
||||
batch_size = inputs_embeds.shape[0]
|
||||
device = inputs_embeds.device
|
||||
|
||||
# Project inputs
|
||||
inputs_embeds = self.small_to_mtp_projection(inputs_embeds)
|
||||
|
||||
# Prefill: process initial embeddings
|
||||
outputs = self.model(
|
||||
input_ids=None,
|
||||
inputs_embeds=inputs_embeds,
|
||||
use_cache=True,
|
||||
output_hidden_states=False,
|
||||
)
|
||||
past_key_values = outputs.past_key_values
|
||||
hidden_states = outputs.last_hidden_state
|
||||
|
||||
# Generate tokens for each codebook
|
||||
generated_tokens = []
|
||||
generation_step = 0 # Start from codebook 1 (index 0 in lm_head)
|
||||
|
||||
for step in range(num_codebooks):
|
||||
# Get logits for current codebook
|
||||
logits = self.lm_head[step](hidden_states[:, -1, :]) # [B, vocab_size]
|
||||
|
||||
# Sample or greedy
|
||||
if do_sample and temperature > 0:
|
||||
logits = logits / temperature
|
||||
|
||||
# Top-k filtering
|
||||
if top_k > 0:
|
||||
top_k_val = min(top_k, logits.size(-1))
|
||||
indices_to_remove = logits < torch.topk(logits, top_k_val)[0][..., -1, None]
|
||||
logits = logits.masked_fill(indices_to_remove, float('-inf'))
|
||||
|
||||
# Top-p filtering
|
||||
if top_p < 1.0:
|
||||
sorted_logits, sorted_indices = torch.sort(logits, descending=True)
|
||||
cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
|
||||
sorted_indices_to_remove = cumulative_probs > top_p
|
||||
sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
|
||||
sorted_indices_to_remove[..., 0] = 0
|
||||
indices_to_remove = sorted_indices_to_remove.scatter(
|
||||
1, sorted_indices, sorted_indices_to_remove
|
||||
)
|
||||
logits = logits.masked_fill(indices_to_remove, float('-inf'))
|
||||
|
||||
probs = F.softmax(logits, dim=-1)
|
||||
next_token = torch.multinomial(probs, num_samples=1)
|
||||
else:
|
||||
next_token = torch.argmax(logits, dim=-1, keepdim=True)
|
||||
|
||||
generated_tokens.append(next_token)
|
||||
|
||||
# Stop if we've generated all codebooks
|
||||
if step == num_codebooks - 1:
|
||||
break
|
||||
|
||||
# Get embedding for next step
|
||||
next_embeds = self.model.get_input_embeddings()[step](next_token)
|
||||
next_embeds = self.small_to_mtp_projection(next_embeds)
|
||||
|
||||
# Forward pass for next position
|
||||
outputs = self.model(
|
||||
input_ids=None,
|
||||
inputs_embeds=next_embeds,
|
||||
past_key_values=past_key_values,
|
||||
use_cache=True,
|
||||
output_hidden_states=False,
|
||||
)
|
||||
past_key_values = outputs.past_key_values
|
||||
hidden_states = outputs.last_hidden_state
|
||||
|
||||
# Concatenate all generated tokens
|
||||
return torch.cat(generated_tokens, dim=-1) # [B, num_codebooks]
|
||||
|
||||
def enable_compile(self, mode: str = "reduce-overhead"):
|
||||
"""
|
||||
Enable torch.compile for the code predictor model.
|
||||
|
||||
This compiles the inner model forward pass for faster execution.
|
||||
Should be called once after model loading.
|
||||
"""
|
||||
self.model.forward = torch.compile(
|
||||
self.model.forward,
|
||||
mode=mode,
|
||||
fullgraph=False, # Allow graph breaks for flexibility
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Qwen3TTSTalkerOutputWithPast(ModelOutput):
|
||||
@@ -1682,6 +1798,10 @@ class Qwen3TTSTalkerForConditionalGeneration(Qwen3TTSTalkerTextPreTrainedModel,
|
||||
sub_talker_loss = sub_talker_outputs.loss
|
||||
return sub_talker_logits, sub_talker_loss
|
||||
|
||||
def enable_fast_codebook_gen(self, enable: bool = True):
|
||||
"""Enable fast codebook generation (bypasses HuggingFace generate() overhead)."""
|
||||
self._use_fast_codebook_gen = enable
|
||||
|
||||
@can_return_tuple
|
||||
def forward(
|
||||
self,
|
||||
@@ -1718,22 +1838,49 @@ class Qwen3TTSTalkerForConditionalGeneration(Qwen3TTSTalkerTextPreTrainedModel,
|
||||
# Generate
|
||||
else:
|
||||
last_id_hidden = self.get_input_embeddings()(input_ids)
|
||||
predictor_result = self.code_predictor.generate(
|
||||
inputs_embeds=torch.cat((past_hidden, last_id_hidden), dim=1),
|
||||
max_new_tokens=self.config.num_code_groups - 1,
|
||||
do_sample=subtalker_dosample,
|
||||
top_p=subtalker_top_p,
|
||||
top_k=subtalker_top_k,
|
||||
temperature=subtalker_temperature,
|
||||
output_hidden_states=True,
|
||||
return_dict_in_generate=True,
|
||||
)
|
||||
codec_ids = torch.cat((input_ids, predictor_result.sequences), dim=-1)
|
||||
codec_hiddens = torch.cat(
|
||||
[last_id_hidden]
|
||||
+ [self.code_predictor.get_input_embeddings()[i](predictor_result.sequences[..., i:i+1]) for i in range(self.config.num_code_groups - 1)],
|
||||
dim=1,
|
||||
)
|
||||
|
||||
# Use fast path if enabled (bypasses HuggingFace generate() overhead)
|
||||
use_fast = getattr(self, '_use_fast_codebook_gen', False)
|
||||
|
||||
if use_fast:
|
||||
# Fast path: direct forward loop (~2-3x faster)
|
||||
# Mark step begin for CUDA graphs to avoid tensor overwrite errors
|
||||
torch.compiler.cudagraph_mark_step_begin()
|
||||
codebook_tokens = self.code_predictor.generate_fast(
|
||||
inputs_embeds=torch.cat((past_hidden, last_id_hidden), dim=1),
|
||||
num_codebooks=self.config.num_code_groups - 1,
|
||||
do_sample=subtalker_dosample if subtalker_dosample is not None else True,
|
||||
temperature=subtalker_temperature if subtalker_temperature is not None else 1.0,
|
||||
top_k=subtalker_top_k if subtalker_top_k is not None else 50,
|
||||
top_p=subtalker_top_p if subtalker_top_p is not None else 1.0,
|
||||
)
|
||||
codec_ids = torch.cat((input_ids, codebook_tokens), dim=-1)
|
||||
codec_hiddens = torch.cat(
|
||||
[last_id_hidden]
|
||||
+ [self.code_predictor.get_input_embeddings()[i](codebook_tokens[..., i:i+1]) for i in range(self.config.num_code_groups - 1)],
|
||||
dim=1,
|
||||
)
|
||||
else:
|
||||
# Original path: HuggingFace generate()
|
||||
# Mark step begin for CUDA graphs to avoid tensor overwrite errors
|
||||
torch.compiler.cudagraph_mark_step_begin()
|
||||
predictor_result = self.code_predictor.generate(
|
||||
inputs_embeds=torch.cat((past_hidden, last_id_hidden), dim=1),
|
||||
max_new_tokens=self.config.num_code_groups - 1,
|
||||
do_sample=subtalker_dosample,
|
||||
top_p=subtalker_top_p,
|
||||
top_k=subtalker_top_k,
|
||||
temperature=subtalker_temperature,
|
||||
output_hidden_states=True,
|
||||
return_dict_in_generate=True,
|
||||
)
|
||||
codec_ids = torch.cat((input_ids, predictor_result.sequences), dim=-1)
|
||||
codec_hiddens = torch.cat(
|
||||
[last_id_hidden]
|
||||
+ [self.code_predictor.get_input_embeddings()[i](predictor_result.sequences[..., i:i+1]) for i in range(self.config.num_code_groups - 1)],
|
||||
dim=1,
|
||||
)
|
||||
|
||||
inputs_embeds = codec_hiddens.sum(1, keepdim=True)
|
||||
|
||||
if generation_step < trailing_text_hidden.shape[1]:
|
||||
@@ -1901,7 +2048,61 @@ class Qwen3TTSForConditionalGeneration(Qwen3TTSPreTrainedModel, GenerationMixin)
|
||||
|
||||
def get_supported_languages(self):
|
||||
return self.supported_languages
|
||||
|
||||
|
||||
def enable_streaming_optimizations(
|
||||
self,
|
||||
decode_window_frames: int = 80,
|
||||
use_compile: bool = True,
|
||||
use_cuda_graphs: bool = True,
|
||||
compile_mode: str = "reduce-overhead",
|
||||
use_fast_codebook: bool = False, # Disabled: needs debugging, currently slower
|
||||
compile_codebook_predictor: bool = True,
|
||||
):
|
||||
"""
|
||||
Enable torch.compile and CUDA graphs optimizations for streaming decode.
|
||||
|
||||
Call this after model loading to speed up streaming generation.
|
||||
The optimizations apply to the speech tokenizer's decoder and talker.
|
||||
|
||||
Args:
|
||||
decode_window_frames: Fixed window size for streaming (must match
|
||||
decode_window_frames parameter in stream_generate_pcm)
|
||||
use_compile: Apply torch.compile to the decoder
|
||||
use_cuda_graphs: Capture CUDA graphs for the fixed window size
|
||||
compile_mode: torch.compile mode ("reduce-overhead" recommended)
|
||||
use_fast_codebook: Use fast codebook generation (bypasses HF generate() overhead)
|
||||
compile_codebook_predictor: Apply torch.compile to codebook predictor (default True)
|
||||
|
||||
Returns:
|
||||
self for method chaining
|
||||
|
||||
Example:
|
||||
model = Qwen3TTSForConditionalGeneration.from_pretrained(...)
|
||||
model.enable_streaming_optimizations(decode_window_frames=80)
|
||||
"""
|
||||
if self.speech_tokenizer is None:
|
||||
raise ValueError("Speech tokenizer not loaded. Call from_pretrained() first.")
|
||||
|
||||
# Enable decoder optimizations
|
||||
self.speech_tokenizer.enable_streaming_optimizations(
|
||||
decode_window_frames=decode_window_frames,
|
||||
use_compile=use_compile,
|
||||
use_cuda_graphs=use_cuda_graphs,
|
||||
compile_mode=compile_mode,
|
||||
)
|
||||
|
||||
# Enable fast codebook generation (bypasses HuggingFace generate() overhead)
|
||||
if use_fast_codebook:
|
||||
print("[Talker] Enabling fast codebook generation...")
|
||||
self.talker.enable_fast_codebook_gen(True)
|
||||
|
||||
# Compile codebook predictor for faster inference
|
||||
if compile_codebook_predictor and use_compile:
|
||||
print(f"[CodePredictor] Compiling model with mode={compile_mode}...")
|
||||
self.talker.code_predictor.enable_compile(mode=compile_mode)
|
||||
|
||||
return self
|
||||
|
||||
@classmethod
|
||||
def from_pretrained(
|
||||
cls,
|
||||
@@ -2314,8 +2515,8 @@ class Qwen3TTSForConditionalGeneration(Qwen3TTSPreTrainedModel, GenerationMixin)
|
||||
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,)
|
||||
],
|
||||
"output_hidden_states": getattr(kwargs, "output_hidden_states", True),
|
||||
"return_dict_in_generate": getattr(kwargs, "return_dict_in_generate", True)
|
||||
"output_hidden_states": kwargs.get("output_hidden_states", True),
|
||||
"return_dict_in_generate": kwargs.get("return_dict_in_generate", True)
|
||||
}
|
||||
|
||||
# Build talker inputs using shared method
|
||||
@@ -2378,6 +2579,8 @@ class Qwen3TTSForConditionalGeneration(Qwen3TTSPreTrainedModel, GenerationMixin)
|
||||
decode_window_frames: int = 80,
|
||||
overlap_samples: int = 512,
|
||||
max_frames: int = 10000,
|
||||
# Optimization flags
|
||||
use_optimized_decode: bool = True,
|
||||
) -> Generator[tuple[np.ndarray, int], None, None]:
|
||||
"""
|
||||
Stream audio generation, yielding PCM chunks as they are generated.
|
||||
@@ -2399,6 +2602,7 @@ class Qwen3TTSForConditionalGeneration(Qwen3TTSPreTrainedModel, GenerationMixin)
|
||||
decode_window_frames: Window size for decoding (longer = better quality, more latency)
|
||||
overlap_samples: Overlap samples for crossfade between chunks
|
||||
max_frames: Maximum number of codec frames to generate
|
||||
use_optimized_decode: Use CUDA graph optimized decode when available (default True)
|
||||
|
||||
Yields:
|
||||
tuple[np.ndarray, int]: (pcm_chunk as float32 array, sample_rate)
|
||||
@@ -2424,6 +2628,9 @@ class Qwen3TTSForConditionalGeneration(Qwen3TTSPreTrainedModel, GenerationMixin)
|
||||
if i != eos_id
|
||||
]
|
||||
|
||||
# Mark step begin for CUDA graphs (required for torch.compile with reduce-overhead)
|
||||
torch.compiler.cudagraph_mark_step_begin()
|
||||
|
||||
# Prefill: single forward pass to initialize KV cache
|
||||
out = self.talker.forward(
|
||||
inputs_embeds=talker_input_embeds,
|
||||
@@ -2446,7 +2653,7 @@ class Qwen3TTSForConditionalGeneration(Qwen3TTSPreTrainedModel, GenerationMixin)
|
||||
past_hidden = out.past_hidden
|
||||
generation_step = out.generation_step
|
||||
|
||||
print(f"[Prefill] done, generation_step={generation_step}, hidden_states type={type(out.hidden_states)}")
|
||||
# Debug removed for performance: prefill done
|
||||
|
||||
# Sample first token from prefill logits
|
||||
last_logits = out.logits[:, -1, :]
|
||||
@@ -2454,25 +2661,25 @@ class Qwen3TTSForConditionalGeneration(Qwen3TTSPreTrainedModel, GenerationMixin)
|
||||
token = _sample_next_token(last_logits, temperature, top_k, top_p, suppress_tokens)
|
||||
else:
|
||||
token = torch.argmax(last_logits, dim=-1)
|
||||
print(f"[Prefill] first token sampled: {token.item()}, eos_id={eos_id}")
|
||||
# Debug removed for performance: first token sampled
|
||||
|
||||
# Decode loop
|
||||
codes_buffer: list[torch.Tensor] = []
|
||||
decoded_tail: Optional[np.ndarray] = None
|
||||
frames_since_emit = 0
|
||||
total_frames_emitted = 0 # Track how many frames we've already emitted audio for
|
||||
import time as _time
|
||||
_last_emit_time = _time.time()
|
||||
|
||||
for step_idx in range(max_frames):
|
||||
_step_start = _time.time()
|
||||
# Mark step begin for CUDA graphs to avoid tensor overwrite errors
|
||||
# This is required when using torch.compile with reduce-overhead mode
|
||||
torch.compiler.cudagraph_mark_step_begin()
|
||||
|
||||
# Single-step forward
|
||||
step_out = self.talker.forward(
|
||||
input_ids=token.unsqueeze(1),
|
||||
use_cache=True,
|
||||
return_dict=True,
|
||||
output_hidden_states=True,
|
||||
output_hidden_states=False, # Disabled: codec_ids accessed via hidden_states[1] still works
|
||||
past_key_values=past_key_values,
|
||||
past_hidden=past_hidden,
|
||||
generation_step=generation_step,
|
||||
@@ -2483,7 +2690,6 @@ class Qwen3TTSForConditionalGeneration(Qwen3TTSPreTrainedModel, GenerationMixin)
|
||||
subtalker_top_p=subtalker_top_p,
|
||||
subtalker_temperature=subtalker_temperature,
|
||||
)
|
||||
_forward_time = _time.time() - _step_start
|
||||
|
||||
# Update state for next iteration
|
||||
past_key_values = step_out.past_key_values
|
||||
@@ -2492,19 +2698,14 @@ 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]
|
||||
codec_ids_cpu = codec_ids[0].detach().cpu()
|
||||
|
||||
# Debug: periodic status
|
||||
if step_idx % 10 == 0:
|
||||
print(f"[Step {step_idx}] forward={_forward_time*1000:.1f}ms, codec[0]={int(codec_ids_cpu[0].item())}, gen_step={generation_step}")
|
||||
|
||||
# Check for EOS in first codebook BEFORE adding to buffer
|
||||
# Check for EOS in first codebook ON GPU (avoids CPU sync bottleneck)
|
||||
# EOS token is out of range for speech tokenizer, so we must not include it
|
||||
if int(codec_ids_cpu[0].item()) == int(eos_id):
|
||||
print(f"[Step {step_idx}] EOS reached")
|
||||
if codec_ids[0, 0] == eos_id:
|
||||
break
|
||||
|
||||
codes_buffer.append(codec_ids_cpu)
|
||||
# CPU transfer AFTER EOS check (not before) to avoid sync on every step
|
||||
codes_buffer.append(codec_ids[0].detach().cpu())
|
||||
|
||||
# Sample next token for first codebook
|
||||
step_logits = step_out.logits[:, -1, :]
|
||||
@@ -2519,23 +2720,29 @@ class Qwen3TTSForConditionalGeneration(Qwen3TTSPreTrainedModel, GenerationMixin)
|
||||
frames_since_emit = 0
|
||||
|
||||
# Decode window of codec frames to PCM
|
||||
_decode_start = _time.time()
|
||||
start = max(0, len(codes_buffer) - decode_window_frames)
|
||||
window = torch.stack(codes_buffer[start:], dim=0) # [T, num_code_groups]
|
||||
|
||||
print(f"[Emit {len(codes_buffer)//emit_every_frames}] decoding window size={window.shape[0]}, time_since_last={_time.time()-_last_emit_time:.2f}s")
|
||||
# Debug removed for performance: emit info
|
||||
|
||||
wavs, sr = self.speech_tokenizer.decode([{"audio_codes": window.to(self.talker.device)}])
|
||||
_decode_time = _time.time() - _decode_start
|
||||
print(f"[Emit] decode took {_decode_time*1000:.1f}ms")
|
||||
_last_emit_time = _time.time()
|
||||
# Use optimized decode path when available
|
||||
# Pass pad_to_size to ensure fixed tensor size for torch.compile
|
||||
if use_optimized_decode and hasattr(self.speech_tokenizer, 'decode_streaming'):
|
||||
wavs, sr = self.speech_tokenizer.decode_streaming(
|
||||
window.to(self.talker.device),
|
||||
use_optimized=True,
|
||||
pad_to_size=decode_window_frames,
|
||||
)
|
||||
else:
|
||||
wavs, sr = self.speech_tokenizer.decode([{"audio_codes": window.to(self.talker.device)}])
|
||||
# Debug removed for performance: decode time tracking
|
||||
|
||||
wav = wavs[0].astype(np.float32)
|
||||
|
||||
# Extract only new samples (tail of decoded window)
|
||||
T = window.shape[0]
|
||||
samples_per_frame = wav.shape[0] / float(T) if T > 0 else 0
|
||||
step_samples = int(round(samples_per_frame * emit_every_frames))
|
||||
# Use fixed upsample rate to avoid floating-point drift
|
||||
samples_per_frame = self.speech_tokenizer.get_decode_upsample_rate()
|
||||
step_samples = samples_per_frame * emit_every_frames
|
||||
chunk = wav[-step_samples:] if step_samples > 0 else wav
|
||||
|
||||
# Crossfade with previous chunk tail for smooth transition
|
||||
@@ -2552,7 +2759,7 @@ class Qwen3TTSForConditionalGeneration(Qwen3TTSPreTrainedModel, GenerationMixin)
|
||||
# Flush: decode only remaining frames that haven't been emitted yet
|
||||
remaining_frames = len(codes_buffer) - total_frames_emitted
|
||||
if remaining_frames > 0:
|
||||
print(f"[Flush] decoding remaining {remaining_frames} frames (total={len(codes_buffer)}, emitted={total_frames_emitted})")
|
||||
# Debug removed for performance: flush info
|
||||
# Decode a window that includes some context for quality
|
||||
context_frames = min(total_frames_emitted, decode_window_frames - remaining_frames)
|
||||
start_idx = total_frames_emitted - context_frames
|
||||
@@ -2574,7 +2781,7 @@ class Qwen3TTSForConditionalGeneration(Qwen3TTSPreTrainedModel, GenerationMixin)
|
||||
head = _crossfade(decoded_tail[-ov:], wav[:ov])
|
||||
wav = np.concatenate([head, wav[ov:]], axis=0)
|
||||
|
||||
print(f"[Flush] done, yielding {len(wav)} samples")
|
||||
# Debug removed for performance: flush done
|
||||
yield wav, sr
|
||||
|
||||
|
||||
|
||||
@@ -16,7 +16,7 @@
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from typing import Callable, Optional, Union, List
|
||||
from typing import Callable, Optional, Union, List, Any
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
@@ -825,6 +825,14 @@ class Qwen3TTSTokenizerV2Decoder(Qwen3TTSTokenizerV2DecoderPreTrainedModel):
|
||||
super().__init__(config)
|
||||
self.total_upsample = np.prod(config.upsample_rates + config.upsampling_ratios)
|
||||
self.pre_transformer = Qwen3TTSTokenizerV2DecoderTransformerModel._from_config(config)
|
||||
|
||||
# Optimization state
|
||||
self._compiled_forward: Optional[Callable] = None
|
||||
self._compile_mode: Optional[str] = None
|
||||
self._cuda_graph: Optional[torch.cuda.CUDAGraph] = None
|
||||
self._static_input: Optional[torch.Tensor] = None
|
||||
self._static_output: Optional[torch.Tensor] = None
|
||||
self._graph_window_size: Optional[int] = None
|
||||
|
||||
self.quantizer = SplitResidualVectorQuantizer(
|
||||
dimension=config.codebook_dim // 2,
|
||||
@@ -894,6 +902,178 @@ class Qwen3TTSTokenizerV2Decoder(Qwen3TTSTokenizerV2DecoderPreTrainedModel):
|
||||
start_index = end_index
|
||||
return torch.cat(wavs, dim=-1)
|
||||
|
||||
def compile_for_streaming(self, mode: str = "reduce-overhead", backend: str = "inductor"):
|
||||
"""
|
||||
Apply torch.compile to the forward pass for faster streaming decode.
|
||||
|
||||
Note: "reduce-overhead" mode already includes CUDA graph optimizations internally,
|
||||
so you should NOT use capture_cuda_graph() when using this mode.
|
||||
|
||||
Args:
|
||||
mode: Compilation mode:
|
||||
- "reduce-overhead" (recommended): Uses CUDA graphs internally, best for streaming
|
||||
- "max-autotune": Maximum optimization, longer compile time
|
||||
- "default": Good balance, no internal CUDA graphs
|
||||
backend: Compilation backend ("inductor" recommended)
|
||||
"""
|
||||
if not hasattr(torch, 'compile'):
|
||||
print("[Decoder] torch.compile not available (requires PyTorch 2.0+)")
|
||||
return self
|
||||
|
||||
print(f"[Decoder] Compiling forward with mode={mode}, backend={backend}...")
|
||||
print(f"[Decoder] Note: mode='reduce-overhead' includes CUDA graphs automatically")
|
||||
self._compiled_forward = torch.compile(
|
||||
self._forward_impl,
|
||||
mode=mode,
|
||||
fullgraph=False,
|
||||
dynamic=False,
|
||||
backend=backend,
|
||||
)
|
||||
self._compile_mode = mode
|
||||
print("[Decoder] Compilation complete")
|
||||
return self
|
||||
|
||||
def _forward_impl(self, codes):
|
||||
"""Internal forward implementation for compilation."""
|
||||
hidden = self.quantizer.decode(codes)
|
||||
hidden = self.pre_conv(hidden).transpose(1, 2)
|
||||
hidden = self.pre_transformer(inputs_embeds=hidden).last_hidden_state
|
||||
hidden = hidden.permute(0, 2, 1)
|
||||
for blocks in self.upsample:
|
||||
for block in blocks:
|
||||
hidden = block(hidden)
|
||||
wav = hidden
|
||||
for block in self.decoder:
|
||||
wav = block(wav)
|
||||
return wav.clamp(min=-1, max=1)
|
||||
|
||||
def capture_cuda_graph(self, window_size: int = 80, warmup_runs: int = 3):
|
||||
"""
|
||||
Capture CUDA graph for a fixed window size.
|
||||
|
||||
CUDA graphs eliminate CPU overhead by capturing and replaying
|
||||
GPU operations. Best for streaming with fixed decode_window_frames.
|
||||
|
||||
WARNING: Do NOT use this with torch.compile mode='reduce-overhead',
|
||||
as that mode already uses CUDA graphs internally and will conflict.
|
||||
Use this only with mode='default' or without torch.compile.
|
||||
|
||||
Args:
|
||||
window_size: Fixed number of codec frames (must match decode_window_frames)
|
||||
warmup_runs: Number of warmup iterations before capture
|
||||
"""
|
||||
if not torch.cuda.is_available():
|
||||
print("[Decoder] CUDA not available, skipping graph capture")
|
||||
return self
|
||||
|
||||
# Check for conflict with torch.compile reduce-overhead mode
|
||||
if self._compiled_forward is not None and getattr(self, '_compile_mode', None) == 'reduce-overhead':
|
||||
print("[Decoder] WARNING: torch.compile with mode='reduce-overhead' already uses CUDA graphs internally.")
|
||||
print("[Decoder] Skipping manual CUDA graph capture to avoid conflicts.")
|
||||
print("[Decoder] The compiled forward will be used instead (already optimized).")
|
||||
return self
|
||||
|
||||
device = next(self.parameters()).device
|
||||
num_quantizers = self.config.num_quantizers
|
||||
|
||||
# Create static input buffer
|
||||
self._static_input = torch.zeros(
|
||||
1, num_quantizers, window_size,
|
||||
dtype=torch.long,
|
||||
device=device
|
||||
)
|
||||
self._graph_window_size = window_size
|
||||
|
||||
# Use non-compiled forward for manual CUDA graph capture
|
||||
forward_fn = self._forward_impl
|
||||
|
||||
# Warmup
|
||||
print(f"[Decoder] Warming up CUDA graph (window_size={window_size})...")
|
||||
s = torch.cuda.Stream()
|
||||
s.wait_stream(torch.cuda.current_stream())
|
||||
|
||||
with torch.cuda.stream(s):
|
||||
for _ in range(warmup_runs):
|
||||
_ = forward_fn(self._static_input)
|
||||
torch.cuda.current_stream().wait_stream(s)
|
||||
|
||||
# Capture
|
||||
print("[Decoder] Capturing CUDA graph...")
|
||||
self._cuda_graph = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(self._cuda_graph):
|
||||
self._static_output = forward_fn(self._static_input)
|
||||
|
||||
print("[Decoder] CUDA graph captured successfully")
|
||||
return self
|
||||
|
||||
def forward_optimized(self, codes):
|
||||
"""
|
||||
Forward pass with optimizations if available.
|
||||
|
||||
Priority:
|
||||
1. Manual CUDA graph (if captured and input matches size)
|
||||
2. Compiled forward (torch.compile, may include internal CUDA graphs)
|
||||
3. Regular forward
|
||||
|
||||
Note: With compile_mode="reduce-overhead", the compiled forward
|
||||
already includes CUDA graph optimizations internally.
|
||||
"""
|
||||
B, Q, T = codes.shape
|
||||
|
||||
# Try manual CUDA graph path (only if we captured one)
|
||||
if (self._cuda_graph is not None
|
||||
and B == 1
|
||||
and T == self._graph_window_size):
|
||||
self._static_input.copy_(codes)
|
||||
self._cuda_graph.replay()
|
||||
return self._static_output.clone()
|
||||
|
||||
# Use compiled forward if available (includes CUDA graphs with reduce-overhead)
|
||||
if self._compiled_forward is not None:
|
||||
# Mark step begin for CUDA graphs to avoid tensor overwrite errors
|
||||
torch.compiler.cudagraph_mark_step_begin()
|
||||
return self._compiled_forward(codes)
|
||||
|
||||
# Fallback to regular forward
|
||||
return self._forward_impl(codes)
|
||||
|
||||
def decode_padded(self, codes: torch.Tensor, target_length: int) -> torch.Tensor:
|
||||
"""
|
||||
Decode with left-padding to fixed size for torch.compile optimization.
|
||||
|
||||
When using torch.compile with dynamic=False, the model recompiles for each
|
||||
new input size. By padding all inputs to target_length, we ensure a single
|
||||
compilation that can be reused for all streaming decode calls.
|
||||
|
||||
Args:
|
||||
codes: Input tensor of shape [B, Q, T] where T <= target_length
|
||||
target_length: Fixed size to pad to (should be decode_window_frames)
|
||||
|
||||
Returns:
|
||||
Waveform tensor with padding samples trimmed from the left
|
||||
"""
|
||||
B, Q, T = codes.shape
|
||||
|
||||
if T < target_length:
|
||||
# Pad with zeros on the left
|
||||
pad = torch.zeros(B, Q, target_length - T, dtype=codes.dtype, device=codes.device)
|
||||
codes_padded = torch.cat([pad, codes], dim=-1)
|
||||
else:
|
||||
codes_padded = codes.contiguous() # Ensure uniform tensor format to avoid recompilation
|
||||
|
||||
# Run forward (uses compiled path if available)
|
||||
wav = self.forward_optimized(codes_padded)
|
||||
|
||||
# Trim padding from output
|
||||
if T < target_length:
|
||||
# Calculate how many samples correspond to the padding frames
|
||||
total_samples = wav.shape[-1]
|
||||
samples_per_frame = total_samples / target_length
|
||||
trim_samples = int((target_length - T) * samples_per_frame)
|
||||
wav = wav[..., trim_samples:]
|
||||
|
||||
return wav
|
||||
|
||||
|
||||
class Qwen3TTSTokenizerV2Encoder(MimiModel):
|
||||
def __init__(self, config: MimiConfig):
|
||||
@@ -1021,5 +1201,95 @@ class Qwen3TTSTokenizerV2Model(Qwen3TTSTokenizerV2PreTrainedModel):
|
||||
|
||||
return Qwen3TTSTokenizerV2DecoderOutput(audio_values)
|
||||
|
||||
def enable_streaming_optimizations(
|
||||
self,
|
||||
decode_window_frames: int = 80,
|
||||
use_compile: bool = True,
|
||||
use_cuda_graphs: bool = False, # Changed default: not needed with reduce-overhead
|
||||
compile_mode: str = "reduce-overhead",
|
||||
):
|
||||
"""
|
||||
Enable optimizations for streaming decode.
|
||||
|
||||
This method applies torch.compile to the decoder for faster streaming generation.
|
||||
|
||||
IMPORTANT: compile_mode="reduce-overhead" (default) already includes CUDA graph
|
||||
optimizations internally. You do NOT need to set use_cuda_graphs=True with this mode.
|
||||
Manual CUDA graphs are only useful with compile_mode="default".
|
||||
|
||||
Args:
|
||||
decode_window_frames: Window size for streaming decode (used for manual CUDA graphs)
|
||||
use_compile: Apply torch.compile to the decoder (recommended)
|
||||
use_cuda_graphs: Capture manual CUDA graphs (only useful with compile_mode="default")
|
||||
compile_mode: Mode for torch.compile:
|
||||
- "reduce-overhead" (recommended): Includes CUDA graphs automatically
|
||||
- "max-autotune": Maximum optimization
|
||||
- "default": Basic compilation, can combine with manual CUDA graphs
|
||||
|
||||
Returns:
|
||||
self for method chaining
|
||||
|
||||
Example:
|
||||
# Recommended: just use torch.compile with reduce-overhead
|
||||
model.speech_tokenizer.model.enable_streaming_optimizations(
|
||||
use_compile=True,
|
||||
compile_mode="reduce-overhead",
|
||||
)
|
||||
"""
|
||||
print(f"[Tokenizer] Enabling streaming optimizations...")
|
||||
print(f" use_compile={use_compile}, compile_mode={compile_mode}")
|
||||
print(f" use_cuda_graphs={use_cuda_graphs} (manual)")
|
||||
|
||||
if use_compile:
|
||||
self.decoder.compile_for_streaming(mode=compile_mode)
|
||||
|
||||
# Only capture manual CUDA graphs if explicitly requested AND not using reduce-overhead
|
||||
if use_cuda_graphs:
|
||||
if compile_mode == "reduce-overhead":
|
||||
print(f"[Tokenizer] Note: compile_mode='reduce-overhead' already includes CUDA graphs")
|
||||
print(f"[Tokenizer] Manual CUDA graph capture skipped (not needed)")
|
||||
else:
|
||||
self.decoder.capture_cuda_graph(window_size=decode_window_frames)
|
||||
|
||||
return self
|
||||
|
||||
def decode_streaming(
|
||||
self,
|
||||
audio_codes: torch.Tensor,
|
||||
use_optimized: bool = True,
|
||||
pad_to_size: Optional[int] = None,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Decode audio codes optimized for streaming (single window).
|
||||
|
||||
Unlike the regular decode(), this method:
|
||||
- Does not use chunked decoding (assumes small window)
|
||||
- Uses CUDA graphs if available
|
||||
- Returns raw tensor (not list)
|
||||
- Optionally pads to fixed size for torch.compile optimization
|
||||
|
||||
Args:
|
||||
audio_codes: [B, T, num_quantizers] tensor of codec indices
|
||||
use_optimized: If True, use CUDA graph path when available
|
||||
pad_to_size: If specified, pad input to this size (in frames) for
|
||||
consistent torch.compile behavior. Should match
|
||||
decode_window_frames for streaming.
|
||||
|
||||
Returns:
|
||||
Waveform tensor [B, samples]
|
||||
"""
|
||||
# Transpose to [B, num_quantizers, T] for decoder
|
||||
codes = audio_codes.transpose(1, 2)
|
||||
|
||||
if use_optimized:
|
||||
if pad_to_size is not None:
|
||||
wav = self.decoder.decode_padded(codes, pad_to_size)
|
||||
else:
|
||||
wav = self.decoder.forward_optimized(codes)
|
||||
else:
|
||||
wav = self.decoder(codes)
|
||||
|
||||
return wav.squeeze(1)
|
||||
|
||||
|
||||
__all__ = ["Qwen3TTSTokenizerV2Model", "Qwen3TTSTokenizerV2PreTrainedModel"]
|
||||
|
||||
@@ -0,0 +1,279 @@
|
||||
# coding=utf-8
|
||||
# Optimized decoder with torch.compile and CUDA graphs support
|
||||
"""
|
||||
Provides optimized decoding for streaming TTS with:
|
||||
1. torch.compile for the decoder forward pass
|
||||
2. CUDA graphs for static-shape decode operations
|
||||
"""
|
||||
|
||||
from typing import Optional, Tuple
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class CUDAGraphDecoder:
|
||||
"""
|
||||
Wraps a decoder with CUDA graph capture for static-shape inference.
|
||||
|
||||
CUDA graphs capture a sequence of GPU operations and replay them
|
||||
without CPU overhead, giving significant speedup for repeated
|
||||
fixed-size operations like streaming decode windows.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
decoder: nn.Module,
|
||||
static_window_size: int = 80,
|
||||
num_quantizers: int = 8,
|
||||
device: torch.device = None,
|
||||
dtype: torch.dtype = torch.bfloat16,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
decoder: The Qwen3TTSTokenizerV2Decoder module
|
||||
static_window_size: Fixed window size for CUDA graph capture
|
||||
num_quantizers: Number of codec quantizers (usually 8)
|
||||
device: CUDA device
|
||||
dtype: Data type for the decoder
|
||||
"""
|
||||
self.decoder = decoder
|
||||
self.static_window_size = static_window_size
|
||||
self.num_quantizers = num_quantizers
|
||||
self.device = device or next(decoder.parameters()).device
|
||||
self.dtype = dtype
|
||||
|
||||
self._graph: Optional[torch.cuda.CUDAGraph] = None
|
||||
self._static_input: Optional[torch.Tensor] = None
|
||||
self._static_output: Optional[torch.Tensor] = None
|
||||
self._is_captured = False
|
||||
|
||||
def warmup_and_capture(self, warmup_runs: int = 3):
|
||||
"""
|
||||
Warm up the decoder and capture CUDA graph.
|
||||
|
||||
Call this once before streaming starts for optimal performance.
|
||||
"""
|
||||
if not torch.cuda.is_available():
|
||||
print("[CUDAGraphDecoder] CUDA not available, skipping graph capture")
|
||||
return
|
||||
|
||||
# Create static input tensor (will be reused for all captures)
|
||||
self._static_input = torch.zeros(
|
||||
1, self.num_quantizers, self.static_window_size,
|
||||
dtype=torch.long,
|
||||
device=self.device
|
||||
)
|
||||
|
||||
# Warmup runs to stabilize CUDA state
|
||||
print(f"[CUDAGraphDecoder] Warming up with {warmup_runs} runs...")
|
||||
s = torch.cuda.Stream()
|
||||
s.wait_stream(torch.cuda.current_stream())
|
||||
|
||||
with torch.cuda.stream(s):
|
||||
for i in range(warmup_runs):
|
||||
_ = self.decoder(self._static_input)
|
||||
|
||||
torch.cuda.current_stream().wait_stream(s)
|
||||
|
||||
# Capture the graph
|
||||
print(f"[CUDAGraphDecoder] Capturing CUDA graph for window_size={self.static_window_size}...")
|
||||
self._graph = torch.cuda.CUDAGraph()
|
||||
|
||||
with torch.cuda.graph(self._graph):
|
||||
self._static_output = self.decoder(self._static_input)
|
||||
|
||||
self._is_captured = True
|
||||
print(f"[CUDAGraphDecoder] Graph captured successfully")
|
||||
|
||||
def decode(self, codes: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Decode codes to audio waveform.
|
||||
|
||||
If codes match the static window size and graph is captured,
|
||||
uses CUDA graph replay for faster execution.
|
||||
Otherwise falls back to regular forward pass.
|
||||
|
||||
Args:
|
||||
codes: [B, num_quantizers, T] tensor of codec indices
|
||||
|
||||
Returns:
|
||||
Waveform tensor [B, 1, samples]
|
||||
"""
|
||||
B, Q, T = codes.shape
|
||||
|
||||
# Check if we can use the captured graph
|
||||
if (self._is_captured
|
||||
and B == 1
|
||||
and Q == self.num_quantizers
|
||||
and T == self.static_window_size):
|
||||
# Copy input to static buffer and replay graph
|
||||
self._static_input.copy_(codes)
|
||||
self._graph.replay()
|
||||
return self._static_output.clone()
|
||||
else:
|
||||
# Fallback to regular forward
|
||||
return self.decoder(codes)
|
||||
|
||||
|
||||
def compile_decoder(
|
||||
decoder: nn.Module,
|
||||
mode: str = "reduce-overhead",
|
||||
fullgraph: bool = False,
|
||||
dynamic: bool = True,
|
||||
) -> nn.Module:
|
||||
"""
|
||||
Apply torch.compile to the decoder for optimized execution.
|
||||
|
||||
Args:
|
||||
decoder: The Qwen3TTSTokenizerV2Decoder module
|
||||
mode: Compilation mode:
|
||||
- "default": Good balance of compile time and speedup
|
||||
- "reduce-overhead": Minimize framework overhead (good for streaming)
|
||||
- "max-autotune": Maximum optimization (longer compile time)
|
||||
fullgraph: If True, requires the entire model to be traceable as single graph.
|
||||
Set False if there are dynamic control flows.
|
||||
dynamic: If True, enables dynamic shape support (slower but more flexible)
|
||||
|
||||
Returns:
|
||||
Compiled decoder module
|
||||
"""
|
||||
if not hasattr(torch, 'compile'):
|
||||
print("[compile_decoder] torch.compile not available (requires PyTorch 2.0+)")
|
||||
return decoder
|
||||
|
||||
print(f"[compile_decoder] Compiling decoder with mode={mode}, fullgraph={fullgraph}, dynamic={dynamic}")
|
||||
|
||||
compiled = torch.compile(
|
||||
decoder,
|
||||
mode=mode,
|
||||
fullgraph=fullgraph,
|
||||
dynamic=dynamic,
|
||||
)
|
||||
|
||||
return compiled
|
||||
|
||||
|
||||
class OptimizedStreamingDecoder:
|
||||
"""
|
||||
Combines torch.compile and CUDA graphs for optimal streaming decode performance.
|
||||
|
||||
Usage:
|
||||
# During model initialization:
|
||||
opt_decoder = OptimizedStreamingDecoder(
|
||||
decoder=model.speech_tokenizer.model.decoder,
|
||||
static_window_size=80, # matches decode_window_frames
|
||||
)
|
||||
opt_decoder.warmup()
|
||||
|
||||
# During streaming:
|
||||
wav = opt_decoder.decode(codes)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
decoder: nn.Module,
|
||||
static_window_size: int = 80,
|
||||
num_quantizers: int = 8,
|
||||
use_compile: bool = True,
|
||||
use_cuda_graphs: bool = True,
|
||||
compile_mode: str = "reduce-overhead",
|
||||
):
|
||||
self.original_decoder = decoder
|
||||
self.static_window_size = static_window_size
|
||||
self.num_quantizers = num_quantizers
|
||||
self.use_compile = use_compile
|
||||
self.use_cuda_graphs = use_cuda_graphs and torch.cuda.is_available()
|
||||
self.compile_mode = compile_mode
|
||||
|
||||
self.device = next(decoder.parameters()).device
|
||||
self.dtype = next(decoder.parameters()).dtype
|
||||
|
||||
self._compiled_decoder: Optional[nn.Module] = None
|
||||
self._cuda_graph_decoder: Optional[CUDAGraphDecoder] = None
|
||||
self._is_warmed_up = False
|
||||
|
||||
def warmup(self, warmup_runs: int = 3):
|
||||
"""
|
||||
Initialize optimizations. Call once before streaming starts.
|
||||
"""
|
||||
if self._is_warmed_up:
|
||||
return
|
||||
|
||||
print(f"[OptimizedStreamingDecoder] Starting warmup...")
|
||||
|
||||
# Step 1: Apply torch.compile
|
||||
if self.use_compile:
|
||||
self._compiled_decoder = compile_decoder(
|
||||
self.original_decoder,
|
||||
mode=self.compile_mode,
|
||||
fullgraph=False, # Decoder has some dynamic ops
|
||||
dynamic=False, # We use static shapes for streaming
|
||||
)
|
||||
else:
|
||||
self._compiled_decoder = self.original_decoder
|
||||
|
||||
# Step 2: Setup CUDA graphs
|
||||
if self.use_cuda_graphs:
|
||||
self._cuda_graph_decoder = CUDAGraphDecoder(
|
||||
decoder=self._compiled_decoder,
|
||||
static_window_size=self.static_window_size,
|
||||
num_quantizers=self.num_quantizers,
|
||||
device=self.device,
|
||||
dtype=self.dtype,
|
||||
)
|
||||
self._cuda_graph_decoder.warmup_and_capture(warmup_runs=warmup_runs)
|
||||
|
||||
self._is_warmed_up = True
|
||||
print(f"[OptimizedStreamingDecoder] Warmup complete")
|
||||
|
||||
def decode(self, codes: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Decode codes to waveform with optimizations.
|
||||
|
||||
Args:
|
||||
codes: [B, num_quantizers, T] tensor
|
||||
|
||||
Returns:
|
||||
Waveform tensor
|
||||
"""
|
||||
if not self._is_warmed_up:
|
||||
# Auto-warmup on first call
|
||||
self.warmup()
|
||||
|
||||
if self._cuda_graph_decoder is not None:
|
||||
return self._cuda_graph_decoder.decode(codes)
|
||||
elif self._compiled_decoder is not None:
|
||||
return self._compiled_decoder(codes)
|
||||
else:
|
||||
return self.original_decoder(codes)
|
||||
|
||||
|
||||
def create_optimized_tokenizer_decode(tokenizer, static_window_size: int = 80):
|
||||
"""
|
||||
Patch a Qwen3TTSTokenizer instance to use optimized decoding.
|
||||
|
||||
Args:
|
||||
tokenizer: Qwen3TTSTokenizer instance
|
||||
static_window_size: Window size for CUDA graph optimization
|
||||
|
||||
Returns:
|
||||
The patched tokenizer (same instance, modified in place)
|
||||
"""
|
||||
decoder = tokenizer.model.decoder
|
||||
num_quantizers = tokenizer.config.decoder_config.num_quantizers
|
||||
|
||||
opt = OptimizedStreamingDecoder(
|
||||
decoder=decoder,
|
||||
static_window_size=static_window_size,
|
||||
num_quantizers=num_quantizers,
|
||||
use_compile=True,
|
||||
use_cuda_graphs=True,
|
||||
)
|
||||
|
||||
# Store reference to prevent garbage collection
|
||||
tokenizer._optimized_decoder = opt
|
||||
|
||||
# Warmup
|
||||
opt.warmup()
|
||||
|
||||
return tokenizer
|
||||
@@ -120,6 +120,57 @@ class Qwen3TTSModel:
|
||||
generate_defaults = model.generate_config
|
||||
return cls(model=model, processor=processor, generate_defaults=generate_defaults)
|
||||
|
||||
def enable_streaming_optimizations(
|
||||
self,
|
||||
decode_window_frames: int = 80,
|
||||
use_compile: bool = True,
|
||||
use_cuda_graphs: bool = True,
|
||||
compile_mode: str = "reduce-overhead",
|
||||
use_fast_codebook: bool = False, # Disabled: needs debugging, currently slower
|
||||
compile_codebook_predictor: bool = True,
|
||||
):
|
||||
"""
|
||||
Enable torch.compile and CUDA graphs optimizations for streaming decode.
|
||||
|
||||
Significantly speeds up streaming generation by:
|
||||
1. Compiling the decoder with torch.compile (reduces Python overhead)
|
||||
2. Capturing CUDA graphs for fixed-size decode windows (eliminates GPU launch overhead)
|
||||
3. Fast codebook generation (bypasses HuggingFace generate() overhead)
|
||||
|
||||
Call this method after loading the model, before starting streaming generation.
|
||||
|
||||
Args:
|
||||
decode_window_frames: Fixed window size for streaming (must match the
|
||||
decode_window_frames parameter in stream_generate_voice_clone)
|
||||
use_compile: Apply torch.compile to the decoder (default True)
|
||||
use_cuda_graphs: Capture CUDA graphs for the fixed window size (default True)
|
||||
compile_mode: torch.compile mode - "reduce-overhead" (recommended for streaming),
|
||||
"max-autotune", or "default"
|
||||
use_fast_codebook: Use fast codebook generation that bypasses HuggingFace's
|
||||
generate() overhead (default True, ~2x faster per step)
|
||||
compile_codebook_predictor: Apply torch.compile to codebook predictor (experimental)
|
||||
|
||||
Returns:
|
||||
self for method chaining
|
||||
|
||||
Example:
|
||||
model = Qwen3TTSModel.from_pretrained("Qwen/Qwen3-TTS-12Hz-1.7B-Base", ...)
|
||||
model.enable_streaming_optimizations(decode_window_frames=80)
|
||||
|
||||
# Now streaming will be faster
|
||||
for chunk, sr in model.stream_generate_voice_clone(..., decode_window_frames=80):
|
||||
...
|
||||
"""
|
||||
self.model.enable_streaming_optimizations(
|
||||
decode_window_frames=decode_window_frames,
|
||||
use_compile=use_compile,
|
||||
use_cuda_graphs=use_cuda_graphs,
|
||||
compile_mode=compile_mode,
|
||||
use_fast_codebook=use_fast_codebook,
|
||||
compile_codebook_predictor=compile_codebook_predictor,
|
||||
)
|
||||
return self
|
||||
|
||||
def _supported_languages_set(self) -> Optional[set]:
|
||||
langs = getattr(self.model, "get_supported_languages", None)
|
||||
if callable(langs):
|
||||
@@ -647,6 +698,8 @@ class Qwen3TTSModel:
|
||||
decode_window_frames: int = 80,
|
||||
overlap_samples: int = 512,
|
||||
max_frames: int = 10000,
|
||||
# Optimization
|
||||
use_optimized_decode: bool = True,
|
||||
**kwargs,
|
||||
) -> Generator[Tuple[np.ndarray, int], None, None]:
|
||||
"""
|
||||
@@ -666,6 +719,8 @@ class Qwen3TTSModel:
|
||||
decode_window_frames: Window size for decoding (longer = better quality, more latency).
|
||||
overlap_samples: Overlap samples for crossfade between chunks.
|
||||
max_frames: Maximum codec frames to generate.
|
||||
use_optimized_decode: Use CUDA graph optimized decode when available (default True).
|
||||
Call enable_streaming_optimizations() first for best performance.
|
||||
**kwargs: Generation parameters (do_sample, top_k, top_p, temperature, etc.)
|
||||
|
||||
Yields:
|
||||
@@ -742,6 +797,7 @@ class Qwen3TTSModel:
|
||||
decode_window_frames=decode_window_frames,
|
||||
overlap_samples=overlap_samples,
|
||||
max_frames=max_frames,
|
||||
use_optimized_decode=use_optimized_decode,
|
||||
**gen_kwargs,
|
||||
):
|
||||
yield chunk, sr
|
||||
|
||||
@@ -408,4 +408,79 @@ class Qwen3TTSTokenizer:
|
||||
Returns:
|
||||
int: Decode upsample rate.
|
||||
"""
|
||||
return int(self.model.get_decode_upsample_rate())
|
||||
return int(self.model.get_decode_upsample_rate())
|
||||
|
||||
def enable_streaming_optimizations(
|
||||
self,
|
||||
decode_window_frames: int = 80,
|
||||
use_compile: bool = True,
|
||||
use_cuda_graphs: bool = True,
|
||||
compile_mode: str = "reduce-overhead",
|
||||
):
|
||||
"""
|
||||
Enable torch.compile and CUDA graphs optimizations for streaming decode.
|
||||
|
||||
Args:
|
||||
decode_window_frames: Fixed window size for streaming (must match streaming params)
|
||||
use_compile: Apply torch.compile to decoder
|
||||
use_cuda_graphs: Capture CUDA graphs for fixed-size operations
|
||||
compile_mode: torch.compile mode ("reduce-overhead" recommended for streaming)
|
||||
|
||||
Returns:
|
||||
self for method chaining
|
||||
"""
|
||||
model_type = self.model.get_model_type()
|
||||
if model_type != "qwen3_tts_tokenizer_12hz":
|
||||
print(f"[Tokenizer] Optimizations only supported for 12Hz tokenizer, got {model_type}")
|
||||
return self
|
||||
|
||||
return self.model.enable_streaming_optimizations(
|
||||
decode_window_frames=decode_window_frames,
|
||||
use_compile=use_compile,
|
||||
use_cuda_graphs=use_cuda_graphs,
|
||||
compile_mode=compile_mode,
|
||||
)
|
||||
|
||||
def decode_streaming(
|
||||
self,
|
||||
audio_codes: torch.Tensor,
|
||||
use_optimized: bool = True,
|
||||
pad_to_size: Optional[int] = None,
|
||||
) -> Tuple[List[np.ndarray], int]:
|
||||
"""
|
||||
Decode audio codes optimized for streaming windows.
|
||||
|
||||
Unlike regular decode(), this method:
|
||||
- Uses CUDA graphs when available and window matches captured size
|
||||
- Does not use chunked decode (assumes single streaming window)
|
||||
- Is optimized for repeated calls with same window size
|
||||
- Optionally pads to fixed size for torch.compile optimization
|
||||
|
||||
Args:
|
||||
audio_codes: [T, num_quantizers] tensor (single window, no batch dim)
|
||||
use_optimized: Whether to use CUDA graph path when available
|
||||
pad_to_size: If specified, pad input to this size (in frames) for
|
||||
consistent torch.compile behavior. Should match
|
||||
decode_window_frames for streaming.
|
||||
|
||||
Returns:
|
||||
Tuple[List[np.ndarray], int]: (list with single waveform, sample_rate)
|
||||
"""
|
||||
model_type = self.model.get_model_type()
|
||||
if model_type != "qwen3_tts_tokenizer_12hz":
|
||||
# Fallback to regular decode for other tokenizers
|
||||
return self.decode({"audio_codes": audio_codes})
|
||||
|
||||
# Add batch dimension if needed: [T, Q] -> [1, T, Q]
|
||||
if audio_codes.dim() == 2:
|
||||
audio_codes = audio_codes.unsqueeze(0)
|
||||
|
||||
wav_tensor = self.model.decode_streaming(
|
||||
audio_codes,
|
||||
use_optimized=use_optimized,
|
||||
pad_to_size=pad_to_size,
|
||||
)
|
||||
|
||||
# Convert to numpy and return
|
||||
wav = wav_tensor[0].to(torch.float32).detach().cpu().numpy()
|
||||
return [wav], int(self.model.get_output_sample_rate())
|
||||
Reference in New Issue
Block a user