Benchmark ParlerTTS + streaming time to first audio.

benchmark_parler_streaming.py

import os
import torch
import time
from parler_tts import ParlerTTSForConditionalGeneration, ParlerTTSStreamer
from transformers import AutoTokenizer
from threading import Thread


# caching allows ~50% compilation time reduction
# see https://docs.google.com/document/d/1y5CRfMLdwEoF1nTk9q8qEu1mgMUuUtvhklPKJ2emLU8/edit#heading=h.o2asbxsrp1ma
CURRENT_DIR = os.path.dirname(os.path.realpath(__file__))
os.environ["TORCHINDUCTOR_CACHE_DIR"] = os.path.join(CURRENT_DIR, "tmp") 

torch._inductor.config.fx_graph_cache = True
torch._logging.set_logs(graph_breaks=True, recompiles=True, cudagraphs=True)

# mind about this parameter ! should be >= 2 * number of compiled models
torch._dynamo.config.cache_size_limit = 15

# reproducibility 
torch.manual_seed(42)


# =============================
# model args
model_name = "parler-tts/parler-tts-mini-v1" 
torch_device = "cuda:0"
torch_dtype = torch.float16
attn_implementation = "sdpa"
compile_mode = "default"
# =============================


# =============================
# load model
model = ParlerTTSForConditionalGeneration.from_pretrained(
    model_name,
    attn_implementation=attn_implementation
).to(torch_device, dtype=torch_dtype)
model.generation_config.cache_implementation = "static"
model.forward = torch.compile(model.forward, mode=compile_mode, fullgraph=True)
# =============================


# =============================
# tokenizer prompt & description
tokenizer = AutoTokenizer.from_pretrained(model_name)

prompt = "This is the second iteration of the model."
description = "A male speaker with a slightly low-pitched voice." 

tokenized_description = tokenizer(description, return_tensors="pt")
input_ids = tokenized_description.input_ids.to(torch_device)

tokenized_prompt = tokenizer(prompt, return_tensors="pt")
prompt_input_ids = tokenized_prompt.input_ids.to(torch_device)
# =============================


# =============================
# warmup 
n_steps = 1 if compile_mode == "default" else 2

print("Warming up...")

start_event = torch.cuda.Event(enable_timing=True)
end_event = torch.cuda.Event(enable_timing=True)

torch.cuda.synchronize()
start_event.record()
for _ in range(n_steps):
    _ = model.generate(
        input_ids=input_ids, 
        prompt_input_ids=prompt_input_ids
    )
end_event.record()
torch.cuda.synchronize()

print(f"Warmed up! Compilation time: {start_event.elapsed_time(end_event) * 1e-3:.3f} s")
# =============================


# =============================
# benchmark time to first audio

sampling_rate = model.audio_encoder.config.sampling_rate
frame_rate = model.audio_encoder.config.frame_rate

play_steps = 10
streamer = ParlerTTSStreamer(model, device=torch_device, play_steps=play_steps)

generation_kwargs = {
    "input_ids": input_ids, 
    "prompt_input_ids": prompt_input_ids, 
    "streamer": streamer, 
    "min_new_tokens": 10
}

thread = Thread(target=model.generate, kwargs=generation_kwargs)
thread.start()

# iterate over chunks of audio
start = time.perf_counter()
for new_audio in streamer:
    if new_audio.shape[0] == 0:
        break
    print(f"time to fist audio: {time.perf_counter() - start:.3f} s, shape {new_audio.shape}")
    break