ACE-Step/Ace-Step-v1.5

1#!/usr/bin/env python32"""3Enhanced profiling script for ACE-Step inference with deep LLM analysis4 5This script helps diagnose why LLM generation is slow by tracking:61. Total tokens generated vs expected throughput (200 tokens/sec baseline)72. Per-iteration timing to detect compilation overhead or slow operations83. Constrained decoding overhead94. CFG overhead (2x forward passes)105. Model forward time vs sampling/processing time11 12Usage:13    python profile_inference.py                    # Standard profiling with warmup14    python profile_inference.py --no-warmup        # Profile first run (includes compilation)15    python profile_inference.py --llm-debug        # Deep LLM performance debugging16    python profile_inference.py --detailed         # Add cProfile function-level analysis17    18    Inference mode options:19    python profile_inference.py --thinking                        # Enable CoT for code generation20    python profile_inference.py --use-constrained-decoding        # Use FSM constrained decoding21    python profile_inference.py --use-cot-metas                  # Enable LM to generate metadata via CoT22"""23 24import time25import argparse26import sys27import os28from contextlib import contextmanager29from collections import defaultdict30import json31from typing import Tuple, Dict, Any, List32from functools import wraps33 34# Add project root to path35project_root = os.path.abspath(os.path.dirname(__file__))36if project_root not in sys.path:37    sys.path.insert(0, project_root)38 39 40def load_env_config():41    """从 .env 文件加载配置"""42    env_config = {43        'ACESTEP_CONFIG_PATH': 'acestep-v15-turbo',44        'ACESTEP_LM_MODEL_PATH': 'acestep-5Hz-lm-0.6B',45        'ACESTEP_DEVICE': 'auto',46        'ACESTEP_LM_BACKEND': 'vllm',47    }48    49    env_file = os.path.join(project_root, '.env')50    if os.path.exists(env_file):51        with open(env_file, 'r', encoding='utf-8') as f:52            for line in f:53                line = line.strip()54                # 跳过空行和注释55                if not line or line.startswith('#'):56                    continue57                # 解析键值对58                if '=' in line:59                    key, value = line.split('=', 1)60                    key = key.strip()61                    value = value.strip()62                    if key in env_config and value:63                        env_config[key] = value64    65    return env_config66 67import torch68from acestep.inference import generate_music, GenerationParams, GenerationConfig69from acestep.handler import AceStepHandler70from acestep.llm_inference import LLMHandler71 72 73class PreciseTimer:74    """High-precision timer with CUDA synchronization for accurate GPU timing"""75    76    def __init__(self, device="cuda"):77        self.device = device78        self.timings = defaultdict(list)79        self.enabled = True80        81    def sync(self):82        """Synchronize CUDA operations for accurate timing"""83        if self.enabled and self.device.startswith("cuda") and torch.cuda.is_available():84            torch.cuda.synchronize()85    86    @contextmanager87    def time(self, name: str):88        """Time a code section with CUDA synchronization"""89        if not self.enabled:90            yield91            return92            93        self.sync()94        start = time.perf_counter()95        try:96            yield97        finally:98            self.sync()99            elapsed = time.perf_counter() - start100            self.timings[name].append(elapsed)101    102    def get_total(self, name: str) -> float:103        """Get total accumulated time for a section"""104        return sum(self.timings.get(name, []))105    106    def get_mean(self, name: str) -> float:107        """Get mean time per call for a section"""108        times = self.timings.get(name, [])109        return sum(times) / len(times) if times else 0.0110    111    def get_count(self, name: str) -> int:112        """Get number of calls for a section"""113        return len(self.timings.get(name, []))114    115    def get_all(self, name: str) -> List[float]:116        """Get all timing samples for a section"""117        return self.timings.get(name, [])118 119 120class LLMDebugger:121    """Track detailed LLM performance metrics to diagnose slow generation"""122    123    def __init__(self):124        self.reset()125    126    def reset(self):127        """Reset all metrics"""128        self.total_tokens = 0129        self.generation_start = None130        self.generation_end = None131        self.output_text = ""132        self.prompt_length = 0133        134    def start(self, prompt_length: int = 0):135        """Mark generation start"""136        self.generation_start = time.perf_counter()137        self.prompt_length = prompt_length138        139    def end(self, output_text: str = ""):140        """Mark generation end and store output"""141        self.generation_end = time.perf_counter()142        self.output_text = output_text143        144    def set_token_count(self, count: int):145        """Set total token count"""146        self.total_tokens = count147        148    def get_throughput(self) -> float:149        """Calculate actual tokens per second"""150        if self.generation_start and self.generation_end and self.total_tokens > 0:151            total_time = self.generation_end - self.generation_start152            if total_time > 0:153                return self.total_tokens / total_time154        return 0.0155    156    def print_analysis(self):157        """Print detailed LLM performance analysis"""158        if not self.generation_start or not self.generation_end:159            return160            161        print("\n" + "=" * 100)162        print("🔍 LLM PERFORMANCE DEEP DIVE")163        print("=" * 100)164        165        total_time = self.generation_end - self.generation_start166        throughput = self.get_throughput()167        168        # Basic metrics table169        print(f"\n{'Metric':<40} {'Value':<20} {'Notes'}")170        print("-" * 100)171        print(f"{'Total Tokens Generated:':<40} {self.total_tokens:<20} (new tokens only)")172        print(f"{'Prompt Length (estimate):':<40} {self.prompt_length:<20} (input tokens)")173        print(f"{'Total Generation Time:':<40} {total_time:<20.3f} seconds")174        print(f"{'Measured Throughput:':<40} {throughput:<20.1f} tokens/sec")175        print(f"{'Expected Throughput:':<40} {'200':<20} tokens/sec (baseline)")176        177        # Calculate performance gap178        if throughput > 0:179            slowdown = 200.0 / throughput180            efficiency = (throughput / 200.0) * 100181            print(f"{'Performance vs Baseline:':<40} {efficiency:<20.1f}% of expected")182            print(f"{'Slowdown Factor:':<40} {slowdown:<20.2f}x slower")183        184        # Analyze generated output185        if self.output_text:186            print(f"\n{'Output Analysis:':<40}")187            print(f"{'  Output length:':<40} {len(self.output_text):<20} characters")188            189            # Count audio codes190            import re191            code_pattern = r'<\|audio_code_\d+\|>'192            codes = re.findall(code_pattern, self.output_text)193            if codes:194                print(f"{'  Audio codes generated:':<40} {len(codes):<20} codes")195                print(f"{'  Expected audio duration:':<40} {f'~{len(codes)/5:.1f}s':<20} (5 codes per second)")196                if total_time > 0:197                    print(f"{'  Time per audio code:':<40} {f'{total_time/len(codes)*1000:.1f}ms':<20}")198            199            # Check for CoT section200            if '<think>' in self.output_text and '</think>' in self.output_text:201                cot_start = self.output_text.find('<think>')202                cot_end = self.output_text.find('</think>') + 8203                cot_section = self.output_text[cot_start:cot_end]204                cot_token_est = len(cot_section) // 4205                print(f"{'  CoT section tokens (estimate):':<40} {f'~{cot_token_est}':<20}")206        207        # Diagnostic guidance208        print("\n" + "=" * 100)209        print("🔧 DIAGNOSTIC GUIDANCE")210        print("=" * 100)211        212        if throughput < 50:213            print("\n⚠️  CRITICAL: Throughput is extremely low (<50 tokens/sec)")214            print("\nThis is ~4x slower than expected. Likely causes:")215            print("  1. ❗ Constrained decoding FSM overhead")216            print("     → Each token triggers FSM state machine validation")217            print("     → Try: set use_constrained_decoding=False in config")218            print("  2. ❗ CFG with double forward passes")219            print("     → cfg_scale > 1.0 means running model twice per token")220            print("     → Check: params.lm_cfg_scale value")221            print("  3. ❗ Running in eager mode without compilation")222            print("     → PyTorch should compile kernels after warmup")223            print("     → Check: torch._dynamo.config settings")224            225        elif throughput < 100:226            print("\n⚠️  WARNING: Throughput is low (50-100 tokens/sec)")227            print("\nLikely causes:")228            print("  1. Constrained decoding overhead (~30-50% slowdown expected)")229            print("  2. CFG enabled (2x compute per token if cfg_scale > 1.0)")230            print("  3. Small model or inefficient GPU utilization")231            232        elif throughput < 150:233            print("\n⚠️  Throughput is below baseline but acceptable (100-150 tokens/sec)")234            print("\nMinor overhead from:")235            print("  - Constrained decoding: ~20-30% overhead")236            print("  - Profiling instrumentation: ~5-10% overhead")237            238        else:239            print(f"\n✓ Throughput is good ({throughput:.1f} tokens/sec)")240            print("  Performance is within acceptable range")241 242 243# Global instances244timer = None245llm_debugger = None246 247 248def wrap_method_with_timing(obj, method_name: str, timing_key: str):249    """Wrap a method with timing instrumentation"""250    original_method = getattr(obj, method_name)251    252    @wraps(original_method)253    def timed_wrapper(*args, **kwargs):254        with timer.time(timing_key):255            return original_method(*args, **kwargs)256    257    setattr(obj, method_name, timed_wrapper)258    return original_method259 260 261def wrap_llm_with_debug_tracking(llm_handler):262    """Wrap LLM generation with detailed performance tracking"""263    original_method = llm_handler.generate_with_stop_condition264    265    @wraps(original_method)266    def debug_wrapper(*args, **kwargs):267        # Estimate prompt length268        caption = kwargs.get('caption', args[0] if len(args) > 0 else "")269        lyrics = kwargs.get('lyrics', args[1] if len(args) > 1 else "")270        prompt_estimate = len(caption) + len(lyrics)271        prompt_tokens_estimate = prompt_estimate // 4272        273        # Start tracking274        llm_debugger.reset()275        llm_debugger.start(prompt_length=prompt_tokens_estimate)276        277        # Call original with timing278        with timer.time('llm_inference'):279            result = original_method(*args, **kwargs)280        281        # Extract and analyze output282        output_text = ""283        if isinstance(result, tuple) and len(result) >= 2:284            if isinstance(result[1], list):285                # Batch mode286                output_text = "".join(result[1])287            else:288                # Single mode289                cot_output = ""290                if isinstance(result[0], dict):291                    for v in result[0].values():292                        if isinstance(v, str):293                            cot_output += v294                output_text = cot_output + str(result[1])295        296        # Count tokens297        import re298        code_pattern = r'<\|audio_code_\d+\|>'299        codes = re.findall(code_pattern, output_text)300        remaining_text = re.sub(code_pattern, '', output_text)301        cot_tokens_estimate = len(remaining_text) // 4302        total_tokens = len(codes) + cot_tokens_estimate303        304        llm_debugger.set_token_count(total_tokens)305        llm_debugger.end(output_text)306        307        return result308    309    llm_handler.generate_with_stop_condition = debug_wrapper310    return original_method311 312 313def instrument_handlers(dit_handler, llm_handler, enable_llm_debug=False):314    """Add timing instrumentation to handler methods"""315    originals = {}316    317    # Instrument LLM318    if llm_handler and llm_handler.llm_initialized:319        if enable_llm_debug:320            originals['llm_generate'] = wrap_llm_with_debug_tracking(llm_handler)321        else:322            originals['llm_generate'] = wrap_method_with_timing(323                llm_handler, 'generate_with_stop_condition', 'llm_inference'324            )325    326    # Instrument DiT handler327    originals['dit_prepare'] = wrap_method_with_timing(328        dit_handler, 'prepare_batch_data', 'prepare_batch_data'329    )330    originals['dit_generate'] = wrap_method_with_timing(331        dit_handler, 'service_generate', 'dit_inference'332    )333    originals['dit_decode'] = wrap_method_with_timing(334        dit_handler, 'tiled_decode', 'vae_decode'335    )336    337    return originals338 339 340def restore_handlers(dit_handler, llm_handler, originals):341    """Restore original handler methods after profiling"""342    if llm_handler and 'llm_generate' in originals:343        llm_handler.generate_with_stop_condition = originals['llm_generate']344    345    dit_handler.prepare_batch_data = originals['dit_prepare']346    dit_handler.service_generate = originals['dit_generate']347    dit_handler.tiled_decode = originals['dit_decode']348 349 350def print_profiling_results(total_time: float, show_llm_debug: bool = False):351    """Print comprehensive profiling results with performance insights"""352    print("\n" + "=" * 100)353    print("🎯 PROFILING RESULTS")354    print("=" * 100)355    356    # Define timing categories357    model_sections = {358        'llm_inference': 'LLM Inference (5Hz Language Model)',359        'dit_inference': 'DiT Inference (Diffusion Transformer)',360        'vae_decode': 'VAE Decode (Audio Decoder)',361    }362    363    non_model_sections = {364        'prepare_batch_data': 'Prepare Batch Data (embedding, formatting)',365    }366    367    # Calculate totals368    model_time = sum(timer.get_total(k) for k in model_sections.keys())369    non_model_time = sum(timer.get_total(k) for k in non_model_sections.keys())370    other_time = total_time - model_time - non_model_time371    372    # Print summary table373    print(f"\n{'CATEGORY':<50} {'TIME (s)':<12} {'%':<8} {'CALLS':<8}")374    print("-" * 100)375    376    # Model time breakdown377    print(f"\n{'🤖 MODEL TIME (Total)':<50} {model_time:<12.3f} {100*model_time/total_time:>6.1f}% {'':<8}")378    for key, desc in model_sections.items():379        t = timer.get_total(key)380        c = timer.get_count(key)381        if c > 0:382            mean = timer.get_mean(key)383            pct = 100 * t / total_time384            print(f"  {'├─ ' + desc:<48} {t:<12.3f} {pct:>6.1f}% {c:<8} (avg: {mean:.3f}s)")385    386    # Non-model time breakdown387    print(f"\n{'⚙️  NON-MODEL TIME (Total)':<50} {non_model_time:<12.3f} {100*non_model_time/total_time:>6.1f}% {'':<8}")388    for key, desc in non_model_sections.items():389        t = timer.get_total(key)390        c = timer.get_count(key)391        if c > 0:392            mean = timer.get_mean(key)393            pct = 100 * t / total_time394            print(f"  {'├─ ' + desc:<48} {t:<12.3f} {pct:>6.1f}% {c:<8} (avg: {mean:.3f}s)")395    396    # Other time397    if other_time > 0.01:398        pct = 100 * other_time / total_time399        print(f"\n{'📦 OTHER TIME (I/O, overhead, audio save)':<50} {other_time:<12.3f} {pct:>6.1f}% {'':<8}")400    401    print(f"\n{'📊 TOTAL TIME':<50} {total_time:<12.3f} {'100.0%':>6} {'':<8}")402    403    # Show LLM detailed analysis if enabled404    if show_llm_debug:405        llm_debugger.print_analysis()406    407    # Performance insights408    print("\n" + "=" * 100)409    print("💡 PERFORMANCE INSIGHTS")410    print("=" * 100)411    412    llm_t = timer.get_total('llm_inference')413    dit_t = timer.get_total('dit_inference')414    vae_t = timer.get_total('vae_decode')415    prep_t = timer.get_total('prepare_batch_data')416    417    # Model time insights418    if model_time > 0:419        print(f"\n✓ Model operations: {model_time:.3f}s ({100*model_time/total_time:.1f}% of total)")420        421        if llm_t > 0:422            print(f"  - LLM: {llm_t:.3f}s ({100*llm_t/model_time:.1f}% of model time)")423        if dit_t > 0:424            print(f"  - DiT: {dit_t:.3f}s ({100*dit_t/model_time:.1f}% of model time)")425        if vae_t > 0:426            print(f"  - VAE: {vae_t:.3f}s ({100*vae_t/model_time:.1f}% of model time)")427    428    # LLM bottleneck analysis429    if llm_t > dit_t and llm_t > 5.0:430        print(f"\n⚠️  LLM IS THE BOTTLENECK: {llm_t:.3f}s ({100*llm_t/total_time:.1f}% of total)")431        print(f"\n   Possible causes:")432        print(f"   1. Generating too many tokens → use --llm-debug to verify")433        print(f"   2. Constrained decoding overhead → FSM validation per token")434        print(f"   3. CFG overhead → cfg_scale > 1.0 = 2x forward passes")435        print(f"   4. First-token latency → warmup should help")436        print(f"   5. KV cache inefficiency → should be ~5-10ms/token")437    438    # Non-model insights439    if non_model_time / total_time > 0.1:440        print(f"\n⚠️  Non-model operations: {non_model_time:.3f}s ({100*non_model_time/total_time:.1f}%)")441        if prep_t > 0.1:442            print(f"   - Batch preparation: {prep_t:.3f}s")443    444    # I/O overhead445    if other_time / total_time > 0.2:446        print(f"\n⚠️  Overhead/I/O: {other_time:.3f}s ({100*other_time/total_time:.1f}%)")447    448    # Recommendations449    print("\n" + "=" * 100)450    print("🚀 OPTIMIZATION RECOMMENDATIONS")451    print("=" * 100)452    453    if llm_t > dit_t * 2:454        print("\n🎯 Priority: Optimize LLM")455        print("  1. Run: python profile_inference.py --llm-debug")456        print("     → Shows exact token count and throughput")457        print("  2. Check constrained decoding overhead")458        print("  3. Check CFG scaling (lm_cfg_scale parameter)")459        print("  4. Profile nanovllm engine step() timing")460        print("  5. Compare vllm vs transformers backends")461 462 463def run_profiled_generation(dit_handler, llm_handler, params, config,464                           enable_cprofile=False, enable_llm_debug=False):465    """Execute generation with full profiling instrumentation"""466    # Instrument handlers467    originals = instrument_handlers(dit_handler, llm_handler, enable_llm_debug)468    469    try:470        print("\n[Profiling] Starting generation...")471        timer.sync()472        total_start = time.perf_counter()473        474        # Optional cProfile475        prof = None476        if enable_cprofile:477            import cProfile478            prof = cProfile.Profile()479            prof.enable()480        481        # Run generation482        result = generate_music(dit_handler, llm_handler, params, config, save_dir="./")483        484        # Stop timing485        timer.sync()486        total_time = time.perf_counter() - total_start487        488        # Save cProfile if enabled489        if enable_cprofile and prof:490            prof.disable()491            492            import pstats493            import io494            495            output_file = "profile_cprofile_detailed.txt"496            with open(output_file, 'w') as f:497                ps = pstats.Stats(prof, stream=f)498                ps.sort_stats('cumulative')499                ps.print_stats(100)500            501            # Print top functions502            print("\n" + "=" * 100)503            print("📊 TOP 20 FUNCTIONS BY CUMULATIVE TIME (cProfile)")504            print("=" * 100)505            s = io.StringIO()506            ps = pstats.Stats(prof, stream=s)507            ps.sort_stats('cumulative')508            ps.print_stats(20)509            print(s.getvalue())510            511            print(f"\nFull report: {output_file}")512        513        # Print results514        print_profiling_results(total_time, show_llm_debug=enable_llm_debug)515        516        return result, total_time517        518    finally:519        restore_handlers(dit_handler, llm_handler, originals)520 521 522def load_example_config(example_file: str) -> Tuple[GenerationParams, GenerationConfig]:523    """Load configuration from example JSON file"""524    try:525        with open(example_file, 'r', encoding='utf-8') as f:526            data = json.load(f)527        528        params = GenerationParams(529            caption=data.get('caption', ''),530            lyrics=data.get('lyrics', ''),531            bpm=data.get('bpm'),532            keyscale=data.get('keyscale', ''),533            timesignature=data.get('timesignature', ''),534            vocal_language=data.get('language', 'unknown'),535            duration=data.get('duration'),536            thinking=data.get('think', False),537            inference_steps=data.get('inference_steps', 8),538            seed=data.get('seed', 42),539        )540        541        config = GenerationConfig(batch_size=data.get('batch_size', 1), seeds=[42])542        543        return params, config544        545    except Exception as e:546        print(f"  ❌ Failed to load: {e}")547        return None, None548 549 550def main():551    global timer, llm_debugger552    553    # 从 .env 文件加载默认配置554    env_config = load_env_config()555    556    parser = argparse.ArgumentParser(557        description="Profile ACE-Step inference with LLM debugging"558    )559    parser.add_argument("--checkpoint-dir", type=str, default="./checkpoints")560    parser.add_argument("--config-path", type=str, default=env_config['ACESTEP_CONFIG_PATH'],561                       help=f"模型配置路径 (默认从 .env: {env_config['ACESTEP_CONFIG_PATH']})")562    parser.add_argument("--device", type=str, default=env_config['ACESTEP_DEVICE'],563                       help=f"设备 (默认从 .env: {env_config['ACESTEP_DEVICE']})")564    parser.add_argument("--lm-model", type=str, default=env_config['ACESTEP_LM_MODEL_PATH'],565                       help=f"LLM 模型路径 (默认从 .env: {env_config['ACESTEP_LM_MODEL_PATH']})")566    parser.add_argument("--lm-backend", type=str, default=env_config['ACESTEP_LM_BACKEND'],567                       help=f"LLM 后端 (默认从 .env: {env_config['ACESTEP_LM_BACKEND']})")568    parser.add_argument("--no-warmup", action="store_true")569    parser.add_argument("--detailed", action="store_true")570    parser.add_argument("--llm-debug", action="store_true",571                       help="Enable deep LLM debugging (token count, throughput)")572    parser.add_argument("--example", type=str, default="example_05.json")573    574    # Inference mode parameters575    parser.add_argument("--thinking", action="store_true",576                       help="Enable CoT reasoning for LM to generate audio codes")577    parser.add_argument("--use-constrained-decoding", action="store_true",578                       help="Use FSM-based constrained decoding for meta generation")579    parser.add_argument("--use-cot-metas", action="store_true",580                       help="Enable LLM to generate music metadata via CoT reasoning")581    582    args = parser.parse_args()583    584    # Initialize585    timer = PreciseTimer(device=args.device)586    llm_debugger = LLMDebugger()587    588    print("=" * 100)589    print("🎵 ACE-Step Inference Profiler (LLM Performance Analysis)")590    print("=" * 100)591    print(f"\n模型配置 (从 .env 加载):")592    print(f"  DiT 模型: {args.config_path}")593    print(f"  LLM 模型: {args.lm_model}")594    print(f"\n运行配置:")595    print(f"  Device: {args.device}")596    print(f"  LLM Backend: {args.lm_backend}")597    print(f"  LLM Debug: {'Enabled' if args.llm_debug else 'Disabled'}")598    print(f"  Warmup: {'Disabled' if args.no_warmup else 'Enabled'}")599    print(f"\nInference Mode:")600    print(f"  Thinking (CoT): {'Enabled' if args.thinking else 'Disabled'}")601    print(f"  Constrained Decoding: {'Enabled' if args.use_constrained_decoding else 'Disabled'}")602    print(f"  Use CoT for Metas: {'Enabled' if args.use_cot_metas else 'Disabled'}")603    604    # Initialize models605    print(f"\nInitializing models...")606    607    dit_handler = AceStepHandler()608    llm_handler = LLMHandler()609    610    print("  🎹 Initializing DiT...")611    status_dit, success_dit = dit_handler.initialize_service(612        project_root=project_root,613        config_path=args.config_path,614        device=args.device,615        use_flash_attention=True,616    )617    if not success_dit:618        print(f"  ❌ Failed: {status_dit}")619        sys.exit(1)620    print(f"     ✓ DiT ready")621    622    print("  🧠 Initializing LLM...")623    if args.thinking or args.use_cot_metas:624        status_llm, success_llm = llm_handler.initialize(625            checkpoint_dir=args.checkpoint_dir,626            lm_model_path=args.lm_model,627            backend=args.lm_backend,628            device=args.device,629        )630        if success_llm:631            print(f"     ✓ LLM ready ({args.lm_backend})")632        else:633            print(f"     ⚠ Failed: {status_llm}")634    else:635        print(f"     ✓ LLM not initialized (thinking or use_cot_metas is disabled)")636    637    # Load example638    example_file = os.path.join(project_root, "examples", "text2music", args.example)639    if not os.path.exists(example_file):640        print(f"\n❌ Not found: {example_file}")641        sys.exit(1)642    643    print(f"\n📄 Loading: {args.example}")644    params, config = load_example_config(example_file)645    646    if not params or not config:647        print("❌ Failed to load config")648        sys.exit(1)649    650    print(f"   Caption: {params.caption[:60]}...")651    print(f"   Batch: {config.batch_size}, Steps: {params.inference_steps}, LLM: {params.thinking}")652    653    # Warmup654    if not args.no_warmup:655        print("\n" + "=" * 100)656        print("🔥 WARMUP RUN")657        print("=" * 100)658        659        warmup_params = GenerationParams(660            caption=params.caption,661            lyrics=params.lyrics,662            bpm=params.bpm,663            keyscale=params.keyscale,664            timesignature=params.timesignature,665            vocal_language=params.vocal_language,666            duration=params.duration,667            thinking=args.thinking,668            use_cot_metas=args.use_cot_metas,669            inference_steps=params.inference_steps,670            seed=params.seed,671        )672        warmup_config = GenerationConfig(batch_size=1, seeds=[42])673        warmup_config.use_constrained_decoding = args.use_constrained_decoding674        675        warmup_start = time.perf_counter()676        warmup_result = generate_music(dit_handler, llm_handler, warmup_params, warmup_config, save_dir="./")677        warmup_time = time.perf_counter() - warmup_start678        679        print(f"\n✓ Warmup: {warmup_time:.2f}s")680        if not warmup_result.success:681            print(f"⚠️  Warning: {warmup_result.error}")682        683        # Reset684        timer = PreciseTimer(device=args.device)685        llm_debugger = LLMDebugger()686    687    # Profiling run688    print("\n" + "=" * 100)689    print("⏱️  PROFILING RUN")690    print("=" * 100)691    692    # Apply inference mode settings693    config.use_constrained_decoding = args.use_constrained_decoding694    # Override thinking and use_cot_metas parameters if specified via CLI695    if args.thinking:696        params.thinking = True697    if args.use_cot_metas:698        params.use_cot_metas = True699    700    result, total_time = run_profiled_generation(701        dit_handler, llm_handler, params, config,702        enable_cprofile=args.detailed,703        enable_llm_debug=args.llm_debug704    )705    706    if not result.success:707        print(f"\n❌ Failed: {result.error}")708        sys.exit(1)709    710    print(f"\n✅ Success! Generated {len(result.audios)} audio file(s)")711    712    # Final tips713    if args.detailed:714        print("\n💡 Check profile_cprofile_detailed.txt for function-level analysis")715    elif not args.llm_debug:716        print("\n💡 Run with --llm-debug to see LLM token count and throughput analysis")717 718 719if __name__ == "__main__":720    main()721