Write

"""
Benchmark LuxTTS with optimizations: fp16, torch.compile, k2 swoosh
"""
import torch
import time
import numpy as np
from zipvoice.luxvoice import LuxTTS
import soundfile as sf

REF_AUDIO = "/home/ubuntu/LuxTTS/ref_audio.wav"
SHORT_TEXT = "Hello, this is a concurrency test."
LONG_TEXT = "The quick brown fox jumps over the lazy dog. This is a test of the text to speech system for benchmarking purposes on the A100 GPU."

def profile_generate(tts, text, encoded, label="", n=5):
    """Run n generations and report stats."""
    times = []
    audio_durs = []
    for i in range(n):
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        wav = tts.generate_speech(text, encoded, num_steps=4)
        torch.cuda.synchronize()
        elapsed = time.perf_counter() - t0
        audio_dur = len(wav[0]) / 48000
        times.append(elapsed)
        audio_durs.append(audio_dur)

    avg_time = np.mean(times[1:])  # skip first (warmup)
    avg_audio = np.mean(audio_durs[1:])
    rtf = avg_time / avg_audio
    print(f"  {label}: {avg_time:.4f}s wall, {avg_audio:.2f}s audio, RTF={rtf:.5f}, {1/rtf:.0f}x realtime")
    return avg_time, avg_audio

def main():
    print("=" * 60)
    print("LuxTTS Optimized Benchmark — A100 80GB")
    print("=" * 60)

    # Check k2
    try:
        import k2
        print(f"k2 version: {k2.__version__} (swoosh ENABLED)")
    except:
        print("k2 NOT available (swoosh DISABLED — slower)")

    # === BASELINE: fp32, no compile ===
    print("\n--- [A] BASELINE: fp32, no compile ---")
    tts = LuxTTS(device='cuda')
    encoded = tts.encode_prompt(REF_AUDIO, duration=5)

    mem = torch.cuda.memory_allocated() / 1024**3
    print(f"  Model VRAM: {mem:.2f} GB")

    profile_generate(tts, SHORT_TEXT, encoded, "short", n=5)
    profile_generate(tts, LONG_TEXT, encoded, "long", n=5)

    # === FP16 AUTOCAST ===
    print("\n--- [B] FP16 autocast ---")

    # Monkey-patch generate to use autocast
    from zipvoice.modeling_utils import generate as orig_generate
    def generate_fp16(*args, **kwargs):
        with torch.amp.autocast('cuda', dtype=torch.float16):
            return orig_generate(*args, **kwargs)

    import zipvoice.luxvoice as lux_module
    import zipvoice.modeling_utils as mod_utils
    # Patch at module level
    orig_gen = mod_utils.generate
    mod_utils.generate = generate_fp16

    # Need to re-import or patch the reference in luxvoice
    tts2 = LuxTTS(device='cuda')
    encoded2 = tts2.encode_prompt(REF_AUDIO, duration=5)

    # Direct call with autocast
    times_short = []
    times_long = []
    for i in range(5):
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        with torch.amp.autocast('cuda', dtype=torch.float16):
            wav = orig_gen(
                encoded2['prompt_tokens'], encoded2['prompt_features_lens'],
                encoded2['prompt_features'], encoded2['prompt_rms'],
                SHORT_TEXT, tts2.model, tts2.vocos, tts2.tokenizer,
                num_step=4, guidance_scale=3.0, t_shift=0.5, speed=1.0
            )
        torch.cuda.synchronize()
        elapsed = time.perf_counter() - t0
        audio_dur = len(wav[0]) / 48000
        times_short.append((elapsed, audio_dur))

    avg = np.mean([t[0] for t in times_short[1:]])
    avg_a = np.mean([t[1] for t in times_short[1:]])
    print(f"  short fp16: {avg:.4f}s wall, {avg_a:.2f}s audio, {avg_a/avg:.0f}x realtime")

    for i in range(5):
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        with torch.amp.autocast('cuda', dtype=torch.float16):
            wav = orig_gen(
                encoded2['prompt_tokens'], encoded2['prompt_features_lens'],
                encoded2['prompt_features'], encoded2['prompt_rms'],
                LONG_TEXT, tts2.model, tts2.vocos, tts2.tokenizer,
                num_step=4, guidance_scale=3.0, t_shift=0.5, speed=1.0
            )
        torch.cuda.synchronize()
        elapsed = time.perf_counter() - t0
        audio_dur = len(wav[0]) / 48000
        times_long.append((elapsed, audio_dur))

    avg = np.mean([t[0] for t in times_long[1:]])
    avg_a = np.mean([t[1] for t in times_long[1:]])
    print(f"  long fp16:  {avg:.4f}s wall, {avg_a:.2f}s audio, {avg_a/avg:.0f}x realtime")

    peak = torch.cuda.max_memory_allocated() / 1024**3
    print(f"  Peak VRAM: {peak:.2f} GB")

    # === TORCH.COMPILE on fm_decoder ===
    print("\n--- [C] torch.compile(fm_decoder) + fp16 ---")
    try:
        tts3 = LuxTTS(device='cuda')
        encoded3 = tts3.encode_prompt(REF_AUDIO, duration=5)

        print("  Compiling fm_decoder...")
        tts3.model.fm_decoder = torch.compile(tts3.model.fm_decoder, mode="reduce-overhead")

        # Warmup compile
        print("  Warmup (first call triggers compilation)...")
        for i in range(3):
            with torch.amp.autocast('cuda', dtype=torch.float16):
                wav = orig_gen(
                    encoded3['prompt_tokens'], encoded3['prompt_features_lens'],
                    encoded3['prompt_features'], encoded3['prompt_rms'],
                    SHORT_TEXT, tts3.model, tts3.vocos, tts3.tokenizer,
                    num_step=4, guidance_scale=3.0, t_shift=0.5, speed=1.0
                )

        # Benchmark
        for text, label in [(SHORT_TEXT, "short"), (LONG_TEXT, "long")]:
            times = []
            for i in range(5):
                torch.cuda.synchronize()
                t0 = time.perf_counter()
                with torch.amp.autocast('cuda', dtype=torch.float16):
                    wav = orig_gen(
                        encoded3['prompt_tokens'], encoded3['prompt_features_lens'],
                        encoded3['prompt_features'], encoded3['prompt_rms'],
                        text, tts3.model, tts3.vocos, tts3.tokenizer,
                        num_step=4, guidance_scale=3.0, t_shift=0.5, speed=1.0
                    )
                torch.cuda.synchronize()
                elapsed = time.perf_counter() - t0
                audio_dur = len(wav[0]) / 48000
                times.append((elapsed, audio_dur))
            avg = np.mean([t[0] for t in times])
            avg_a = np.mean([t[1] for t in times])
            print(f"  {label} compiled+fp16: {avg:.4f}s wall, {avg_a:.2f}s audio, {avg_a/avg:.0f}x realtime")
    except Exception as e:
        print(f"  torch.compile failed: {e}")

    # === Model parameter count ===
    print("\n--- Model Info ---")
    total_params = sum(p.numel() for p in tts.model.parameters())
    print(f"  Total params: {total_params/1e6:.1f}M")
    for name, module in tts.model.named_children():
        params = sum(p.numel() for p in module.parameters())
        if params > 0:
            print(f"    {name}: {params/1e6:.1f}M")

if __name__ == "__main__":
    main()


File created successfully at: /home/ubuntu/LuxTTS/bench_optimized.py

"""
Benchmark LuxTTS with optimizations: fp16, torch.compile, k2 swoosh
"""
import torch
import time
import numpy as np
from zipvoice.luxvoice import LuxTTS
import soundfile as sf

REF_AUDIO = "/home/ubuntu/LuxTTS/ref_audio.wav"
SHORT_TEXT = "Hello, this is a concurrency test."
LONG_TEXT = "The quick brown fox jumps over the lazy dog. This is a test of the text to speech system for benchmarking purposes on the A100 GPU."

def profile_generate(tts, text, encoded, label="", n=5):
    """Run n generations and report stats."""
    times = []
    audio_durs = []
    for i in range(n):
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        wav = tts.generate_speech(text, encoded, num_steps=4)
        torch.cuda.synchronize()
        elapsed = time.perf_counter() - t0
        audio_dur = len(wav[0]) / 48000
        times.append(elapsed)
        audio_durs.append(audio_dur)

    avg_time = np.mean(times[1:])  # skip first (warmup)
    avg_audio = np.mean(audio_durs[1:])
    rtf = avg_time / avg_audio
    print(f"  {label}: {avg_time:.4f}s wall, {avg_audio:.2f}s audio, RTF={rtf:.5f}, {1/rtf:.0f}x realtime")
    return avg_time, avg_audio

def main():
    print("=" * 60)
    print("LuxTTS Optimized Benchmark — A100 80GB")
    print("=" * 60)

    # Check k2
    try:
        import k2
        print(f"k2 version: {k2.__version__} (swoosh ENABLED)")
    except:
        print("k2 NOT available (swoosh DISABLED — slower)")

    # === BASELINE: fp32, no compile ===
    print("\n--- [A] BASELINE: fp32, no compile ---")
    tts = LuxTTS(device='cuda')
    encoded = tts.encode_prompt(REF_AUDIO, duration=5)

    mem = torch.cuda.memory_allocated() / 1024**3
    print(f"  Model VRAM: {mem:.2f} GB")

    profile_generate(tts, SHORT_TEXT, encoded, "short", n=5)
    profile_generate(tts, LONG_TEXT, encoded, "long", n=5)

    # === FP16 AUTOCAST ===
    print("\n--- [B] FP16 autocast ---")

    # Monkey-patch generate to use autocast
    from zipvoice.modeling_utils import generate as orig_generate
    def generate_fp16(*args, **kwargs):
        with torch.amp.autocast('cuda', dtype=torch.float16):
            return orig_generate(*args, **kwargs)

    import zipvoice.luxvoice as lux_module
    import zipvoice.modeling_utils as mod_utils
    # Patch at module level
    orig_gen = mod_utils.generate
    mod_utils.generate = generate_fp16

    # Need to re-import or patch the reference in luxvoice
    tts2 = LuxTTS(device='cuda')
    encoded2 = tts2.encode_prompt(REF_AUDIO, duration=5)

    # Direct call with autocast
    times_short = []
    times_long = []
    for i in range(5):
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        with torch.amp.autocast('cuda', dtype=torch.float16):
            wav = orig_gen(
                encoded2['prompt_tokens'], encoded2['prompt_features_lens'],
                encoded2['prompt_features'], encoded2['prompt_rms'],
                SHORT_TEXT, tts2.model, tts2.vocos, tts2.tokenizer,
                num_step=4, guidance_scale=3.0, t_shift=0.5, speed=1.0
            )
        torch.cuda.synchronize()
        elapsed = time.perf_counter() - t0
        audio_dur = len(wav[0]) / 48000
        times_short.append((elapsed, audio_dur))

    avg = np.mean([t[0] for t in times_short[1:]])
    avg_a = np.mean([t[1] for t in times_short[1:]])
    print(f"  short fp16: {avg:.4f}s wall, {avg_a:.2f}s audio, {avg_a/avg:.0f}x realtime")

    for i in range(5):
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        with torch.amp.autocast('cuda', dtype=torch.float16):
            wav = orig_gen(
                encoded2['prompt_tokens'], encoded2['prompt_features_lens'],
                encoded2['prompt_features'], encoded2['prompt_rms'],
                LONG_TEXT, tts2.model, tts2.vocos, tts2.tokenizer,
                num_step=4, guidance_scale=3.0, t_shift=0.5, speed=1.0
            )
        torch.cuda.synchronize()
        elapsed = time.perf_counter() - t0
        audio_dur = len(wav[0]) / 48000
        times_long.append((elapsed, audio_dur))

    avg = np.mean([t[0] for t in times_long[1:]])
    avg_a = np.mean([t[1] for t in times_long[1:]])
    print(f"  long fp16:  {avg:.4f}s wall, {avg_a:.2f}s audio, {avg_a/avg:.0f}x realtime")

    peak = torch.cuda.max_memory_allocated() / 1024**3
    print(f"  Peak VRAM: {peak:.2f} GB")

    # === TORCH.COMPILE on fm_decoder ===
    print("\n--- [C] torch.compile(fm_decoder) + fp16 ---")
    try:
        tts3 = LuxTTS(device='cuda')
        encoded3 = tts3.encode_prompt(REF_AUDIO, duration=5)

        print("  Compiling fm_decoder...")
        tts3.model.fm_decoder = torch.compile(tts3.model.fm_decoder, mode="reduce-overhead")

        # Warmup compile
        print("  Warmup (first call triggers compilation)...")
        for i in range(3):
            with torch.amp.autocast('cuda', dtype=torch.float16):
                wav = orig_gen(
                    encoded3['prompt_tokens'], encoded3['prompt_features_lens'],
                    encoded3['prompt_features'], encoded3['prompt_rms'],
                    SHORT_TEXT, tts3.model, tts3.vocos, tts3.tokenizer,
                    num_step=4, guidance_scale=3.0, t_shift=0.5, speed=1.0
                )

        # Benchmark
        for text, label in [(SHORT_TEXT, "short"), (LONG_TEXT, "long")]:
            times = []
            for i in range(5):
                torch.cuda.synchronize()
                t0 = time.perf_counter()
                with torch.amp.autocast('cuda', dtype=torch.float16):
                    wav = orig_gen(
                        encoded3['prompt_tokens'], encoded3['prompt_features_lens'],
                        encoded3['prompt_features'], encoded3['prompt_rms'],
                        text, tts3.model, tts3.vocos, tts3.tokenizer,
                        num_step=4, guidance_scale=3.0, t_shift=0.5, speed=1.0
                    )
                torch.cuda.synchronize()
                elapsed = time.perf_counter() - t0
                audio_dur = len(wav[0]) / 48000
                times.append((elapsed, audio_dur))
            avg = np.mean([t[0] for t in times])
            avg_a = np.mean([t[1] for t in times])
            print(f"  {label} compiled+fp16: {avg:.4f}s wall, {avg_a:.2f}s audio, {avg_a/avg:.0f}x realtime")
    except Exception as e:
        print(f"  torch.compile failed: {e}")

    # === Model parameter count ===
    print("\n--- Model Info ---")
    total_params = sum(p.numel() for p in tts.model.parameters())
    print(f"  Total params: {total_params/1e6:.1f}M")
    for name, module in tts.model.named_children():
        params = sum(p.numel() for p in module.parameters())
        if params > 0:
            print(f"    {name}: {params/1e6:.1f}M")

