moe2382/dramabox-dit-int8

DramaBox DiT INT8 — Selective Weight-Only Quantization

A selectively quantized version of the DramaBox TTS 3.3B DiT (Diffusion Transformer) model from Resemble AI. Reduces VRAM by 20% and checkpoint size by 45% while preserving audio quality.

Base model: ResembleAI/Dramabox | Code: resemble-ai/DramaBox | Architecture: LTX-2.3 DiT + Gemma 3 12B

What's included

You still need the other components from ResembleAI/Dramabox:

• —dramabox-audio-components.safetensors (1.9 GB) — VAE + vocoder
• —unsloth/gemma-3-12b-it-bnb-4bit (~8 GB) — text encoder

Results

MCD (Mel-Cepstral Distortion) measures spectral distance from the BF16 baseline. Lower is better. Scores below 5.0 dB are perceptually near-identical for speech.

Quantization details

Method: Selective INT8 weight-only quantization via torchao Int8WeightOnlyConfig. Weights are stored as INT8 with per-channel BF16 scales and dequantized at runtime during matrix multiplication.

What's quantized (562 layers, ~81.5% of DiT parameters):

• —All attention projections (to_q, to_k, to_v, to_out) across all 48 transformer blocks
• —All gate_logits layers
• —All FFN GELU projections (audio_ff.net.0.proj) across all 48 blocks
• —FFN output projections (audio_ff.net.2) in blocks 15–47, excluding block 17
• —Input/output projections (audio_patchify_proj, audio_proj_out)

What's NOT quantized (kept in BF16):

• —All normalization layers — extremely sensitive to precision changes
• —AdaLN conditioning layers — controls the diffusion process globally
• —Timestep embedder — conditioning pathway, highly sensitive
• —FFN output projections in blocks 0–14 — early blocks are most sensitive to quantization
• —FFN output projection in block 17 — anomalously sensitive individual block

This layer map was discovered through 80+ automated experiments using Andrej Karpathy's auto-research methodology, systematically testing each layer type and block index.

Usage

Option 1: Runtime quantization (simplest, no extra downloads)

If you just want VRAM savings without downloading this checkpoint, you can apply quantization at load time to the original DramaBox model:

import torch, re
from torchao.quantization import quantize_, Int8WeightOnlyConfig

# After loading the standard DramaBox TTSServer:
attn_proj_keys = ("to_q", "to_k", "to_v", "to_out")

def dit_filter(mod, fqn):
    if not isinstance(mod, torch.nn.Linear): return False
    if "norm" in fqn: return False
    if "gate_logits" in fqn: return True
    if any(k in fqn for k in attn_proj_keys): return True
    if "audio_ff" in fqn:
        m = re.search(r'transformer_blocks\.(\d+)\.', fqn)
        if m:
            idx = int(m.group(1))
            if "net.2" in fqn and idx >= 15 and idx != 17: return True
            if "net.0.proj" in fqn: return True
    return False

def io_filter(mod, fqn):
    return fqn in ("audio_patchify_proj", "audio_proj_out") and isinstance(mod, torch.nn.Linear)

quantize_(tts._velocity_model, Int8WeightOnlyConfig(), filter_fn=dit_filter)
quantize_(tts._velocity_model, Int8WeightOnlyConfig(), filter_fn=io_filter)

Option 2: Load pre-quantized weights (faster startup)

from load_int8 import load_int8_dit

# Loads the INT8 safetensors and reconstructs quantized Linear layers
load_int8_dit(tts._velocity_model, "dramabox-dit-int8-selective.safetensors")

Option 3: Full optimized pipeline with Gemma offload

For maximum VRAM savings (5.93 GB during denoising), use the included inference_optimized.py which also offloads Gemma 12B to CPU between text encoding and audio generation.

Requirements

• —PyTorch >= 2.4
• —torchao >= 0.15.0
• —CUDA GPU with >= 16 GB VRAM (14 GB with Gemma offload)
• —The original DramaBox model and its dependencies

How this was made

We ran 80+ experiments using an automated loop inspired by Karpathy's auto-research methodology:

1. 1.Start from the BF16 baseline
2. 2.Modify quantization config (which layers, which precision, which blocks)
3. 3.Generate 3 evaluation audio samples with fixed prompts/seeds
4. 4.Measure peak VRAM, generation time, and MCD vs baseline
5. 5.Keep the change if MCD < 5.0 dB, discard otherwise
6. 6.Repeat

Key findings from the search:

• —Flow-matching diffusion models are far more precision-sensitive than autoregressive LLMs. All 4-bit approaches (NF4, NVFP4, FP4, Int4) produced unacceptable quality (MCD 17–32 dB).
• —FP8 is worse than INT8 for weight representation in this model (MCD 11.8 vs 4.35).
• —`torch.compile` breaks audio output even on the unquantized baseline (MCD 24–32 dB). The iterative denoising loop is numerically sensitive to graph optimizations.
• —Early transformer blocks (0–14) are most sensitive in their FFN output projections. Block 17 is an outlier.
• —Attention projections and GELU gates are universally robust to INT8 across all 48 blocks.

Citation

If you use this work, please cite the original DramaBox model:

@misc{dramabox2025,
  title={DramaBox: Expressive Text to Speech Model},
  author={Resemble AI},
  year={2025},
  url={https://github.com/resemble-ai/DramaBox}
}

License

Same as the base DramaBox model — LTX-2 Community License.