optimize_cpu.py

"""
AICoverGen CPU Performance Optimizer
Place this file in the AICoverGen root directory and run it instead of main.py

Usage:
    python optimize_cpu.py -i "song_link_or_path" -dir "model_folder_name" -p 0

This script automatically applies all CPU optimizations including:
- PyTorch threading optimization
- Faster processing parameters
- Intel MKL optimizations (if available)
"""

import os
import sys
import argparse
import subprocess
import multiprocessing

# ========== CPU OPTIMIZATION SETTINGS ==========
def setup_cpu_optimizations():
    """Apply CPU-specific optimizations before running"""
    
    # Get optimal thread count (number of physical cores)
    cpu_count = multiprocessing.cpu_count()
    
    # Set PyTorch threading for better CPU performance
    # These reduce overhead and improve throughput
    os.environ['OMP_NUM_THREADS'] = str(cpu_count)
    os.environ['MKL_NUM_THREADS'] = str(cpu_count)
    os.environ['OPENBLAS_NUM_THREADS'] = str(cpu_count)
    os.environ['VECLIB_MAXIMUM_THREADS'] = str(cpu_count)
    os.environ['NUMEXPR_NUM_THREADS'] = str(cpu_count)
    
    # Enable Intel MKL optimizations if available
    try:
        import intel_extension_for_pytorch as ipex
        print("✓ Intel Extension for PyTorch detected - extra optimizations enabled!")
    except ImportError:
        print("ℹ Intel Extension not installed (optional). Install with: pip install intel-extension-for-pytorch")
    
    print(f"✓ CPU optimizations applied (using {cpu_count} threads)")

# ========== OPTIMIZED PARAMETER DEFAULTS ==========
def get_optimized_args():
    """Parse arguments with CPU-optimized defaults"""
    
    parser = argparse.ArgumentParser(
        description='AICoverGen with CPU optimizations',
        formatter_class=argparse.RawDescriptionHelpFormatter
    )
    
    # Required arguments
    parser.add_argument('-i', '--song-input', type=str, required=True,
                        help='YouTube link or local audio file path')
    parser.add_argument('-dir', '--rvc-dirname', type=str, required=True,
                        help='Name of folder in rvc_models directory')
    parser.add_argument('-p', '--pitch-change', type=int, required=True,
                        help='Pitch change in octaves: 1(male->female), -1(female->male), 0(no change)')
    
    # Optimized optional arguments (defaults set for speed)
    parser.add_argument('-k', '--keep-files', action='store_true',
                        help='Keep intermediate files (default: False for speed)')
    parser.add_argument('-ir', '--index-rate', type=float, default=0.3,
                        help='Index rate (default: 0.3, lower=faster)')
    parser.add_argument('-fr', '--filter-radius', type=int, default=0,
                        help='Median filtering (default: 0 for speed, normally 3)')
    parser.add_argument('-rms', '--rms-mix-rate', type=float, default=0.25,
                        help='RMS mix rate (default: 0.25)')
    parser.add_argument('-palgo', '--pitch-detection-algo', type=str, default='rmvpe',
                        choices=['rmvpe', 'mangio-crepe'],
                        help='Pitch detection (default: rmvpe - fastest)')
    parser.add_argument('-hop', '--crepe-hop-length', type=int, default=128,
                        help='Crepe hop length (default: 128, higher=faster)')
    parser.add_argument('-pro', '--protect', type=float, default=0.15,
                        help='Protect consonants (default: 0.15, lower=faster)')
    parser.add_argument('-mv', '--main-vol', type=int, default=0,
                        help='Main vocals volume change in dB (default: 0)')
    parser.add_argument('-bv', '--backup-vol', type=int, default=0,
                        help='Backup vocals volume change in dB (default: 0)')
    parser.add_argument('-iv', '--inst-vol', type=int, default=0,
                        help='Instrumental volume change in dB (default: 0)')
    parser.add_argument('-pall', '--pitch-change-all', type=int, default=0,
                        help='Pitch change for all audio in semitones (default: 0)')
    parser.add_argument('-rsize', '--reverb-size', type=float, default=0.0,
                        help='Reverb room size (default: 0.0 for speed)')
    parser.add_argument('-rwet', '--reverb-wetness', type=float, default=0.0,
                        help='Reverb wetness (default: 0.0 for speed)')
    parser.add_argument('-rdry', '--reverb-dryness', type=float, default=0.8,
                        help='Reverb dryness (default: 0.8)')
    parser.add_argument('-rdamp', '--reverb-damping', type=float, default=0.7,
                        help='Reverb damping (default: 0.7)')
    parser.add_argument('-oformat', '--output-format', type=str, default='mp3',
                        choices=['mp3', 'wav'],
                        help='Output format (default: mp3 for smaller size)')
    
    return parser.parse_args()

# ========== MAIN EXECUTION ==========
def main():
    print("=" * 60)
    print("AICoverGen CPU Performance Optimizer")
    print("=" * 60)
    
    # Apply CPU optimizations first
    setup_cpu_optimizations()
    
    # Parse arguments
    args = get_optimized_args()
    
    # Build command for main.py
    cmd = [
        sys.executable,
        'src/main.py',
        '-i', args.song_input,
        '-dir', args.rvc_dirname,
        '-p', str(args.pitch_change),
        '-ir', str(args.index_rate),
        '-fr', str(args.filter_radius),
        '-rms', str(args.rms_mix_rate),
        '-palgo', args.pitch_detection_algo,
        '-hop', str(args.crepe_hop_length),
        '-pro', str(args.protect),
        '-mv', str(args.main_vol),
        '-bv', str(args.backup_vol),
        '-iv', str(args.inst_vol),
        '-pall', str(args.pitch_change_all),
        '-rsize', str(args.reverb_size),
        '-rwet', str(args.reverb_wetness),
        '-rdry', str(args.reverb_dryness),
        '-rdamp', str(args.reverb_damping),
        '-oformat', args.output_format
    ]
    
    if args.keep_files:
        cmd.append('-k')
    
    print("\n📊 Optimized Settings Applied:")
    print(f"  • Threading: {os.environ['OMP_NUM_THREADS']} threads")
    print(f"  • Pitch Detection: {args.pitch_detection_algo} (fastest)")
    print(f"  • Filter Radius: {args.filter_radius} (0=fastest)")
    print(f"  • Protect: {args.protect} (lower=faster)")
    print(f"  • Reverb: Disabled for speed")
    print(f"  • Output Format: {args.output_format}")
    
    print("\n🚀 Starting optimized processing...")
    print("=" * 60)
    
    # Run the main script with optimizations
    try:
        subprocess.run(cmd, check=True)
        print("\n" + "=" * 60)
        print("✅ Processing complete!")
        print("=" * 60)
    except subprocess.CalledProcessError as e:
        print(f"\n❌ Error during processing: {e}")
        sys.exit(1)
    except KeyboardInterrupt:
        print("\n\n⚠️ Processing interrupted by user")
        sys.exit(0)

if __name__ == '__main__':
    main()