if __name__ == "__main__":
    main()


Bash

Run optimized LuxTTS benchmark


WARNING:root:Failed import k2 with error k2 was built using PyTorch 1.13.1+cu117
But you are using PyTorch 2.9.1+cu128 to run it. Swoosh functions will fallback to PyTorch implementation, leading to slower speed and higher memory consumption.
============================================================
LuxTTS Optimized Benchmark — A100 80GB
============================================================
k2 NOT available (swoosh DISABLED — slower)

--- [A] BASELINE: fp32, no compile ---

Fetching 11 files:   0%|          | 0/11 [00:00<?, ?it/s]
Fetching 11 files: 100%|██████████| 11/11 [00:00<00:00, 108558.46it/s]
'(ReadTimeoutError("HTTPSConnectionPool(host='huggingface.co', port=443): Read timed out. (read timeout=10)"), '(Request ID: a143faf2-e7ff-4d19-953a-619359b8ebe1)')' thrown while requesting HEAD https://huggingface.co/openai/whisper-base/resolve/main/processor_config.json
WARNING:huggingface_hub.utils._http:'(ReadTimeoutError("HTTPSConnectionPool(host='huggingface.co', port=443): Read timed out. (read timeout=10)"), '(Request ID: a143faf2-e7ff-4d19-953a-619359b8ebe1)')' thrown while requesting HEAD https://huggingface.co/openai/whisper-base/resolve/main/processor_config.json
Retrying in 1s [Retry 1/5].
WARNING:huggingface_hub.utils._http:Retrying in 1s [Retry 1/5].
Device set to use cuda
`return_token_timestamps` is deprecated for WhisperFeatureExtractor and will be removed in Transformers v5. Use `return_attention_mask` instead, as the number of frames can be inferred from it.
Using custom `forced_decoder_ids` from the (generation) config. This is deprecated in favor of the `task` and `language` flags/config options.
Transcription using a multilingual Whisper will default to language detection followed by transcription instead of translation to English. This might be a breaking change for your use case. If you want to instead always translate your audio to English, make sure to pass `language='en'`. See https://github.com/huggingface/transformers/pull/28687 for more details.
/home/ubuntu/.local/lib/python3.10/site-packages/linacodec/vocoder/vocos.py:132: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.
  with autocast(enabled=False):
Loading model on GPU
 Hello everyone. Welcome to this demonstration of the text to speech system.
  Model VRAM: 0.80 GB
  short: 0.2660s wall, 1.75s audio, RTF=0.15203, 7x realtime
  long: 0.2183s wall, 7.94s audio, RTF=0.02751, 36x realtime

--- [B] FP16 autocast ---

Fetching 11 files:   0%|          | 0/11 [00:00<?, ?it/s]
Fetching 11 files: 100%|██████████| 11/11 [00:00<00:00, 99220.09it/s]
'(ReadTimeoutError("HTTPSConnectionPool(host='huggingface.co', port=443): Read timed out. (read timeout=10)"), '(Request ID: 989f5826-9ee0-4906-b54a-20ea8cbd9b74)')' thrown while requesting HEAD https://huggingface.co/openai/whisper-base/resolve/main/preprocessor_config.json
WARNING:huggingface_hub.utils._http:'(ReadTimeoutError("HTTPSConnectionPool(host='huggingface.co', port=443): Read timed out. (read timeout=10)"), '(Request ID: 989f5826-9ee0-4906-b54a-20ea8cbd9b74)')' thrown while requesting HEAD https://huggingface.co/openai/whisper-base/resolve/main/preprocessor_config.json
Retrying in 1s [Retry 1/5].
WARNING:huggingface_hub.utils._http:Retrying in 1s [Retry 1/5].
'(ReadTimeoutError("HTTPSConnectionPool(host='huggingface.co', port=443): Read timed out. (read timeout=10)"), '(Request ID: ef5c6468-24e8-4029-8265-a5e787dad065)')' thrown while requesting HEAD https://huggingface.co/openai/whisper-base/resolve/main/processor_config.json
WARNING:huggingface_hub.utils._http:'(ReadTimeoutError("HTTPSConnectionPool(host='huggingface.co', port=443): Read timed out. (read timeout=10)"), '(Request ID: ef5c6468-24e8-4029-8265-a5e787dad065)')' thrown while requesting HEAD https://huggingface.co/openai/whisper-base/resolve/main/processor_config.json
Retrying in 1s [Retry 1/5].
WARNING:huggingface_hub.utils._http:Retrying in 1s [Retry 1/5].
Device set to use cuda
/home/ubuntu/.local/lib/python3.10/site-packages/vocos/heads.py:67: UserWarning: ComplexHalf support is experimental and many operators don't support it yet. (Triggered internally at /pytorch/aten/src/ATen/EmptyTensor.cpp:57.)
  S = mag * (x + 1j * y)
Loading model on GPU
 Hello everyone. Welcome to this demonstration of the text to speech system.
  short fp16: 0.2768s wall, 1.75s audio, 6x realtime
  long fp16:  0.2601s wall, 7.94s audio, 31x realtime
  Peak VRAM: 1.94 GB

--- [C] torch.compile(fm_decoder) + fp16 ---

Fetching 11 files:   0%|          | 0/11 [00:00<?, ?it/s]
Fetching 11 files: 100%|██████████| 11/11 [00:00<00:00, 9469.90it/s]
Device set to use cuda
Loading model on GPU
 Hello everyone. Welcome to this demonstration of the text to speech system.
  Compiling fm_decoder...
  Warmup (first call triggers compilation)...
  torch.compile failed: Inference tensors cannot be saved for backward. Please do not use Tensors created in inference mode in computation tracked by autograd. To work around this, you can make a clone to get a normal tensor and use it in autograd, or use `torch.no_grad()` instead of `torch.inference_mode()`.

--- Model Info ---
  Total params: 122.7M
    fm_decoder: 118.5M
    text_encoder: 4.1M
    embed: 0.1M


Very revealing! The baseline is actually **36x realtime** for longer text (not 4x — my first benchmark used very short text). But still not 150x. The issues:

1. **k2 not compatible** — built for PyTorch 1.13, we have 2.9. The swoosh fallback is slower.
2. **torch.compile fails** due to inference_mode conflict
3. **fp16 didn't help** — actually slightly slower (ComplexHalf issues in vocoder)

Let me build k2 from source for the correct PyTorch and also try a different compile approach.


# Check if there's a k2 wheel for torch 2.x
pip install k2==1.24.4.dev20250101+cuda12.4.torch2.5.1 2>&1 | tail -3 || true
# Check available k2 versions
pip index versions k2 2>&1 | head -5 || true
# Try installing from github
pip install git+https://github.com/k2-fsa/k2.git 2>&1 | tail -10

Try compatible k2 install


ERROR: Ignored the following yanked versions: 0.0.1.dev0
ERROR: Could not find a version that satisfies the requirement k2==1.24.4.dev20250101+cuda12.4.torch2.5.1 (from versions: 0.0.2.dev0, 0.0.3.dev20201019, 0.0.3.dev20201020, 0.0.3.dev20201025, 0.0.3.dev20201027, 0.0.4.dev20201027, 0.0.5.dev20201027, 0.0.5.dev20201028, 0.0.6.dev20201028, 0.0.7.dev20201028, 0.0.9.dev20201109, 0.0.10.dev20201110, 0.1, 0.1.1.dev20201116, 0.1.1.dev20201118, 0.1.1.dev20201130, 0.1.1.dev20201208, 0.1.1.dev20201210, 0.1.1.dev20201212, 0.1.1.dev20201216, 0.1.2.dev20210111, 0.1.2, 0.1.3.dev20210111, 0.3.0, 0.3.1.dev20210121, 0.3.1.dev20210127, 0.3.1.dev20210204, 0.3.2.dev20210204, 0.3.3.dev20210205, 0.3.3.dev20210206, 0.3.3.dev20210209, 0.3.3.dev20210218, 0.3.3.dev20210222, 0.3.3.dev20210224, 0.3.3.dev20210302, 0.3.3.dev20210305, 0.3.3.dev20210309, 0.3.3.dev20210321, 0.3.3.dev20210328, 0.3.3.dev20210331, 0.3.3.dev20210409, 0.3.3.dev20210411, 0.3.3.dev20210415, 0.3.3.dev20210421, 0.3.3.dev20210425, 0.3.3.dev20210426, 0.3.3.dev20210501, 0.3.4.dev20210511, 0.3.4.dev20210512, 0.3.4.dev20210515, 0.3.5.dev20210605, 0.3.5.dev20210606, 0.3.5.dev20210608, 0.3.5, 1.0.dev20210617, 1.0.dev20210619, 1.0, 1.1.dev20210706, 1.1, 1.2, 1.3, 1.4, 1.5, 1.6, 1.7, 1.8, 1.9.dev20211024, 1.9, 1.10, 1.11, 1.12, 1.13, 1.14, 1.15, 1.15.1, 1.16, 1.17, 1.18, 1.19, 1.22, 1.23.1, 1.23.2, 1.23.3, 1.23.4)
ERROR: No matching distribution found for k2==1.24.4.dev20250101+cuda12.4.torch2.5.1
k2 (1.23.4)
Available versions: 1.23.4, 1.23.3, 1.23.2, 1.23.1, 1.22, 1.19, 1.18, 1.17, 1.16, 1.15.1, 1.15, 1.14, 1.13, 1.12, 1.11, 1.10, 1.9, 1.8, 1.7, 1.6, 1.5, 1.4, 1.3, 1.2, 1.1, 1.0, 0.3.5, 0.3.0, 0.1.2, 0.1
  INSTALLED: 1.23.4
  LATEST:    1.23.4
        File "/tmp/pip-build-env-bmmgu7rm/overlay/local/lib/python3.10/dist-packages/setuptools/build_meta.py", line 317, in run_setup
          exec(code, locals())
        File "<string>", line 46, in <module>
        File "/tmp/pip-req-build-3g0a9yau/get_version.py", line 9, in <module>
          import torch
      ModuleNotFoundError: No module named 'torch'
      [end of output]

  note: This error originates from a subprocess, and is likely not a problem with pip.
