inference_pytorch.py

# coding: utf-8
__author__ = 'PyTorch Optimized Inference Implementation'

import argparse
import time
import librosa
from tqdm.auto import tqdm
import sys
import os
import glob
import torch
import soundfile as sf
import torch.nn as nn
import numpy as np
import pickle
from assets.i18n.i18n import I18nAuto

# Set inference path for compatibility
INFERENCE_PATH = os.path.abspath(__file__)

i18n = I18nAuto()

current_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.append(current_dir)

from utils import get_model_from_config, normalize_audio, denormalize_audio
from utils import prefer_target_instrument, load_start_checkpoint, apply_tta, demix
from pytorch_backend import PyTorchBackend, PyTorchOptimizer, create_inference_session

import warnings
warnings.filterwarnings("ignore")


def shorten_filename(filename, max_length=30):
    """Dosya adını belirtilen maksimum uzunluğa kısaltır."""
    base, ext = os.path.splitext(filename)
    if len(base) <= max_length:
        return filename
    shortened = base[:15] + "..." + base[-10:] + ext
    return shortened


def get_soundfile_subtype(pcm_type, is_float=False):
    """PCM türüne göre uygun soundfile alt türünü belirler."""
    if is_float:
        return 'FLOAT'
    subtype_map = {
        'PCM_16': 'PCM_16',
        'PCM_24': 'PCM_24',
        'FLOAT': 'FLOAT'
    }
    return subtype_map.get(pcm_type, 'FLOAT')


def demix_pytorch_optimized(
    config,
    backend: PyTorchBackend,
    mix: np.ndarray,
    device: torch.device,
    pbar: bool = False
) -> dict:
    """
    Optimized PyTorch backend ile audio source separation.
    
    Parameters:
    ----------
    config : ConfigDict
        Configuration object
    backend : PyTorchBackend
        PyTorch backend with optimized model
    mix : np.ndarray
        Input audio array
    device : torch.device
        Computation device
    pbar : bool
        Show progress bar
        
    Returns:
    -------
    dict
        Dictionary of separated sources
    """
    mix = torch.tensor(mix, dtype=torch.float32)
    
    chunk_size = config.audio.chunk_size
    num_instruments = len(prefer_target_instrument(config))
    num_overlap = config.inference.num_overlap
    
    fade_size = chunk_size // 10
    step = chunk_size // num_overlap
    border = chunk_size - step
    length_init = mix.shape[-1]
    
    # Windowing array
    fadein = torch.linspace(0, 1, fade_size)
    fadeout = torch.linspace(1, 0, fade_size)
    windowing_array = torch.ones(chunk_size)
    windowing_array[-fade_size:] = fadeout
    windowing_array[:fade_size] = fadein
    
    # Add padding
    if length_init > 2 * border and border > 0:
        mix = nn.functional.pad(mix, (border, border), mode="reflect")
    
    batch_size = config.inference.batch_size
    use_amp = getattr(config.training, 'use_amp', True)
    
    with torch.cuda.amp.autocast(enabled=use_amp):
        with torch.inference_mode():
            # Initialize result and counter tensors
            req_shape = (num_instruments,) + mix.shape
            result = torch.zeros(req_shape, dtype=torch.float32)
            counter = torch.zeros(req_shape, dtype=torch.float32)
            
            i = 0
            batch_data = []
            batch_locations = []
            
            # Progress reporting for GUI (no terminal tqdm)
            total_samples = mix.shape[1]
            last_reported_percent = -1
            
            while i < mix.shape[1]:
                # Extract chunk
                part = mix[:, i:i + chunk_size].to(device)
                chunk_len = part.shape[-1]
                
                if chunk_len > chunk_size // 2:
                    pad_mode = "reflect"
                else:
                    pad_mode = "constant"
                    
                part = nn.functional.pad(
                    part, 
                    (0, chunk_size - chunk_len), 
                    mode=pad_mode, 
                    value=0
                )
                
                batch_data.append(part)
                batch_locations.append((i, chunk_len))
                i += step
                
                # Process batch
                if len(batch_data) >= batch_size or i >= mix.shape[1]:
                    arr = torch.stack(batch_data, dim=0)
                    
                    # Use optimized PyTorch backend for inference
                    x = backend(arr)
                    
                    window = windowing_array.clone()
                    if i - step == 0:  # First chunk
                        window[:fade_size] = 1
                    elif i >= mix.shape[1]:  # Last chunk
                        window[-fade_size:] = 1
                    
                    for j, (start, seg_len) in enumerate(batch_locations):
                        result[..., start:start + seg_len] += x[j, ..., :seg_len].cpu() * window[..., :seg_len]
                        counter[..., start:start + seg_len] += window[..., :seg_len]
                    
                    batch_data.clear()
                    batch_locations.clear()
                
                # Report real progress percentage for GUI capture (every 1% for smooth updates)
                # Use unique prefix [SESA_PROGRESS] to avoid confusion with other log messages
                current_percent = int((i / total_samples) * 100)
                if current_percent > last_reported_percent:
                    last_reported_percent = current_percent
                    print(f"[SESA_PROGRESS]{current_percent}", flush=True)
            
            print("[SESA_PROGRESS]100", flush=True)
            
            # Compute final estimated sources
            estimated_sources = result / counter
            estimated_sources = estimated_sources.cpu().numpy()
            np.nan_to_num(estimated_sources, copy=False, nan=0.0)
            
            # Remove padding
            if length_init > 2 * border and border > 0:
                estimated_sources = estimated_sources[..., border:-border]
    
    # Return as dictionary
    instruments = prefer_target_instrument(config)
    ret_data = {k: v for k, v in zip(instruments, estimated_sources)}
    
    return ret_data


