ASesYusuf1/SESA_Audio_Separation
14
1import os2import shutil3import glob4import re5import subprocess6import random7import yaml8from pathlib import Path9import torch10import gradio as gr11import threading12import time13import librosa14import soundfile as sf15import numpy as np16import requests17import json18import locale19from datetime import datetime20import yt_dlp21import validators22from pytube import YouTube23 24# Google API imports (optional - for Colab/Google Drive support)25try:26 from googleapiclient.discovery import build27 from googleapiclient.http import MediaIoBaseDownload28 from google.oauth2.credentials import Credentials29 GOOGLE_API_AVAILABLE = True30except ImportError:31 GOOGLE_API_AVAILABLE = False32 build = None33 MediaIoBaseDownload = None34 Credentials = None35 36import io37import math38import hashlib39import gc40import psutil41import concurrent.futures42from tqdm import tqdm43import tempfile44from urllib.parse import urlparse, quote45import argparse46from tqdm.auto import tqdm47import torch.nn as nn48from model import get_model_config, MODEL_CONFIGS, get_all_model_configs_with_custom, load_custom_models49from assets.i18n.i18n import I18nAuto50import matchering as mg51from scipy.signal import find_peaks52 53i18n = I18nAuto()54 55# Temel dizinler56BASE_DIR = os.path.dirname(os.path.abspath(__file__))57INPUT_DIR = os.path.join(BASE_DIR, "input")58OUTPUT_DIR = os.path.join(BASE_DIR, "output")59OLD_OUTPUT_DIR = os.path.join(BASE_DIR, "old_output")60AUTO_ENSEMBLE_TEMP = os.path.join(BASE_DIR, "auto_ensemble_temp")61AUTO_ENSEMBLE_OUTPUT = os.path.join(BASE_DIR, "ensemble_folder")62VIDEO_TEMP = os.path.join(BASE_DIR, "video_temp")63ENSEMBLE_DIR = os.path.join(BASE_DIR, "ensemble")64COOKIE_PATH = os.path.join(BASE_DIR, "cookies.txt")65INFERENCE_SCRIPT_PATH = os.path.join(BASE_DIR, "inference.py")66 67def extract_model_name_from_checkpoint(checkpoint_path):68 if not checkpoint_path:69 return "Unknown"70 base_name = os.path.basename(checkpoint_path)71 model_name = os.path.splitext(base_name)[0]72 return model_name.strip()73 74for directory in [BASE_DIR, INPUT_DIR, OUTPUT_DIR, OLD_OUTPUT_DIR, AUTO_ENSEMBLE_TEMP, AUTO_ENSEMBLE_OUTPUT, VIDEO_TEMP, ENSEMBLE_DIR]:75 os.makedirs(directory, exist_ok=True)76 77class IndentDumper(yaml.Dumper):78 def increase_indent(self, flow=False, indentless=False):79 return super(IndentDumper, self).increase_indent(flow, False)80 81def tuple_constructor(loader, node):82 """YAML'dan bir tuple yükler."""83 values = loader.construct_sequence(node)84 return tuple(values)85 86yaml.SafeLoader.add_constructor('tag:yaml.org,2002:python/tuple', tuple_constructor)87 88def clean_model(model):89 """90 Cleans a model name by removing unwanted characters like ⭐ and extra whitespace.91 92 Args:93 model (str): The model name to clean.94 95 Returns:96 str: The cleaned model name, or None if input is invalid.97 """98 if not model or not isinstance(model, str):99 return None100 # Remove ⭐ and extra whitespace101 cleaned = model.replace("⭐", "").strip()102 # Remove any additional unwanted characters if needed103 cleaned = cleaned.replace("\t", " ").replace("\n", " ")104 return cleaned105 106def get_original_category(translated_category):107 all_configs = get_all_model_configs_with_custom()108 for original_cat in all_configs.keys():109 if i18n(original_cat) == translated_category:110 return original_cat111 return None112 113def clamp_percentage(value):114 """Clamp percentage values to the 0-100 range."""115 try:116 return min(max(float(value), 0), 100)117 except (ValueError, TypeError):118 print(f"Warning: Invalid percentage value {value}, defaulting to 0")119 return 0 120 121def update_model_dropdown(category, favorites=None):122 # Get all configs including custom models123 all_configs = get_all_model_configs_with_custom()124 # Map translated category back to English125 eng_cat = next((k for k in all_configs.keys() if i18n(k) == category), list(all_configs.keys())[0])126 models = all_configs.get(eng_cat, {})127 choices = []128 favorite_models = []129 non_favorite_models = []130 131 for model in models:132 model_name = f"{model} ⭐" if favorites and model in favorites else model133 if favorites and model in favorites:134 favorite_models.append(model_name)135 else:136 non_favorite_models.append(model_name)137 138 choices = favorite_models + non_favorite_models139 return {"choices": choices}140 141def get_model_categories():142 """Get all model categories including Custom Models if any exist."""143 all_configs = get_all_model_configs_with_custom()144 return list(all_configs.keys())145 146def handle_file_upload(uploaded_file, file_path, is_auto_ensemble=False):147 clear_temp_folder("/tmp", exclude_items=["gradio", "config.json"])148 clear_directory(INPUT_DIR)149 os.makedirs(INPUT_DIR, exist_ok=True)150 clear_directory(INPUT_DIR)151 if uploaded_file:152 target_path = save_uploaded_file(uploaded_file, is_input=True)153 return target_path, target_path154 elif file_path and os.path.exists(file_path):155 target_path = os.path.join(INPUT_DIR, os.path.basename(file_path))156 shutil.copy(file_path, target_path)157 return target_path, target_path158 return None, None159 160 if torch.cuda.is_available():161 torch.cuda.empty_cache()162 163def clear_directory(directory):164 """Verilen dizindeki tüm dosyaları siler."""165 files = glob.glob(os.path.join(directory, '*'))166 for f in files:167 try:168 os.remove(f)169 except Exception as e:170 print(i18n("file_deletion_error").format(f, e))171 172def clear_temp_folder(folder_path, exclude_items=None):173 """Dizinin içeriğini güvenli bir şekilde temizler ve belirtilen öğeleri korur."""174 try:175 if not os.path.exists(folder_path):176 print(i18n("directory_not_exist_warning").format(folder_path))177 return False178 if not os.path.isdir(folder_path):179 print(i18n("not_a_directory_warning").format(folder_path))180 return False181 exclude_items = exclude_items or []182 for item_name in os.listdir(folder_path):183 item_path = os.path.join(folder_path, item_name)184 if item_name in exclude_items:185 continue186 try:187 if os.path.isfile(item_path) or os.path.islink(item_path):188 os.unlink(item_path)189 elif os.path.isdir(item_path):190 shutil.rmtree(item_path)191 except Exception as e:192 print(i18n("item_deletion_error").format(item_path, e))193 return True194 except Exception as e:195 print(i18n("critical_error").format(e))196 return False197 198def clear_old_output():199 old_output_folder = os.path.join(BASE_DIR, 'old_output')200 try:201 if not os.path.exists(old_output_folder):202 return i18n("old_output_not_exist")203 shutil.rmtree(old_output_folder)204 os.makedirs(old_output_folder, exist_ok=True)205 return i18n("old_outputs_cleared")206 except Exception as e:207 return i18n("error").format(e)208 209def shorten_filename(filename, max_length=30):210 """Dosya adını belirtilen maksimum uzunluğa kısaltır."""211 base, ext = os.path.splitext(filename)212 if len(base) <= max_length:213 return filename214 return base[:15] + "..." + base[-10:] + ext215 216def clean_filename(title):217 """Dosya adından özel karakterleri kaldırır."""218 return re.sub(r'[^\w\-_\. ]', '', title).strip()219 220def sanitize_filename(filename):221 base, ext = os.path.splitext(filename)222 base = re.sub(r'\.+', '_', base)223 base = re.sub(r'[#<>:"/\\|?*]', '_', base)224 base = re.sub(r'\s+', '_', base)225 base = re.sub(r'_+', '_', base)226 base = base.strip('_')227 return f"{base}{ext}"228 229def convert_to_wav(file_path):230 """Ses dosyasını WAV formatına dönüştürür."""231 original_filename = os.path.basename(file_path)232 filename, ext = os.path.splitext(original_filename)233 if ext.lower() == '.wav':234 return file_path235 wav_output = os.path.join(ENSEMBLE_DIR, f"{filename}.wav")236 try:237 command = [238 'ffmpeg', '-y', '-i', file_path,239 '-acodec', 'pcm_s16le', '-ar', '44100', wav_output240 ]241 subprocess.run(command, check=True, capture_output=True)242 return wav_output243 except subprocess.CalledProcessError as e:244 print(i18n("ffmpeg_error").format(e.returncode, e.stderr.decode()))245 return None246 247def generate_random_port():248 """Rastgele bir port numarası oluşturur."""249 return random.randint(1000, 9000)250 251def save_segment(audio, sr, path):252 """253 Save audio segment to a file.254 255 Args:256 audio (np.ndarray): Audio data.257 sr (int): Sample rate.258 path (str): Output file path.259 """260 sf.write(path, audio, sr)261 262def run_matchering(reference_path, target_path, output_path, passes=1, bit_depth=24):263 """264 Run Matchering to master the target audio using the reference audio.265 266 Args:267 reference_path (str): Path to the reference audio (clear segment).268 target_path (str): Path to the target audio to be mastered.269 output_path (str): Path for the mastered output.270 passes (int): Number of Matchering passes (1-4).271 bit_depth (int): Output bit depth (16 or 24).272 273 Returns:274 str: Path to the mastered output file.275 """276 # Ensure inputs are WAV files277 ref_audio, sr = librosa.load(reference_path, sr=44100, mono=False)278 tgt_audio, sr = librosa.load(target_path, sr=44100, mono=False)279 280 # Save temporary WAV files281 temp_ref = os.path.join(tempfile.gettempdir(), "matchering_ref.wav")282 temp_tgt = os.path.join(tempfile.gettempdir(), "matchering_tgt.wav")283 save_segment(ref_audio.T if ref_audio.ndim > 1 else ref_audio, sr, temp_ref)284 save_segment(tgt_audio.T if tgt_audio.ndim > 1 else tgt_audio, sr, temp_tgt)285 286 # Configure Matchering with default settings287 config = mg.Config() # No parameters, use defaults288 289 # Select bit depth for output290 result_format = mg.pcm24 if bit_depth == 24 else mg.pcm16291 292 # Run Matchering for multiple passes293 current_tgt = temp_tgt294 for i in range(passes):295 temp_out = os.path.join(tempfile.gettempdir(), f"matchering_out_pass_{i}.wav")296 mg.process(297 reference=temp_ref,298 target=current_tgt,299 results=[result_format(temp_out)], # Bit depth control300 config=config301 )302 current_tgt = temp_out303 304 # Move final output to desired path305 shutil.move(current_tgt, output_path)306 307 # Clean up temporary files308 for temp_file in [temp_ref, temp_tgt] + [os.path.join(tempfile.gettempdir(), f"matchering_out_pass_{i}.wav") for i in range(passes-1)]:309 if os.path.exists(temp_file):310 os.remove(temp_file)311 312 return output_path 313 314def find_clear_segment(audio_path, segment_duration=15, sr=44100):315 """316 Find the clearest (high-energy, low-noise) segment in an audio file.317 318 Args:319 audio_path (str): Path to the original audio file.320 segment_duration (float): Duration of the segment to extract (in seconds).321 sr (int): Sample rate for loading audio.322 323 Returns:324 tuple: (start_time, end_time, segment_audio) of the clearest segment.325 """326 # Load audio327 audio, sr = librosa.load(audio_path, sr=sr, mono=True)328 329 # Compute RMS energy in windows330 window_size = int(5 * sr) # 5-second windows331 hop_length = window_size // 2332 rms = librosa.feature.rms(y=audio, frame_length=window_size, hop_length=hop_length)[0]333 334 # Compute spectral flatness for noise detection335 flatness = librosa.feature.spectral_flatness(y=audio, n_fft=window_size, hop_length=hop_length)[0]336 337 # Combine metrics: high RMS and low flatness indicate clear, high-energy segments338 score = rms / (flatness + 1e-6) # Avoid division by zero339 340 # Find peaks in the score341 peaks, _ = find_peaks(score, height=np.mean(score), distance=5)342 if len(peaks) == 0:343 # Fallback: Use the middle of the track344 peak_idx = len(score) // 2345 else:346 peak_idx = peaks[np.argmax(score[peaks])]347 348 # Calculate start and end times349 start_sample = peak_idx * hop_length350 end_sample = start_sample + int(segment_duration * sr)351 352 # Ensure the segment fits within the audio353 if end_sample > len(audio):354 end_sample = len(audio)355 start_sample = max(0, end_sample - int(segment_duration * sr))356 357 start_time = start_sample / sr358 end_time = end_sample / sr359 segment_audio = audio[start_sample:end_sample]360 361 return start_time, end_time, segment_audio362 363def update_file_list():364 output_files = glob.glob(os.path.join(OUTPUT_DIR, "*.wav"))365 old_output_files = glob.glob(os.path.join(OLD_OUTPUT_DIR, "*.wav"))366 files = output_files + old_output_files367 return gr.Dropdown(choices=files)368 369def save_uploaded_file(uploaded_file, is_input=False, target_dir=None):370 """Yüklenen dosyayı belirtilen dizine kaydeder."""371 media_extensions = ['.mp3', '.wav', '.flac', '.aac', '.ogg', '.m4a', '.mp4']372 target_dir = target_dir or (INPUT_DIR if is_input else OUTPUT_DIR)373 timestamp_patterns = [374 r'_\d{8}_\d{6}_\d{6}$', r'_\d{14}$', r'_\d{10}$', r'_\d+$'375 ]376 377 if hasattr(uploaded_file, 'name'):378 original_filename = os.path.basename(uploaded_file.name)379 else:380 original_filename = os.path.basename(str(uploaded_file))381 382 if is_input:383 base_filename = original_filename384 for pattern in timestamp_patterns:385 base_filename = re.sub(pattern, '', base_filename)386 for ext in media_extensions:387 base_filename = base_filename.replace(ext, '')388 file_ext = next(389 (ext for ext in media_extensions if original_filename.lower().endswith(ext)),390 '.wav'391 )392 clean_filename = f"{base_filename.strip('_- ')}{file_ext}"393 else:394 clean_filename = original_filename395 396 target_path = os.path.join(target_dir, clean_filename)397 os.makedirs(target_dir, exist_ok=True)398 399 if os.path.exists(target_path):400 os.remove(target_path)401 402 if hasattr(uploaded_file, 'read'):403 with open(target_path, "wb") as f:404 f.write(uploaded_file.read())405 else:406 shutil.copy(uploaded_file, target_path)407 408 print(i18n("file_saved_successfully").format(os.path.basename(target_path)))409 return target_path410 411def move_old_files(output_folder):412 """Eski dosyaları old_output dizinine taşır."""413 os.makedirs(OLD_OUTPUT_DIR, exist_ok=True)414 for filename in os.listdir(output_folder):415 file_path = os.path.join(output_folder, filename)416 if os.path.isfile(file_path):417 new_filename = f"{os.path.splitext(filename)[0]}_old{os.path.splitext(filename)[1]}"418 new_file_path = os.path.join(OLD_OUTPUT_DIR, new_filename)419 shutil.move(file_path, new_file_path)420