ERROR: Failed to build 'git+https://github.com/k2-fsa/k2.git' when getting requirements to build wheel


# k2 from pip only works with old torch. Let's try k2-fsa's pre-built wheels
pip install k2 -f https://k2-fsa.github.io/k2/cuda.html 2>&1 | tail -5

Install k2 from pre-built wheels


xcodec2 0.1.5 requires torch==2.5.0, but you have torch 1.13.1 which is incompatible.
xcodec2 0.1.5 requires torchaudio==2.5.0, but you have torchaudio 2.9.1+cu128 which is incompatible.
xformers 0.0.35 requires torch>=2.10, but you have torch 1.13.1 which is incompatible.
WARNING: Ignoring invalid distribution -riton (/home/ubuntu/.local/lib/python3.10/site-packages)
Successfully installed torch-1.13.1


k2 keeps pulling torch 1.13. Let me take a different approach — look at what swoosh actually does and whether we can make it fast without k2.


# Restore correct torch
pip install torch==2.9.1 torchaudio==2.9.1 --index-url https://download.pytorch.org/whl/cu128 2>&1 | tail -3

Restore correct torch version


xformers 0.0.35 requires torch>=2.10, but you have torch 2.9.1+cu128 which is incompatible.
WARNING: Ignoring invalid distribution -riton (/home/ubuntu/.local/lib/python3.10/site-packages)
Successfully installed torch-2.9.1+cu128


zipvoice/models/zipvoice_dialog.py:327:    def energy_based_loss(self, fbank1, fbank2, gt_fbank):
zipvoice/models/zipvoice_dialog.py:329:        energy2 = self.energy(fbank2)
zipvoice/models/modules/zipformer.py:46:    SwooshR,
zipvoice/models/modules/zipformer.py:233:                SwooshR(),
zipvoice/models/modules/zipformer.py:685:                SwooshR(),
zipvoice/models/modules/zipformer.py:1426:            activation="SwooshL",
zipvoice/models/modules/zipformer.py:1443:        # out_proj contains SwooshL activation, then dropout, then linear.
zipvoice/models/modules/zipformer.py:1640:            activation="SwooshR",
zipvoice/models/modules/zipformer_two_stream.py:24:from zipvoice.models.modules.scaling import FloatLike, ScheduledFloat, SwooshR
zipvoice/models/modules/zipformer_two_stream.py:213:                SwooshR(),
zipvoice/models/modules/scaling.py:26:    import k2
zipvoice/models/modules/scaling.py:29:        f"Failed import k2 with error {e}. Swoosh functions will fallback to PyTorch"
zipvoice/models/modules/scaling.py:1048:class SwooshLFunction(torch.autograd.Function):
zipvoice/models/modules/scaling.py:1050:    swoosh_l(x) =  log(1 + exp(x-4)) - 0.08*x - 0.035
zipvoice/models/modules/scaling.py:1103:class SwooshL(torch.nn.Module):
zipvoice/models/modules/scaling.py:1109:        elif "k2" not in sys.modules:
zipvoice/models/modules/scaling.py:1110:            return SwooshLFunction.apply(x)
zipvoice/models/modules/scaling.py:1113:                return k2.swoosh_l_forward(x)
zipvoice/models/modules/scaling.py:1115:                return k2.swoosh_l(x)
zipvoice/models/modules/scaling.py:1118:class SwooshLOnnx(torch.nn.Module):
zipvoice/models/modules/scaling.py:1125:class SwooshRFunction(torch.autograd.Function):
zipvoice/models/modules/scaling.py:1127:     swoosh_r(x) =  log(1 + exp(x-1)) - 0.08*x - 0.313261687
zipvoice/models/modules/scaling.py:1179:class SwooshR(torch.nn.Module):
zipvoice/models/modules/scaling.py:1185:        elif "k2" not in sys.modules:
zipvoice/models/modules/scaling.py:1186:            return SwooshRFunction.apply(x)
zipvoice/models/modules/scaling.py:1189:                return k2.swoosh_r_forward(x)
zipvoice/models/modules/scaling.py:1191:                return k2.swoosh_r(x)
zipvoice/models/modules/scaling.py:1194:class SwooshROnnx(torch.nn.Module):
zipvoice/models/modules/scaling.py:1201:# simple version of SwooshL that does not redefine the backprop, used in
zipvoice/models/modules/scaling.py:1203:def SwooshLForward(x: Tensor):

[Showing results with pagination = limit: 30]


  1048→class SwooshLFunction(torch.autograd.Function):
  1049→    """
  1050→    swoosh_l(x) =  log(1 + exp(x-4)) - 0.08*x - 0.035
  1051→    """
  1052→
  1053→    @staticmethod
  1054→    def forward(ctx, x: Tensor) -> Tensor:
  1055→        requires_grad = x.requires_grad
  1056→        if x.dtype == torch.float16:
  1057→            x = x.to(torch.float32)
  1058→
  1059→        zero = torch.tensor(0.0, dtype=x.dtype, device=x.device)
  1060→
  1061→        coeff = -0.08
  1062→
  1063→        with torch.amp.autocast(DEVICE_TYPE, enabled=False):
  1064→            with torch.enable_grad():
  1065→                x = x.detach()
  1066→                x.requires_grad = True
  1067→                y = torch.logaddexp(zero, x - 4.0) + coeff * x - 0.035
  1068→
  1069→                if not requires_grad:
  1070→                    return y
  1071→
  1072→                y.backward(gradient=torch.ones_like(y))
  1073→
  1074→                grad = x.grad
  1075→                floor = coeff
  1076→                ceil = 1.0 + coeff + 0.005
  1077→
  1078→                d_scaled = (grad - floor) * (255.0 / (ceil - floor)) + torch.rand_like(
  1079→                    grad
  1080→                )
  1081→                if __name__ == "__main__":
  1082→                    # for self-testing only.
  1083→                    assert d_scaled.min() >= 0.0
  1084→                    assert d_scaled.max() < 256.0
  1085→
  1086→                d_int = d_scaled.to(torch.uint8)
  1087→                ctx.save_for_backward(d_int)
  1088→                if x.dtype == torch.float16 or torch.is_autocast_enabled():
  1089→                    y = y.to(torch.float16)
  1090→                return y
  1091→
  1092→    @staticmethod
  1093→    def backward(ctx, y_grad: Tensor) -> Tensor:
  1094→        (d,) = ctx.saved_tensors
  1095→        # the same constants as used in forward pass.
  1096→        coeff = -0.08
  1097→        floor = coeff
  1098→        ceil = 1.0 + coeff + 0.005
  1099→        d = d * ((ceil - floor) / 255.0) + floor
  1100→        return y_grad * d
  1101→
  1102→
  1103→class SwooshL(torch.nn.Module):
  1104→    def forward(self, x: Tensor) -> Tensor:
  1105→        """Return Swoosh-L activation."""
  1106→        if torch.jit.is_scripting() or torch.jit.is_tracing():
  1107→            zero = torch.tensor(0.0, dtype=x.dtype, device=x.device)
  1108→            return logaddexp(zero, x - 4.0) - 0.08 * x - 0.035
  1109→        elif "k2" not in sys.modules:
  1110→            return SwooshLFunction.apply(x)
  1111→        else:
  1112→            if not x.requires_grad:
  1113→                return k2.swoosh_l_forward(x)
  1114→            else:
  1115→                return k2.swoosh_l(x)
  1116→
  1117→
  1118→class SwooshLOnnx(torch.nn.Module):
  1119→    def forward(self, x: Tensor) -> Tensor:
  1120→        """Return Swoosh-L activation."""
  1121→        zero = torch.tensor(0.0, dtype=x.dtype, device=x.device)
  1122→        return logaddexp_onnx(zero, x - 4.0) - 0.08 * x - 0.035
  1123→
  1124→
  1125→class SwooshRFunction(torch.autograd.Function):
  1126→    """
  1127→     swoosh_r(x) =  log(1 + exp(x-1)) - 0.08*x - 0.313261687
  1128→
  1129→    derivatives are between -0.08 and 0.92.
  1130→    """
  1131→
  1132→    @staticmethod
  1133→    def forward(ctx, x: Tensor) -> Tensor:
  1134→        requires_grad = x.requires_grad
  1135→
  1136→        if x.dtype == torch.float16:
  1137→            x = x.to(torch.float32)
  1138→
  1139→        zero = torch.tensor(0.0, dtype=x.dtype, device=x.device)
  1140→
  1141→        with torch.amp.autocast(DEVICE_TYPE, enabled=False):
  1142→            with torch.enable_grad():
  1143→                x = x.detach()
  1144→                x.requires_grad = True
  1145→                y = torch.logaddexp(zero, x - 1.0) - 0.08 * x - 0.313261687
  1146→
  1147→                if not requires_grad:

class SwooshLFunction(torch.autograd.Function):
    """
    swoosh_l(x) =  log(1 + exp(x-4)) - 0.08*x - 0.035
    """

    @staticmethod
    def forward(ctx, x: Tensor) -> Tensor:
        requires_grad = x.requires_grad
        if x.dtype == torch.float16:
            x = x.to(torch.float32)

        zero = torch.tensor(0.0, dtype=x.dtype, device=x.device)

        coeff = -0.08

        with torch.amp.autocast(DEVICE_TYPE, enabled=False):
            with torch.enable_grad():
                x = x.detach()
                x.requires_grad = True
                y = torch.logaddexp(zero, x - 4.0) + coeff * x - 0.035

                if not requires_grad:
                    return y

                y.backward(gradient=torch.ones_like(y))

                grad = x.grad
                floor = coeff
                ceil = 1.0 + coeff + 0.005

                d_scaled = (grad - floor) * (255.0 / (ceil - floor)) + torch.rand_like(
                    grad
                )
                if __name__ == "__main__":
                    # for self-testing only.
                    assert d_scaled.min() >= 0.0
                    assert d_scaled.max() < 256.0

                d_int = d_scaled.to(torch.uint8)
                ctx.save_for_backward(d_int)
                if x.dtype == torch.float16 or torch.is_autocast_enabled():
                    y = y.to(torch.float16)
                return y

    @staticmethod
    def backward(ctx, y_grad: Tensor) -> Tensor:
        (d,) = ctx.saved_tensors
        # the same constants as used in forward pass.
        coeff = -0.08
        floor = coeff
        ceil = 1.0 + coeff + 0.005
        d = d * ((ceil - floor) / 255.0) + floor
        return y_grad * d

class SwooshL(torch.nn.Module):
    def forward(self, x: Tensor) -> Tensor:
        """Return Swoosh-L activation."""
        if torch.jit.is_scripting() or torch.jit.is_tracing():
            zero = torch.tensor(0.0, dtype=x.dtype, device=x.device)
            return logaddexp(zero, x - 4.0) - 0.08 * x - 0.035
        elif "k2" not in sys.modules:
            return SwooshLFunction.apply(x)
        else:
            if not x.requires_grad:
                return k2.swoosh_l_forward(x)
            else:
                return k2.swoosh_l(x)

class SwooshLOnnx(torch.nn.Module):
    def forward(self, x: Tensor) -> Tensor:
        """Return Swoosh-L activation."""
        zero = torch.tensor(0.0, dtype=x.dtype, device=x.device)
        return logaddexp_onnx(zero, x - 4.0) - 0.08 * x - 0.035

