ASesYusuf1/SESA_Audio_Separation

1# coding: utf-82__author__ = 'Roman Solovyev (ZFTurbo): https://github.com/ZFTurbo/'3 4import argparse5import time6import librosa7from tqdm.auto import tqdm8import sys9import os10import glob11import torch12import soundfile as sf13import torch.nn as nn14import numpy as np15from assets.i18n.i18n import I18nAuto16 17# Colab kontrolü18try:19    from google.colab import drive20    IS_COLAB = True21except ImportError:22    IS_COLAB = False23 24i18n = I18nAuto()25 26current_dir = os.path.dirname(os.path.abspath(__file__))27sys.path.append(current_dir)28 29from utils import demix, get_model_from_config, normalize_audio, denormalize_audio30from utils import prefer_target_instrument, apply_tta, load_start_checkpoint, load_lora_weights31 32# PyTorch optimized backend (always available)33try:34    from pytorch_backend import PyTorchBackend35    PYTORCH_OPTIMIZED_AVAILABLE = True36except ImportError:37    PYTORCH_OPTIMIZED_AVAILABLE = False38 39import warnings40warnings.filterwarnings("ignore")41 42def shorten_filename(filename, max_length=30):43    """Dosya adını belirtilen maksimum uzunluğa kısaltır."""44    base, ext = os.path.splitext(filename)45    if len(base) <= max_length:46        return filename47    shortened = base[:15] + "..." + base[-10:] + ext48    return shortened49 50def get_soundfile_subtype(pcm_type, is_float=False):51    """PCM türüne göre uygun soundfile alt türünü belirler."""52    if is_float:53        return 'FLOAT'54    subtype_map = {55        'PCM_16': 'PCM_16',56        'PCM_24': 'PCM_24',57        'FLOAT': 'FLOAT'58    }59    return subtype_map.get(pcm_type, 'FLOAT')60 61def run_folder(model, args, config, device, verbose: bool = False):62    start_time = time.time()63    model.eval()64 65    mixture_paths = sorted(glob.glob(os.path.join(args.input_folder, '*.*')))66    sample_rate = getattr(config.audio, 'sample_rate', 44100)67 68    print(i18n("total_files_found").format(len(mixture_paths), sample_rate))69 70    instruments = prefer_target_instrument(config)[:]71 72    # Çıktı klasörünü kullan (processing.py tarafından ayarlandı)73    store_dir = args.store_dir74    os.makedirs(store_dir, exist_ok=True)75 76    if not verbose:77        mixture_paths = tqdm(mixture_paths, desc=i18n("total_progress"))78    else:79        mixture_paths = mixture_paths80 81    detailed_pbar = not args.disable_detailed_pbar82    print(i18n("detailed_pbar_enabled").format(detailed_pbar))83 84    for path in mixture_paths:85        try:86            mix, sr = librosa.load(path, sr=sample_rate, mono=False)87            print(i18n("loaded_audio").format(path, mix.shape))88        except Exception as e:89            print(i18n("cannot_read_track").format(path))90            print(i18n("error_message").format(str(e)))91            continue92 93        mix_orig = mix.copy()94        if 'normalize' in config.inference:95            if config.inference['normalize'] is True:96                mix, norm_params = normalize_audio(mix)97 98        waveforms_orig = demix(config, model, mix, device, model_type=args.model_type, pbar=detailed_pbar)99 100        if args.use_tta:101            waveforms_orig = apply_tta(config, model, mix, waveforms_orig, device, args.model_type)102 103        if args.demud_phaseremix_inst:104            print(i18n("demudding_track").format(path))105            instr = 'vocals' if 'vocals' in instruments else instruments[0]106            instruments.append('instrumental_phaseremix')107            if 'instrumental' not in instruments and 'Instrumental' not in instruments:108                mix_modified = mix_orig - 2*waveforms_orig[instr]109                mix_modified_ = mix_modified.copy()110                waveforms_modified = demix(config, model, mix_modified, device, model_type=args.model_type, pbar=detailed_pbar)111                if args.use_tta:112                    waveforms_modified = apply_tta(config, model, mix_modified, waveforms_modified, device, args.model_type)113                waveforms_orig['instrumental_phaseremix'] = mix_orig + waveforms_modified[instr]114            else:115                mix_modified = 2*waveforms_orig[instr] - mix_orig116                mix_modified_ = mix_modified.copy()117                waveforms_modified = demix(config, model, mix_modified, device, model_type=args.model_type, pbar=detailed_pbar)118                if args.use_tta:119                    waveforms_modified = apply_tta(config, model, mix_modified, waveforms_orig, device, args.model_type)120                waveforms_orig['instrumental_phaseremix'] = mix_orig + mix_modified_ - waveforms_modified[instr]121 122        if args.extract_instrumental:123            instr = 'vocals' if 'vocals' in instruments else instruments[0]124            waveforms_orig['instrumental'] = mix_orig - waveforms_orig[instr]125            if 'instrumental' not in instruments:126                instruments.append('instrumental')127 128        for instr in instruments:129            estimates = waveforms_orig[instr]130            if 'normalize' in config.inference:131                if config.inference['normalize'] is True:132                    estimates = denormalize_audio(estimates, norm_params)133 134            is_float = getattr(args, 'export_format', '').startswith('wav FLOAT')135            codec = 'flac' if getattr(args, 'flac_file', False) else 'wav'136            if codec == 'flac':137                subtype = get_soundfile_subtype(args.pcm_type, is_float)138            else:139                subtype = get_soundfile_subtype('FLOAT', is_float)140 141            shortened_filename = shorten_filename(os.path.basename(path))142            output_filename = f"{shortened_filename}_{instr}.{codec}"143            output_path = os.path.join(store_dir, output_filename)144            sf.write(output_path, estimates.T, sr, subtype=subtype)145 146    print(i18n("elapsed_time").format(time.time() - start_time))147 148def proc_folder(args, use_tensorrt=False):149    """150    Process folder with optional TensorRT backend.151    152    Parameters:153    ----------154    args : list or None155        Command line arguments156    use_tensorrt : bool157        Use TensorRT backend if available158    """159    parser = argparse.ArgumentParser(description=i18n("proc_folder_description"))160    parser.add_argument("--model_type", type=str, default='mdx23c', help=i18n("model_type_help"))161    parser.add_argument("--config_path", type=str, help=i18n("config_path_help"))162    parser.add_argument("--demud_phaseremix_inst", action='store_true', help=i18n("demud_phaseremix_help"))163    parser.add_argument("--start_check_point", type=str, default='', help=i18n("start_checkpoint_help"))164    parser.add_argument("--input_folder", type=str, help=i18n("input_folder_help"))165    parser.add_argument("--audio_path", type=str, help=i18n("audio_path_help"))166    parser.add_argument("--store_dir", type=str, default="", help=i18n("store_dir_help"))167    parser.add_argument("--device_ids", nargs='+', type=int, default=0, help=i18n("device_ids_help"))168    parser.add_argument("--extract_instrumental", action='store_true', help=i18n("extract_instrumental_help"))169    parser.add_argument("--disable_detailed_pbar", action='store_true', help=i18n("disable_detailed_pbar_help"))170    parser.add_argument("--force_cpu", action='store_true', help=i18n("force_cpu_help"))171    parser.add_argument("--flac_file", action='store_true', help=i18n("flac_file_help"))172    parser.add_argument("--export_format", type=str, choices=['wav FLOAT', 'flac PCM_16', 'flac PCM_24'], default='flac PCM_24', help=i18n("export_format_help"))173    parser.add_argument("--pcm_type", type=str, choices=['PCM_16', 'PCM_24'], default='PCM_24', help=i18n("pcm_type_help"))174    parser.add_argument("--use_tta", action='store_true', help=i18n("use_tta_help"))175    parser.add_argument("--lora_checkpoint", type=str, default='', help=i18n("lora_checkpoint_help"))176    parser.add_argument("--chunk_size", type=int, default=1000000, help="Inference chunk size")177    parser.add_argument("--overlap", type=int, default=4, help="Inference overlap factor")178    parser.add_argument("--optimize_mode", type=str, choices=['default', 'compile', 'jit', 'channels_last'], default='channels_last', help="PyTorch optimization mode (always enabled)")179    parser.add_argument("--enable_amp", action='store_true', default=True, help="Enable automatic mixed precision")180    parser.add_argument("--enable_tf32", action='store_true', default=True, help="Enable TF32 (Ampere GPUs)")181    parser.add_argument("--enable_cudnn_benchmark", action='store_true', default=True, help="Enable cuDNN benchmark")182 183    if args is None:184        args = parser.parse_args()185    else:186        args = parser.parse_args(args)187 188    device = "cpu"189    if args.force_cpu:190        device = "cpu"191    elif torch.cuda.is_available():192        print(i18n("cuda_available"))193        device = f'cuda:{args.device_ids[0]}' if type(args.device_ids) == list else f'cuda:{args.device_ids}'194    elif torch.backends.mps.is_available():195         device = "mps"196 197    print(i18n("using_device").format(device))198 199    model_load_start_time = time.time()200    torch.backends.cudnn.benchmark = True201 202    model, config = get_model_from_config(args.model_type, args.config_path)203 204    if args.start_check_point != '':205        load_start_checkpoint(args, model, type_='inference')206 207    print(i18n("instruments_print").format(config.training.instruments))208 209    if type(args.device_ids) == list and len(args.device_ids) > 1 and not args.force_cpu:210        model = nn.DataParallel(model, device_ids=args.device_ids)211 212    model = model.to(device)213 214    print(i18n("model_load_time").format(time.time() - model_load_start_time))215 216    # Always use optimized PyTorch backend if available217    if PYTORCH_OPTIMIZED_AVAILABLE:218        print(f"Using optimized PyTorch backend")219        print(f"   Mode: {args.optimize_mode}")220        print(f"   AMP: {args.enable_amp} | TF32: {args.enable_tf32} | cuDNN: {args.enable_cudnn_benchmark}")221        from inference_pytorch import proc_folder_pytorch_optimized222        # Recreate args for optimized PyTorch inference223        sys.argv = sys.argv[:1]  # Keep only script name224        for key, value in vars(args).items():225            if value is not None and value is not False:226                if isinstance(value, bool):227                    sys.argv.append(f"--{key}")228                elif isinstance(value, list):229                    sys.argv.append(f"--{key}")230                    sys.argv.extend(map(str, value))231                else:232                    sys.argv.extend([f"--{key}", str(value)])233        proc_folder_pytorch_optimized(None)234    else:235        print("Warning: PyTorch optimized backend not available, using standard inference")236        run_folder(model, args, config, device, verbose=False)237 238if __name__ == "__main__":239    proc_folder(None)240