build_engine.py

#!/usr/bin/env python3
"""
Orpheus TTS TensorRT-LLM Engine Builder for RTX 4090

This script downloads the Orpheus TTS model and builds an optimized
TensorRT-LLM engine with FP8 quantization for RTX 4090.

Orpheus TTS is based on Llama 3.2 3B architecture.

Follow the TensorRT-LLM Llama guide:
1. Download/convert HF checkpoint
2. Quantize to FP8 (optimal for 4090)
3. Build TensorRT engine
4. Serve with separate LLM + SNAC decoder processes

Usage:
    python build_engine.py --config config/model_config.yaml
    
    # Step by step:
    python build_engine.py --step download
    python build_engine.py --step convert
    python build_engine.py --step build
    python build_engine.py --step verify
"""

import os
import sys
import yaml
import argparse
import subprocess
import shutil
import logging
from pathlib import Path
from typing import Optional, Dict, Any

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


def load_config(config_path: str) -> dict:
    """Load configuration from YAML file."""
    with open(config_path, 'r') as f:
        return yaml.safe_load(f)


def run_command(cmd: list, cwd: str = None, env: dict = None) -> subprocess.CompletedProcess:
    """Run a command and handle errors."""
    logger.info(f"Running: {' '.join(cmd)}")
    
    full_env = os.environ.copy()
    if env:
        full_env.update(env)
    
    result = subprocess.run(
        cmd,
        cwd=cwd,
        env=full_env,
        capture_output=True,
        text=True
    )
    
    if result.returncode != 0:
        logger.error(f"Command failed with code {result.returncode}")
        logger.error(f"STDOUT: {result.stdout}")
        logger.error(f"STDERR: {result.stderr}")
        raise RuntimeError(f"Command failed: {' '.join(cmd)}")
    
    if result.stdout:
        logger.debug(result.stdout)
    
    return result


def check_tensorrt_llm():
    """Check TensorRT-LLM installation."""
    try:
        import tensorrt_llm
        logger.info(f"TensorRT-LLM version: {tensorrt_llm.__version__}")
        return True
    except ImportError:
        logger.error("TensorRT-LLM not installed!")
        logger.error("Install with: pip install tensorrt_llm -U --pre --extra-index-url https://pypi.nvidia.com")
        return False


def check_gpu():
    """Check GPU availability and capabilities."""
    try:
        import torch
        if not torch.cuda.is_available():
            logger.error("CUDA not available!")
            return False
        
        gpu_name = torch.cuda.get_device_name(0)
        gpu_memory = torch.cuda.get_device_properties(0).total_memory / (1024**3)
        compute_cap = torch.cuda.get_device_capability(0)
        
        logger.info(f"GPU: {gpu_name}")
        logger.info(f"Memory: {gpu_memory:.1f} GB")
        logger.info(f"Compute Capability: {compute_cap[0]}.{compute_cap[1]}")
        
        # Check for FP8 support (SM 89+ for Ada Lovelace / RTX 4090)
        if compute_cap[0] >= 8 and compute_cap[1] >= 9:
            logger.info("✓ FP8 support: YES (Ada Lovelace - RTX 4090)")
        elif compute_cap[0] >= 9:
            logger.info("✓ FP8 support: YES (Hopper+)")
        else:
            logger.warning("⚠ FP8 support: Limited - will use FP16 fallback")
        
        return True
    except Exception as e:
        logger.error(f"GPU check failed: {e}")
        return False


def step_download(config: dict, base_dir: Path):
    """Step 1: Download Orpheus TTS model from HuggingFace."""
    from huggingface_hub import snapshot_download
    
    model_config = config['model']
    model_name = model_config['name']
    output_dir = base_dir / model_config['hf_model_path']
    
    logger.info(f"Downloading model: {model_name}")
    
    if output_dir.exists() and list(output_dir.glob("*.safetensors")):
        logger.info(f"Model already exists at {output_dir}")
        return str(output_dir)
    
    output_dir.mkdir(parents=True, exist_ok=True)
    
    local_dir = snapshot_download(
        repo_id=model_name,
        local_dir=str(output_dir),
        local_dir_use_symlinks=False,
    )
    
    logger.info(f"Model downloaded to: {local_dir}")
    return local_dir