def run_folder_pytorch_optimized(backend, args, config, device, model=None, verbose: bool = False):
    """
    PyTorch backend ile klasör işleme.
    """
    start_time = time.time()
    
    mixture_paths = sorted(glob.glob(os.path.join(args.input_folder, '*.*')))
    sample_rate = getattr(config.audio, 'sample_rate', 44100)
    
    print(f"PyTorch Backend | {len(mixture_paths)} dosya | SR: {sample_rate}")
    
    instruments = prefer_target_instrument(config)[:]
    
    # Çıktı klasörünü kullan
    store_dir = args.store_dir
    os.makedirs(store_dir, exist_ok=True)
    
    # Progress is reported via print statements for GUI capture (no terminal tqdm)
    total_files = len(mixture_paths)
    detailed_pbar = not args.disable_detailed_pbar
    print(i18n("detailed_pbar_enabled").format(detailed_pbar))
    
    for file_idx, path in enumerate(mixture_paths):
        try:
            mix, sr = librosa.load(path, sr=sample_rate, mono=False)
            print(i18n("loaded_audio").format(path, mix.shape))
        except Exception as e:
            print(i18n("cannot_read_track").format(path))
            print(i18n("error_message").format(str(e)))
            continue
        
        mix_orig = mix.copy()
        if 'normalize' in config.inference:
            if config.inference['normalize'] is True:
                mix, norm_params = normalize_audio(mix)
        
        # Use optimized PyTorch backend
        waveforms_orig = demix_pytorch_optimized(config, backend, mix, device, pbar=detailed_pbar)
        
        if args.use_tta and model is not None:
            waveforms_orig = apply_tta(config, model, mix, waveforms_orig, device, args.model_type)

        if args.demud_phaseremix_inst and model is not None:
            print(f"DemudPhaseRemix: {path}")
            instr = 'vocals' if 'vocals' in instruments else instruments[0]
            instruments.append('instrumental_phaseremix')
            if 'instrumental' not in instruments and 'Instrumental' not in instruments:
                mix_modified = mix_orig - 2 * waveforms_orig[instr]
                mix_modified_ = mix_modified.copy()
                waveforms_modified = demix(config, model, mix_modified, device, model_type=args.model_type)
                if args.use_tta:
                    waveforms_modified = apply_tta(config, model, mix_modified, waveforms_modified, device, args.model_type)
                waveforms_orig['instrumental_phaseremix'] = mix_orig + waveforms_modified[instr]
            else:
                mix_modified = 2 * waveforms_orig[instr] - mix_orig
                mix_modified_ = mix_modified.copy()
                waveforms_modified = demix(config, model, mix_modified, device, model_type=args.model_type)
                if args.use_tta:
                    waveforms_modified = apply_tta(config, model, mix_modified, waveforms_orig, device, args.model_type)
                waveforms_orig['instrumental_phaseremix'] = mix_orig + mix_modified_ - waveforms_modified[instr]

        if args.extract_instrumental:
            instr = 'vocals' if 'vocals' in instruments else instruments[0]
            waveforms_orig['instrumental'] = mix_orig - waveforms_orig[instr]
            if 'instrumental' not in instruments:
                instruments.append('instrumental')
        
        for instr in instruments:
            estimates = waveforms_orig[instr]
            if 'normalize' in config.inference:
                if config.inference['normalize'] is True:
                    estimates = denormalize_audio(estimates, norm_params)
            
            is_float = getattr(args, 'export_format', '').startswith('wav FLOAT')
            codec = 'flac' if getattr(args, 'flac_file', False) else 'wav'
            if codec == 'flac':
                subtype = get_soundfile_subtype(args.pcm_type, is_float)
            else:
                subtype = get_soundfile_subtype('FLOAT', is_float)
            
            shortened_filename = shorten_filename(os.path.basename(path))
            output_filename = f"{shortened_filename}_{instr}.{codec}"
            output_path = os.path.join(store_dir, output_filename)
            sf.write(output_path, estimates.T, sr, subtype=subtype)
    
    print(i18n("elapsed_time").format(time.time() - start_time))


