profile_inference.py

#!/usr/bin/env python3
"""
ACE-Step 1.5 Inference Profiler & Benchmark

Comprehensive profiling tool that supports all features, devices, and backends.
Uses the high-level inference API and built-in time_costs for accurate timing.

Modes:
    profile         - Profile a single generation run with detailed timing breakdown
    benchmark       - Run a matrix of configurations and produce a summary table
    tier-test       - Auto-test across simulated GPU tiers (4/6/8/12/16/24/48 GB)
    understand      - Profile the understand_music() API (audio codes -> metadata)
    create_sample   - Profile the create_sample() API (inspiration/simple mode)
    format_sample   - Profile the format_sample() API (caption+lyrics -> metadata)

Usage:
    # Profile text2music with default settings
    python profile_inference.py

    # Profile with thinking enabled on MPS
    python profile_inference.py --device mps --thinking

    # Benchmark across configurations
    python profile_inference.py --mode benchmark

    # Test all GPU tiers automatically (the key feature!)
    python profile_inference.py --mode tier-test

    # Test specific tiers only
    python profile_inference.py --mode tier-test --tiers 6 8 16

    # Test tiers with LM enabled (where supported)
    python profile_inference.py --mode tier-test --tier-with-lm

    # Profile create_sample (inspiration mode)
    python profile_inference.py --mode create_sample --sample-query "a soft Bengali love song"

    # Profile understand mode
    python profile_inference.py --mode understand

    # Full profiling with cProfile
    python profile_inference.py --detailed --llm-debug
"""

import time
import argparse
import sys
import os
import json
import tempfile
import traceback
from contextlib import contextmanager
from collections import defaultdict
from typing import Tuple, Dict, Any, List, Optional

# Add project root to path
PROJECT_ROOT = os.path.abspath(os.path.dirname(__file__))
if PROJECT_ROOT not in sys.path:
    sys.path.insert(0, PROJECT_ROOT)

import torch

from acestep.inference import (
    generate_music,
    understand_music,
    create_sample,
    format_sample,
    GenerationParams,
    GenerationConfig,
    GenerationResult,
)
from acestep.handler import AceStepHandler
from acestep.llm_inference import LLMHandler
from acestep.gpu_config import (
    get_gpu_config,
    set_global_gpu_config,
    get_gpu_tier,
    find_best_lm_model_on_disk,
    is_lm_model_size_allowed,
    GPUConfig,
    VRAM_AUTO_OFFLOAD_THRESHOLD_GB,
)


# =============================================================================
# Device / Backend helpers
# =============================================================================


def resolve_device(device: str) -> str:
    """Resolve 'auto' device to the best available device."""
    if device == "auto":
        if hasattr(torch, "xpu") and torch.xpu.is_available():
            return "xpu"
        if torch.cuda.is_available():
            return "cuda"
        if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
            return "mps"
        return "cpu"
    return device


def auto_detect_backend(device: str) -> str:
    """Auto-detect the best LLM backend for the resolved device."""
    if device == "mps":
        try:
            import mlx.core  # noqa: F401
            return "mlx"
        except ImportError:
            return "pt"
    if device.startswith("cuda"):
        return "vllm"
    return "pt"