def step_convert_checkpoint(config: dict, base_dir: Path):
    """
    Step 2: Convert HuggingFace checkpoint to TensorRT-LLM format with FP8 quantization.
    
    This uses TensorRT-LLM's convert_checkpoint.py for Llama models.
    For FP8, we use quantization flags to enable FP8 weights.
    """
    model_config = config['model']
    engine_config = config['engine']
    
    hf_model_path = base_dir / model_config['hf_model_path']
    checkpoint_path = base_dir / model_config['checkpoint_path']
    
    logger.info("Converting checkpoint to TensorRT-LLM format...")
    
    if checkpoint_path.exists() and (list(checkpoint_path.glob("*.safetensors")) or list(checkpoint_path.glob("*.bin"))):
        logger.info(f"Checkpoint already exists at {checkpoint_path}")
        return str(checkpoint_path)
    
    checkpoint_path.mkdir(parents=True, exist_ok=True)
    
    # Build conversion command
    cmd = [
        sys.executable, "-m", "tensorrt_llm.commands.convert_checkpoint",
        "--model_dir", str(hf_model_path),
        "--output_dir", str(checkpoint_path),
        "--dtype", engine_config['dtype'],
        "--tp_size", str(engine_config['tp_size']),
        "--pp_size", str(engine_config['pp_size']),
    ]
    
    # Add FP8 quantization flags for RTX 4090
    if engine_config.get('quantization') == 'fp8':
        logger.info("Enabling FP8 quantization (optimal for RTX 4090)")
        # Use FP8 rowwise quantization
        cmd.append("--use_fp8_rowwise")
        # Enable FP8 KV cache
        if engine_config.get('use_fp8_kv_cache', True):
            cmd.append("--fp8_kv_cache")
    
    try:
        run_command(cmd, cwd=str(base_dir))
    except RuntimeError as e:
        logger.warning(f"Standard conversion failed, trying alternative method: {e}")
        # Fallback: try without FP8 flags if they're not supported
        cmd_fallback = [
            sys.executable, "-m", "tensorrt_llm.commands.convert_checkpoint",
            "--model_dir", str(hf_model_path),
            "--output_dir", str(checkpoint_path),
            "--dtype", engine_config['dtype'],
            "--tp_size", str(engine_config['tp_size']),
            "--pp_size", str(engine_config['pp_size']),
        ]
        run_command(cmd_fallback, cwd=str(base_dir))
    
    logger.info(f"Checkpoint converted to: {checkpoint_path}")
    return str(checkpoint_path)