def proc_folder_pytorch_optimized(args):
    """
    PyTorch ile inference işleme fonksiyonu.
    """
    parser = argparse.ArgumentParser(description="PyTorch Inference for Music Source Separation")
    parser.add_argument("--model_type", type=str, default='mdx23c', help="Model type")
    parser.add_argument("--config_path", type=str, help="Config path")
    parser.add_argument("--start_check_point", type=str, default='', help="Checkpoint path (.ckpt)")
    parser.add_argument("--input_folder", type=str, help="Input folder path")
    parser.add_argument("--store_dir", type=str, default="", help="Output directory")
    parser.add_argument("--device_ids", nargs='+', type=int, default=0, help="Device IDs")
    parser.add_argument("--extract_instrumental", action='store_true', help="Extract instrumental")
    parser.add_argument("--disable_detailed_pbar", action='store_true', help="Disable detailed progress bar")
    parser.add_argument("--flac_file", action='store_true', help="Output as FLAC")
    parser.add_argument("--export_format", type=str, choices=['wav FLOAT', 'flac PCM_16', 'flac PCM_24'], 
                        default='flac PCM_24', help="Export format")
    parser.add_argument("--pcm_type", type=str, choices=['PCM_16', 'PCM_24'], default='PCM_24', help="PCM type")
    parser.add_argument("--chunk_size", type=int, default=1000000, help="Inference chunk size")
    parser.add_argument("--overlap", type=int, default=4, help="Inference overlap factor")
    parser.add_argument("--optimize_mode", type=str, choices=['channels_last', 'compile', 'jit', 'default'], 
                        default='channels_last', help="PyTorch optimization mode (channels_last recommended)")
    parser.add_argument("--enable_amp", action='store_true', help="Enable automatic mixed precision (2x faster)")
    parser.add_argument("--enable_tf32", action='store_true', help="Enable TF32 for RTX 30xx+ (faster)")
    parser.add_argument("--enable_cudnn_benchmark", action='store_true', help="Enable cuDNN benchmark (faster after warmup)")
    parser.add_argument("--lora_checkpoint", type=str, default='', help="Initial checkpoint to LoRA weights")
    parser.add_argument("--use_tta", action='store_true', help="Test Time Augmentation (flips + polarity)")
    parser.add_argument("--demud_phaseremix_inst", action='store_true', help="DemudPhaseRemix instrumental extraction")
    
    if args is None:
        args = parser.parse_args()
    else:
        args = parser.parse_args(args)
    
    # Device setup
    device = "cpu"
    if torch.cuda.is_available():
        print(i18n("cuda_available"))
        device = f'cuda:{args.device_ids[0]}' if type(args.device_ids) == list else f'cuda:{args.device_ids}'
    elif torch.backends.mps.is_available():
        device = "mps"
        print("Using MPS (Metal) backend")
    
    print(i18n("using_device").format(device))
    
    # Load model
    model_load_start_time = time.time()
    
    model, config = get_model_from_config(args.model_type, args.config_path)
    
    if args.start_check_point != '':
        try:
            checkpoint = torch.load(args.start_check_point, map_location=device, weights_only=False)
        except (pickle.UnpicklingError, RuntimeError, EOFError) as e:
            error_details = f"""
CHECKPOINT FILE CORRUPTED

Error: {str(e)}

The checkpoint file appears to be corrupted or was not downloaded correctly.
File: {args.start_check_point}

Common causes:
  - File is an HTML page (wrong download URL, e.g., HuggingFace /blob/ instead of /resolve/)
  - Incomplete or interrupted download
  - Network issues during download
  - File system corruption

Solution:
  1. Delete the corrupted checkpoint file:
     {args.start_check_point}
  2. Re-run the application - it will automatically re-download the model
  3. If the problem persists, check that your model URL uses /resolve/ not /blob/
     Example: https://huggingface.co/user/repo/resolve/main/model.ckpt
"""
            print(error_details)
            import sys
            sys.exit(1)
        
        # Handle different checkpoint formats
        if isinstance(checkpoint, dict):
            if 'state_dict' in checkpoint:
                state_dict = checkpoint['state_dict']
            elif 'model' in checkpoint:
                state_dict = checkpoint['model']
            elif 'state' in checkpoint:
                state_dict = checkpoint['state']
            else:
                state_dict = checkpoint
        else:
            state_dict = checkpoint
        
        model.load_state_dict(state_dict, strict=False)
        model = model.eval().to(device)
    
    print(i18n("instruments_print").format(config.training.instruments))
    
    # Create optimized PyTorch backend
    backend = create_inference_session(
        model=model,
        device=device,
        optimize_mode=args.optimize_mode,
        enable_amp=args.enable_amp,
        enable_tf32=args.enable_tf32,
        enable_cudnn_benchmark=args.enable_cudnn_benchmark
    )
    
    print(i18n("model_load_time").format(time.time() - model_load_start_time))
    
    # Run inference (pass raw model for TTA/demud support)
    run_folder_pytorch_optimized(backend, args, config, device, model=model, verbose=False)


if __name__ == "__main__":
    proc_folder_pytorch_optimized(None)