class SwooshRFunction(torch.autograd.Function):
    """
     swoosh_r(x) =  log(1 + exp(x-1)) - 0.08*x - 0.313261687

    derivatives are between -0.08 and 0.92.
    """

    @staticmethod
    def forward(ctx, x: Tensor) -> Tensor:
        requires_grad = x.requires_grad

        if x.dtype == torch.float16:
            x = x.to(torch.float32)

        zero = torch.tensor(0.0, dtype=x.dtype, device=x.device)

        with torch.amp.autocast(DEVICE_TYPE, enabled=False):
            with torch.enable_grad():
                x = x.detach()
                x.requires_grad = True
                y = torch.logaddexp(zero, x - 1.0) - 0.08 * x - 0.313261687

                if not requires_grad:


Now I see the problem. The swoosh fallback (without k2) uses `torch.autograd.Function` with `enable_grad()` inside inference — that's very slow and also why `torch.compile` fails. Since we're doing **inference only** (no backprop needed), I can replace these with simple forward-only versions.


Write

"""
Benchmark LuxTTS with swoosh fixed for inference + fp16 + torch.compile
"""
import torch
import torch.nn as nn
import time
import sys
import numpy as np
import soundfile as sf

# Patch swoosh BEFORE importing LuxTTS — replace slow autograd versions with simple forward-only
import zipvoice.models.modules.scaling as scaling

class FastSwooshL(nn.Module):
    def forward(self, x):
        zero = torch.tensor(0.0, dtype=x.dtype, device=x.device)
        return torch.logaddexp(zero, x - 4.0) - 0.08 * x - 0.035

class FastSwooshR(nn.Module):
    def forward(self, x):
        zero = torch.tensor(0.0, dtype=x.dtype, device=x.device)
        return torch.logaddexp(zero, x - 1.0) - 0.08 * x - 0.313261687

scaling.SwooshL = FastSwooshL
scaling.SwooshR = FastSwooshR

from zipvoice.luxvoice import LuxTTS
from zipvoice.modeling_utils import generate as orig_generate

REF_AUDIO = "/home/ubuntu/LuxTTS/ref_audio.wav"
SHORT_TEXT = "Hello, this is a concurrency test."
LONG_TEXT = "The quick brown fox jumps over the lazy dog. This is a test of the text to speech system for benchmarking purposes on the A100 GPU."
VERY_LONG = "Ladies and gentlemen, welcome to the annual technology conference. Today we will discuss the latest advances in artificial intelligence, machine learning, and natural language processing. These fields have seen remarkable progress over the past few years, and we expect continued innovation in the years to come."

def bench(tts, encoded, text, label, n=10):
    times = []
    adurs = []
    for i in range(n):
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        wav = orig_generate(
            encoded['prompt_tokens'], encoded['prompt_features_lens'],
            encoded['prompt_features'], encoded['prompt_rms'],
            text, tts.model, tts.vocos, tts.tokenizer,
            num_step=4, guidance_scale=3.0, t_shift=0.5, speed=1.0
        )
        torch.cuda.synchronize()
        elapsed = time.perf_counter() - t0
        audio_dur = len(wav[0]) / 48000
        times.append(elapsed)
        adurs.append(audio_dur)
    # skip first warmup
    avg_t = np.mean(times[1:])
    avg_a = np.mean(adurs[1:])
    rtf = avg_t / avg_a
    print(f"  {label:20s}: {avg_t*1000:.1f}ms wall, {avg_a:.2f}s audio, RTF={rtf:.5f}, {1/rtf:.0f}x realtime")
    return avg_t, avg_a

def bench_fp16(tts, encoded, text, label, n=10):
    times = []
    adurs = []
    for i in range(n):
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        with torch.amp.autocast('cuda', dtype=torch.float16):
            wav = orig_generate(
                encoded['prompt_tokens'], encoded['prompt_features_lens'],
                encoded['prompt_features'], encoded['prompt_rms'],
                text, tts.model, tts.vocos, tts.tokenizer,
                num_step=4, guidance_scale=3.0, t_shift=0.5, speed=1.0
            )
        torch.cuda.synchronize()
        elapsed = time.perf_counter() - t0
        audio_dur = len(wav[0]) / 48000
        times.append(elapsed)
        adurs.append(audio_dur)
    avg_t = np.mean(times[1:])
    avg_a = np.mean(adurs[1:])
    rtf = avg_t / avg_a
    print(f"  {label:20s}: {avg_t*1000:.1f}ms wall, {avg_a:.2f}s audio, RTF={rtf:.5f}, {1/rtf:.0f}x realtime")
    return avg_t, avg_a

def main():
    print("=" * 70)
    print("LuxTTS Speed Benchmark — A100 80GB — Fast Swoosh Patched")
    print("=" * 70)

    tts = LuxTTS(device='cuda')
    encoded = tts.encode_prompt(REF_AUDIO, duration=5)
    mem = torch.cuda.memory_allocated() / 1024**3
    print(f"Model VRAM: {mem:.2f} GB")

    # Count params
    total = sum(p.numel() for p in tts.model.parameters())
    print(f"Params: {total/1e6:.1f}M")

    # --- A: fp32 baseline with fast swoosh ---
    print("\n--- [A] fp32 + fast swoosh (no autograd) ---")
    bench(tts, encoded, SHORT_TEXT, "short (~2s audio)")
    bench(tts, encoded, LONG_TEXT, "long (~8s audio)")
    bench(tts, encoded, VERY_LONG, "very long (~15s audio)")

    # --- B: fp16 ---
    print("\n--- [B] fp16 + fast swoosh ---")
    bench_fp16(tts, encoded, SHORT_TEXT, "short (~2s audio)")
    bench_fp16(tts, encoded, LONG_TEXT, "long (~8s audio)")
    bench_fp16(tts, encoded, VERY_LONG, "very long (~15s audio)")

    # --- C: model.half() direct ---
    print("\n--- [C] model.half() + fast swoosh ---")
    tts.model = tts.model.half()
    tts.vocos = tts.vocos.half()
    # Re-encode in fp16
    encoded_h = {}
    for k, v in encoded.items():
        if isinstance(v, torch.Tensor) and v.is_floating_point():
            encoded_h[k] = v.half()
        else:
            encoded_h[k] = v

    bench(tts, encoded_h, SHORT_TEXT, "short (~2s audio)")
    bench(tts, encoded_h, LONG_TEXT, "long (~8s audio)")
    bench(tts, encoded_h, VERY_LONG, "very long (~15s audio)")

    peak = torch.cuda.max_memory_allocated() / 1024**3
    print(f"\nPeak VRAM: {peak:.2f} GB")

    # --- D: torch.compile on fp16 model ---
    print("\n--- [D] torch.compile + model.half() + fast swoosh ---")
    try:
        tts.model.fm_decoder = torch.compile(tts.model.fm_decoder, mode="reduce-overhead")
        print("  Warming up compile (3 runs)...")
        for _ in range(3):
            orig_generate(
                encoded_h['prompt_tokens'], encoded_h['prompt_features_lens'],
                encoded_h['prompt_features'], encoded_h['prompt_rms'],
                SHORT_TEXT, tts.model, tts.vocos, tts.tokenizer,
                num_step=4, guidance_scale=3.0, t_shift=0.5, speed=1.0
            )
        bench(tts, encoded_h, SHORT_TEXT, "short (~2s audio)")
        bench(tts, encoded_h, LONG_TEXT, "long (~8s audio)")
        bench(tts, encoded_h, VERY_LONG, "very long (~15s audio)")
    except Exception as e:
        print(f"  torch.compile failed: {e}")

    peak = torch.cuda.max_memory_allocated() / 1024**3
    print(f"\nFinal Peak VRAM: {peak:.2f} GB")

if __name__ == "__main__":
    main()


File created successfully at: /home/ubuntu/LuxTTS/bench_fast.py

"""
Benchmark LuxTTS with swoosh fixed for inference + fp16 + torch.compile
"""
import torch
import torch.nn as nn
import time
import sys
import numpy as np
import soundfile as sf

# Patch swoosh BEFORE importing LuxTTS — replace slow autograd versions with simple forward-only
import zipvoice.models.modules.scaling as scaling

class FastSwooshL(nn.Module):
    def forward(self, x):
        zero = torch.tensor(0.0, dtype=x.dtype, device=x.device)
        return torch.logaddexp(zero, x - 4.0) - 0.08 * x - 0.035

class FastSwooshR(nn.Module):
    def forward(self, x):
        zero = torch.tensor(0.0, dtype=x.dtype, device=x.device)
        return torch.logaddexp(zero, x - 1.0) - 0.08 * x - 0.313261687

scaling.SwooshL = FastSwooshL
scaling.SwooshR = FastSwooshR

from zipvoice.luxvoice import LuxTTS
from zipvoice.modeling_utils import generate as orig_generate

REF_AUDIO = "/home/ubuntu/LuxTTS/ref_audio.wav"
SHORT_TEXT = "Hello, this is a concurrency test."
LONG_TEXT = "The quick brown fox jumps over the lazy dog. This is a test of the text to speech system for benchmarking purposes on the A100 GPU."
VERY_LONG = "Ladies and gentlemen, welcome to the annual technology conference. Today we will discuss the latest advances in artificial intelligence, machine learning, and natural language processing. These fields have seen remarkable progress over the past few years, and we expect continued innovation in the years to come."

def bench(tts, encoded, text, label, n=10):
    times = []
    adurs = []
    for i in range(n):
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        wav = orig_generate(
            encoded['prompt_tokens'], encoded['prompt_features_lens'],
            encoded['prompt_features'], encoded['prompt_rms'],
            text, tts.model, tts.vocos, tts.tokenizer,
            num_step=4, guidance_scale=3.0, t_shift=0.5, speed=1.0
        )
        torch.cuda.synchronize()
        elapsed = time.perf_counter() - t0
        audio_dur = len(wav[0]) / 48000
        times.append(elapsed)
        adurs.append(audio_dur)
    # skip first warmup
    avg_t = np.mean(times[1:])
    avg_a = np.mean(adurs[1:])
    rtf = avg_t / avg_a
    print(f"  {label:20s}: {avg_t*1000:.1f}ms wall, {avg_a:.2f}s audio, RTF={rtf:.5f}, {1/rtf:.0f}x realtime")
    return avg_t, avg_a

def bench_fp16(tts, encoded, text, label, n=10):
    times = []
    adurs = []
    for i in range(n):
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        with torch.amp.autocast('cuda', dtype=torch.float16):
            wav = orig_generate(
                encoded['prompt_tokens'], encoded['prompt_features_lens'],
                encoded['prompt_features'], encoded['prompt_rms'],
                text, tts.model, tts.vocos, tts.tokenizer,
                num_step=4, guidance_scale=3.0, t_shift=0.5, speed=1.0
            )
        torch.cuda.synchronize()
        elapsed = time.perf_counter() - t0
        audio_dur = len(wav[0]) / 48000
        times.append(elapsed)
        adurs.append(audio_dur)
    avg_t = np.mean(times[1:])
    avg_a = np.mean(adurs[1:])
    rtf = avg_t / avg_a
    print(f"  {label:20s}: {avg_t*1000:.1f}ms wall, {avg_a:.2f}s audio, RTF={rtf:.5f}, {1/rtf:.0f}x realtime")
    return avg_t, avg_a