def step_build_engine(config: dict, base_dir: Path):
    """
    Step 3: Build TensorRT-LLM engine from checkpoint.
    
    Uses trtllm-build CLI for Llama models with FP8 optimizations.
    """
    model_config = config['model']
    engine_config = config['engine']
    
    checkpoint_path = base_dir / model_config['checkpoint_path']
    engine_path = base_dir / model_config['trt_engine_path']
    
    logger.info("Building TensorRT-LLM engine with FP8 optimizations...")
    
    if engine_path.exists() and list(engine_path.glob("*.engine")):
        logger.info(f"Engine already exists at {engine_path}")
        return str(engine_path)
    
    engine_path.mkdir(parents=True, exist_ok=True)
    
    # Build command
    cmd = [
        "trtllm-build",
        "--checkpoint_dir", str(checkpoint_path),
        "--output_dir", str(engine_path),
        "--max_batch_size", str(engine_config['max_batch_size']),
        "--max_input_len", str(engine_config['max_input_len']),
        "--max_seq_len", str(engine_config['max_input_len'] + engine_config['max_output_len']),
        "--max_num_tokens", str(engine_config.get('max_num_tokens', 8192)),
    ]
    
    # Plugin configurations
    gpt_plugin = engine_config.get('use_gpt_attention_plugin', 'float16')
    gemm_plugin = engine_config.get('use_gemm_plugin', 'float16')
    
    cmd.extend(["--gpt_attention_plugin", gpt_plugin])
    cmd.extend(["--gemm_plugin", gemm_plugin])
    
    # Memory optimizations
    if engine_config.get('paged_kv_cache', True):
        cmd.extend(["--paged_kv_cache", "enable"])
        if engine_config.get('tokens_per_block'):
            cmd.extend(["--tokens_per_block", str(engine_config['tokens_per_block'])])
    
    if engine_config.get('remove_input_padding', True):
        cmd.extend(["--remove_input_padding", "enable"])
    
    if engine_config.get('context_fmha', True):
        cmd.extend(["--context_fmha", "enable"])
    
    if engine_config.get('multi_block_mode', True):
        cmd.extend(["--multi_block_mode", "enable"])
    
    # FP8 context FMHA for better performance
    if engine_config.get('use_fp8_kv_cache', False):
        cmd.extend(["--use_fp8_context_fmha", "enable"])
    
    # Builder optimization level
    if engine_config.get('builder_opt_level'):
        cmd.extend(["--builder_opt", str(engine_config['builder_opt_level'])])
    
    # Strongly typed for better performance
    if engine_config.get('strongly_typed', True):
        cmd.append("--strongly_typed")
    
    try:
        run_command(cmd, cwd=str(base_dir))
    except RuntimeError as e:
        logger.warning(f"Build with all options failed, trying minimal build: {e}")
        # Fallback: minimal build
        cmd_minimal = [
            "trtllm-build",
            "--checkpoint_dir", str(checkpoint_path),
            "--output_dir", str(engine_path),
            "--max_batch_size", str(engine_config['max_batch_size']),
            "--max_input_len", str(engine_config['max_input_len']),
            "--max_seq_len", str(engine_config['max_input_len'] + engine_config['max_output_len']),
            "--gpt_attention_plugin", "float16",
            "--gemm_plugin", "float16",
        ]
        run_command(cmd_minimal, cwd=str(base_dir))
    
    logger.info(f"Engine built at: {engine_path}")
    return str(engine_path)


def step_verify_engine(config: dict, base_dir: Path):
    """Step 4: Verify the built engine works correctly."""
    model_config = config['model']
    
    engine_path = base_dir / model_config['trt_engine_path']
    tokenizer_path = base_dir / model_config['hf_model_path']
    
    logger.info("Verifying engine...")
    
    try:
        import torch
        from transformers import AutoTokenizer
        from tensorrt_llm.runtime import ModelRunner
        
        # Load tokenizer
        tokenizer = AutoTokenizer.from_pretrained(
            str(tokenizer_path),
            trust_remote_code=True
        )
        
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token
        
        # Load engine
        runner = ModelRunner.from_dir(
            engine_dir=str(engine_path),
            rank=0,
        )
        
        # Test prompt with Orpheus format
        test_text = "<|audio|>tara<|/audio|>Hello, this is a test of the Orpheus TTS system."
        input_ids = tokenizer.encode(test_text, return_tensors="pt")
        input_ids = input_ids.cuda()
        
        # Generate
        with torch.no_grad():
            outputs = runner.generate(
                batch_input_ids=input_ids,
                max_new_tokens=100,
                end_id=tokenizer.eos_token_id,
                pad_id=tokenizer.pad_token_id,
                temperature=0.6,
                top_p=0.95,
            )
        
        output_ids = outputs[0][0].tolist()
        num_generated = len(output_ids) - input_ids.shape[1]
        
        logger.info(f"Generated {num_generated} tokens")
        logger.info("✓ Engine verification PASSED!")
        
        return True
        
    except Exception as e:
        logger.error(f"✗ Engine verification FAILED: {e}")
        import traceback
        traceback.print_exc()
        return False