def load_env_config() -> Dict[str, str]:
    """Load configuration defaults from .env file."""
    env_config = {
        "ACESTEP_CONFIG_PATH": "acestep-v15-turbo",
        "ACESTEP_LM_MODEL_PATH": "acestep-5Hz-lm-0.6B",
        "ACESTEP_DEVICE": "auto",
        "ACESTEP_LM_BACKEND": "auto",
    }
    env_file = os.path.join(PROJECT_ROOT, ".env")
    if os.path.exists(env_file):
        with open(env_file, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line or line.startswith("#"):
                    continue
                if "=" in line:
                    key, value = line.split("=", 1)
                    key = key.strip()
                    value = value.strip()
                    if key in env_config and value:
                        env_config[key] = value
    return env_config


# =============================================================================
# Timer utilities
# =============================================================================


class PreciseTimer:
    """High-precision timer with GPU synchronization for accurate timing."""

    def __init__(self, device: str = "cpu"):
        self.device = device
        self.timings: Dict[str, List[float]] = defaultdict(list)
        self.enabled = True

    def sync(self):
        """Synchronize GPU operations for accurate timing."""
        if not self.enabled:
            return
        if self.device.startswith("cuda") and torch.cuda.is_available():
            torch.cuda.synchronize()
        elif self.device == "mps" and hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
            if hasattr(torch, "mps"):
                torch.mps.synchronize()
        elif self.device.startswith("xpu") and hasattr(torch, "xpu"):
            torch.xpu.synchronize()

    @contextmanager
    def time(self, name: str):
        """Time a code section with GPU synchronization."""
        if not self.enabled:
            yield
            return
        self.sync()
        start = time.perf_counter()
        try:
            yield
        finally:
            self.sync()
            elapsed = time.perf_counter() - start
            self.timings[name].append(elapsed)

    def get_total(self, name: str) -> float:
        return sum(self.timings.get(name, []))

    def get_mean(self, name: str) -> float:
        times = self.timings.get(name, [])
        return sum(times) / len(times) if times else 0.0

    def get_count(self, name: str) -> int:
        return len(self.timings.get(name, []))

    def reset(self):
        self.timings.clear()


# =============================================================================
# Example config loader
# =============================================================================


def load_example_config(
    example_file: str, cli_overrides: argparse.Namespace
) -> Tuple[Optional[GenerationParams], Optional[GenerationConfig]]:
    """Load configuration from example JSON file, applying CLI overrides."""
    try:
        with open(example_file, "r", encoding="utf-8") as f:
            data = json.load(f)

        params = GenerationParams(
            caption=data.get("caption", ""),
            lyrics=data.get("lyrics", ""),
            bpm=data.get("bpm"),
            keyscale=data.get("keyscale", ""),
            timesignature=data.get("timesignature", ""),
            vocal_language=data.get("language", "unknown"),
            duration=(
                cli_overrides.duration
                if cli_overrides.duration is not None
                else data.get("duration", -1.0)
            ),
            thinking=cli_overrides.thinking,
            use_cot_metas=cli_overrides.use_cot_metas,
            use_cot_caption=cli_overrides.use_cot_caption,
            use_cot_language=cli_overrides.use_cot_language,
            use_constrained_decoding=cli_overrides.use_constrained_decoding,
            inference_steps=(
                cli_overrides.inference_steps
                if cli_overrides.inference_steps is not None
                else data.get("inference_steps", 8)
            ),
            seed=(
                cli_overrides.seed
                if cli_overrides.seed is not None
                else data.get("seed", 42)
            ),
            task_type=cli_overrides.task_type,
            lm_temperature=cli_overrides.lm_temperature,
            lm_cfg_scale=cli_overrides.lm_cfg_scale,
            guidance_scale=cli_overrides.guidance_scale,
            reference_audio=cli_overrides.reference_audio,
            src_audio=cli_overrides.src_audio,
        )

        config = GenerationConfig(
            batch_size=(
                cli_overrides.batch_size
                if cli_overrides.batch_size is not None
                else data.get("batch_size", 1)
            ),
            seeds=[params.seed] if params.seed >= 0 else None,
            use_random_seed=(params.seed < 0),
            audio_format="flac",
        )

        return params, config

    except Exception as e:
        print(f"  Failed to load example: {e}")
        return None, None


# =============================================================================
# Printing helpers
# =============================================================================


def print_time_costs_breakdown(
    time_costs: Dict[str, float], total_wall_time: float
):
    """Print a detailed timing breakdown from result.extra_outputs['time_costs']."""
    print("\n" + "=" * 100)
    print("PROFILING RESULTS")
    print("=" * 100)

    if not time_costs:
        print("\n  (No time_costs data available from the pipeline)")
        print(f"\n  Total wall time: {total_wall_time:.3f}s")
        return

    # Categorize keys
    lm_keys = {
        k: v
        for k, v in time_costs.items()
        if k.startswith("lm_") and isinstance(v, (int, float))
    }
    dit_keys = {
        k: v
        for k, v in time_costs.items()
        if k.startswith("dit_") and isinstance(v, (int, float))
    }
    pipeline_keys = {
        k: v
        for k, v in time_costs.items()
        if k.startswith("pipeline_") and isinstance(v, (int, float))
    }
    other_keys = {
        k: v
        for k, v in time_costs.items()
        if not k.startswith(("lm_", "dit_", "pipeline_"))
        and isinstance(v, (int, float))
    }

    print(f"\n{'COMPONENT':<50} {'TIME (s)':<12} {'% of wall':<10}")
    print("-" * 72)

    # LM timing
    lm_total = lm_keys.get("lm_total_time", 0.0)
    if lm_keys:
        print(
            f"\n{'LLM (5Hz Language Model)':<50} "
            f"{lm_total:<12.3f} {100 * lm_total / total_wall_time:>6.1f}%"
        )
        for k, v in sorted(lm_keys.items()):
            if k != "lm_total_time":
                label = k.replace("lm_", "  ")
                print(
                    f"  {label:<48} "
                    f"{v:<12.3f} {100 * v / total_wall_time:>6.1f}%"
                )

    # DiT timing
    dit_total = dit_keys.get("dit_total_time_cost", 0.0)
    if dit_keys:
        print(
            f"\n{'DiT (Diffusion Transformer)':<50} "
            f"{dit_total:<12.3f} {100 * dit_total / total_wall_time:>6.1f}%"
        )
        for k, v in sorted(dit_keys.items()):
            if k != "dit_total_time_cost":
                label = k.replace("dit_", "  ")
                print(
                    f"  {label:<48} "
                    f"{v:<12.3f} {100 * v / total_wall_time:>6.1f}%"
                )

    # Pipeline total
    if pipeline_keys:
        for k, v in sorted(pipeline_keys.items()):
            print(
                f"\n{'Pipeline: ' + k:<50} "
                f"{v:<12.3f} {100 * v / total_wall_time:>6.1f}%"
            )

    # Other keys
    if other_keys:
        print(f"\n{'Other:':<50}")
        for k, v in sorted(other_keys.items()):
            print(
                f"  {k:<48} "
                f"{v:<12.3f} {100 * v / total_wall_time:>6.1f}%"
            )

    # Overhead (wall time minus accounted time)
    accounted = lm_total + dit_total
    overhead = total_wall_time - accounted
    if overhead > 0.01:
        print(
            f"\n{'Overhead (I/O, audio save, etc.)':<50} "
            f"{overhead:<12.3f} {100 * overhead / total_wall_time:>6.1f}%"
        )

    print(f"\n{'TOTAL WALL TIME':<50} {total_wall_time:<12.3f} {'100.0%':>6}")

    # Performance insights
    print("\n" + "=" * 100)
    print("PERFORMANCE INSIGHTS")
    print("=" * 100)

    if lm_total > 0 and dit_total > 0:
        if lm_total > dit_total * 2:
            print(
                f"\n  LLM is the bottleneck: {lm_total:.1f}s "
                f"({100 * lm_total / total_wall_time:.0f}% of total)"
            )
            print("  Suggestions:")
            print("    1. Run with --llm-debug for token-level throughput analysis")
            print("    2. Try --no-constrained-decoding to reduce FSM overhead")
            print("    3. Compare backends: --lm-backend vllm vs pt vs mlx")
            print(
                "    4. Reduce lm_cfg_scale "
                "(currently doubles forward passes if > 1.0)"
            )
        elif dit_total > lm_total * 2:
            print(
                f"\n  DiT is the bottleneck: {dit_total:.1f}s "
                f"({100 * dit_total / total_wall_time:.0f}% of total)"
            )
            print("  Suggestions:")
            print("    1. Reduce --inference-steps (turbo model supports 4-8)")
            print("    2. Reduce --duration")
            print("    3. Try --quantization int8_weight_only")
        else:
            print(
                f"\n  Balanced pipeline: LLM={lm_total:.1f}s, DiT={dit_total:.1f}s"
            )
    elif dit_total > 0:
        print(f"\n  DiT only (no LLM): {dit_total:.1f}s")
        vae_time = dit_keys.get("dit_vae_decode_time_cost", 0.0)
        diffusion_time = dit_keys.get(
            "dit_diffusion_time_cost", dit_total - vae_time
        )
        if vae_time > 0:
            print(
                f"    Diffusion: {diffusion_time:.1f}s, "
                f"VAE decode: {vae_time:.1f}s"
            )


def print_result_summary(result: GenerationResult, mode: str = "profile"):
    """Print a short summary of the generation result."""
    if result.success:
        n_audios = len(result.audios)
        silent_count = sum(1 for a in result.audios if a.get("silent", False))
        print(f"\n  Success! Generated {n_audios} audio(s)", end="")
        if silent_count:
            print(f" ({silent_count} silent)", end="")
        print()
    else:
        print(f"\n  FAILED: {result.error}")


# =============================================================================
# Mode: profile (text2music and other task types)
# =============================================================================


def run_profile_mode(dit_handler, llm_handler, args, timer: PreciseTimer):
    """Run a single profiled generation."""
    example_dir = "text2music"
    example_file = os.path.join(
        PROJECT_ROOT, "examples", example_dir, args.example
    )
    if not os.path.exists(example_file):
        print(f"\n  Example not found: {example_file}")
        sys.exit(1)

    print(f"\n  Loading example: {args.example}")
    params, config = load_example_config(example_file, args)
    if not params or not config:
        print("  Failed to load example config")
        sys.exit(1)

    caption_preview = (
        params.caption[:80] + "..."
        if len(params.caption) > 80
        else params.caption
    )
    print(f"  Caption: {caption_preview}")
    print(
        f"  Task: {params.task_type}, Batch: {config.batch_size}, "
        f"Steps: {params.inference_steps}"
    )
    print(
        f"  Thinking: {params.thinking}, CoT Metas: {params.use_cot_metas}, "
        f"CoT Caption: {params.use_cot_caption}"
    )

    # Use a temporary directory for output (don't pollute project root)
    save_dir = tempfile.mkdtemp(prefix="acestep_profile_")

    # Warmup
    if not args.no_warmup:
        print("\n" + "-" * 100)
        print("WARMUP RUN")
        print("-" * 100)
        warmup_params = GenerationParams(
            caption=params.caption,
            lyrics=params.lyrics,
            bpm=params.bpm,
            keyscale=params.keyscale,
            timesignature=params.timesignature,
            vocal_language=params.vocal_language,
            duration=params.duration,
            thinking=params.thinking,
            use_cot_metas=params.use_cot_metas,
            use_cot_caption=params.use_cot_caption,
            use_cot_language=params.use_cot_language,
            use_constrained_decoding=params.use_constrained_decoding,
            inference_steps=params.inference_steps,
            seed=42,
            task_type=params.task_type,
            lm_temperature=params.lm_temperature,
            lm_cfg_scale=params.lm_cfg_scale,
            guidance_scale=params.guidance_scale,
        )
        warmup_config = GenerationConfig(
            batch_size=1, seeds=[42], use_random_seed=False, audio_format="flac"
        )
        warmup_start = time.perf_counter()
        warmup_result = generate_music(
            dit_handler, llm_handler, warmup_params, warmup_config,
            save_dir=save_dir,
        )
        warmup_time = time.perf_counter() - warmup_start
        print(f"  Warmup completed: {warmup_time:.2f}s")
        if not warmup_result.success:
            print(f"  Warning: warmup failed: {warmup_result.error}")
        timer.reset()

    # Profiling run
    print("\n" + "=" * 100)
    print("PROFILING RUN")
    print("=" * 100)

    # Optional cProfile
    prof = None
    if args.detailed:
        import cProfile
        prof = cProfile.Profile()
        prof.enable()

    timer.sync()
    total_start = time.perf_counter()

    result = generate_music(
        dit_handler, llm_handler, params, config, save_dir=save_dir
    )

    timer.sync()
    total_wall_time = time.perf_counter() - total_start

    if args.detailed and prof:
        prof.disable()
        _print_cprofile(prof)

    # Print results
    print_result_summary(result, "profile")

    time_costs = (
        result.extra_outputs.get("time_costs", {}) if result.success else {}
    )
    print_time_costs_breakdown(time_costs, total_wall_time)

    # Cleanup temp dir
    _cleanup_dir(save_dir)

    return result, total_wall_time


# =============================================================================
# Mode: benchmark
# =============================================================================


def run_benchmark_mode(dit_handler, llm_handler, args, timer: PreciseTimer):
    """Run a matrix of configurations and produce a summary table."""
    example_file = os.path.join(
        PROJECT_ROOT, "examples", "text2music", args.example
    )
    if not os.path.exists(example_file):
        print(f"\n  Example not found: {example_file}")
        sys.exit(1)

    with open(example_file, "r", encoding="utf-8") as f:
        example_data = json.load(f)

    save_dir = tempfile.mkdtemp(prefix="acestep_bench_")

    # Define benchmark matrix
    durations = [30, 60, 120]
    batch_sizes = [1, 2]
    thinking_options = (
        [False, True] if llm_handler.llm_initialized else [False]
    )
    inference_steps_options = [8]

    # Clamp to GPU limits
    gpu_config = get_gpu_config()
    max_dur = gpu_config.max_duration_without_lm
    max_batch = gpu_config.max_batch_size_without_lm
    durations = [d for d in durations if d <= max_dur]
    batch_sizes = [b for b in batch_sizes if b <= max_batch]

    if not durations:
        durations = [30]
    if not batch_sizes:
        batch_sizes = [1]

    configs = []
    for dur in durations:
        for bs in batch_sizes:
            for think in thinking_options:
                for steps in inference_steps_options:
                    configs.append(
                        {
                            "duration": dur,
                            "batch_size": bs,
                            "thinking": think,
                            "inference_steps": steps,
                        }
                    )

    print(f"\n  Running {len(configs)} benchmark configurations...")
    print(f"  Durations: {durations}, Batch sizes: {batch_sizes}")
    print(f"  Thinking: {thinking_options}, Steps: {inference_steps_options}")

    # Warmup
    if not args.no_warmup:
        print("\n  Warmup run...")
        warmup_params = GenerationParams(
            caption=example_data.get("caption", ""),
            lyrics=example_data.get("lyrics", ""),
            duration=30,
            thinking=False,
            inference_steps=8,
            seed=42,
        )
        warmup_config = GenerationConfig(
            batch_size=1, seeds=[42], use_random_seed=False, audio_format="flac"
        )
        generate_music(
            dit_handler, llm_handler, warmup_params, warmup_config,
            save_dir=save_dir,
        )
        print("  Warmup done.")

    # Run benchmark
    results = []
    for i, cfg in enumerate(configs):
        label = (
            f"dur={cfg['duration']}s, bs={cfg['batch_size']}, "
            f"think={cfg['thinking']}, steps={cfg['inference_steps']}"
        )
        print(f"\n  [{i + 1}/{len(configs)}] {label}")

        params = GenerationParams(
            caption=example_data.get("caption", ""),
            lyrics=example_data.get("lyrics", ""),
            bpm=example_data.get("bpm"),
            keyscale=example_data.get("keyscale", ""),
            timesignature=example_data.get("timesignature", ""),
            vocal_language=example_data.get("language", "unknown"),
            duration=cfg["duration"],
            thinking=cfg["thinking"],
            use_cot_metas=cfg["thinking"],
            use_cot_caption=cfg["thinking"],
            use_cot_language=cfg["thinking"],
            use_constrained_decoding=args.use_constrained_decoding,
            inference_steps=cfg["inference_steps"],
            seed=42,
            lm_temperature=args.lm_temperature,
            lm_cfg_scale=args.lm_cfg_scale,
            guidance_scale=args.guidance_scale,
        )
        config = GenerationConfig(
            batch_size=cfg["batch_size"],
            seeds=[42 + j for j in range(cfg["batch_size"])],
            use_random_seed=False,
            audio_format="flac",
        )

        timer.sync()
        t0 = time.perf_counter()
        result = generate_music(
            dit_handler, llm_handler, params, config, save_dir=save_dir
        )
        timer.sync()
        wall_time = time.perf_counter() - t0

        tc = (
            result.extra_outputs.get("time_costs", {})
            if result.success
            else {}
        )
        entry = {
            "config": cfg,
            "wall_time": wall_time,
            "success": result.success,
            "error": result.error,
            "lm_time": tc.get("lm_total_time", 0.0),
            "dit_time": tc.get("dit_total_time_cost", 0.0),
            "vae_time": tc.get("dit_vae_decode_time_cost", 0.0),
            "n_audios": len(result.audios) if result.success else 0,
        }
        results.append(entry)

        status = "OK" if result.success else f"FAIL: {result.error}"
        print(
            f"    {status} | wall={wall_time:.1f}s, "
            f"lm={entry['lm_time']:.1f}s, dit={entry['dit_time']:.1f}s"
        )

    # Print summary table
    print("\n" + "=" * 120)
    print("BENCHMARK SUMMARY")
    print("=" * 120)

    header = (
        f"{'Duration':<10} {'Batch':<7} {'Think':<7} {'Steps':<7} "
        f"{'Wall(s)':<10} {'LM(s)':<10} {'DiT(s)':<10} "
        f"{'VAE(s)':<10} {'Status':<10}"
    )
    print(header)
    print("-" * 120)

    for entry in results:
        cfg = entry["config"]
        status = "OK" if entry["success"] else "FAIL"
        print(
            f"{cfg['duration']:<10} {cfg['batch_size']:<7} "
            f"{str(cfg['thinking']):<7} {cfg['inference_steps']:<7} "
            f"{entry['wall_time']:<10.2f} {entry['lm_time']:<10.2f} "
            f"{entry['dit_time']:<10.2f} {entry['vae_time']:<10.2f} "
            f"{status:<10}"
        )

    # Save benchmark results as JSON
    if args.benchmark_output:
        output_path = args.benchmark_output
        with open(output_path, "w", encoding="utf-8") as f:
            json.dump(results, f, indent=2, default=str)
        print(f"\n  Benchmark results saved to: {output_path}")

    _cleanup_dir(save_dir)
    return results


# =============================================================================
# Mode: tier-test  (THE KEY FEATURE)
# =============================================================================


def _get_vram_info_str() -> str:
    """Get current VRAM usage string for logging."""
    if not torch.cuda.is_available():
        return "N/A"
    allocated = torch.cuda.memory_allocated() / (1024 ** 3)
    reserved = torch.cuda.memory_reserved() / (1024 ** 3)
    return f"alloc={allocated:.2f}GB, reserved={reserved:.2f}GB"


def _run_single_tier_test(
    sim_gb: float,
    gpu_config: GPUConfig,
    args,
    example_data: Dict,
    checkpoint_dir: str,
    disk_lm_models: List[str],
    *,
    offload_override: Optional[bool] = None,
    offload_dit_override: Optional[bool] = None,
    quantization_override: Optional[str] = "USE_DEFAULT",
    test_variant: str = "default",
    batch_size_override: Optional[int] = None,
    use_lm_override: Optional[bool] = None,
) -> Dict[str, Any]:
    """
    Run a single tier test with the given configuration.

    Args:
        sim_gb: Simulated VRAM in GB
        gpu_config: GPU configuration for this tier
        args: CLI arguments
        example_data: Example JSON data for generation
        checkpoint_dir: Path to checkpoints directory
        disk_lm_models: List of LM models found on disk
        offload_override: If not None, override offload_to_cpu setting
        offload_dit_override: If not None, override offload_dit_to_cpu setting
        quantization_override: If not "USE_DEFAULT", override quantization setting
                               (None means no quantization, "int8_weight_only" etc.)
        test_variant: Label for this test variant ("default", "no-quant", "no-offload")
        batch_size_override: If not None, override batch size (used by batch boundary tests)
        use_lm_override: If not None, force LM on (True) or off (False)

    Returns:
        Result dictionary for this test
    """
    tier = gpu_config.tier

    # Determine test configuration
    if use_lm_override is not None:
        use_lm = use_lm_override and gpu_config.init_lm_default and bool(gpu_config.available_lm_models)
    else:
        use_lm = args.tier_with_lm and gpu_config.init_lm_default and bool(gpu_config.available_lm_models)

    if hasattr(torch, "xpu") and torch.xpu.is_available():
        offload_override = False

    if offload_override is not None:
        offload = offload_override
    else:
        offload = gpu_config.offload_to_cpu_default

    if offload_dit_override is not None:
        offload_dit = offload_dit_override
    else:
        offload_dit = gpu_config.offload_dit_to_cpu_default

    if quantization_override != "USE_DEFAULT":
        quantization = quantization_override
    else:
        quantization = "int8_weight_only" if gpu_config.quantization_default else None

    # Find LM model on disk
    lm_model = None
    lm_backend = gpu_config.recommended_backend
    if use_lm:
        lm_model = find_best_lm_model_on_disk(
            gpu_config.recommended_lm_model, disk_lm_models
        )
        if not lm_model:
            print(f"  ⚠️ No compatible LM model on disk for tier {tier}, skipping LM")
            use_lm = False

    # Clamp duration to tier limit
    test_duration = args.tier_duration
    max_dur = gpu_config.max_duration_with_lm if use_lm else gpu_config.max_duration_without_lm
    if test_duration > max_dur:
        test_duration = max_dur
        print(f"  Duration clamped to {test_duration}s (tier limit)")

    batch_size = batch_size_override if batch_size_override is not None else 1

    print(f"\n  Test config [{test_variant}]: duration={test_duration}s, batch={batch_size}, LM={use_lm}")
    if use_lm:
        print(f"    LM model: {lm_model}, backend: {lm_backend}")
    print(f"    offload={offload}, offload_dit={offload_dit}, quant={quantization}")

    # Enforce VRAM cap
    if torch.cuda.is_available():
        total_bytes = torch.cuda.get_device_properties(0).total_memory
        total_gb = total_bytes / (1024 ** 3)
        if sim_gb < total_gb:
            reference_context_gb = 0.5
            allocator_budget_gb = max(0.5, sim_gb - reference_context_gb)
            fraction = max(0.01, min(1.0, allocator_budget_gb / total_gb))
            torch.cuda.set_per_process_memory_fraction(fraction)

        torch.cuda.empty_cache()
        torch.cuda.reset_peak_memory_stats()

    # Initialize result entry
    result_entry = {
        "tier_gb": sim_gb,
        "tier": tier,
        "test_variant": test_variant,
        "use_lm": use_lm,
        "lm_model": lm_model,
        "lm_backend": lm_backend,
        "offload": offload,
        "offload_dit": offload_dit,
        "quantization": quantization,
        "duration": test_duration,
        "batch_size": batch_size,
        "init_success": False,
        "gen_success": False,
        "wall_time": 0.0,
        "error": None,
        "peak_vram_gb": 0.0,
    }

    dit_handler = None
    llm_handler = None

    try:
        print(f"\n  Initializing DiT handler... ({_get_vram_info_str()})")
        dit_handler = AceStepHandler()

        # Determine flash attention availability
        use_flash_attention = False
        try:
            import flash_attn  # noqa: F401
            use_flash_attention = True
        except ImportError:
            pass

        # compile_model must be True when quantization is used;
        # --tier-skip-compile can skip it for non-quantized tiers to save time
        if quantization:
            compile_model = True
        elif args.tier_skip_compile:
            compile_model = False
        else:
            compile_model = gpu_config.compile_model_default

        status_dit, success_dit = dit_handler.initialize_service(
            project_root=PROJECT_ROOT,
            config_path=args.config_path,
            device="auto",
            use_flash_attention=use_flash_attention,
            compile_model=compile_model,
            offload_to_cpu=offload,
            offload_dit_to_cpu=offload_dit,
            quantization=quantization,
        )

        if not success_dit:
            result_entry["error"] = f"DiT init failed: {status_dit}"
            print(f"  ❌ DiT init failed: {status_dit}")
            _cleanup_handlers(dit_handler, None)
            return result_entry

        print(f"  ✅ DiT ready ({_get_vram_info_str()})")

        llm_handler = LLMHandler()

        if use_lm:
            print(f"  Initializing LLM handler (backend={lm_backend})... ({_get_vram_info_str()})")
            status_llm, success_llm = llm_handler.initialize(
                checkpoint_dir=checkpoint_dir,
                lm_model_path=lm_model,
                backend=lm_backend,
                device="auto",
                offload_to_cpu=offload,
                dtype=None,
            )
            if success_llm:
                print(f"  ✅ LLM ready ({_get_vram_info_str()})")
            else:
                print(f"  ⚠️ LLM init failed: {status_llm}")
                use_lm = False
                result_entry["use_lm"] = False
                result_entry["error"] = f"LM init failed (non-fatal): {status_llm}"

        result_entry["init_success"] = True

    except torch.cuda.OutOfMemoryError as e:
        result_entry["error"] = f"Init OOM: {e}"
        print(f"  ❌ Init OOM: {e}")
        _cleanup_handlers(dit_handler, llm_handler)
        return result_entry
    except Exception as e:
        result_entry["error"] = f"Init exception: {e}"
        print(f"  ❌ Init exception: {e}")
        traceback.print_exc()
        _cleanup_handlers(dit_handler, llm_handler)
        return result_entry

    # Run generation
    try:
        print(f"\n  Running generation... ({_get_vram_info_str()})")
        save_dir = tempfile.mkdtemp(prefix=f"acestep_tier{int(sim_gb)}_{test_variant}_")

        params = GenerationParams(
            caption=example_data.get("caption", ""),
            lyrics=example_data.get("lyrics", ""),
            bpm=example_data.get("bpm"),
            keyscale=example_data.get("keyscale", ""),
            timesignature=example_data.get("timesignature", ""),
            vocal_language=example_data.get("language", "unknown"),
            duration=test_duration,
            thinking=use_lm,
            use_cot_metas=use_lm,
            use_cot_caption=False,
            use_cot_language=False,
            use_constrained_decoding=True,
            inference_steps=8,
            seed=42,
            lm_temperature=0.85,
            lm_cfg_scale=2.0,
            guidance_scale=7.0,
        )
        config = GenerationConfig(
            batch_size=batch_size,
            seeds=[42 + j for j in range(batch_size)],
            use_random_seed=False,
            audio_format="flac",
        )

        # When testing batch boundaries, temporarily override the GPU tier config's
        # max_batch limits so that inference.py's clamping doesn't reduce our test
        # batch size. We restore the original values after the test.
        _patched_tier_config = False
        _orig_batch_with_lm = None
        _orig_batch_without_lm = None
        if batch_size_override is not None and batch_size_override > 1:
            from acestep.gpu_config import GPU_TIER_CONFIGS as _tier_configs
            tier = gpu_config.tier
            if tier in _tier_configs:
                _patched_tier_config = True
                _orig_batch_with_lm = _tier_configs[tier]["max_batch_size_with_lm"]
                _orig_batch_without_lm = _tier_configs[tier]["max_batch_size_without_lm"]
                _tier_configs[tier]["max_batch_size_with_lm"] = max(batch_size_override, _orig_batch_with_lm)
                _tier_configs[tier]["max_batch_size_without_lm"] = max(batch_size_override, _orig_batch_without_lm)

        t0 = time.perf_counter()
        try:
            result = generate_music(
                dit_handler, llm_handler, params, config, save_dir=save_dir
            )
        finally:
            # Restore original tier config values
            if _patched_tier_config:
                _tier_configs[tier]["max_batch_size_with_lm"] = _orig_batch_with_lm
                _tier_configs[tier]["max_batch_size_without_lm"] = _orig_batch_without_lm
        wall_time = time.perf_counter() - t0

        result_entry["wall_time"] = wall_time
        result_entry["gen_success"] = result.success

        if result.success:
            tc = result.extra_outputs.get("time_costs", {})
            result_entry["lm_time"] = tc.get("lm_total_time", 0.0)
            result_entry["dit_time"] = tc.get("dit_total_time_cost", 0.0)
            result_entry["vae_time"] = tc.get("dit_vae_decode_time_cost", 0.0)
            n_audios = len(result.audios)
            print(f"  ✅ [{test_variant}] Generation OK: {n_audios} audio(s) in {wall_time:.1f}s")
        else:
            result_entry["error"] = result.error
            print(f"  ❌ [{test_variant}] Generation FAILED: {result.error}")

        _cleanup_dir(save_dir)

    except torch.cuda.OutOfMemoryError as e:
        result_entry["error"] = f"OOM: {e}"
        print(f"  ❌ [{test_variant}] OOM ERROR: {e}")
    except Exception as e:
        result_entry["error"] = f"Generation exception: {e}"
        print(f"  ❌ [{test_variant}] Exception: {e}")
        traceback.print_exc()

    # Record peak VRAM
    if torch.cuda.is_available():
        peak_bytes = torch.cuda.max_memory_allocated()
        result_entry["peak_vram_gb"] = peak_bytes / (1024 ** 3)
        print(f"  Peak VRAM: {result_entry['peak_vram_gb']:.2f}GB")

    # Cleanup
    _cleanup_handlers(dit_handler, llm_handler)

    return result_entry


def run_tier_test_mode(args):
    """
    Automatically test inference across multiple simulated GPU tiers.

    For each tier:
      1. Set MAX_CUDA_VRAM to simulate the VRAM limit
      2. Initialize gpu_config for that tier
      3. Initialize DiT + (optionally) LLM handlers with tier-appropriate settings
      4. Run a short generation and verify it completes without OOM
      5. Report results

    When --tier-boundary is enabled, each tier is tested with up to 3 configurations:
      - default: tier's default settings (quantization + offload as configured)
      - no-quant: same as default but with quantization disabled
      - no-offload: no quantization AND no CPU offload (all models on GPU)

    This replaces the manual workflow of:
      MAX_CUDA_VRAM=8 uv run acestep → click UI → wait → check
    """
    # Determine which tiers to test
    default_tiers = [4, 6, 8, 12, 16, 24, 48]
    tiers_to_test = args.tiers if args.tiers else default_tiers

    # Load example for generation
    example_file = os.path.join(
        PROJECT_ROOT, "examples", "text2music", args.example
    )
    if not os.path.exists(example_file):
        print(f"\n  Example not found: {example_file}")
        sys.exit(1)

    with open(example_file, "r", encoding="utf-8") as f:
        example_data = json.load(f)

    # Scan available LM models on disk
    checkpoint_dir = os.path.join(PROJECT_ROOT, "checkpoints")
    disk_lm_models = []
    if os.path.exists(checkpoint_dir):
        for item in sorted(os.listdir(checkpoint_dir)):
            if os.path.isdir(os.path.join(checkpoint_dir, item)) and item.startswith("acestep-5Hz-lm-"):
                disk_lm_models.append(item)

    boundary_mode = getattr(args, "tier_boundary", False)
    batch_boundary_mode = getattr(args, "tier_batch_boundary", False)

    print(f"\n  Tiers to test: {tiers_to_test}")
    print(f"  LM models on disk: {disk_lm_models}")
    print(f"  Test with LM: {args.tier_with_lm}")
    print(f"  Test duration: {args.tier_duration}s")
    print(f"  Boundary testing: {boundary_mode}")
    print(f"  Batch boundary testing: {batch_boundary_mode}")
    print(f"  Example: {args.example}")

    # Results collector
    all_results = []

    for sim_gb in tiers_to_test:
        print("\n" + "=" * 120)
        print(f"  TIER TEST: {sim_gb}GB simulated VRAM")
        print("=" * 120)

        # Configure GPU simulation
        os.environ["MAX_CUDA_VRAM"] = str(sim_gb)

        # Force re-detection of GPU config
        gpu_config = get_gpu_config(gpu_memory_gb=float(sim_gb))
        set_global_gpu_config(gpu_config)

        tier = gpu_config.tier
        print(f"  Tier: {tier}")
        print(f"  init_lm_default: {gpu_config.init_lm_default}")
        print(f"  available_lm_models: {gpu_config.available_lm_models}")
        print(f"  recommended_lm_model: {gpu_config.recommended_lm_model}")
        print(f"  recommended_backend: {gpu_config.recommended_backend}")
        print(f"  lm_backend_restriction: {gpu_config.lm_backend_restriction}")
        print(f"  offload_to_cpu: {gpu_config.offload_to_cpu_default}")
        print(f"  offload_dit_to_cpu: {gpu_config.offload_dit_to_cpu_default}")
        print(f"  quantization: {gpu_config.quantization_default}")
        print(f"  max_duration_with_lm: {gpu_config.max_duration_with_lm}s")
        print(f"  max_duration_without_lm: {gpu_config.max_duration_without_lm}s")
        print(f"  max_batch_with_lm: {gpu_config.max_batch_size_with_lm}")
        print(f"  max_batch_without_lm: {gpu_config.max_batch_size_without_lm}")

        # ---- Test 1: Default configuration ----
        print(f"\n  --- Variant: default ---")
        result_default = _run_single_tier_test(
            sim_gb, gpu_config, args, example_data,
            checkpoint_dir, disk_lm_models,
            test_variant="default",
        )
        all_results.append(result_default)

        if boundary_mode:
            # ---- Test 2: No quantization (keep offload as default) ----
            # Skip if the tier already doesn't use quantization (no point re-testing)
            if gpu_config.quantization_default:
                print(f"\n  --- Variant: no-quant (offload={gpu_config.offload_to_cpu_default}) ---")
                result_no_quant = _run_single_tier_test(
                    sim_gb, gpu_config, args, example_data,
                    checkpoint_dir, disk_lm_models,
                    quantization_override=None,
                    test_variant="no-quant",
                )
                all_results.append(result_no_quant)
            else:
                print(f"\n  --- Variant: no-quant — SKIPPED (tier already has quantization=False) ---")

            # ---- Test 3: No quantization AND no offload ----
            # Skip if the tier already has both disabled
            # Also skip if simulated VRAM is too small — the unquantized DiT model
            # alone needs ~6GB; without offload there is no room left for VAE decode,
            # which causes a fallback to CPU VAE with tiny chunk_size and 20+ hour runs.
            MIN_VRAM_FOR_NO_OFFLOAD = 8  # GB — DiT (~6GB) + VAE headroom (~2GB)
            if sim_gb < MIN_VRAM_FOR_NO_OFFLOAD:
                print(f"\n  --- Variant: no-offload — SKIPPED (simulated {sim_gb}GB < {MIN_VRAM_FOR_NO_OFFLOAD}GB minimum for no-offload) ---")
            elif gpu_config.quantization_default or gpu_config.offload_to_cpu_default:
                print(f"\n  --- Variant: no-offload (quant=None, offload=False) ---")
                result_no_offload = _run_single_tier_test(
                    sim_gb, gpu_config, args, example_data,
                    checkpoint_dir, disk_lm_models,
                    offload_override=False,
                    offload_dit_override=False,
                    quantization_override=None,
                    test_variant="no-offload",
                )
                all_results.append(result_no_offload)
            else:
                print(f"\n  --- Variant: no-offload — SKIPPED (tier already has offload=False, quant=False) ---")

        if batch_boundary_mode:
            # ---- Batch boundary tests: escalate batch size until OOM ----
            BATCH_SIZES_TO_TEST = [1, 2, 4, 8]

            # Test WITHOUT LM
            print(f"\n  --- Batch boundary: without LM ---")
            for bs in BATCH_SIZES_TO_TEST:
                print(f"\n  --- Variant: batch-noLM-{bs} (batch_size={bs}, no LM) ---")
                result_batch = _run_single_tier_test(
                    sim_gb, gpu_config, args, example_data,
                    checkpoint_dir, disk_lm_models,
                    test_variant=f"batch-noLM-{bs}",
                    batch_size_override=bs,
                    use_lm_override=False,
                )
                all_results.append(result_batch)
                if not result_batch["gen_success"]:
                    print(f"  ⚠️ Batch size {bs} failed without LM — stopping escalation")
                    break

            # Test WITH LM (if tier supports it)
            if gpu_config.init_lm_default and bool(gpu_config.available_lm_models):
                print(f"\n  --- Batch boundary: with LM ---")
                for bs in BATCH_SIZES_TO_TEST:
                    print(f"\n  --- Variant: batch-LM-{bs} (batch_size={bs}, with LM) ---")
                    result_batch_lm = _run_single_tier_test(
                        sim_gb, gpu_config, args, example_data,
                        checkpoint_dir, disk_lm_models,
                        test_variant=f"batch-LM-{bs}",
                        batch_size_override=bs,
                        use_lm_override=True,
                    )
                    all_results.append(result_batch_lm)
                    if not result_batch_lm["gen_success"]:
                        print(f"  ⚠️ Batch size {bs} failed with LM — stopping escalation")
                        break

    # ---- Print summary ----
    _print_tier_test_summary(all_results)

    if boundary_mode:
        _print_boundary_summary(all_results)

    if batch_boundary_mode:
        _print_batch_boundary_summary(all_results)

    # Save results
    if args.benchmark_output:
        with open(args.benchmark_output, "w", encoding="utf-8") as f:
            json.dump(all_results, f, indent=2, default=str)
        print(f"\n  Results saved to: {args.benchmark_output}")

    return all_results


def _cleanup_handlers(dit_handler, llm_handler):
    """Clean up handlers and free GPU memory."""
    try:
        if dit_handler is not None:
            if hasattr(dit_handler, 'model') and dit_handler.model is not None:
                dit_handler.model = None
            if hasattr(dit_handler, 'vae') and dit_handler.vae is not None:
                dit_handler.vae = None
            if hasattr(dit_handler, 'text_encoder') and dit_handler.text_encoder is not None:
                dit_handler.text_encoder = None
            del dit_handler
    except Exception:
        pass

    try:
        if llm_handler is not None:
            if hasattr(llm_handler, 'llm') and llm_handler.llm is not None:
                llm_handler.llm = None
            del llm_handler
    except Exception:
        pass

    import gc
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.reset_peak_memory_stats()


def _print_tier_test_summary(results: List[Dict]):
    """Print a summary table of all tier test results."""
    # Detect if any result has a test_variant (boundary mode)
    has_variants = any(r.get("test_variant", "default") != "default" for r in results)

    print("\n" + "=" * 160)
    print("TIER TEST SUMMARY")
    print("=" * 160)

    if has_variants:
        header = (
            f"{'VRAM':>6} {'Tier':<10} {'Variant':<12} {'LM':>4} {'LM Model':<24} {'Backend':<8} "
            f"{'Offload':<8} {'Quant':<6} {'Init':>5} {'Gen':>5} "
            f"{'Wall(s)':>8} {'Peak(GB)':>9} {'Status':<30}"
        )
    else:
        header = (
            f"{'VRAM':>6} {'Tier':<10} {'LM':>4} {'LM Model':<28} {'Backend':<8} "
            f"{'Offload':<8} {'Quant':<6} {'Init':>5} {'Gen':>5} "
            f"{'Wall(s)':>8} {'Peak(GB)':>9} {'Status':<30}"
        )
    print(header)
    print("-" * 160)

    pass_count = 0
    fail_count = 0

    for r in results:
        lm_model_short = (r.get("lm_model") or "-")
        max_lm_len = 22 if has_variants else 26
        if len(lm_model_short) > max_lm_len:
            lm_model_short = lm_model_short[:max_lm_len] + ".."

        init_ok = "✅" if r["init_success"] else "❌"
        gen_ok = "✅" if r["gen_success"] else "❌"
        status = "PASS" if r["gen_success"] else (r.get("error", "FAIL") or "FAIL")
        if len(status) > 28:
            status = status[:28] + ".."

        if r["gen_success"]:
            pass_count += 1
        else:
            fail_count += 1

        quant = "int8" if r.get("quantization") else "-"
        variant = r.get("test_variant", "default")

        if has_variants:
            print(
                f"{r['tier_gb']:5d}GB {r['tier']:<10} {variant:<12} "
                f"{'Y' if r['use_lm'] else 'N':>4} {lm_model_short:<24} "
                f"{r.get('lm_backend', '-'):<8} "
                f"{'Y' if r['offload'] else 'N':<8} {quant:<6} "
                f"{init_ok:>5} {gen_ok:>5} "
                f"{r['wall_time']:>8.1f} {r.get('peak_vram_gb', 0):>9.2f} "
                f"{status:<30}"
            )
        else:
            print(
                f"{r['tier_gb']:5d}GB {r['tier']:<10} "
                f"{'Y' if r['use_lm'] else 'N':>4} {lm_model_short:<28} "
                f"{r.get('lm_backend', '-'):<8} "
                f"{'Y' if r['offload'] else 'N':<8} {quant:<6} "
                f"{init_ok:>5} {gen_ok:>5} "
                f"{r['wall_time']:>8.1f} {r.get('peak_vram_gb', 0):>9.2f} "
                f"{status:<30}"
            )

    print("-" * 160)
    print(f"  Total: {len(results)} tests run, {pass_count} PASSED, {fail_count} FAILED")


def _print_boundary_summary(results: List[Dict]):
    """
    Print a boundary analysis summary showing the minimum tier for each capability.

    Analyzes results from boundary testing to determine:
    - Minimum tier that works WITHOUT INT8 quantization
    - Minimum tier that works WITHOUT CPU offload (and without quantization)
    """
    print("\n" + "=" * 100)
    print("BOUNDARY ANALYSIS")
    print("=" * 100)
    print()
    print("  This analysis shows the minimum VRAM tier at which each optimization")
    print("  can be safely disabled while still completing inference successfully.")
    print()

    # Collect results by variant
    no_quant_results = [r for r in results if r.get("test_variant") == "no-quant"]
    no_offload_results = [r for r in results if r.get("test_variant") == "no-offload"]
    default_results = [r for r in results if r.get("test_variant") == "default"]

    # Also consider default results where the tier already has quant/offload disabled
    # (e.g., tier6b default already has quantization=False)
    for r in default_results:
        if not r.get("quantization") and r not in no_quant_results:
            # This tier's default already runs without quantization
            no_quant_results.append(r)
        if not r.get("offload") and not r.get("quantization") and r not in no_offload_results:
            # This tier's default already runs without offload and without quantization
            no_offload_results.append(r)

    # Sort by VRAM
    no_quant_results.sort(key=lambda r: r["tier_gb"])
    no_offload_results.sort(key=lambda r: r["tier_gb"])

    # Find minimum passing tier for each capability
    def _find_min_passing(result_list, capability_name):
        passing = [r for r in result_list if r.get("gen_success")]
        failing = [r for r in result_list if not r.get("gen_success")]

        if passing:
            min_pass = passing[0]
            print(f"  {capability_name}:")
            print(f"    Minimum tier:  {min_pass['tier']} ({min_pass['tier_gb']}GB)")
            print(f"    Peak VRAM:     {min_pass.get('peak_vram_gb', 0):.2f}GB")
            if failing:
                max_fail = failing[-1]
                print(f"    Last failure:  {max_fail['tier']} ({max_fail['tier_gb']}GB) — {max_fail.get('error', 'unknown')[:60]}")
        else:
            if failing:
                print(f"  {capability_name}:")
                print(f"    ❌ No tier passed this test. All tested tiers failed.")
                for r in failing:
                    err = (r.get("error") or "unknown")[:50]
                    print(f"       {r['tier_gb']}GB ({r['tier']}): {err}")
            else:
                print(f"  {capability_name}:")
                print(f"    ⚠️ No test results available for this capability.")
        print()
        return passing[0] if passing else None

    min_no_quant = _find_min_passing(no_quant_results, "Without INT8 Quantization")
    min_no_offload = _find_min_passing(no_offload_results, "Without CPU Offload (and no quantization)")

    # Print compact summary table
    print("  " + "-" * 60)
    print(f"  {'Capability':<45} {'Min Tier':<10} {'VRAM':>6}")
    print("  " + "-" * 60)

    if min_no_quant:
        print(f"  {'No INT8 Quantization':<45} {min_no_quant['tier']:<10} {min_no_quant['tier_gb']:>5}GB")
    else:
        print(f"  {'No INT8 Quantization':<45} {'N/A':<10} {'N/A':>6}")

    if min_no_offload:
        print(f"  {'No CPU Offload (all models on GPU)':<45} {min_no_offload['tier']:<10} {min_no_offload['tier_gb']:>5}GB")
    else:
        print(f"  {'No CPU Offload (all models on GPU)':<45} {'N/A':<10} {'N/A':>6}")

    print("  " + "-" * 60)
    print()
    print("  Note: These boundaries are empirical and may vary based on:")
    print("    - DiT model variant (turbo vs base)")
    print("    - Whether LM is enabled (--tier-with-lm)")
    print("    - Generation duration and batch size")
    print("    - Flash attention availability")


def _print_batch_boundary_summary(results: List[Dict]):
    """
    Print a batch boundary analysis summary showing the maximum safe batch size per tier.

    Analyzes results from batch boundary testing to determine:
    - Maximum batch size WITHOUT LM for each tier
    - Maximum batch size WITH LM for each tier
    """
    print("\n" + "=" * 120)
    print("BATCH BOUNDARY ANALYSIS")
    print("=" * 120)
    print()
    print("  This analysis shows the maximum batch size that completed successfully")
    print("  for each simulated VRAM tier.")
    print()

    # Collect batch boundary results
    batch_no_lm = [r for r in results if r.get("test_variant", "").startswith("batch-noLM-")]
    batch_with_lm = [r for r in results if r.get("test_variant", "").startswith("batch-LM-")]

    # Group by tier_gb
    def _group_by_tier(result_list):
        groups = {}
        for r in result_list:
            tier_gb = r["tier_gb"]
            if tier_gb not in groups:
                groups[tier_gb] = {"tier": r["tier"], "results": []}
            groups[tier_gb]["results"].append(r)
        return groups

    no_lm_groups = _group_by_tier(batch_no_lm)
    with_lm_groups = _group_by_tier(batch_with_lm)

    # Find max passing batch per tier
    def _max_passing_batch(group_results):
        max_bs = 0
        peak_vram = 0.0
        for r in group_results:
            if r.get("gen_success"):
                bs = r.get("batch_size", 1)
                if bs > max_bs:
                    max_bs = bs
                    peak_vram = r.get("peak_vram_gb", 0)
        return max_bs, peak_vram

    # Collect all tier_gb values
    all_tier_gbs = sorted(set(list(no_lm_groups.keys()) + list(with_lm_groups.keys())))

    # Print table
    print(f"  {'VRAM':>6}  {'Tier':<12}  {'Max Batch (no LM)':>18}  {'Peak VRAM':>10}  {'Max Batch (with LM)':>20}  {'Peak VRAM':>10}")
    print("  " + "-" * 90)

    summary_rows = []
    for tier_gb in all_tier_gbs:
        tier_name = no_lm_groups.get(tier_gb, with_lm_groups.get(tier_gb, {})).get("tier", "?")

        no_lm_max, no_lm_peak = (0, 0.0)
        if tier_gb in no_lm_groups:
            no_lm_max, no_lm_peak = _max_passing_batch(no_lm_groups[tier_gb]["results"])

        with_lm_max, with_lm_peak = (0, 0.0)
        if tier_gb in with_lm_groups:
            with_lm_max, with_lm_peak = _max_passing_batch(with_lm_groups[tier_gb]["results"])

        no_lm_str = str(no_lm_max) if no_lm_max > 0 else "FAIL"
        with_lm_str = str(with_lm_max) if with_lm_max > 0 else ("N/A" if tier_gb not in with_lm_groups else "FAIL")

        no_lm_peak_str = f"{no_lm_peak:.2f}GB" if no_lm_max > 0 else "-"
        with_lm_peak_str = f"{with_lm_peak:.2f}GB" if with_lm_max > 0 else "-"

        print(
            f"  {tier_gb:5d}GB  {tier_name:<12}  {no_lm_str:>18}  {no_lm_peak_str:>10}  "
            f"{with_lm_str:>20}  {with_lm_peak_str:>10}"
        )

        summary_rows.append({
            "tier_gb": tier_gb,
            "tier": tier_name,
            "max_batch_no_lm": no_lm_max,
            "max_batch_with_lm": with_lm_max if tier_gb in with_lm_groups else None,
        })

    print("  " + "-" * 90)
    print()

    # Print comparison with current GPU_TIER_CONFIGS
    print("  Comparison with current GPU_TIER_CONFIGS:")
    print(f"  {'VRAM':>6}  {'Tier':<12}  {'Config (no LM)':>15}  {'Tested (no LM)':>15}  {'Config (LM)':>12}  {'Tested (LM)':>12}  {'Recommendation':<30}")
    print("  " + "-" * 110)

    for row in summary_rows:
        tier_gb = row["tier_gb"]
        tier_name = row["tier"]
        cfg = get_gpu_config(gpu_memory_gb=float(tier_gb))

        cfg_no_lm = cfg.max_batch_size_without_lm
        cfg_with_lm = cfg.max_batch_size_with_lm
        tested_no_lm = row["max_batch_no_lm"]
        tested_with_lm = row["max_batch_with_lm"]

        tested_no_lm_str = str(tested_no_lm) if tested_no_lm > 0 else "FAIL"
        tested_with_lm_str = str(tested_with_lm) if tested_with_lm is not None and tested_with_lm > 0 else ("N/A" if tested_with_lm is None else "FAIL")

        # Recommendation
        rec_parts = []
        if tested_no_lm > 0 and tested_no_lm != cfg_no_lm:
            rec_parts.append(f"no_lm: {cfg_no_lm}→{tested_no_lm}")
        if tested_with_lm is not None and tested_with_lm > 0 and tested_with_lm != cfg_with_lm:
            rec_parts.append(f"lm: {cfg_with_lm}→{tested_with_lm}")
        recommendation = ", ".join(rec_parts) if rec_parts else "OK"

        print(
            f"  {tier_gb:5d}GB  {tier_name:<12}  {cfg_no_lm:>15}  {tested_no_lm_str:>15}  "
            f"{cfg_with_lm:>12}  {tested_with_lm_str:>12}  {recommendation:<30}"
        )

    print("  " + "-" * 110)
    print()
    print("  Note: Batch boundary results are empirical and depend on:")
    print("    - DiT model variant (turbo vs base)")
    print("    - Generation duration (longer = more VRAM per batch)")
    print("    - Flash attention availability")
    print("    - LM model size (0.6B vs 1.7B vs 4B)")
    print("    - Quantization and offload settings")


# =============================================================================
# Mode: understand
# =============================================================================


def run_understand_mode(dit_handler, llm_handler, args, timer: PreciseTimer):
    """Profile the understand_music() API."""
    if not llm_handler.llm_initialized:
        print("\n  LLM not initialized. understand mode requires LLM.")
        print("  Re-run with --thinking or ensure LLM is available.")
        sys.exit(1)

    audio_codes = args.audio_codes if args.audio_codes else ""

    print(
        f"\n  Audio codes: "
        f"{'<provided>' if audio_codes else '<empty - will generate sample>'}"
    )

    timer.sync()
    t0 = time.perf_counter()

    result = understand_music(
        llm_handler=llm_handler,
        audio_codes=audio_codes,
        temperature=args.lm_temperature,
        use_constrained_decoding=args.use_constrained_decoding,
    )

    timer.sync()
    wall_time = time.perf_counter() - t0

    print(f"\n  Wall time: {wall_time:.3f}s")
    print(f"  Success: {result.success}")
    if result.success:
        print(f"  Caption: {result.caption[:100]}...")
        print(
            f"  BPM: {result.bpm}, Duration: {result.duration}, "
            f"Key: {result.keyscale}"
        )
        print(
            f"  Language: {result.language}, Time Sig: {result.timesignature}"
        )
        if result.lyrics:
            print(f"  Lyrics: {result.lyrics[:100]}...")
    else:
        print(f"  Error: {result.error}")

    return result, wall_time


# =============================================================================
# Mode: create_sample
# =============================================================================


def run_create_sample_mode(
    dit_handler, llm_handler, args, timer: PreciseTimer
):
    """Profile the create_sample() API (inspiration/simple mode)."""
    if not llm_handler.llm_initialized:
        print("\n  LLM not initialized. create_sample mode requires LLM.")
        sys.exit(1)

    query = args.sample_query or "a soft love song for a quiet evening"
    print(f"\n  Query: {query}")
    print(f"  Instrumental: {args.instrumental}")

    timer.sync()
    t0 = time.perf_counter()

    result = create_sample(
        llm_handler=llm_handler,
        query=query,
        instrumental=args.instrumental,
        temperature=args.lm_temperature,
        use_constrained_decoding=args.use_constrained_decoding,
    )

    timer.sync()
    wall_time = time.perf_counter() - t0

    print(f"\n  Wall time: {wall_time:.3f}s")
    print(f"  Success: {result.success}")
    if result.success:
        print(f"  Caption: {result.caption[:100]}...")
        print(
            f"  BPM: {result.bpm}, Duration: {result.duration}, "
            f"Key: {result.keyscale}"
        )
        print(
            f"  Language: {result.language}, Time Sig: {result.timesignature}"
        )
        print(f"  Instrumental: {result.instrumental}")
        if result.lyrics:
            print(f"  Lyrics: {result.lyrics[:100]}...")
    else:
        print(f"  Error: {result.error}")

    return result, wall_time


# =============================================================================
# Mode: format_sample
# =============================================================================


def run_format_sample_mode(
    dit_handler, llm_handler, args, timer: PreciseTimer
):
    """Profile the format_sample() API."""
    if not llm_handler.llm_initialized:
        print("\n  LLM not initialized. format_sample mode requires LLM.")
        sys.exit(1)

    example_file = os.path.join(
        PROJECT_ROOT, "examples", "text2music", args.example
    )
    if not os.path.exists(example_file):
        print(f"\n  Example not found: {example_file}")
        sys.exit(1)

    with open(example_file, "r", encoding="utf-8") as f:
        data = json.load(f)

    caption = data.get("caption", "Latin pop, reggaeton")
    lyrics = data.get("lyrics", "[Verse 1]\nHola mundo")

    print(f"\n  Caption: {caption[:80]}...")
    print(f"  Lyrics: {lyrics[:80]}...")

    timer.sync()
    t0 = time.perf_counter()

    result = format_sample(
        llm_handler=llm_handler,
        caption=caption,
        lyrics=lyrics,
        temperature=args.lm_temperature,
        use_constrained_decoding=args.use_constrained_decoding,
    )

    timer.sync()
    wall_time = time.perf_counter() - t0

    print(f"\n  Wall time: {wall_time:.3f}s")
    print(f"  Success: {result.success}")
    if result.success:
        print(f"  Caption: {result.caption[:100]}...")
        print(
            f"  BPM: {result.bpm}, Duration: {result.duration}, "
            f"Key: {result.keyscale}"
        )
        print(
            f"  Language: {result.language}, Time Sig: {result.timesignature}"
        )
    else:
        print(f"  Error: {result.error}")

    return result, wall_time


# =============================================================================
# cProfile helper
# =============================================================================


def _print_cprofile(prof):
    """Print cProfile results and save to file."""
    import pstats
    import io

    output_file = "profile_cprofile_detailed.txt"
    with open(output_file, "w") as f:
        ps = pstats.Stats(prof, stream=f)
        ps.sort_stats("cumulative")
        ps.print_stats(100)

    print("\n" + "=" * 100)
    print("TOP 20 FUNCTIONS BY CUMULATIVE TIME (cProfile)")
    print("=" * 100)
    s = io.StringIO()
    ps = pstats.Stats(prof, stream=s)
    ps.sort_stats("cumulative")
    ps.print_stats(20)
    print(s.getvalue())
    print(f"Full report saved to: {output_file}")


def _cleanup_dir(path: str):
    """Remove temporary directory silently."""
    try:
        import shutil
        shutil.rmtree(path, ignore_errors=True)
    except Exception:
        pass


# =============================================================================
# Handler initialization (for non-tier-test modes)
# =============================================================================


def initialize_handlers(
    args, device: str
) -> Tuple[AceStepHandler, LLMHandler]:
    """Initialize DiT and LLM handlers with current API."""
    dit_handler = AceStepHandler()
    llm_handler = LLMHandler()

    # Determine flash attention availability
    use_flash_attention = False
    if device.startswith("cuda"):
        try:
            import flash_attn  # noqa: F401
            use_flash_attention = True
        except ImportError:
            pass

    compile_model = os.environ.get(
        "ACESTEP_COMPILE_MODEL", ""
    ).strip().lower() in {"1", "true", "yes", "y", "on"}

    print("  Initializing DiT handler...")
    status_dit, success_dit = dit_handler.initialize_service(
        project_root=PROJECT_ROOT,
        config_path=args.config_path,
        device=args.device,  # Pass original device string (handler resolves "auto")
        use_flash_attention=use_flash_attention,
        compile_model=compile_model,
        offload_to_cpu=args.offload_to_cpu,
        offload_dit_to_cpu=args.offload_dit_to_cpu,
        quantization=args.quantization,
    )
    if not success_dit:
        print(f"  DiT initialization failed: {status_dit}")
        sys.exit(1)
    print(f"  DiT ready (device={dit_handler.device})")

    # Determine if LLM should be initialized
    need_llm = (
        args.thinking
        or args.use_cot_metas
        or args.use_cot_caption
        or args.use_cot_language
        or args.mode in ("understand", "create_sample", "format_sample")
    )

    if need_llm:
        print(f"  Initializing LLM handler (backend={args.lm_backend})...")
        status_llm, success_llm = llm_handler.initialize(
            checkpoint_dir=os.path.join(PROJECT_ROOT, "checkpoints"),
            lm_model_path=args.lm_model,
            backend=args.lm_backend,
            device=args.device,
            offload_to_cpu=args.offload_to_cpu,
            dtype=None,
        )
        if success_llm:
            print(f"  LLM ready (backend={llm_handler.llm_backend})")
        else:
            print(f"  LLM initialization failed: {status_llm}")
            if args.mode in ("understand", "create_sample", "format_sample"):
                sys.exit(1)
    else:
        print(
            "  LLM not needed for current configuration "
            "(thinking/CoT disabled)"
        )

    return dit_handler, llm_handler


# =============================================================================
# CLI argument parser
# =============================================================================


def build_parser() -> argparse.ArgumentParser:
    """Build the argument parser with all options."""
    env_config = load_env_config()

    parser = argparse.ArgumentParser(
        description="ACE-Step 1.5 Inference Profiler & Benchmark",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python profile_inference.py                                    # Profile text2music
  python profile_inference.py --thinking --llm-debug             # With LLM analysis
  python profile_inference.py --mode benchmark                   # Benchmark matrix
  python profile_inference.py --mode tier-test                   # Test all GPU tiers
  python profile_inference.py --mode tier-test --tiers 6 8 16    # Test specific tiers
  python profile_inference.py --mode tier-test --tier-with-lm    # Test tiers with LM
  python profile_inference.py --mode understand                  # Profile understand API
  python profile_inference.py --mode create_sample --sample-query "jazz ballad"
  python profile_inference.py --device mps --lm-backend mlx      # Apple Silicon
  python profile_inference.py --device cuda --lm-backend vllm    # NVIDIA GPU
""",
    )

    # Mode
    parser.add_argument(
        "--mode",
        type=str,
        default="profile",
        choices=[
            "profile",
            "benchmark",
            "tier-test",
            "understand",
            "create_sample",
            "format_sample",
        ],
        help="Profiling mode (default: profile)",
    )

    # Device & backend
    parser.add_argument(
        "--device",
        type=str,
        default=env_config["ACESTEP_DEVICE"],
        help=(
            f"Device: auto/cuda/mps/cpu "
            f"(default: {env_config['ACESTEP_DEVICE']})"
        ),
    )
    parser.add_argument(
        "--lm-backend",
        type=str,
        default=env_config["ACESTEP_LM_BACKEND"],
        choices=["auto", "vllm", "pt", "mlx"],
        help=(
            f"LLM backend "
            f"(default: {env_config['ACESTEP_LM_BACKEND']})"
        ),
    )

    # Model paths
    parser.add_argument(
        "--config-path",
        type=str,
        default=env_config["ACESTEP_CONFIG_PATH"],
        help=(
            f"DiT model config "
            f"(default: {env_config['ACESTEP_CONFIG_PATH']})"
        ),
    )
    parser.add_argument(
        "--lm-model",
        type=str,
        default=env_config["ACESTEP_LM_MODEL_PATH"],
        help=(
            f"LLM model path "
            f"(default: {env_config['ACESTEP_LM_MODEL_PATH']})"
        ),
    )

    # Hardware options
    parser.add_argument(
        "--offload-to-cpu",
        action="store_true",
        help="Offload models to CPU when not in use",
    )
    parser.add_argument(
        "--offload-dit-to-cpu",
        action="store_true",
        help="Offload DiT to CPU when not in use",
    )
    parser.add_argument(
        "--quantization",
        type=str,
        default=None,
        choices=["int8_weight_only", "fp8_weight_only", "w8a8_dynamic"],
        help="Quantization mode for DiT model",
    )

    # Example & input
    parser.add_argument(
        "--example",
        type=str,
        default="example_05.json",
        help="Example JSON file from examples/text2music/",
    )

    # Task type
    parser.add_argument(
        "--task-type",
        type=str,
        default="text2music",
        choices=[
            "text2music",
            "cover",
            "repaint",
            "lego",
            "extract",
            "complete",
        ],
        help="Generation task type (default: text2music)",
    )
    parser.add_argument(
        "--reference-audio",
        type=str,
        default=None,
        help="Reference audio path (for cover/style transfer)",
    )
    parser.add_argument(
        "--src-audio",
        type=str,
        default=None,
        help="Source audio path (for audio-to-audio tasks)",
    )

    # Generation parameters
    parser.add_argument(
        "--duration",
        type=float,
        default=None,
        help="Audio duration in seconds (overrides example)",
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=None,
        help="Batch size (overrides example)",
    )
    parser.add_argument(
        "--inference-steps",
        type=int,
        default=None,
        help="Diffusion inference steps (overrides example)",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=None,
        help="Random seed (overrides example)",
    )
    parser.add_argument(
        "--guidance-scale",
        type=float,
        default=7.0,
        help="CFG guidance scale for DiT (default: 7.0)",
    )

    # LLM / CoT parameters
    parser.add_argument(
        "--thinking",
        action="store_true",
        help="Enable 5Hz LM Chain-of-Thought reasoning",
    )
    parser.add_argument(
        "--use-cot-metas",
        action="store_true",
        help="Enable LLM to generate music metadata via CoT",
    )
    parser.add_argument(
        "--use-cot-caption",
        action="store_true",
        help="Enable LLM to rewrite/format caption via CoT",
    )
    parser.add_argument(
        "--use-cot-language",
        action="store_true",
        help="Enable LLM to detect vocal language via CoT",
    )
    parser.add_argument(
        "--use-constrained-decoding",
        action="store_true",
        default=True,
        help="Use FSM-based constrained decoding (default: True)",
    )
    parser.add_argument(
        "--no-constrained-decoding",
        action="store_true",
        help="Disable constrained decoding",
    )
    parser.add_argument(
        "--lm-temperature",
        type=float,
        default=0.85,
        help="LLM sampling temperature (default: 0.85)",
    )
    parser.add_argument(
        "--lm-cfg-scale",
        type=float,
        default=2.0,
        help="LLM CFG scale (default: 2.0)",
    )

    # Profiling options
    parser.add_argument(
        "--no-warmup",
        action="store_true",
        help="Skip warmup run (includes compilation overhead)",
    )
    parser.add_argument(
        "--detailed",
        action="store_true",
        help="Enable cProfile function-level analysis",
    )
    parser.add_argument(
        "--llm-debug",
        action="store_true",
        help="Enable deep LLM debugging (token count, throughput)",
    )

    # Benchmark options
    parser.add_argument(
        "--benchmark-output",
        type=str,
        default=None,
        help="Save benchmark results to JSON file",
    )

    # Tier-test options
    parser.add_argument(
        "--tiers",
        type=int,
        nargs="+",
        default=None,
        help="Specific VRAM tiers to test (e.g., --tiers 6 8 16). Default: all tiers",
    )
    parser.add_argument(
        "--tier-with-lm",
        action="store_true",
        help="Enable LM for tiers that support it (default: DiT-only test)",
    )
    parser.add_argument(
        "--tier-duration",
        type=float,
        default=240,
        help="Test generation duration in seconds for tier-test (default: 240)",
    )
    parser.add_argument(
        "--tier-skip-compile",
        action="store_true",
        help="Skip torch.compile for non-quantized tiers (faster testing, less realistic)",
    )
    parser.add_argument(
        "--tier-boundary",
        action="store_true",
        help="Enable boundary testing: for each tier, also test without INT8 quantization "
             "and without CPU offload to find the minimum VRAM tier for each capability",
    )
    parser.add_argument(
        "--tier-batch-boundary",
        action="store_true",
        help="Enable batch size boundary testing: for each tier, progressively test "
             "batch sizes 1, 2, 4, 8 (stop at first OOM) to find the maximum safe batch "
             "size. Tests both with-LM and without-LM configurations.",
    )

    # create_sample / understand options
    parser.add_argument(
        "--sample-query",
        type=str,
        default=None,
        help="Query for create_sample mode",
    )
    parser.add_argument(
        "--instrumental",
        action="store_true",
        help="Generate instrumental music (for create_sample)",
    )
    parser.add_argument(
        "--audio-codes",
        type=str,
        default=None,
        help="Audio codes string (for understand mode)",
    )

    return parser


# =============================================================================
# Main
# =============================================================================


def main():
    parser = build_parser()
    args = parser.parse_args()

    # Handle --no-constrained-decoding
    if args.no_constrained_decoding:
        args.use_constrained_decoding = False

    # Tier-test mode has its own initialization flow
    if args.mode == "tier-test":
        print("=" * 120)
        print("ACE-Step 1.5 Tier Compatibility Test")
        print("=" * 120)
        run_tier_test_mode(args)
        print("\n" + "=" * 120)
        print("DONE")
        print("=" * 120)
        return

    # Resolve device
    device = resolve_device(args.device)

    # Auto-detect backend
    if args.lm_backend == "auto":
        args.lm_backend = auto_detect_backend(device)

    # Setup GPU config
    gpu_config = get_gpu_config()
    set_global_gpu_config(gpu_config)

    # Auto-enable offload for small GPUs
    if (
        gpu_config.gpu_memory_gb > 0
        and gpu_config.gpu_memory_gb < VRAM_AUTO_OFFLOAD_THRESHOLD_GB
        and not args.offload_to_cpu
    ):
        args.offload_to_cpu = True

    # Print header
    print("=" * 100)
    print("ACE-Step 1.5 Inference Profiler")
    print("=" * 100)
    print(f"\n  Mode:           {args.mode}")
    print(f"  Device:         {device} (requested: {args.device})")
    print(f"  LLM Backend:    {args.lm_backend}")
    print(f"  DiT Config:     {args.config_path}")
    print(f"  LLM Model:      {args.lm_model}")
    print(
        f"  GPU Memory:     {gpu_config.gpu_memory_gb:.1f} GB "
        f"(tier: {gpu_config.tier})"
    )
    if args.quantization:
        print(f"  Quantization:   {args.quantization}")
    if args.offload_to_cpu:
        print("  CPU Offload:    enabled")
    print(f"\n  Thinking:       {args.thinking}")
    print(f"  CoT Metas:      {args.use_cot_metas}")
    print(f"  CoT Caption:    {args.use_cot_caption}")
    print(f"  CoT Language:   {args.use_cot_language}")
    print(f"  Constrained:    {args.use_constrained_decoding}")
    print(f"  Warmup:         {'disabled' if args.no_warmup else 'enabled'}")

    # Initialize handlers
    print("\n" + "-" * 100)
    print("INITIALIZING MODELS")
    print("-" * 100)

    dit_handler, llm_handler = initialize_handlers(args, device)

    # Create timer with resolved device
    actual_device = getattr(dit_handler, "device", device)
    timer = PreciseTimer(device=actual_device)

    # Dispatch to mode
    print("\n" + "=" * 100)
    print(f"RUNNING MODE: {args.mode.upper()}")
    print("=" * 100)

    if args.mode == "profile":
        run_profile_mode(dit_handler, llm_handler, args, timer)
    elif args.mode == "benchmark":
        run_benchmark_mode(dit_handler, llm_handler, args, timer)
    elif args.mode == "understand":
        run_understand_mode(dit_handler, llm_handler, args, timer)
    elif args.mode == "create_sample":
        run_create_sample_mode(dit_handler, llm_handler, args, timer)
    elif args.mode == "format_sample":
        run_format_sample_mode(dit_handler, llm_handler, args, timer)

    print("\n" + "=" * 100)
    print("DONE")
    print("=" * 100)


if __name__ == "__main__":
    main()