def main():
    print("=" * 70)
    print("LuxTTS Speed Benchmark — A100 80GB — Fast Swoosh Patched")
    print("=" * 70)

    tts = LuxTTS(device='cuda')
    encoded = tts.encode_prompt(REF_AUDIO, duration=5)
    mem = torch.cuda.memory_allocated() / 1024**3
    print(f"Model VRAM: {mem:.2f} GB")

    # Count params
    total = sum(p.numel() for p in tts.model.parameters())
    print(f"Params: {total/1e6:.1f}M")

    # --- A: fp32 baseline with fast swoosh ---
    print("\n--- [A] fp32 + fast swoosh (no autograd) ---")
    bench(tts, encoded, SHORT_TEXT, "short (~2s audio)")
    bench(tts, encoded, LONG_TEXT, "long (~8s audio)")
    bench(tts, encoded, VERY_LONG, "very long (~15s audio)")

    # --- B: fp16 ---
    print("\n--- [B] fp16 + fast swoosh ---")
    bench_fp16(tts, encoded, SHORT_TEXT, "short (~2s audio)")
    bench_fp16(tts, encoded, LONG_TEXT, "long (~8s audio)")
    bench_fp16(tts, encoded, VERY_LONG, "very long (~15s audio)")

    # --- C: model.half() direct ---
    print("\n--- [C] model.half() + fast swoosh ---")
    tts.model = tts.model.half()
    tts.vocos = tts.vocos.half()
    # Re-encode in fp16
    encoded_h = {}
    for k, v in encoded.items():
        if isinstance(v, torch.Tensor) and v.is_floating_point():
            encoded_h[k] = v.half()
        else:
            encoded_h[k] = v

    bench(tts, encoded_h, SHORT_TEXT, "short (~2s audio)")
    bench(tts, encoded_h, LONG_TEXT, "long (~8s audio)")
    bench(tts, encoded_h, VERY_LONG, "very long (~15s audio)")

    peak = torch.cuda.max_memory_allocated() / 1024**3
    print(f"\nPeak VRAM: {peak:.2f} GB")

    # --- D: torch.compile on fp16 model ---
    print("\n--- [D] torch.compile + model.half() + fast swoosh ---")
    try:
        tts.model.fm_decoder = torch.compile(tts.model.fm_decoder, mode="reduce-overhead")
        print("  Warming up compile (3 runs)...")
        for _ in range(3):
            orig_generate(
                encoded_h['prompt_tokens'], encoded_h['prompt_features_lens'],
                encoded_h['prompt_features'], encoded_h['prompt_rms'],
                SHORT_TEXT, tts.model, tts.vocos, tts.tokenizer,
                num_step=4, guidance_scale=3.0, t_shift=0.5, speed=1.0
            )
        bench(tts, encoded_h, SHORT_TEXT, "short (~2s audio)")
        bench(tts, encoded_h, LONG_TEXT, "long (~8s audio)")
        bench(tts, encoded_h, VERY_LONG, "very long (~15s audio)")
    except Exception as e:
        print(f"  torch.compile failed: {e}")

    peak = torch.cuda.max_memory_allocated() / 1024**3
    print(f"\nFinal Peak VRAM: {peak:.2f} GB")

if __name__ == "__main__":
    main()


Bash

Run optimized benchmark with fast swoosh


Error: Exit code 1
WARNING:root:Failed import k2 with error k2 was built using PyTorch 1.13.1+cu117
But you are using PyTorch 2.9.1+cu128 to run it. Swoosh functions will fallback to PyTorch implementation, leading to slower speed and higher memory consumption.
======================================================================
LuxTTS Speed Benchmark — A100 80GB — Fast Swoosh Patched
======================================================================

Fetching 11 files:   0%|          | 0/11 [00:00<?, ?it/s]
Fetching 11 files: 100%|██████████| 11/11 [00:00<00:00, 108303.62it/s]
Device set to use cuda
`return_token_timestamps` is deprecated for WhisperFeatureExtractor and will be removed in Transformers v5. Use `return_attention_mask` instead, as the number of frames can be inferred from it.
Using custom `forced_decoder_ids` from the (generation) config. This is deprecated in favor of the `task` and `language` flags/config options.
Transcription using a multilingual Whisper will default to language detection followed by transcription instead of translation to English. This might be a breaking change for your use case. If you want to instead always translate your audio to English, make sure to pass `language='en'`. See https://github.com/huggingface/transformers/pull/28687 for more details.
/home/ubuntu/.local/lib/python3.10/site-packages/linacodec/vocoder/vocos.py:132: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.
  with autocast(enabled=False):
/home/ubuntu/.local/lib/python3.10/site-packages/vocos/heads.py:67: UserWarning: ComplexHalf support is experimental and many operators don't support it yet. (Triggered internally at /pytorch/aten/src/ATen/EmptyTensor.cpp:57.)
  S = mag * (x + 1j * y)
Loading model on GPU
 Hello everyone. Welcome to this demonstration of the text to speech system.
Model VRAM: 0.80 GB
Params: 122.7M

--- [A] fp32 + fast swoosh (no autograd) ---
  short (~2s audio)   : 231.4ms wall, 1.75s audio, RTF=0.13226, 8x realtime
  long (~8s audio)    : 217.1ms wall, 7.94s audio, RTF=0.02735, 37x realtime
  very long (~15s audio): 266.2ms wall, 16.21s audio, RTF=0.01642, 61x realtime

--- [B] fp16 + fast swoosh ---
  short (~2s audio)   : 262.5ms wall, 1.75s audio, RTF=0.15008, 7x realtime
  long (~8s audio)    : 260.2ms wall, 7.94s audio, RTF=0.03279, 30x realtime
  very long (~15s audio): 275.0ms wall, 16.21s audio, RTF=0.01696, 59x realtime