def download_snac_model(base_dir: Path):
    """Download SNAC audio decoder model for Orpheus."""
    from huggingface_hub import snapshot_download
    
    snac_path = base_dir / "models" / "snac_24khz"
    
    if snac_path.exists() and list(snac_path.glob("*.pt")):
        logger.info(f"SNAC model already exists at {snac_path}")
        return str(snac_path)
    
    logger.info("Downloading SNAC audio decoder (24kHz)...")
    
    snac_path.mkdir(parents=True, exist_ok=True)
    
    local_dir = snapshot_download(
        repo_id="hubertsiuzdak/snac_24khz",
        local_dir=str(snac_path),
        local_dir_use_symlinks=False,
    )
    
    logger.info(f"SNAC model downloaded to: {local_dir}")
    return local_dir


def main():
    parser = argparse.ArgumentParser(
        description="Build TensorRT-LLM engine for Orpheus TTS (FP8 on RTX 4090)"
    )
    parser.add_argument(
        "--config",
        type=str,
        default="config/model_config.yaml",
        help="Path to config file"
    )
    parser.add_argument(
        "--step",
        type=str,
        choices=["all", "download", "convert", "build", "verify", "snac"],
        default="all",
        help="Which step to run"
    )
    parser.add_argument(
        "--skip-checks",
        action="store_true",
        help="Skip GPU and TensorRT-LLM checks"
    )
    parser.add_argument(
        "--force",
        action="store_true",
        help="Force rebuild even if files exist"
    )
    
    args = parser.parse_args()
    
    # Load configuration
    base_dir = Path(__file__).parent.resolve()
    config = load_config(args.config)
    
    logger.info("=" * 60)
    logger.info("Orpheus TTS TensorRT-LLM Engine Builder")
    logger.info("Optimized for RTX 4090 with FP8 quantization")
    logger.info("=" * 60)
    
    # Pre-flight checks
    if not args.skip_checks:
        if not check_tensorrt_llm():
            logger.error("Please install TensorRT-LLM first:")
            logger.error("  pip install tensorrt_llm -U --pre --extra-index-url https://pypi.nvidia.com")
            sys.exit(1)
        if not check_gpu():
            sys.exit(1)
    
    # Create directories
    for dir_name in ["models", "checkpoints", "engines", "logs"]:
        (base_dir / dir_name).mkdir(parents=True, exist_ok=True)
    
    # Force rebuild if requested
    if args.force:
        model_config = config['model']
        for path_key in ['checkpoint_path', 'trt_engine_path']:
            path = base_dir / model_config[path_key]
            if path.exists():
                logger.info(f"Removing existing: {path}")
                shutil.rmtree(path)
    
    # Run steps
    steps_to_run = []
    if args.step == "all":
        steps_to_run = ["download", "snac", "convert", "build", "verify"]
    else:
        steps_to_run = [args.step]
    
    for step in steps_to_run:
        logger.info(f"\n{'='*60}")
        logger.info(f"Step: {step.upper()}")
        logger.info("=" * 60)
        
        try:
            if step == "download":
                step_download(config, base_dir)
            elif step == "snac":
                download_snac_model(base_dir)
            elif step == "convert":
                step_convert_checkpoint(config, base_dir)
            elif step == "build":
                step_build_engine(config, base_dir)
            elif step == "verify":
                if not step_verify_engine(config, base_dir):
                    logger.warning("Verification failed but continuing...")
        except Exception as e:
            logger.error(f"Step {step} failed: {e}")
            if step in ["convert", "build"]:
                logger.error("This is a critical step. Please check the error above.")
                sys.exit(1)
    
    logger.info("\n" + "=" * 60)
    logger.info("BUILD COMPLETE!")
    logger.info("=" * 60)
    logger.info(f"Engine path: {base_dir / config['model']['trt_engine_path']}")
    logger.info("")
    logger.info("Next steps:")
    logger.info("  1. Start the streaming server (LLM + SNAC in separate processes):")
    logger.info("     python streaming_server.py --config config/model_config.yaml")
    logger.info("")
    logger.info("  2. Or start individual components:")
    logger.info("     python llm_process.py --config config/model_config.yaml")
    logger.info("     python snac_decoder.py --config config/model_config.yaml")
    logger.info("")
    logger.info("  3. Test with client:")
    logger.info("     python client.py --text 'Hello world!' --output test.wav")


if __name__ == "__main__":
    main()