ASesYusuf1/SESA_Audio_Separation
14
1import os2import glob3import subprocess4import time5import gc6import shutil7import sys8from assets.i18n.i18n import I18nAuto9 10i18n = I18nAuto()11current_dir = os.path.dirname(os.path.abspath(__file__))12sys.path.append(current_dir)13 14from datetime import datetime15from helpers import INPUT_DIR, OLD_OUTPUT_DIR, ENSEMBLE_DIR, AUTO_ENSEMBLE_TEMP, move_old_files, clear_directory, BASE_DIR, clean_model, extract_model_name_from_checkpoint, sanitize_filename, find_clear_segment, save_segment, run_matchering, clamp_percentage16from model import get_model_config, get_model_chunk_size17from apollo_processing import process_with_apollo # Import Apollo processing18import torch19 20# PyTorch optimized backend (always available)21try:22 from pytorch_backend import PyTorchBackend23 PYTORCH_OPTIMIZED_AVAILABLE = True24except ImportError:25 PYTORCH_OPTIMIZED_AVAILABLE = False26import yaml27import gradio as gr28import threading29import random30import librosa31import soundfile as sf32import numpy as np33import requests34import json35import locale36import re37import psutil38import concurrent.futures39from tqdm import tqdm40 41# Google OAuth imports (optional - for Colab/Google Drive support)42try:43 from google.oauth2.credentials import Credentials44 GOOGLE_OAUTH_AVAILABLE = True45except ImportError:46 GOOGLE_OAUTH_AVAILABLE = False47 Credentials = None48 49import tempfile50from urllib.parse import urlparse, quote51try:52 from google.colab import drive53 # Verify we're actually in a working Colab environment54 IS_COLAB = True55except ImportError:56 IS_COLAB = False57 drive = None58import matchering as mg59 60import warnings61warnings.filterwarnings("ignore")62 63BASE_DIR = os.path.dirname(os.path.abspath(__file__))64INFERENCE_PATH = os.path.join(BASE_DIR, "inference.py")65ENSEMBLE_PATH = os.path.join(BASE_DIR, "ensemble.py")66 67if IS_COLAB:68 AUTO_ENSEMBLE_OUTPUT = "/content/drive/MyDrive/ensemble_output"69 OUTPUT_DIR = "/content/drive/MyDrive/!output_file"70else:71 AUTO_ENSEMBLE_OUTPUT = os.path.join(BASE_DIR, "ensemble_output")72 OUTPUT_DIR = os.path.join(BASE_DIR, "output")73 74os.makedirs(AUTO_ENSEMBLE_OUTPUT, exist_ok=True)75os.makedirs(OUTPUT_DIR, exist_ok=True)76 77def setup_directories():78 """Create necessary directories and check Google Drive access."""79 if IS_COLAB:80 try:81 # Check if Google Drive is already mounted82 if os.path.exists('/content/drive/MyDrive'):83 pass # Already mounted, no action needed84 else:85 print("Mounting Google Drive...")86 try:87 from google.colab import drive88 drive.mount('/content/drive', force_remount=True)89 except AttributeError as ae:90 # Handle 'NoneType' object has no attribute 'kernel' error91 print(f"Warning: Google Drive mount skipped (Colab kernel issue): {str(ae)}")92 print("Continuing with local storage...")93 except Exception as mount_error:94 print(f"Warning: Google Drive mount failed: {str(mount_error)}")95 print("Continuing with local storage...")96 except Exception as e:97 print(f"Warning: Google Drive setup error: {str(e)}")98 print("Continuing without Google Drive...")99 os.makedirs(OUTPUT_DIR, exist_ok=True)100 os.makedirs(INPUT_DIR, exist_ok=True)101 os.makedirs(OLD_OUTPUT_DIR, exist_ok=True)102 os.makedirs(AUTO_ENSEMBLE_OUTPUT, exist_ok=True)103 104def refresh_auto_output():105 try:106 output_files = glob.glob(os.path.join(AUTO_ENSEMBLE_OUTPUT, "*.wav"))107 if not output_files:108 return None, "No output files found"109 110 latest_file = max(output_files, key=os.path.getctime)111 return latest_file, "Output refreshed successfully"112 except Exception as e:113 return None, f"Error refreshing output: {str(e)}"114 115def update_progress_html(progress_label, progress_percent, download_info=None):116 """Generate progress HTML with smooth animations and optional download percentage.117 118 Args:119 progress_label: Text label to show above the progress bar120 progress_percent: Overall progress percentage (0-100)121 download_info: Optional dict with 'filename' and 'percent' for download progress122 """123 progress_percent = clamp_percentage(progress_percent)124 125 # Determine if processing is active for pulse animation126 is_active = 0 < progress_percent < 100127 pulse_style = "animation: progress-pulse 1.5s ease-in-out infinite;" if is_active else ""128 129 # Build download sub-bar if downloading130 download_html = ""131 if download_info and isinstance(download_info, dict):132 dl_filename = download_info.get('filename', '')133 dl_percent = clamp_percentage(download_info.get('percent', 0))134 download_html = f"""135 <div style="margin-top: 8px; padding: 8px; background: rgba(0,0,0,0.3); border-radius: 5px;">136 <div style="font-size: 0.85rem; color: #a0a0a0; margin-bottom: 4px;">{dl_filename} - %{int(dl_percent)}</div>137 <div style="width: 100%; background-color: #333; border-radius: 4px; overflow: hidden;">138 <div style="width: {dl_percent}%; height: 14px; background: linear-gradient(90deg, #4ade80, #22d3ee); transition: width 0.3s ease-out; border-radius: 4px;"></div>139 </div>140 </div>141 """142 143 return f"""144 <style>145 @keyframes progress-pulse {{146 0%, 100% {{ opacity: 1; }}147 50% {{ opacity: 0.85; }}148 }}149 </style>150 <div id="custom-progress" style="margin-top: 10px;">151 <div style="font-size: 1rem; color: #C0C0C0; margin-bottom: 5px;" id="progress-label">{progress_label}</div>152 <div style="width: 100%; background-color: #444; border-radius: 5px; overflow: hidden;">153 <div id="progress-bar" style="width: {progress_percent}%; height: 20px; background: linear-gradient(90deg, #6e8efb, #a855f7); transition: width 0.5s ease-out; max-width: 100%; {pulse_style}"></div>154 </div>155 {download_html}156 </div>157 """158 159def extract_model_name_from_checkpoint(checkpoint_path):160 if not checkpoint_path:161 return "Unknown"162 base_name = os.path.basename(checkpoint_path)163 model_name = os.path.splitext(base_name)[0]164 return model_name.strip()165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187def run_command_and_process_files(188 model_type,189 config_path,190 start_check_point,191 INPUT_DIR,192 OUTPUT_DIR,193 extract_instrumental,194 use_tta,195 demud_phaseremix_inst,196 progress=None,197 use_apollo=True,198 apollo_normal_model="Apollo Universal Model",199 inference_chunk_size=352800,200 inference_overlap=2,201 apollo_chunk_size=19,202 apollo_overlap=2,203 apollo_method="normal_method",204 apollo_midside_model=None,205 output_format="wav",206 optimize_mode='channels_last',207 enable_amp=True,208 enable_tf32=True,209 enable_cudnn_benchmark=True210):211 """212 Run inference.py with specified parameters and process output files.213 This is a generator function that yields progress updates for real-time UI feedback.214 """215 try:216 # Create directories and check Google Drive access217 setup_directories()218 219 if not config_path:220 raise ValueError(f"Configuration path is empty: model_type: {model_type}")221 if not os.path.exists(config_path):222 raise FileNotFoundError(f"Configuration file not found: {config_path}")223 if not start_check_point or not os.path.exists(start_check_point):224 raise FileNotFoundError(f"Checkpoint file not found: {start_check_point}")225 226 # Validate inference parameters227 try:228 inference_chunk_size = int(inference_chunk_size)229 inference_overlap = int(inference_overlap)230 except (TypeError, ValueError) as e:231 print(f"Invalid inference_chunk_size or inference_overlap: {e}. Defaulting to: inference_chunk_size=352800, inference_overlap=2")232 inference_chunk_size = 352800233 inference_overlap = 2234 235 # Validate Apollo parameters236 try:237 apollo_chunk_size = int(apollo_chunk_size)238 apollo_overlap = int(apollo_overlap)239 except (TypeError, ValueError) as e:240 print(f"Invalid apollo_chunk_size or apollo_overlap: {e}. Defaulting to: apollo_chunk_size=19, apollo_overlap=2")241 apollo_chunk_size = 19242 apollo_overlap = 2243 244 # Initial progress yield245 yield {"progress": 0, "status": "Starting audio separation...", "outputs": None}246 247 # Always use optimized PyTorch backend248 python_exe = "python"249 250 if PYTORCH_OPTIMIZED_AVAILABLE:251 from inference_pytorch import INFERENCE_PATH as PYTORCH_INFERENCE_PATH252 inference_script = PYTORCH_INFERENCE_PATH if os.path.exists(PYTORCH_INFERENCE_PATH) else INFERENCE_PATH253 print(f"Using PyTorch backend (mode: {optimize_mode})")254 print(f" AMP: {enable_amp} | TF32: {enable_tf32} | cuDNN: {enable_cudnn_benchmark}")255 else:256 inference_script = INFERENCE_PATH257 print("Warning: PyTorch optimized backend not available, using standard inference")258 259 260 261 262 263 264 265 266 cmd_parts = [267 python_exe, inference_script,268 "--model_type", model_type,269 "--config_path", config_path,270 "--start_check_point", start_check_point,271 "--input_folder", INPUT_DIR,272 "--store_dir", OUTPUT_DIR,273 "--chunk_size", str(inference_chunk_size),274 "--overlap", str(inference_overlap),275 "--export_format", f"{output_format} FLOAT"276 ]277 278 279 280 281 282 283 284 285 286 287 # Add optimized backend arguments (always enabled)288 if PYTORCH_OPTIMIZED_AVAILABLE:289 cmd_parts.extend([290 "--optimize_mode", optimize_mode291 ])292 if enable_amp:293 cmd_parts.append("--enable_amp")294 if enable_tf32:295 cmd_parts.append("--enable_tf32")296 if enable_cudnn_benchmark:297 cmd_parts.append("--enable_cudnn_benchmark")298 299 if extract_instrumental:300 cmd_parts.append("--extract_instrumental")301 if use_tta:302 cmd_parts.append("--use_tta")303 if demud_phaseremix_inst:304 cmd_parts.append("--demud_phaseremix_inst")305 306 print(f"Running command: {' '.join(cmd_parts)}")307 308 # Use subprocess.Popen for real-time progress capture309 process = subprocess.Popen(310 cmd_parts,311 cwd=BASE_DIR,312 stdout=subprocess.PIPE,313 stderr=subprocess.PIPE,314 text=True,315 bufsize=1,316 universal_newlines=True317 )318 319 stderr_output = ""320 last_yield_percent = -1321 downloading_file = None322 323 # Read stdout line-by-line for real-time progress updates324 for line in process.stdout:325 line_stripped = line.strip()326 327 # Check for download progress [SESA_DOWNLOAD]328 if line_stripped.startswith("[SESA_DOWNLOAD]"):329 try:330 dl_info = line_stripped.replace("[SESA_DOWNLOAD]", "")331 if dl_info.startswith("START:"):332 downloading_file = dl_info.replace("START:", "")333 yield {"progress": 0, "status": i18n("downloading_model_file").format(downloading_file), "outputs": None}334 elif dl_info.startswith("END:"):335 downloading_file = None336 elif ":" in dl_info:337 parts = dl_info.rsplit(":", 1)338 if len(parts) == 2:339 filename, percent_str = parts340 download_percent = int(percent_str)341 yield {"progress": 0, "status": i18n("downloading_file_progress").format(filename, download_percent), "outputs": None}342 except (ValueError, TypeError):343 pass344 # Check for [SESA_PROGRESS] prefix from inference script345 elif line_stripped.startswith("[SESA_PROGRESS]"):346 try:347 percentage_str = line_stripped.replace("[SESA_PROGRESS]", "").strip()348 percentage = float(percentage_str) if percentage_str else 0349 percentage = min(max(percentage, 0), 100)350 351 # Scale progress to 0-80% range (saving 80-100% for Apollo)352 scaled_progress = int(percentage * 0.8)353 354 # Yield on every percent change for smooth updates355 if int(percentage) != last_yield_percent:356 last_yield_percent = int(percentage)357 yield {"progress": scaled_progress, "status": f"Separating audio... {int(percentage)}%", "outputs": None}358 except (ValueError, TypeError):359 pass360 else:361 # Only print important non-progress lines (errors, warnings, key info)362 if line_stripped and not line_stripped.startswith((" ", " ")):363 print(line_stripped)364 365 # Capture stderr (only print errors)366 for line in process.stderr:367 stderr_output += line368 line_s = line.strip()369 if line_s and ("error" in line_s.lower() or "warning" in line_s.lower() or "traceback" in line_s.lower()):370 print(f"Warning: {line_s}")371 372 process.wait()373 374 if process.returncode != 0:375 raise subprocess.CalledProcessError(process.returncode, cmd_parts, stderr=stderr_output)376 377 yield {"progress": 80, "status": "Separation complete, processing outputs...", "outputs": None}378 379 # Check if output files were created380 filename_model = extract_model_name_from_checkpoint(start_check_point)381 output_files = os.listdir(OUTPUT_DIR)382 if not output_files:383 raise FileNotFoundError("No output files created in OUTPUT_DIR")384 385 def rename_files_with_model(folder, filename_model):386 timestamp = datetime.now().strftime("%d-%m-%Y_%H-%M")387 for filename in sorted(os.listdir(folder)):388 file_path = os.path.join(folder, filename)389 if not any(filename.lower().endswith(ext) for ext in ['.mp3', '.wav', '.flac', '.aac', '.ogg', '.m4a']):390 continue391 base, ext = os.path.splitext(filename)392 detected_type = None393 for type_key in ['vocals', 'instrumental', 'instrument', 'phaseremix', 'drum', 'bass', 'other', 'effects', 'speech', 'music', 'dry', 'male', 'female', 'bleed', 'karaoke', 'mid', 'side']:394 if type_key in base.lower():395 detected_type = type_key396 break397 # Normalize 'instrument' to 'Instrumental' for consistency398 type_suffix = 'Instrumental' if detected_type == 'instrument' else (detected_type.capitalize() if detected_type else "Processed")399 clean_base = sanitize_filename(base.split('_')[0]).rsplit('.', 1)[0]400 new_filename = f"{timestamp}_{clean_base}_{type_suffix}_{filename_model}{ext}"401 new_file_path = os.path.join(folder, new_filename)402 try:403 os.rename(file_path, new_file_path)404 except Exception as e:405 print(f"Could not rename file: {os.path.basename(file_path)} -> {os.path.basename(new_file_path)}: {str(e)}")406 407 rename_files_with_model(OUTPUT_DIR, filename_model)408 409 output_files = os.listdir(OUTPUT_DIR)410 if not output_files:411 raise FileNotFoundError("No output files in OUTPUT_DIR after renaming")412 413 def find_file(keywords):414 """Find file matching any of the keywords (can be single keyword or list)."""415 if isinstance(keywords, str):416 keywords = [keywords]417 matching_files = [418 os.path.join(OUTPUT_DIR, f) for f in output_files 419 if any(kw in f.lower() for kw in keywords)420 ]421 return matching_files[0] if matching_files else None422 423 output_list = [424 find_file('vocals'), find_file(['instrumental', 'instrument']), find_file('phaseremix'),425 find_file('drum'), find_file('bass'), find_file('other'), find_file('effects'),426 find_file('speech'), find_file('music'), find_file('dry'), find_file('male'),427 find_file('female'), find_file('bleed'), find_file('karaoke'),428 find_file('mid'), find_file('side')429 ]430 431 normalized_outputs = []432 for output_file in output_list:433 if output_file and os.path.exists(output_file):434 normalized_file = os.path.join(OUTPUT_DIR, f"{sanitize_filename(os.path.splitext(os.path.basename(output_file))[0])}.{output_format}")435 if output_file.endswith(f".{output_format}") and output_file != normalized_file:436 shutil.copy(output_file, normalized_file)437 elif output_file != normalized_file:438 audio, sr = librosa.load(output_file, sr=None, mono=False)439 sf.write(normalized_file, audio.T if audio.ndim > 1 else audio, sr)440 else:441 normalized_file = output_file442 normalized_outputs.append(normalized_file)443 else:444 normalized_outputs.append(output_file)445 446 # Apollo processing447 if use_apollo:448 yield {"progress": 80, "status": "Enhancing with Apollo...", "outputs": None}449 normalized_outputs = process_with_apollo(450 output_files=normalized_outputs,451 output_dir=OUTPUT_DIR,452 apollo_chunk_size=apollo_chunk_size,453 apollo_overlap=apollo_overlap,454 apollo_method=apollo_method,455 apollo_normal_model=apollo_normal_model,456 apollo_midside_model=apollo_midside_model,457 output_format=output_format,458 progress=progress,459 total_progress_start=80,460 total_progress_end=100461 )462 463 # Final yield with outputs464 yield {"progress": 100, "status": "Separation complete", "outputs": tuple(normalized_outputs)}465 466 except subprocess.CalledProcessError as e:467 print(f"Subprocess failed, code: {e.returncode}: {e.stderr}")468 yield {"progress": 0, "status": f"Error: {e.stderr}", "outputs": (None,) * 16}469 except Exception as e:470 print(f"run_command_and_process_files error: {str(e)}")471 import traceback472 traceback.print_exc()473 yield {"progress": 0, "status": f"Error: {str(e)}", "outputs": (None,) * 16}474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502def process_audio(503 input_audio_file,504 model,505 chunk_size,506 overlap,507 export_format,508 optimize_mode,509 enable_amp,510 enable_tf32,511 enable_cudnn_benchmark,512 use_tta,513 demud_phaseremix_inst,514 extract_instrumental,515 use_apollo,516 apollo_chunk_size,517 apollo_overlap,518 apollo_method,519 apollo_normal_model,520 apollo_midside_model,521 use_matchering,522 matchering_passes,523 progress=gr.Progress(track_tqdm=True),524 *args,525 **kwargs526):527 """528 Process audio with the selected model. This is a generator function that yields529 progress updates for real-time UI feedback.530 """531 try:532 # Check Google Drive connection533 setup_directories()534 535 if input_audio_file is not None:536 audio_path = input_audio_file.name if hasattr(input_audio_file, 'name') else input_audio_file537 else:538 yield (539 None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None,540 "No audio file provided",541 update_progress_html("No input provided", 0)542 )543 return544 545 os.makedirs(OUTPUT_DIR, exist_ok=True)546 os.makedirs(OLD_OUTPUT_DIR, exist_ok=True)547 move_old_files(OUTPUT_DIR)548 549 # Clean model name, remove ⭐ and other unwanted characters550 clean_model_name = clean_model(model) if not model.startswith("/") else extract_model_name_from_checkpoint(model)551 print(f"Processing: {os.path.basename(audio_path)} | Model: {clean_model_name}")552 553 # Validate inference parameters554 _use_yaml_chunk = (chunk_size == "yaml")555 try:556 inference_chunk_size = 352800 if _use_yaml_chunk else int(chunk_size)557 except (TypeError, ValueError):558 print(f"Invalid chunk_size: {chunk_size}. Defaulting to: 352800.")559 inference_chunk_size = 352800560 _use_yaml_chunk = True # fallback: read from YAML561 562 try:563 inference_overlap = int(overlap)564 except (TypeError, ValueError):565 print(f"Invalid overlap: {overlap}. Defaulting to: 2.")566 inference_overlap = 2567 568 # Validate Apollo parameters569 try:570 apollo_chunk_size = int(apollo_chunk_size)571 except (TypeError, ValueError):572 print(f"Invalid apollo_chunk_size: {apollo_chunk_size}. Defaulting to: 19.")573 apollo_chunk_size = 19574 575 try:576 apollo_overlap = int(apollo_overlap)577 except (TypeError, ValueError):578 print(f"Invalid apollo_overlap: {apollo_overlap}. Defaulting to: 2.")579 apollo_overlap = 2580 581 # Map apollo_method to backend values582 if apollo_method in ["Mid-side method", "2", 2, "mid_side_method"]:583 apollo_method = "mid_side_method"584 elif apollo_method in ["Normal method", "1", 1, "normal_method"]:585 apollo_method = "normal_method"586 else:587 print(f"Invalid apollo_method: {apollo_method}. Defaulting to: normal_method.")588 apollo_method = "normal_method"589 # Copy input file to INPUT_DIR590 input_filename = os.path.basename(audio_path)591 dest_path = os.path.join(INPUT_DIR, input_filename)592 shutil.copy(audio_path, dest_path)593 594 # Yield status for model loading595 yield (596 None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None,597 f"Loading model: {clean_model_name}...",598 update_progress_html(f"Loading model: {clean_model_name}", 0)599 )600 601 # Get model configuration with cleaned model name (downloads if needed)602 model_type, config_path, start_check_point = get_model_config(clean_model_name, inference_chunk_size, inference_overlap)603 604 # Read the model's native chunk_size from its YAML (now guaranteed to be downloaded)605 native_chunk = get_model_chunk_size(clean_model_name)606 if _use_yaml_chunk and native_chunk:607 print(f"Using model's native chunk_size from YAML: {native_chunk}")608 inference_chunk_size = native_chunk609 elif not _use_yaml_chunk:610 print(f"Using user-selected chunk_size: {inference_chunk_size}")611 612 # Iterate over the generator and yield progress updates613 outputs = None614 for update in run_command_and_process_files(615 model_type=model_type,616 config_path=config_path,617 start_check_point=start_check_point,618 INPUT_DIR=INPUT_DIR,619 OUTPUT_DIR=OUTPUT_DIR,620 extract_instrumental=extract_instrumental,621 use_tta=use_tta,622 demud_phaseremix_inst=demud_phaseremix_inst,623 progress=progress,624 use_apollo=use_apollo,625 apollo_normal_model=apollo_normal_model,626 inference_chunk_size=inference_chunk_size,627 inference_overlap=inference_overlap,628 apollo_chunk_size=apollo_chunk_size,629 apollo_overlap=apollo_overlap,630 apollo_method=apollo_method,631 apollo_midside_model=apollo_midside_model,632 output_format=export_format.split()[0].lower(),633 optimize_mode=optimize_mode,634 enable_amp=enable_amp,635 enable_tf32=enable_tf32,636 enable_cudnn_benchmark=enable_cudnn_benchmark637 ):638 if update.get("outputs") is not None:639 outputs = update["outputs"]640 # Yield progress update to Gradio641 yield (642 None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None,643 update["status"],644 update_progress_html(update["status"], update["progress"])645 )646 647 if outputs is None or all(output is None for output in outputs):648 raise ValueError("run_command_and_process_files returned None or all None outputs")649 650 # Apply Matchering (if enabled)651 if use_matchering:652 # Yield progress update for Matchering653 yield (654 None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None,655 "Applying Matchering...",656 update_progress_html("Applying Matchering...", 90)657 )658 659 # Find clean segment from original audio660 segment_start, segment_end, segment_audio = find_clear_segment(audio_path)661 segment_path = os.path.join(tempfile.gettempdir(), "matchering_segment.wav")662 save_segment(segment_audio, 44100, segment_path)663 664 # Process each output with Matchering665 mastered_outputs = []666 for output in outputs:667 if output and os.path.exists(output):668 output_base = sanitize_filename(os.path.splitext(os.path.basename(output))[0])669 mastered_path = os.path.join(OUTPUT_DIR, f"{output_base}_mastered.wav")670 mastered_output = run_matchering(671 reference_path=segment_path,672 target_path=output,673 output_path=mastered_path,674 passes=matchering_passes,675 bit_depth=24676 )677 mastered_outputs.append(mastered_path)678 else:679 mastered_outputs.append(output)680 681 # Clean up segment file682 if os.path.exists(segment_path):683 os.remove(segment_path)684 685 outputs = tuple(mastered_outputs)686 687 # Final yield with all outputs688 yield (689 outputs[0], outputs[1], outputs[2], outputs[3], outputs[4], outputs[5], outputs[6],690 outputs[7], outputs[8], outputs[9], outputs[10], outputs[11], outputs[12], outputs[13],691 outputs[14], outputs[15],692 "Audio processing completed",693 update_progress_html("Audio processing completed", 100)694 )695 696 except Exception as e:697 print(f"process_audio error: {str(e)}")698 import traceback699 traceback.print_exc()700 yield (701 None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None,702 f"Error occurred: {str(e)}",703 update_progress_html("Error occurred", 0)704 )705 706def ensemble_audio_fn(files, method, weights, progress=gr.Progress()):707 try:708 if len(files) < 2:709 return None, "Minimum two files required"710 711 valid_files = [f for f in files if os.path.exists(f)]712 if len(valid_files) < 2:713 return None, "Valid files not found"714 715 output_dir = os.path.join(BASE_DIR, "ensembles")716 os.makedirs(output_dir, exist_ok=True)717 718 timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")719 output_path = f"{output_dir}/ensemble_{timestamp}.wav"720 721 ensemble_args = [722 "--files", *valid_files,723 "--type", method.lower().replace(' ', '_'),724 "--output", output_path725 ]726 727 if weights and weights.strip():728 weights_list = [str(w) for w in map(float, weights.split(','))]729 ensemble_args += ["--weights", *weights_list]730 731 progress(0, desc="Starting ensemble process", total=100)732 733 # Run ensemble subprocess with real-time output capture734 process = subprocess.Popen(735 ["python", "ensemble.py"] + ensemble_args,736 stdout=subprocess.PIPE,737 stderr=subprocess.PIPE,738 text=True,739 bufsize=1,740 universal_newlines=True741 )742 743 stdout_output = ""744 stderr_output = ""745 746 # Read output in real-time and capture actual progress747 for line in process.stdout:748 stdout_output += line749 line_stripped = line.strip()750 751 # Capture real progress percentage from ensemble.py with new format752 if line_stripped.startswith("[SESA_PROGRESS]"):753 try:754 percent_str = line_stripped.replace("[SESA_PROGRESS]", "").strip()755 percent = int(float(percent_str)) if percent_str else 0756 percent = min(max(percent, 0), 100)757 progress(percent, desc=f"Ensemble progress: {percent}%")758 except (ValueError, TypeError):759 pass760 # Legacy format support761 elif line_stripped.startswith("Progress:"):762 try:763 percent = int(line_stripped.split(":")[1].strip().replace("%", ""))764 percent = min(max(percent, 0), 100)765 progress(percent, desc=f"Ensemble progress: {percent}%")766 except (ValueError, IndexError):767 pass768 elif "loading" in line.lower():769 print(f"Ensemble: {line_stripped}")770 progress(5, desc="Loading audio files for ensemble...")771 elif "processing ensemble" in line.lower():772 print(f"Ensemble: {line_stripped}")773 progress(10, desc="Starting ensemble processing...")774 elif "saving" in line.lower():775 print(f"Ensemble: {line_stripped}")776 progress(95, desc="Saving ensemble output...")777 elif line_stripped and not line_stripped.startswith("[SESA_PROGRESS]") and not line_stripped.startswith("Progress:"):778 # Only print non-progress messages779 print(f"Ensemble: {line_stripped}")780 781 for line in process.stderr:782 stderr_output += line783 print(f"Ensemble stderr: {line.strip()}")784 785 process.wait()786 result = type('Result', (), {'stdout': stdout_output, 'stderr': stderr_output, 'returncode': process.returncode})()787 788 progress(100, desc="Ensemble complete")789 log = f"Success: {result.stdout}" if not result.stderr else f"Error: {result.stderr}"790 return output_path, log791 792 except Exception as e:793 return None, f"Critical error: {str(e)}"794 finally:795 progress(100, desc="Ensemble process completed")796 797 798def auto_ensemble_process(799 auto_input_audio_file,800 selected_models,801 auto_chunk_size,802 auto_overlap,803 export_format,804 auto_use_tta,805 auto_extract_instrumental,806 auto_ensemble_type,807 _state,808 auto_use_apollo=True,809 auto_apollo_normal_model="Apollo Universal Model",810 auto_apollo_chunk_size=19,811 auto_apollo_overlap=2,812 auto_apollo_method="normal_method",813 auto_use_matchering=False,814 auto_matchering_passes=1,815 apollo_midside_model=None,816 progress=gr.Progress(track_tqdm=True)817):818 """Process audio with multiple models and ensemble the results, saving output to Google Drive."""819 try:820 # Check Google Drive connection and setup directories821 setup_directories()822 823 if not selected_models or len(selected_models) < 1:824 yield None, i18n("no_models_selected"), update_progress_html(i18n("error_occurred"), 0)825 return826 827 if auto_input_audio_file is None:828 existing_files = os.listdir(INPUT_DIR)829 if not existing_files:830 yield None, i18n("no_input_audio_provided"), update_progress_html(i18n("error_occurred"), 0)831 return832 audio_path = os.path.join(INPUT_DIR, existing_files[0])833 else:834 audio_path = auto_input_audio_file.name if hasattr(auto_input_audio_file, 'name') else auto_input_audio_file835 836 # Copy input file to INPUT_DIR837 input_filename = os.path.basename(audio_path)838 dest_path = os.path.join(INPUT_DIR, input_filename)839 shutil.copy(audio_path, dest_path)840 841 # Parse apollo method842 if auto_apollo_method in ["2", 2]:843 auto_apollo_method = "mid_side_method"844 elif auto_apollo_method in ["1", 1]:845 auto_apollo_method = "normal_method"846 847 corrected_auto_chunk_size = int(auto_apollo_chunk_size)848 corrected_auto_overlap = int(auto_apollo_overlap)849 850 # Setup temporary directories851 auto_ensemble_temp = os.path.join(BASE_DIR, "auto_ensemble_temp")852 os.makedirs(auto_ensemble_temp, exist_ok=True)853 clear_directory(auto_ensemble_temp)854 855 all_outputs = []856 total_models = len(selected_models)857 model_progress_range = 60858 model_progress_per_step = model_progress_range / total_models if total_models > 0 else 0859 860 for i, model in enumerate(selected_models):861 clean_model_name = clean_model(model)862 model_output_dir = os.path.join(auto_ensemble_temp, clean_model_name)863 os.makedirs(model_output_dir, exist_ok=True)864 865 current_progress = i * model_progress_per_step866 current_progress = clamp_percentage(current_progress)867 yield None, i18n("loading_model").format(i+1, total_models, clean_model_name), update_progress_html(868 i18n("loading_model_progress").format(i+1, total_models, clean_model_name, current_progress),869 current_progress870 )871 872 model_type, config_path, start_check_point = get_model_config(clean_model_name, auto_chunk_size, auto_overlap)873 874 # Read the model's native chunk_size from its YAML after download875 native_chunk = get_model_chunk_size(clean_model_name)876 effective_chunk_size = native_chunk if native_chunk else auto_chunk_size877 if native_chunk:878 print(f"Using model's native chunk_size from YAML: {native_chunk} (UI value was: {auto_chunk_size})")879 880 cmd = [881 "python", INFERENCE_PATH,882 "--model_type", model_type,883 "--config_path", config_path,884 "--start_check_point", start_check_point,885 "--input_folder", INPUT_DIR,886 "--store_dir", model_output_dir,887 "--chunk_size", str(effective_chunk_size),888 "--overlap", str(auto_overlap),889 "--export_format", f"{export_format.split()[0].lower()} FLOAT"890 ]891 if auto_use_tta:892 cmd.append("--use_tta")893 if auto_extract_instrumental:894 cmd.append("--extract_instrumental")895 896 print(f"Running command: {' '.join(cmd)}")897 process = subprocess.Popen(898 cmd,899 stdout=subprocess.PIPE,900 stderr=subprocess.PIPE,901 text=True,902 bufsize=1,903 universal_newlines=True904 )905 906 stderr_output = ""907 last_yield_percent = -1908 downloading_file = None909 910 for line in process.stdout:911 line_stripped = line.strip()912 913 # Check for download progress [SESA_DOWNLOAD]914 if line_stripped.startswith("[SESA_DOWNLOAD]"):915 try:916 dl_info = line_stripped.replace("[SESA_DOWNLOAD]", "")917 if dl_info.startswith("START:"):918 downloading_file = dl_info.replace("START:", "")919 yield None, i18n("downloading_model_file").format(downloading_file), update_progress_html(920 i18n("downloading_model_file").format(downloading_file),921 i * model_progress_per_step,922 download_info={"filename": downloading_file, "percent": 0}923 )924 elif dl_info.startswith("END:"):925 downloading_file = None926 elif ":" in dl_info:927 parts = dl_info.rsplit(":", 1)928 if len(parts) == 2:929 filename, percent_str = parts930 download_percent = int(percent_str)931 yield None, i18n("downloading_file_progress").format(filename, download_percent), update_progress_html(932 i18n("downloading_model_file").format(filename),933 i * model_progress_per_step,934 download_info={"filename": filename, "percent": download_percent}935 )936 except (ValueError, TypeError):937 pass938 # Check for unique progress prefix [SESA_PROGRESS]939 elif line_stripped.startswith("[SESA_PROGRESS]"):940 try:941 # Extract percentage from [SESA_PROGRESS]XX format942 percentage_str = line_stripped.replace("[SESA_PROGRESS]", "").strip()943 percentage = float(percentage_str) if percentage_str else 0944 percentage = min(max(percentage, 0), 100) # Clamp to 0-100945 946 model_percentage = (percentage / 100) * model_progress_per_step947 current_progress = (i * model_progress_per_step) + model_percentage948 current_progress = clamp_percentage(current_progress)949 950 # Yield on every percent change for smooth updates951 if int(percentage) != last_yield_percent:952 last_yield_percent = int(percentage)953 yield None, i18n("loading_model_progress_label").format(i+1, total_models, clean_model_name, int(percentage)), update_progress_html(954 f"Model {i+1}/{total_models}: {clean_model_name} - {int(percentage)}%",955 current_progress956 )957 except (ValueError, TypeError):958 # Silently ignore parsing errors for progress lines959 pass960 # Also support legacy "Progress: XX%" format for backwards compatibility961 elif line_stripped.startswith("Progress:"):962 try:963 match = re.search(r"Progress:\s*(\d+(?:\.\d+)?)%?", line_stripped)964 if match:965 percentage = float(match.group(1))966 percentage = min(max(percentage, 0), 100)967 968 model_percentage = (percentage / 100) * model_progress_per_step969 current_progress = (i * model_progress_per_step) + model_percentage970 current_progress = clamp_percentage(current_progress)971 972 if int(percentage) != last_yield_percent:973 last_yield_percent = int(percentage)974 yield None, i18n("loading_model_progress_label").format(i+1, total_models, clean_model_name, int(percentage)), update_progress_html(975 f"Model {i+1}/{total_models}: {clean_model_name} - {int(percentage)}%",976 current_progress977 )978 except (ValueError, TypeError):979 pass980 else:981 # Print non-progress lines982 if line_stripped:983 print(line_stripped)984 985 for line in process.stderr:986 stderr_output += line987 print(line.strip())988 989 process.wait()990 if process.returncode != 0:991 print(f"Error: {stderr_output}")992 yield None, i18n("model_failed").format(clean_model_name, stderr_output), update_progress_html(993 i18n("error_occurred"), 0994 )995 return996 997 gc.collect()998 if torch.cuda.is_available():999 torch.cuda.empty_cache()1000 1001 current_progress = (i + 1) * model_progress_per_step1002 current_progress = clamp_percentage(current_progress)1003 yield None, i18n("completed_model").format(i+1, total_models, clean_model_name), update_progress_html(1004 i18n("completed_model_progress").format(i+1, total_models, clean_model_name, current_progress),1005 current_progress1006 )1007 1008 model_outputs = glob.glob(os.path.join(model_output_dir, "*.wav"))1009 if not model_outputs:1010 raise FileNotFoundError(i18n("model_output_failed").format(clean_model_name))1011 all_outputs.extend(model_outputs)1012 1013 # Select compatible files for ensemble1014 preferred_type = 'instrumental' if auto_extract_instrumental else 'vocals'1015 ensemble_files = [output for output in all_outputs if preferred_type.lower() in output.lower()]1016 print(f"Selected ensemble files: {ensemble_files}")1017 if len(ensemble_files) < 2:1018 print(f"Warning: Insufficient {preferred_type} files ({len(ensemble_files)}). Falling back to all outputs.")1019 ensemble_files = all_outputs1020 if len(ensemble_files) < 2:1021 raise ValueError(i18n("insufficient_files_for_ensemble").format(len(ensemble_files)))1022 1023 # Enhanced outputs with Apollo (if enabled)1024 if auto_use_apollo:1025 yield None, i18n("enhancing_with_apollo").format(0, len(all_outputs)), update_progress_html(1026 i18n("waiting_for_files"), 601027 )1028 1029 all_outputs = process_with_apollo(1030 output_files=all_outputs,1031 output_dir=auto_ensemble_temp,1032 apollo_chunk_size=corrected_auto_chunk_size,1033 apollo_overlap=corrected_auto_overlap,1034 apollo_method=auto_apollo_method,1035 apollo_normal_model=auto_apollo_normal_model,1036 apollo_midside_model=apollo_midside_model,1037 output_format=export_format.split()[0].lower(),1038 progress=progress,1039 total_progress_start=60,1040 total_progress_end=901041 )1042 1043 # Perform ensemble1044 yield None, i18n("performing_ensemble"), update_progress_html(1045 i18n("performing_ensemble"), 901046 )1047 1048 quoted_files = [f'"{f}"' for f in ensemble_files]1049 timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")1050 output_path = os.path.join(AUTO_ENSEMBLE_OUTPUT, f"auto_ensemble_output_{timestamp}.wav")1051 1052 ensemble_cmd = [1053 "python", ENSEMBLE_PATH,1054 "--files", *quoted_files,1055 "--type", auto_ensemble_type,1056 "--output", f'"{output_path}"'1057 ]1058 1059 print(f"Running ensemble command: {' '.join(ensemble_cmd)}")1060 try:1061 process = subprocess.Popen(1062 " ".join(ensemble_cmd),1063 shell=True,1064 stdout=subprocess.PIPE,1065 stderr=subprocess.PIPE,1066 text=True,1067 bufsize=1,1068 universal_newlines=True1069 )1070 1071 stdout_output = ""1072 stderr_output = ""1073 for line in process.stdout:1074 stdout_output += line1075 print(f"Ensemble stdout: {line.strip()}")1076 for line in process.stderr:1077 stderr_output += line1078 print(f"Ensemble stderr: {line.strip()}")1079 1080 process.wait()1081 if process.returncode != 0:1082 print(f"Ensemble subprocess failed with code {process.returncode}: {stderr_output}")1083 yield None, i18n("ensemble_error").format(stderr_output), update_progress_html(1084 i18n("error_occurred"), 01085 )1086 return1087 1088 print(f"Checking if output file exists: {output_path}")1089 if not os.path.exists(output_path):1090 raise RuntimeError(f"Ensemble output file not created at {output_path}. Stdout: {stdout_output}, Stderr: {stderr_output}")1091 1092 except Exception as e:1093 print(f"Ensemble command execution failed: {str(e)}")1094 yield None, i18n("ensemble_error").format(str(e)), update_progress_html(1095 i18n("error_occurred"), 01096 )1097 return1098 1099 # Apply Matchering (if enabled)1100 if auto_use_matchering and os.path.exists(output_path):1101 yield None, i18n("applying_matchering"), update_progress_html(1102 i18n("applying_matchering"), 981103 )1104 1105 try:1106 # Find clean segment1107 segment_start, segment_end, segment_audio = find_clear_segment(audio_path)1108 segment_path = os.path.join(tempfile.gettempdir(), "matchering_segment.wav")1109 save_segment(segment_audio, 44100, segment_path)1110 1111 # Master the ensemble output1112 mastered_output_path = os.path.join(AUTO_ENSEMBLE_OUTPUT, f"auto_ensemble_output_{timestamp}_mastered.wav")1113 print(f"Running Matchering: reference={segment_path}, target={output_path}, output={mastered_output_path}")1114 mastered_output = run_matchering(1115 reference_path=segment_path,1116 target_path=output_path,1117 output_path=mastered_output_path,1118 passes=auto_matchering_passes,1119 bit_depth=241120 )1121 1122 # Verify mastered output1123 if not os.path.exists(mastered_output_path):1124 raise RuntimeError(f"Matchering failed to create output at {mastered_output_path}")1125 1126 # Clean up segment file1127 if os.path.exists(segment_path):1128 os.remove(segment_path)1129 1130 output_path = mastered_output_path1131 print(f"Matchering completed: {mastered_output_path}")1132 except Exception as e:1133 print(f"Matchering error: {str(e)}")1134 yield None, i18n("error").format(f"Matchering failed: {str(e)}"), update_progress_html(1135 i18n("error_occurred"), 01136 )1137 return1138 1139 yield None, i18n("finalizing_ensemble_output"), update_progress_html(1140 i18n("finalizing_ensemble_output"), 981141 )1142 1143 if not os.path.exists(output_path):1144 raise RuntimeError(i18n("ensemble_file_creation_failed").format(output_path))1145 1146 # Verify write permissions for Google Drive directory1147 try:1148 print(f"Verifying write permissions for {AUTO_ENSEMBLE_OUTPUT}")1149 test_file = os.path.join(AUTO_ENSEMBLE_OUTPUT, "test_write.txt")1150 with open(test_file, "w") as f:1151 f.write("Test")1152 os.remove(test_file)1153 print(f"Write permissions verified for {AUTO_ENSEMBLE_OUTPUT}")1154 except Exception as e:1155 print(f"Write permission error for {AUTO_ENSEMBLE_OUTPUT}: {str(e)}")1156 yield None, i18n("error").format(f"Write permission error: {str(e)}"), update_progress_html(1157 i18n("error_occurred"), 01158 )1159 return1160 1161 # Verify file in Google Drive1162 print(f"Final output file: {output_path}")1163 if IS_COLAB:1164 drive_output_path = os.path.join("/content/drive/MyDrive/ensemble_output", os.path.basename(output_path))1165 print(f"Checking if file exists in Google Drive: {drive_output_path}")1166 if not os.path.exists(drive_output_path):1167 print(f"File not found in Google Drive, copying from local path: {output_path}")1168 shutil.copy(output_path, drive_output_path)1169 print(f"Copied to Google Drive: {drive_output_path}")1170 else:1171 drive_output_path = output_path1172 1173 yield output_path, i18n("success_output_created") + f" Saved to {drive_output_path if IS_COLAB else output_path}", update_progress_html(1174 i18n("ensemble_completed"), 1001175 )1176 1177 except Exception as e:1178 print(f"auto_ensemble_process error: {str(e)}")1179 import traceback1180 traceback.print_exc()1181 yield None, i18n("error").format(str(e)), update_progress_html(1182 i18n("error_occurred"), 01183 )1184 finally:1185 shutil.rmtree(auto_ensemble_temp, ignore_errors=True)1186 gc.collect()1187 if torch.cuda.is_available():1188 torch.cuda.empty_cache()1189