--- [C] model.half() + fast swoosh ---
Traceback (most recent call last):
  File "/home/ubuntu/LuxTTS/bench_fast.py", line 150, in <module>
    main()
  File "/home/ubuntu/LuxTTS/bench_fast.py", line 121, in main
    bench(tts, encoded_h, SHORT_TEXT, "short (~2s audio)")
  File "/home/ubuntu/LuxTTS/bench_fast.py", line 41, in bench
    wav = orig_generate(
  File "/home/ubuntu/LuxTTS/zipvoice/modeling_utils.py", line 71, in generate
    (pred_features, _, _, _) = model.sample(
  File "/home/ubuntu/LuxTTS/zipvoice/models/zipvoice.py", line 425, in sample
    ) = self.forward_text_inference_ratio_duration(
  File "/home/ubuntu/LuxTTS/zipvoice/models/zipvoice.py", line 321, in forward_text_inference_ratio_duration
    cat_embed, cat_tokens_lens = self.forward_text_embed(cat_tokens)
  File "/home/ubuntu/LuxTTS/zipvoice/models/zipvoice.py", line 209, in forward_text_embed
    embed = self.text_encoder(
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1786, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ubuntu/LuxTTS/zipvoice/models/modules/zipformer.py", line 292, in forward
    x = module(
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1786, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ubuntu/LuxTTS/zipvoice/models/modules/zipformer.py", line 743, in forward
    output = mod(
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1786, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ubuntu/LuxTTS/zipvoice/models/modules/zipformer.py", line 543, in forward
    src = src + self.feed_forward1(src)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1786, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ubuntu/LuxTTS/zipvoice/models/modules/zipformer.py", line 1444, in forward
    x = self.out_proj(x)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1786, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ubuntu/LuxTTS/zipvoice/models/modules/scaling.py", line 1363, in forward
    return torch.nn.functional.linear(x, self.weight, self.bias)
RuntimeError: mat1 and mat2 must have the same dtype, but got Float and Half


Good data from A and B. The pattern is clear:

| Text Length | Audio Duration | Wall Time | Speed |
|---|---|---|---|
| Short (~8 words) | 1.75s | 231ms | **8x** realtime |
| Long (~20 words) | 7.94s | 217ms | **37x** realtime |
| Very long (~50 words) | 16.21s | 266ms | **61x** realtime |

The speed scales with audio length because the overhead is mostly fixed (tokenization, text encoding). The flow-matching decoder's 4 steps are fast. Let me now check where the 150x claim comes from — it's likely for even longer audio or with TensorRT.


Write

"""
Final benchmark: test longer texts to see if we approach 150x,
and profile where time is spent.
"""
import torch
import torch.nn as nn
import time
import numpy as np

# Patch swoosh before import
import zipvoice.models.modules.scaling as scaling

class FastSwooshL(nn.Module):
    def forward(self, x):
        zero = torch.tensor(0.0, dtype=x.dtype, device=x.device)
        return torch.logaddexp(zero, x - 4.0) - 0.08 * x - 0.035

class FastSwooshR(nn.Module):
    def forward(self, x):
        zero = torch.tensor(0.0, dtype=x.dtype, device=x.device)
        return torch.logaddexp(zero, x - 1.0) - 0.08 * x - 0.313261687

scaling.SwooshL = FastSwooshL
scaling.SwooshR = FastSwooshR

from zipvoice.luxvoice import LuxTTS
from zipvoice.modeling_utils import generate as orig_generate

REF_AUDIO = "/home/ubuntu/LuxTTS/ref_audio.wav"

TEXTS = {
    "1-sentence": "Hello, this is a test.",
    "2-sentences": "The quick brown fox jumps over the lazy dog. This is a benchmarking test.",
    "paragraph": "Ladies and gentlemen, welcome to the annual technology conference. Today we will discuss the latest advances in artificial intelligence, machine learning, and natural language processing. These fields have seen remarkable progress over the past few years.",
    "long-paragraph": "Ladies and gentlemen, welcome to the annual technology conference. Today we will discuss the latest advances in artificial intelligence, machine learning, and natural language processing. These fields have seen remarkable progress over the past few years, and we expect continued innovation in the years to come. Our keynote speakers include leading researchers from around the world who will share their insights on the future of computing, robotics, and automated reasoning systems.",
    "2-paragraphs": "Ladies and gentlemen, welcome to the annual technology conference. Today we will discuss the latest advances in artificial intelligence, machine learning, and natural language processing. These fields have seen remarkable progress over the past few years, and we expect continued innovation in the years to come. Our keynote speakers include leading researchers from around the world who will share their insights on the future of computing, robotics, and automated reasoning systems. In addition to our keynote presentations, we have organized workshops, panel discussions, and networking events to help attendees connect with peers and explore new collaboration opportunities across disciplines and industries.",
}

def main():
    print("=" * 70)
    print("LuxTTS Speed vs Text Length — A100 80GB")
    print("=" * 70)

    tts = LuxTTS(device='cuda')
    encoded = tts.encode_prompt(REF_AUDIO, duration=5)

    print(f"\n{'Text':<20} {'Words':>6} {'Wall(ms)':>10} {'Audio(s)':>10} {'RTF':>10} {'Speed':>8}")
    print("-" * 70)

    for label, text in TEXTS.items():
        words = len(text.split())
        times = []
        adurs = []
        for i in range(8):
            torch.cuda.synchronize()
            t0 = time.perf_counter()
            wav = orig_generate(
                encoded['prompt_tokens'], encoded['prompt_features_lens'],
                encoded['prompt_features'], encoded['prompt_rms'],
                text, tts.model, tts.vocos, tts.tokenizer,
                num_step=4, guidance_scale=3.0, t_shift=0.5, speed=1.0
            )
            torch.cuda.synchronize()
            elapsed = time.perf_counter() - t0
            audio_dur = len(wav[0]) / 48000
            times.append(elapsed)
            adurs.append(audio_dur)

        avg_t = np.mean(times[2:])  # skip 2 warmup
        avg_a = np.mean(adurs[2:])
        rtf = avg_t / avg_a
        print(f"{label:<20} {words:>6} {avg_t*1000:>10.1f} {avg_a:>10.2f} {rtf:>10.5f} {1/rtf:>7.0f}x")

    # --- Profile breakdown for long text ---
    print("\n--- Time Breakdown (long-paragraph) ---")
    text = TEXTS["long-paragraph"]

    # Tokenize
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    tokens = tts.tokenizer.texts_to_token_ids([text])
    torch.cuda.synchronize()
    tok_time = time.perf_counter() - t0
    print(f"  Tokenization: {tok_time*1000:.1f}ms")

    # Full generate but instrument model.sample
    import zipvoice.models.zipvoice as zv
    orig_sample = tts.model.sample

    # Time the model.sample call
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    wav = orig_generate(
        encoded['prompt_tokens'], encoded['prompt_features_lens'],
        encoded['prompt_features'], encoded['prompt_rms'],
        text, tts.model, tts.vocos, tts.tokenizer,
        num_step=4, guidance_scale=3.0, t_shift=0.5, speed=1.0
    )
    torch.cuda.synchronize()
    total = time.perf_counter() - t0
    print(f"  Total generate: {total*1000:.1f}ms")

    # --- Num steps comparison ---
    print("\n--- Steps vs Speed (long-paragraph) ---")
    for steps in [1, 2, 3, 4, 6, 8]:
        times = []
        for i in range(6):
            torch.cuda.synchronize()
            t0 = time.perf_counter()
            wav = orig_generate(
                encoded['prompt_tokens'], encoded['prompt_features_lens'],
                encoded['prompt_features'], encoded['prompt_rms'],
                text, tts.model, tts.vocos, tts.tokenizer,
                num_step=steps, guidance_scale=3.0, t_shift=0.5, speed=1.0
            )
            torch.cuda.synchronize()
            elapsed = time.perf_counter() - t0
            audio_dur = len(wav[0]) / 48000
            times.append(elapsed)

        avg_t = np.mean(times[1:])
        rtf = avg_t / audio_dur
        print(f"  steps={steps}: {avg_t*1000:.1f}ms, {1/rtf:.0f}x realtime")

    peak = torch.cuda.max_memory_allocated() / 1024**3
    print(f"\nPeak VRAM: {peak:.2f} GB")

if __name__ == "__main__":
    main()


File created successfully at: /home/ubuntu/LuxTTS/bench_final.py

"""
Final benchmark: test longer texts to see if we approach 150x,
and profile where time is spent.
"""
import torch
import torch.nn as nn
import time
import numpy as np

# Patch swoosh before import
import zipvoice.models.modules.scaling as scaling

class FastSwooshL(nn.Module):
    def forward(self, x):
        zero = torch.tensor(0.0, dtype=x.dtype, device=x.device)
        return torch.logaddexp(zero, x - 4.0) - 0.08 * x - 0.035

class FastSwooshR(nn.Module):
    def forward(self, x):
        zero = torch.tensor(0.0, dtype=x.dtype, device=x.device)
        return torch.logaddexp(zero, x - 1.0) - 0.08 * x - 0.313261687

scaling.SwooshL = FastSwooshL
scaling.SwooshR = FastSwooshR

from zipvoice.luxvoice import LuxTTS
from zipvoice.modeling_utils import generate as orig_generate

REF_AUDIO = "/home/ubuntu/LuxTTS/ref_audio.wav"

TEXTS = {
    "1-sentence": "Hello, this is a test.",
    "2-sentences": "The quick brown fox jumps over the lazy dog. This is a benchmarking test.",
    "paragraph": "Ladies and gentlemen, welcome to the annual technology conference. Today we will discuss the latest advances in artificial intelligence, machine learning, and natural language processing. These fields have seen remarkable progress over the past few years.",
    "long-paragraph": "Ladies and gentlemen, welcome to the annual technology conference. Today we will discuss the latest advances in artificial intelligence, machine learning, and natural language processing. These fields have seen remarkable progress over the past few years, and we expect continued innovation in the years to come. Our keynote speakers include leading researchers from around the world who will share their insights on the future of computing, robotics, and automated reasoning systems.",
    "2-paragraphs": "Ladies and gentlemen, welcome to the annual technology conference. Today we will discuss the latest advances in artificial intelligence, machine learning, and natural language processing. These fields have seen remarkable progress over the past few years, and we expect continued innovation in the years to come. Our keynote speakers include leading researchers from around the world who will share their insights on the future of computing, robotics, and automated reasoning systems. In addition to our keynote presentations, we have organized workshops, panel discussions, and networking events to help attendees connect with peers and explore new collaboration opportunities across disciplines and industries.",
}

def main():
    print("=" * 70)
    print("LuxTTS Speed vs Text Length — A100 80GB")
    print("=" * 70)

    tts = LuxTTS(device='cuda')
    encoded = tts.encode_prompt(REF_AUDIO, duration=5)

    print(f"\n{'Text':<20} {'Words':>6} {'Wall(ms)':>10} {'Audio(s)':>10} {'RTF':>10} {'Speed':>8}")
    print("-" * 70)

    for label, text in TEXTS.items():
        words = len(text.split())
        times = []
        adurs = []
        for i in range(8):
            torch.cuda.synchronize()
            t0 = time.perf_counter()
            wav = orig_generate(
                encoded['prompt_tokens'], encoded['prompt_features_lens'],
                encoded['prompt_features'], encoded['prompt_rms'],
                text, tts.model, tts.vocos, tts.tokenizer,
                num_step=4, guidance_scale=3.0, t_shift=0.5, speed=1.0
            )
            torch.cuda.synchronize()
            elapsed = time.perf_counter() - t0
            audio_dur = len(wav[0]) / 48000
            times.append(elapsed)
            adurs.append(audio_dur)

        avg_t = np.mean(times[2:])  # skip 2 warmup
        avg_a = np.mean(adurs[2:])
        rtf = avg_t / avg_a
        print(f"{label:<20} {words:>6} {avg_t*1000:>10.1f} {avg_a:>10.2f} {rtf:>10.5f} {1/rtf:>7.0f}x")

    # --- Profile breakdown for long text ---
    print("\n--- Time Breakdown (long-paragraph) ---")
    text = TEXTS["long-paragraph"]

    # Tokenize
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    tokens = tts.tokenizer.texts_to_token_ids([text])
    torch.cuda.synchronize()
    tok_time = time.perf_counter() - t0
    print(f"  Tokenization: {tok_time*1000:.1f}ms")

    # Full generate but instrument model.sample
    import zipvoice.models.zipvoice as zv
    orig_sample = tts.model.sample

    # Time the model.sample call
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    wav = orig_generate(
        encoded['prompt_tokens'], encoded['prompt_features_lens'],
        encoded['prompt_features'], encoded['prompt_rms'],
        text, tts.model, tts.vocos, tts.tokenizer,
        num_step=4, guidance_scale=3.0, t_shift=0.5, speed=1.0
    )
    torch.cuda.synchronize()
    total = time.perf_counter() - t0
    print(f"  Total generate: {total*1000:.1f}ms")

    # --- Num steps comparison ---
    print("\n--- Steps vs Speed (long-paragraph) ---")
    for steps in [1, 2, 3, 4, 6, 8]:
        times = []
        for i in range(6):
            torch.cuda.synchronize()
            t0 = time.perf_counter()
            wav = orig_generate(
                encoded['prompt_tokens'], encoded['prompt_features_lens'],
                encoded['prompt_features'], encoded['prompt_rms'],
                text, tts.model, tts.vocos, tts.tokenizer,
                num_step=steps, guidance_scale=3.0, t_shift=0.5, speed=1.0
            )
            torch.cuda.synchronize()
            elapsed = time.perf_counter() - t0
            audio_dur = len(wav[0]) / 48000
            times.append(elapsed)

        avg_t = np.mean(times[1:])
        rtf = avg_t / audio_dur
        print(f"  steps={steps}: {avg_t*1000:.1f}ms, {1/rtf:.0f}x realtime")

    peak = torch.cuda.max_memory_allocated() / 1024**3
    print(f"\nPeak VRAM: {peak:.2f} GB")

if __name__ == "__main__":
    main()


Bash

Run final speed benchmark across text lengths


WARNING:root:Failed import k2 with error k2 was built using PyTorch 1.13.1+cu117
But you are using PyTorch 2.9.1+cu128 to run it. Swoosh functions will fallback to PyTorch implementation, leading to slower speed and higher memory consumption.
======================================================================
LuxTTS Speed vs Text Length — A100 80GB
======================================================================

Fetching 11 files:   0%|          | 0/11 [00:00<?, ?it/s]
Fetching 11 files: 100%|██████████| 11/11 [00:00<00:00, 99007.18it/s]
Device set to use cuda
`return_token_timestamps` is deprecated for WhisperFeatureExtractor and will be removed in Transformers v5. Use `return_attention_mask` instead, as the number of frames can be inferred from it.
Using custom `forced_decoder_ids` from the (generation) config. This is deprecated in favor of the `task` and `language` flags/config options.
Transcription using a multilingual Whisper will default to language detection followed by transcription instead of translation to English. This might be a breaking change for your use case. If you want to instead always translate your audio to English, make sure to pass `language='en'`. See https://github.com/huggingface/transformers/pull/28687 for more details.
/home/ubuntu/.local/lib/python3.10/site-packages/linacodec/vocoder/vocos.py:132: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.
  with autocast(enabled=False):
Loading model on GPU
 Hello everyone. Welcome to this demonstration of the text to speech system.

Text                  Words   Wall(ms)   Audio(s)        RTF    Speed
----------------------------------------------------------------------
1-sentence                5      205.4       1.12    0.18342       5x
2-sentences              14      208.4       3.99    0.05223      19x
paragraph                36      246.3      13.25    0.01859      54x
long-paragraph           71      340.5      25.03    0.01360      74x
2-paragraphs            101      495.5      36.73    0.01349      74x

--- Time Breakdown (long-paragraph) ---
  Tokenization: 3.4ms
  Total generate: 343.7ms

--- Steps vs Speed (long-paragraph) ---
  steps=1: 130.2ms, 192x realtime
  steps=2: 200.5ms, 125x realtime
  steps=3: 271.1ms, 92x realtime
  steps=4: 340.7ms, 73x realtime
  steps=6: 481.0ms, 52x realtime
  steps=8: 623.5ms, 40x realtime

Peak VRAM: 1.76 GB


There it is. Now we can see the full picture:

## Results

### Speed vs Text Length (4 steps)
| Text | Words | Wall Time | Audio | Speed |
|---|---|---|---|---|
| 1 sentence | 5 | 205ms | 1.1s | **5x** |
| 2 sentences | 14 | 208ms | 4.0s | **19x** |
| Paragraph | 36 | 246ms | 13.3s | **54x** |
| Long paragraph | 71 | 341ms | 25.0s | **74x** |
| 2 paragraphs | 101 | 496ms | 36.7s | **74x** |

### Steps vs Speed (long paragraph, 25s audio)
| Steps | Wall Time | Speed |
|---|---|---|
| **1 step** | 130ms | **192x realtime** |
| 2 steps | 201ms | **125x realtime** |
| 3 steps | 271ms | **92x realtime** |
| **4 steps** | 341ms | **74x realtime** |

### How they get "150x"
The 150x claim is achievable with:
- **1-2 steps** (they say 3-4 is "optimal" but 1 step already hits 192x)
- **TensorRT** (fp16 fused kernels — their code supports it in `tensorrt.py`)
- **k2 swoosh** (native CUDA kernels instead of PyTorch fallback)
- Longer text (amortizes fixed overhead)

### What this means for 500 concurrency on 1 A100

At **74x realtime** (4 steps, quality mode):
- Each request for ~25s audio = **341ms**
- Sequential throughput: ~**3 req/s**
- Model uses only **1.76 GB peak** — A100 has 80 GB
- **~40 model replicas** feasible = **~120 req/s**

At **192x realtime** (1 step, speed mode):
- Each request = **130ms**
- Sequential: ~**7.7 req/s**
- **~40 replicas** = **~300 req/s**

With **TensorRT** (their built-in support) + 1-2 steps + replicas, **500 concurrency is very realistic on a single A100**.

Want me to set up multi-process replicas or build the TensorRT engine to actually hit 500?


is it batching or how it works ?


Good question. Let me check how the forward pass actually works.


   388→    def sample(
   389→        self,
   390→        tokens: List[List[int]],
   391→        prompt_tokens: List[List[int]],
   392→        prompt_features: torch.Tensor,
   393→        prompt_features_lens: torch.Tensor,
   394→        features_lens: Optional[torch.Tensor] = None,
   395→        speed: float = 1.0,
   396→        t_shift: float = 1.0,
   397→        duration: str = "predict",
   398→        num_step: int = 5,
   399→        guidance_scale: float = 0.5,
   400→    ) -> torch.Tensor:
   401→        """
   402→        Generate acoustic features, given text tokens, prompts feature
   403→            and prompt transcription's text tokens.
   404→        Args:
   405→            tokens: a list of list of text tokens.
   406→            prompt_tokens: a list of list of prompt tokens.
   407→            prompt_features: the prompt feature with the shape
   408→                (batch_size, seq_len, feat_dim).
   409→            prompt_features_lens: the length of each prompt feature,
   410→                with the shape (batch_size,).
   411→            features_lens: the length of the predicted eature, with the
   412→                shape (batch_size,). It is used only when duration is "real".
   413→            duration: "real" or "predict". If "real", the predicted
   414→                feature length is given by features_lens.
   415→            num_step: the number of steps to use in the ODE solver.
   416→            guidance_scale: the guidance scale for classifier-free guidance.
   417→        """
   418→
   419→        assert duration in ["real", "predict"]
   420→
   421→        if duration == "predict":
   422→            (
   423→                text_condition,
   424→                padding_mask,
   425→            ) = self.forward_text_inference_ratio_duration(
   426→                tokens=tokens,
   427→                prompt_tokens=prompt_tokens,
   428→                prompt_features_lens=prompt_features_lens,
   429→                speed=speed,
   430→            )
   431→        else:
   432→            assert features_lens is not None
   433→            text_condition, padding_mask = self.forward_text_inference_gt_duration(
   434→                tokens=tokens,
   435→                features_lens=features_lens,
   436→                prompt_tokens=prompt_tokens,
   437→                prompt_features_lens=prompt_features_lens,
   438→            )
   439→        batch_size, num_frames, _ = text_condition.shape
   440→
   441→        speech_condition = torch.nn.functional.pad(
   442→            prompt_features, (0, 0, 0, num_frames - prompt_features.size(1))
   443→        )  # (B, T, F)
   444→
   445→        # False means speech condition positions.
   446→        speech_condition_mask = make_pad_mask(prompt_features_lens, num_frames)
   447→        speech_condition = torch.where(

    def sample(
        self,
        tokens: List[List[int]],
        prompt_tokens: List[List[int]],
        prompt_features: torch.Tensor,
        prompt_features_lens: torch.Tensor,
        features_lens: Optional[torch.Tensor] = None,
        speed: float = 1.0,
        t_shift: float = 1.0,
        duration: str = "predict",
        num_step: int = 5,
        guidance_scale: float = 0.5,
    ) -> torch.Tensor:
        """
        Generate acoustic features, given text tokens, prompts feature
            and prompt transcription's text tokens.
        Args:
            tokens: a list of list of text tokens.
            prompt_tokens: a list of list of prompt tokens.
            prompt_features: the prompt feature with the shape
                (batch_size, seq_len, feat_dim).
            prompt_features_lens: the length of each prompt feature,
                with the shape (batch_size,).
            features_lens: the length of the predicted eature, with the
                shape (batch_size,). It is used only when duration is "real".
            duration: "real" or "predict". If "real", the predicted
                feature length is given by features_lens.
            num_step: the number of steps to use in the ODE solver.
            guidance_scale: the guidance scale for classifier-free guidance.
        """

        assert duration in ["real", "predict"]

        if duration == "predict":
            (
                text_condition,
                padding_mask,
            ) = self.forward_text_inference_ratio_duration(
                tokens=tokens,
                prompt_tokens=prompt_tokens,
                prompt_features_lens=prompt_features_lens,
                speed=speed,
            )
        else:
            assert features_lens is not None
            text_condition, padding_mask = self.forward_text_inference_gt_duration(
                tokens=tokens,
                features_lens=features_lens,
                prompt_tokens=prompt_tokens,
                prompt_features_lens=prompt_features_lens,
            )
        batch_size, num_frames, _ = text_condition.shape

        speech_condition = torch.nn.functional.pad(
            prompt_features, (0, 0, 0, num_frames - prompt_features.size(1))
        )  # (B, T, F)

        # False means speech condition positions.
        speech_condition_mask = make_pad_mask(prompt_features_lens, num_frames)
        speech_condition = torch.where(


   447→        speech_condition = torch.where(
   448→            speech_condition_mask.unsqueeze(-1),
   449→            torch.zeros_like(speech_condition),
   450→            speech_condition,
   451→        )
   452→
   453→        x0 = torch.randn(
   454→            batch_size,
   455→            num_frames,
   456→            prompt_features.size(-1),
   457→            device=text_condition.device,
   458→        )
   459→
   460→        x1 = self.solver.sample(
   461→            x=x0,
   462→            text_condition=text_condition,
   463→            speech_condition=speech_condition,
   464→            padding_mask=padding_mask,
   465→            num_step=num_step,
   466→            guidance_scale=guidance_scale,
   467→            t_shift=t_shift,
   468→        )
   469→        x1_wo_prompt_lens = (~padding_mask).sum(-1) - prompt_features_lens
   470→        x1_prompt = torch.zeros(
   471→            x1.size(0), prompt_features_lens.max(), x1.size(2), device=x1.device
   472→        )
   473→        x1_wo_prompt = torch.zeros(
   474→            x1.size(0), x1_wo_prompt_lens.max(), x1.size(2), device=x1.device
   475→        )
   476→        for i in range(x1.size(0)):
   477→            x1_wo_prompt[i, : x1_wo_prompt_lens[i], :] = x1[
   478→                i,
   479→                prompt_features_lens[i] : prompt_features_lens[i]
   480→                + x1_wo_prompt_lens[i],
   481→            ]
   482→            x1_prompt[i, : prompt_features_lens[i], :] = x1[
   483→                i, : prompt_features_lens[i]
   484→            ]
   485→
   486→        return x1_wo_prompt, x1_wo_prompt_lens, x1_prompt, prompt_features_lens
   487→
   488→    def sample_intermediate(
   489→        self,
   490→        tokens: List[List[int]],
   491→        features: torch.Tensor,
   492→        features_lens: torch.Tensor,
   493→        noise: torch.Tensor,
   494→        speech_condition_mask: torch.Tensor,
   495→        t_start: float,
   496→        t_end: float,

        speech_condition = torch.where(
            speech_condition_mask.unsqueeze(-1),
            torch.zeros_like(speech_condition),
            speech_condition,
        )

        x0 = torch.randn(
            batch_size,
            num_frames,
            prompt_features.size(-1),
            device=text_condition.device,
        )

        x1 = self.solver.sample(
            x=x0,
            text_condition=text_condition,
            speech_condition=speech_condition,
            padding_mask=padding_mask,
            num_step=num_step,
            guidance_scale=guidance_scale,
            t_shift=t_shift,
        )
        x1_wo_prompt_lens = (~padding_mask).sum(-1) - prompt_features_lens
        x1_prompt = torch.zeros(
            x1.size(0), prompt_features_lens.max(), x1.size(2), device=x1.device
        )
        x1_wo_prompt = torch.zeros(
            x1.size(0), x1_wo_prompt_lens.max(), x1.size(2), device=x1.device
        )
        for i in range(x1.size(0)):
            x1_wo_prompt[i, : x1_wo_prompt_lens[i], :] = x1[
                i,
                prompt_features_lens[i] : prompt_features_lens[i]
                + x1_wo_prompt_lens[i],
            ]
            x1_prompt[i, : prompt_features_lens[i], :] = x1[
                i, : prompt_features_lens[i]
            ]

        return x1_wo_prompt, x1_wo_prompt_lens, x1_prompt, prompt_features_lens

    def sample_intermediate(
        self,
        tokens: List[List[int]],
        features: torch.Tensor,
        features_lens: torch.Tensor,
        noise: torch.Tensor,
        speech_condition_mask: torch.Tensor,
        t_start: float,
        t_end: float,


   120→    def __init__(
   121→        self,
   122→        model: torch.nn.Module,
   123→        func_name: str = "forward_fm_decoder",
   124→    ):
   125→        super().__init__(model=model, func_name=func_name)
   126→
   127→    def forward(
   128→        self,
   129→        t: torch.Tensor,
   130→        x: torch.Tensor,
   131→        text_condition: torch.Tensor,
   132→        speech_condition: torch.Tensor,
   133→        padding_mask: Optional[torch.Tensor] = None,
   134→        guidance_scale: Union[float, torch.Tensor] = 0.0,
   135→        **kwargs
   136→    ) -> torch.Tensor:
   137→        """
   138→        Forward function that Handles the classifier-free guidance.
   139→        Args:
   140→            t: The current timestep, a tensor of a single float.
   141→            x: The initial value, with the shape (batch, seq_len, emb_dim).
   142→            text_condition: The text_condition of the diffision model, with
   143→                the shape (batch, seq_len, emb_dim).
   144→            speech_condition: The speech_condition of the diffision model, with the
   145→                shape (batch, seq_len, emb_dim).
   146→            padding_mask: The mask for padding; True means masked position, with the
   147→                shape (batch, seq_len).
   148→            guidance_scale: The scale of classifier-free guidance, a float or a tensor
   149→                of shape (batch, 1, 1).
   150→        Retrun:
   151→            The prediction with the shape (batch, seq_len, emb_dim).
   152→        """
   153→        if not torch.is_tensor(guidance_scale):
   154→            guidance_scale = torch.tensor(
   155→                guidance_scale, dtype=t.dtype, device=t.device
   156→            )
   157→        return self.model_func(
   158→            t=t,
   159→            xt=x,
   160→            text_condition=text_condition,
   161→            speech_condition=speech_condition,
   162→            padding_mask=padding_mask,
   163→            guidance_scale=guidance_scale,
   164→            **kwargs
   165→        )
   166→
   167→
   168→class EulerSolver:
   169→    def __init__(
   170→        self,
   171→        model: torch.nn.Module,
   172→        func_name: str = "forward_fm_decoder",
   173→    ):
   174→        """Construct a Euler Solver
   175→        Args:
   176→            model: The diffusion model.
   177→            func_name: The function name to call.
   178→        """
   179→        self.model = DiffusionModel(model, func_name=func_name)

    def __init__(
        self,
        model: torch.nn.Module,
        func_name: str = "forward_fm_decoder",
    ):
        super().__init__(model=model, func_name=func_name)

    def forward(
        self,
        t: torch.Tensor,
        x: torch.Tensor,
        text_condition: torch.Tensor,
        speech_condition: torch.Tensor,
        padding_mask: Optional[torch.Tensor] = None,
        guidance_scale: Union[float, torch.Tensor] = 0.0,
        **kwargs
    ) -> torch.Tensor:
        """
        Forward function that Handles the classifier-free guidance.
        Args:
            t: The current timestep, a tensor of a single float.
            x: The initial value, with the shape (batch, seq_len, emb_dim).
            text_condition: The text_condition of the diffision model, with
                the shape (batch, seq_len, emb_dim).
            speech_condition: The speech_condition of the diffision model, with the
                shape (batch, seq_len, emb_dim).
            padding_mask: The mask for padding; True means masked position, with the
                shape (batch, seq_len).
            guidance_scale: The scale of classifier-free guidance, a float or a tensor
                of shape (batch, 1, 1).
        Retrun:
            The prediction with the shape (batch, seq_len, emb_dim).
        """
        if not torch.is_tensor(guidance_scale):
            guidance_scale = torch.tensor(
                guidance_scale, dtype=t.dtype, device=t.device
            )
        return self.model_func(
            t=t,
            xt=x,
            text_condition=text_condition,
            speech_condition=speech_condition,
            padding_mask=padding_mask,
            guidance_scale=guidance_scale,
            **kwargs
        )

class EulerSolver:
    def __init__(
        self,
        model: torch.nn.Module,
        func_name: str = "forward_fm_decoder",
    ):
        """Construct a Euler Solver
        Args:
            model: The diffusion model.
            func_name: The function name to call.
        """
        self.model = DiffusionModel(model, func_name=func_name)


   180→
   181→    def sample(
   182→        self,
   183→        x: torch.Tensor,
   184→        text_condition: torch.Tensor,
   185→        speech_condition: torch.Tensor,
   186→        padding_mask: torch.Tensor,
   187→        num_step: int = 10,
   188→        guidance_scale: Union[float, torch.Tensor] = 0.0,
   189→        t_start: float = 0.0,
   190→        t_end: float = 1.0,
   191→        t_shift: float = 1.0,
   192→        **kwargs
   193→    ) -> torch.Tensor:
   194→        device = x.device
   195→        assert isinstance(t_start, float) and isinstance(t_end, float)
   196→
   197→        # Generate the schedule of timesteps
   198→        timesteps = get_time_steps(
   199→            t_start=t_start,
   200→            t_end=t_end,
   201→            num_step=num_step,
   202→            t_shift=t_shift,
   203→            device=device,
   204→        )
   205→
   206→        for step in range(num_step):
   207→            t_cur = timesteps[step]
   208→            t_next = timesteps[step + 1]
   209→
   210→            # Predict velocity (v)
   211→            v = self.model(
   212→                t=t_cur,
   213→                x=x,
   214→                text_condition=text_condition,
   215→                speech_condition=speech_condition,
   216→                padding_mask=padding_mask,
   217→                guidance_scale=guidance_scale,
   218→                **kwargs
   219→            )
   220→
   221→            # 1. Predict the clean 'data' (x_1) and 'noise' (x_0)
   222→            # Flow matching formulation: x_t = (1 - t) * x_0 + t * x_1
   223→            # Therefore: v = x_1 - x_0
   224→            x_1_pred = x + (1.0 - t_cur) * v
   225→            x_0_pred = x - t_cur * v
   226→
   227→            if step < num_step - 1:
   228→                # 2. Probability Flow ODE update (Anchor-based)
   229→                # This 'anchors' the next point along the predicted line,
   230→                # making it more robust than simple Euler integration.
   231→                x = (1.0 - t_next) * x_0_pred + t_next * x_1_pred
   232→            else:
   233→                # Final step: Snap directly to the predicted clean data
   234→                x = x_1_pred
   235→
   236→        return x
   237→
   238→
   239→class DistillEulerSolver(EulerSolver):
   240→    def __init__(
   241→        self,
   242→        model: torch.nn.Module,
   243→        func_name: str = "forward_fm_decoder",
   244→    ):
   245→        """Construct a Euler Solver for distilled diffusion models.
   246→        Args:
   247→            model: The diffusion model.
   248→        """
   249→        self.model = DistillDiffusionModel(model, func_name=func_name)
   250→
   251→
   252→def get_time_steps(
   253→    t_start: float = 0.0,
   254→    t_end: float = 1.0,
   255→    num_step: int = 10,
   256→    t_shift: float = 1.0,
   257→    device: torch.device = torch.device("cpu"),
   258→) -> torch.Tensor:
   259→    """Compute the intermediate time steps for sampling.

    def sample(
        self,
        x: torch.Tensor,
        text_condition: torch.Tensor,
        speech_condition: torch.Tensor,
        padding_mask: torch.Tensor,
        num_step: int = 10,
        guidance_scale: Union[float, torch.Tensor] = 0.0,
        t_start: float = 0.0,
        t_end: float = 1.0,
        t_shift: float = 1.0,
        **kwargs
    ) -> torch.Tensor:
        device = x.device
        assert isinstance(t_start, float) and isinstance(t_end, float)

        # Generate the schedule of timesteps
        timesteps = get_time_steps(
            t_start=t_start,
            t_end=t_end,
            num_step=num_step,
            t_shift=t_shift,
            device=device,
        )

        for step in range(num_step):
            t_cur = timesteps[step]
            t_next = timesteps[step + 1]

            # Predict velocity (v)
            v = self.model(
                t=t_cur,
                x=x,
                text_condition=text_condition,
                speech_condition=speech_condition,
                padding_mask=padding_mask,
                guidance_scale=guidance_scale,
                **kwargs
            )

            # 1. Predict the clean 'data' (x_1) and 'noise' (x_0)
            # Flow matching formulation: x_t = (1 - t) * x_0 + t * x_1
            # Therefore: v = x_1 - x_0
            x_1_pred = x + (1.0 - t_cur) * v
            x_0_pred = x - t_cur * v

            if step < num_step - 1:
                # 2. Probability Flow ODE update (Anchor-based)
                # This 'anchors' the next point along the predicted line,
                # making it more robust than simple Euler integration.
                x = (1.0 - t_next) * x_0_pred + t_next * x_1_pred
            else:
                # Final step: Snap directly to the predicted clean data
                x = x_1_pred

        return x

class DistillEulerSolver(EulerSolver):
    def __init__(
        self,
        model: torch.nn.Module,
        func_name: str = "forward_fm_decoder",
    ):
        """Construct a Euler Solver for distilled diffusion models.
        Args:
            model: The diffusion model.
        """
        self.model = DistillDiffusionModel(model, func_name=func_name)

def get_time_steps(
    t_start: float = 0.0,
    t_end: float = 1.0,
    num_step: int = 10,
    t_shift: float = 1.0,
    device: torch.device = torch.device("cpu"),
) -> torch.Tensor:
    """Compute the intermediate time steps for sampling.


Now let me check what `generate()` passes — is it always batch_size=1?


    64→def generate(prompt_tokens, prompt_features_lens, prompt_features, prompt_rms, text, model, vocoder, tokenizer, num_step=4, guidance_scale=3.0, speed=1.0, t_shift=0.5, target_rms=0.1):
    65→    tokens = tokenizer.texts_to_token_ids([text])
    66→    device = next(model.parameters()).device  # Auto-detect device
    67→
    68→    speed = speed * 1.3
    69→
    70→    with torch.inference_mode():
    71→        (pred_features, _, _, _) = model.sample(
    72→            tokens=tokens,
    73→            prompt_tokens=prompt_tokens,
    74→            prompt_features=prompt_features,
    75→            prompt_features_lens=prompt_features_lens,
    76→            speed=speed,
    77→            t_shift=t_shift,
    78→            duration='predict',
    79→            num_step=num_step,
    80→            guidance_scale=guidance_scale,
    81→        )
    82→
    83→    # Convert to waveform
    84→    pred_features = pred_features.permute(0, 2, 1) / 0.1
    85→    wav = vocoder.decode(pred_features).squeeze(1).clamp(-1, 1)
    86→
    87→    # Volume matching
    88→    if prompt_rms < target_rms:
    89→        wav = wav * (prompt_rms / target_rms)
    90→
    91→    return wav

def generate(prompt_tokens, prompt_features_lens, prompt_features, prompt_rms, text, model, vocoder, tokenizer, num_step=4, guidance_scale=3.0, speed=1.0, t_shift=0.5, target_rms=0.1):
    tokens = tokenizer.texts_to_token_ids([text])
    device = next(model.parameters()).device  # Auto-detect device

    speed = speed * 1.3

    with torch.inference_mode():
        (pred_features, _, _, _) = model.sample(
            tokens=tokens,
            prompt_tokens=prompt_tokens,
            prompt_features=prompt_features,
            prompt_features_lens=prompt_features_lens,
            speed=speed,
            t_shift=t_shift,
            duration='predict',
            num_step=num_step,
            guidance_scale=guidance_scale,
        )

    # Convert to waveform
    pred_features = pred_features.permute(0, 2, 1) / 0.1
    wav = vocoder.decode(pred_features).squeeze(1).clamp(-1, 1)

    # Volume matching
    if prompt_rms < target_rms:
        wav = wav * (prompt_rms / target_rms)

    return wav


Here's how it works — **no batching at all**:

## Current Architecture (batch_size=1 always)

```
Request → tokenize([text])     →  batch=1 list
       → model.sample()        →  (B=1, T, 100) mel features
         └─ text_encoder       →  encode text + prompt tokens
         └─ duration_predictor  →  predict # of mel frames
         └─ EulerSolver.sample  →  4 steps of flow matching
            └─ fm_decoder(xt, t, padding_mask)  ← 118.5M params, runs 4x
       → vocoder.decode()      →  mel → 48kHz waveform
```

**Every request = 1 item, processed alone.** The `(B, T, F)` tensors all have B=1. The fm_decoder (118.5M param Zipformer) runs 4 times per request, each time processing a single sequence.

## Why it's "fast" but not 150x

The A100 is **massively underutilized** at B=1. The model is only 122M params, using 1.76 GB of 80 GB. The GPU compute cores are mostly idle waiting for memory transfers.

## How to actually hit 150x+ and 500 concurrency

There are 3 approaches, from easiest to hardest:

| Approach | How | Expected Speed | Concurrency |
|---|---|---|---|
| **Multi-process replicas** | Spawn 40 independent processes, each B=1 | 74x each, ~120 req/s total | ~120 |
| **Batched inference** | Modify `generate()` to accept B>1, pad sequences, run fm_decoder on batch | 150x+ per batch | 500+ on 1 GPU |
| **TensorRT** | Export fm_decoder to TRT (code already exists in `tensorrt.py`), fp16 fused kernels | 200x+ | 500+ with batching |

**Batching is the key.** The fm_decoder already supports `(B, T, F)` tensors — the architecture is batch-ready. The bottleneck is just that `generate()` wraps everything at B=1. A batched version would:

1. Collect N requests into a batch
2. Pad token sequences to same length
3. Run fm_decoder once on `(N, T_max, 100)` — GPU fully utilized
4. Split output back per request

With batch=16-32 on A100, each fm_decoder call processes 16-32 requests simultaneously. That's where 150x+ comes from.

Want me to build the batched inference version?


then concutency is 32 only?


No — concurrency and batch size are different things.

**Batch size = how many requests the GPU processes at once** (e.g., 32)

**Concurrency = how many requests are waiting/in-flight** (e.g., 500)

It works like this:

```
500 concurrent requests
        │
        ▼
   Request Queue
        │
        ▼
  Batch Scheduler (grabs 32 at a time)
        │
        ▼
   GPU: process batch of 32 in ~340ms
        │
        ▼
   Return 32 results, grab next 32
```

**Math:**
- Batch of 32 takes ~340ms (same wall time as B=1 for long text — that's the point of batching)
- Throughput: 32 / 0.34s = **~94 req/s**
- 500 requests drain in ~5.3s
- All 500 complete, just with some queuing latency

**To increase throughput further:**
- Bigger batch (B=64): ~188 req/s → 500 in ~2.7s
- TensorRT + B=64: could hit 300+ req/s
- Multiple CUDA streams processing batches in parallel

So **500 concurrency is fine** — it just means some requests wait in the queue. The p50 latency for 500 concurrent would be ~2-3s, which is totally acceptable for TTS.

Want me to build the batched inference + queue server?


is stremaing ther ?


No streaming in LuxTTS. It generates the **entire audio at once** — the flow matching solver needs all 4 steps to complete before you get any audio.

```
Current: [wait 340ms...] → full audio out at once
Streaming: [50ms] chunk1 → [50ms] chunk2 → [50ms] chunk3 → ...
```

The architecture makes streaming hard because:
1. **Flow matching is non-autoregressive** — it generates all frames simultaneously, not left-to-right
2. **Duration prediction** needs full text upfront to decide total length
3. **Vocoder** (Vocos) also runs on the full mel at once

## Options for streaming-like behavior

| Approach | Latency | Complexity |
|---|---|---|
| **Chunked text splitting** — split text into sentences, generate each separately, stream chunks | First chunk ~200ms | Easy |
| **Overlap-add chunking** — split text into overlapping segments, crossfade outputs | First chunk ~200ms | Medium |
| **Speculative streaming** — start vocoding partial mel while solver still running | ~100ms | Hard, needs model surgery |

The **chunked approach** is the most practical — split input at sentence boundaries, fire each chunk as a separate batch item, and stream results as they complete:

```
"Hello world. How are you. Nice day."
  → chunk1: "Hello world."        → audio1 (200ms) → stream immediately
  → chunk2: "How are you."        → audio2 (200ms) → stream when ready
  → chunk3: "Nice day."           → audio3 (200ms) → stream when ready
```

With batching, all 3 chunks can run in **one batch** — so you get first audio at ~200ms and all audio by ~200ms (same batch).

Want me to build this chunked streaming + batched server?


do setup the bacthing infernce optimized once .. we can tetx it


Let me first understand all the pieces I need to batch, then build it.