ASesYusuf1/SESA_Audio_Separation
14
1# coding: utf-82__author__ = 'PyTorch Optimized Inference Implementation'3 4import argparse5import time6import librosa7from tqdm.auto import tqdm8import sys9import os10import glob11import torch12import soundfile as sf13import torch.nn as nn14import numpy as np15import pickle16from assets.i18n.i18n import I18nAuto17 18# Set inference path for compatibility19INFERENCE_PATH = os.path.abspath(__file__)20 21i18n = I18nAuto()22 23current_dir = os.path.dirname(os.path.abspath(__file__))24sys.path.append(current_dir)25 26from utils import get_model_from_config, normalize_audio, denormalize_audio27from utils import prefer_target_instrument, load_start_checkpoint, apply_tta, demix28from pytorch_backend import PyTorchBackend, PyTorchOptimizer, create_inference_session29 30import warnings31warnings.filterwarnings("ignore")32 33 34def shorten_filename(filename, max_length=30):35 """Dosya adını belirtilen maksimum uzunluğa kısaltır."""36 base, ext = os.path.splitext(filename)37 if len(base) <= max_length:38 return filename39 shortened = base[:15] + "..." + base[-10:] + ext40 return shortened41 42 43def get_soundfile_subtype(pcm_type, is_float=False):44 """PCM türüne göre uygun soundfile alt türünü belirler."""45 if is_float:46 return 'FLOAT'47 subtype_map = {48 'PCM_16': 'PCM_16',49 'PCM_24': 'PCM_24',50 'FLOAT': 'FLOAT'51 }52 return subtype_map.get(pcm_type, 'FLOAT')53 54 55def demix_pytorch_optimized(56 config,57 backend: PyTorchBackend,58 mix: np.ndarray,59 device: torch.device,60 pbar: bool = False61) -> dict:62 """63 Optimized PyTorch backend ile audio source separation.64 65 Parameters:66 ----------67 config : ConfigDict68 Configuration object69 backend : PyTorchBackend70 PyTorch backend with optimized model71 mix : np.ndarray72 Input audio array73 device : torch.device74 Computation device75 pbar : bool76 Show progress bar77 78 Returns:79 -------80 dict81 Dictionary of separated sources82 """83 mix = torch.tensor(mix, dtype=torch.float32)84 85 chunk_size = config.audio.chunk_size86 num_instruments = len(prefer_target_instrument(config))87 num_overlap = config.inference.num_overlap88 89 fade_size = chunk_size // 1090 step = chunk_size // num_overlap91 border = chunk_size - step92 length_init = mix.shape[-1]93 94 # Windowing array95 fadein = torch.linspace(0, 1, fade_size)96 fadeout = torch.linspace(1, 0, fade_size)97 windowing_array = torch.ones(chunk_size)98 windowing_array[-fade_size:] = fadeout99 windowing_array[:fade_size] = fadein100 101 # Add padding102 if length_init > 2 * border and border > 0:103 mix = nn.functional.pad(mix, (border, border), mode="reflect")104 105 batch_size = config.inference.batch_size106 use_amp = getattr(config.training, 'use_amp', True)107 108 with torch.cuda.amp.autocast(enabled=use_amp):109 with torch.inference_mode():110 # Initialize result and counter tensors111 req_shape = (num_instruments,) + mix.shape112 result = torch.zeros(req_shape, dtype=torch.float32)113 counter = torch.zeros(req_shape, dtype=torch.float32)114 115 i = 0116 batch_data = []117 batch_locations = []118 119 # Progress reporting for GUI (no terminal tqdm)120 total_samples = mix.shape[1]121 last_reported_percent = -1122 123 while i < mix.shape[1]:124 # Extract chunk125 part = mix[:, i:i + chunk_size].to(device)126 chunk_len = part.shape[-1]127 128 if chunk_len > chunk_size // 2:129 pad_mode = "reflect"130 else:131 pad_mode = "constant"132 133 part = nn.functional.pad(134 part, 135 (0, chunk_size - chunk_len), 136 mode=pad_mode, 137 value=0138 )139 140 batch_data.append(part)141 batch_locations.append((i, chunk_len))142 i += step143 144 # Process batch145 if len(batch_data) >= batch_size or i >= mix.shape[1]:146 arr = torch.stack(batch_data, dim=0)147 148 # Use optimized PyTorch backend for inference149 x = backend(arr)150 151 window = windowing_array.clone()152 if i - step == 0: # First chunk153 window[:fade_size] = 1154 elif i >= mix.shape[1]: # Last chunk155 window[-fade_size:] = 1156 157 for j, (start, seg_len) in enumerate(batch_locations):158 result[..., start:start + seg_len] += x[j, ..., :seg_len].cpu() * window[..., :seg_len]159 counter[..., start:start + seg_len] += window[..., :seg_len]160 161 batch_data.clear()162 batch_locations.clear()163 164 # Report real progress percentage for GUI capture (every 1% for smooth updates)165 # Use unique prefix [SESA_PROGRESS] to avoid confusion with other log messages166 current_percent = int((i / total_samples) * 100)167 if current_percent > last_reported_percent:168 last_reported_percent = current_percent169 print(f"[SESA_PROGRESS]{current_percent}", flush=True)170 171 print("[SESA_PROGRESS]100", flush=True)172 173 # Compute final estimated sources174 estimated_sources = result / counter175 estimated_sources = estimated_sources.cpu().numpy()176 np.nan_to_num(estimated_sources, copy=False, nan=0.0)177 178 # Remove padding179 if length_init > 2 * border and border > 0:180 estimated_sources = estimated_sources[..., border:-border]181 182 # Return as dictionary183 instruments = prefer_target_instrument(config)184 ret_data = {k: v for k, v in zip(instruments, estimated_sources)}185 186 return ret_data187 188 189def run_folder_pytorch_optimized(backend, args, config, device, model=None, verbose: bool = False):190 """191 PyTorch backend ile klasör işleme.192 """193 start_time = time.time()194 195 mixture_paths = sorted(glob.glob(os.path.join(args.input_folder, '*.*')))196 sample_rate = getattr(config.audio, 'sample_rate', 44100)197 198 print(f"PyTorch Backend | {len(mixture_paths)} dosya | SR: {sample_rate}")199 200 instruments = prefer_target_instrument(config)[:]201 202 # Çıktı klasörünü kullan203 store_dir = args.store_dir204 os.makedirs(store_dir, exist_ok=True)205 206 # Progress is reported via print statements for GUI capture (no terminal tqdm)207 total_files = len(mixture_paths)208 detailed_pbar = not args.disable_detailed_pbar209 print(i18n("detailed_pbar_enabled").format(detailed_pbar))210 211 for file_idx, path in enumerate(mixture_paths):212 try:213 mix, sr = librosa.load(path, sr=sample_rate, mono=False)214 print(i18n("loaded_audio").format(path, mix.shape))215 except Exception as e:216 print(i18n("cannot_read_track").format(path))217 print(i18n("error_message").format(str(e)))218 continue219 220 mix_orig = mix.copy()221 if 'normalize' in config.inference:222 if config.inference['normalize'] is True:223 mix, norm_params = normalize_audio(mix)224 225 # Use optimized PyTorch backend226 waveforms_orig = demix_pytorch_optimized(config, backend, mix, device, pbar=detailed_pbar)227 228 if args.use_tta and model is not None:229 waveforms_orig = apply_tta(config, model, mix, waveforms_orig, device, args.model_type)230 231 if args.demud_phaseremix_inst and model is not None:232 print(f"DemudPhaseRemix: {path}")233 instr = 'vocals' if 'vocals' in instruments else instruments[0]234 instruments.append('instrumental_phaseremix')235 if 'instrumental' not in instruments and 'Instrumental' not in instruments:236 mix_modified = mix_orig - 2 * waveforms_orig[instr]237 mix_modified_ = mix_modified.copy()238 waveforms_modified = demix(config, model, mix_modified, device, model_type=args.model_type)239 if args.use_tta:240 waveforms_modified = apply_tta(config, model, mix_modified, waveforms_modified, device, args.model_type)241 waveforms_orig['instrumental_phaseremix'] = mix_orig + waveforms_modified[instr]242 else:243 mix_modified = 2 * waveforms_orig[instr] - mix_orig244 mix_modified_ = mix_modified.copy()245 waveforms_modified = demix(config, model, mix_modified, device, model_type=args.model_type)246 if args.use_tta:247 waveforms_modified = apply_tta(config, model, mix_modified, waveforms_orig, device, args.model_type)248 waveforms_orig['instrumental_phaseremix'] = mix_orig + mix_modified_ - waveforms_modified[instr]249 250 if args.extract_instrumental:251 instr = 'vocals' if 'vocals' in instruments else instruments[0]252 waveforms_orig['instrumental'] = mix_orig - waveforms_orig[instr]253 if 'instrumental' not in instruments:254 instruments.append('instrumental')255 256 for instr in instruments:257 estimates = waveforms_orig[instr]258 if 'normalize' in config.inference:259 if config.inference['normalize'] is True:260 estimates = denormalize_audio(estimates, norm_params)261 262 is_float = getattr(args, 'export_format', '').startswith('wav FLOAT')263 codec = 'flac' if getattr(args, 'flac_file', False) else 'wav'264 if codec == 'flac':265 subtype = get_soundfile_subtype(args.pcm_type, is_float)266 else:267 subtype = get_soundfile_subtype('FLOAT', is_float)268 269 shortened_filename = shorten_filename(os.path.basename(path))270 output_filename = f"{shortened_filename}_{instr}.{codec}"271 output_path = os.path.join(store_dir, output_filename)272 sf.write(output_path, estimates.T, sr, subtype=subtype)273 274 print(i18n("elapsed_time").format(time.time() - start_time))275 276 277def proc_folder_pytorch_optimized(args):278 """279 PyTorch ile inference işleme fonksiyonu.280 """281 parser = argparse.ArgumentParser(description="PyTorch Inference for Music Source Separation")282 parser.add_argument("--model_type", type=str, default='mdx23c', help="Model type")283 parser.add_argument("--config_path", type=str, help="Config path")284 parser.add_argument("--start_check_point", type=str, default='', help="Checkpoint path (.ckpt)")285 parser.add_argument("--input_folder", type=str, help="Input folder path")286 parser.add_argument("--store_dir", type=str, default="", help="Output directory")287 parser.add_argument("--device_ids", nargs='+', type=int, default=0, help="Device IDs")288 parser.add_argument("--extract_instrumental", action='store_true', help="Extract instrumental")289 parser.add_argument("--disable_detailed_pbar", action='store_true', help="Disable detailed progress bar")290 parser.add_argument("--flac_file", action='store_true', help="Output as FLAC")291 parser.add_argument("--export_format", type=str, choices=['wav FLOAT', 'flac PCM_16', 'flac PCM_24'], 292 default='flac PCM_24', help="Export format")293 parser.add_argument("--pcm_type", type=str, choices=['PCM_16', 'PCM_24'], default='PCM_24', help="PCM type")294 parser.add_argument("--chunk_size", type=int, default=1000000, help="Inference chunk size")295 parser.add_argument("--overlap", type=int, default=4, help="Inference overlap factor")296 parser.add_argument("--optimize_mode", type=str, choices=['channels_last', 'compile', 'jit', 'default'], 297 default='channels_last', help="PyTorch optimization mode (channels_last recommended)")298 parser.add_argument("--enable_amp", action='store_true', help="Enable automatic mixed precision (2x faster)")299 parser.add_argument("--enable_tf32", action='store_true', help="Enable TF32 for RTX 30xx+ (faster)")300 parser.add_argument("--enable_cudnn_benchmark", action='store_true', help="Enable cuDNN benchmark (faster after warmup)")301 parser.add_argument("--lora_checkpoint", type=str, default='', help="Initial checkpoint to LoRA weights")302 parser.add_argument("--use_tta", action='store_true', help="Test Time Augmentation (flips + polarity)")303 parser.add_argument("--demud_phaseremix_inst", action='store_true', help="DemudPhaseRemix instrumental extraction")304 305 if args is None:306 args = parser.parse_args()307 else:308 args = parser.parse_args(args)309 310 # Device setup311 device = "cpu"312 if torch.cuda.is_available():313 print(i18n("cuda_available"))314 device = f'cuda:{args.device_ids[0]}' if type(args.device_ids) == list else f'cuda:{args.device_ids}'315 elif torch.backends.mps.is_available():316 device = "mps"317 print("Using MPS (Metal) backend")318 319 print(i18n("using_device").format(device))320 321 # Load model322 model_load_start_time = time.time()323 324 model, config = get_model_from_config(args.model_type, args.config_path)325 326 if args.start_check_point != '':327 try:328 checkpoint = torch.load(args.start_check_point, map_location=device, weights_only=False)329 except (pickle.UnpicklingError, RuntimeError, EOFError) as e:330 error_details = f"""331CHECKPOINT FILE CORRUPTED332 333Error: {str(e)}334 335The checkpoint file appears to be corrupted or was not downloaded correctly.336File: {args.start_check_point}337 338Common causes:339 - File is an HTML page (wrong download URL, e.g., HuggingFace /blob/ instead of /resolve/)340 - Incomplete or interrupted download341 - Network issues during download342 - File system corruption343 344Solution:345 1. Delete the corrupted checkpoint file:346 {args.start_check_point}347 2. Re-run the application - it will automatically re-download the model348 3. If the problem persists, check that your model URL uses /resolve/ not /blob/349 Example: https://huggingface.co/user/repo/resolve/main/model.ckpt350"""351 print(error_details)352 import sys353 sys.exit(1)354 355 # Handle different checkpoint formats356 if isinstance(checkpoint, dict):357 if 'state_dict' in checkpoint:358 state_dict = checkpoint['state_dict']359 elif 'model' in checkpoint:360 state_dict = checkpoint['model']361 elif 'state' in checkpoint:362 state_dict = checkpoint['state']363 else:364 state_dict = checkpoint365 else:366 state_dict = checkpoint367 368 model.load_state_dict(state_dict, strict=False)369 model = model.eval().to(device)370 371 print(i18n("instruments_print").format(config.training.instruments))372 373 # Create optimized PyTorch backend374 backend = create_inference_session(375 model=model,376 device=device,377 optimize_mode=args.optimize_mode,378 enable_amp=args.enable_amp,379 enable_tf32=args.enable_tf32,380 enable_cudnn_benchmark=args.enable_cudnn_benchmark381 )382 383 print(i18n("model_load_time").format(time.time() - model_load_start_time))384 385 # Run inference (pass raw model for TTA/demud support)386 run_folder_pytorch_optimized(backend, args, config, device, model=model, verbose=False)387 388 389if __name__ == "__main__":390 proc_folder_pytorch_optimized(None)391