Update app.py

app.py

@@ -25,6 +25,13 @@ import scipy.io.wavfile
 import subprocess
 import spaces
 import torchaudio
+from models_config import (
+    EXTENDED_MODELS, get_all_models, get_categories, get_model_choices,
+    find_model_filename, add_custom_model, delete_custom_model, load_custom_models,
+    get_custom_models_list, ensure_model_files_downloaded,
+    get_audio_duration, split_audio_segments, concatenate_segment_outputs,
+    MAX_UNSPLIT_DURATION, SEGMENT_DURATION
+)
 
 # Logging setup
 logging.basicConfig(level=logging.INFO)
@@ -63,97 +70,8 @@ max_retries = 2
 time_budget = 300  # ZeroGPU için işlem sınırı
 gpu_lock = Lock()
 
-# ROFORMER_MODELS and OUTPUT_FORMATS
-ROFORMER_MODELS = {
-    "Vocals": {
-        'MelBand Roformer | Big Beta 6X by unwa': 'melband_roformer_big_beta6x.ckpt',
-        'MelBand Roformer Kim | Big Beta 4 FT by unwa': 'melband_roformer_big_beta4.ckpt',
-        'MelBand Roformer Kim | Big Beta 5e FT by unwa': 'melband_roformer_big_beta5e.ckpt',
-        'MelBand Roformer | Big Beta 6 by unwa': 'melband_roformer_big_beta6.ckpt',
-        'MelBand Roformer | Vocals by Kimberley Jensen': 'vocals_mel_band_roformer.ckpt',
-        'MelBand Roformer Kim | FT 3 by unwa': 'mel_band_roformer_kim_ft3_unwa.ckpt',
-        'MelBand Roformer Kim | FT by unwa': 'mel_band_roformer_kim_ft_unwa.ckpt',
-        'MelBand Roformer Kim | FT 2 by unwa': 'mel_band_roformer_kim_ft2_unwa.ckpt',
-        'MelBand Roformer Kim | FT 2 Bleedless by unwa': 'mel_band_roformer_kim_ft2_bleedless_unwa.ckpt',
-        'MelBand Roformer | Vocals by becruily': 'mel_band_roformer_vocals_becruily.ckpt',
-        'MelBand Roformer | Vocals Fullness by Aname': 'mel_band_roformer_vocal_fullness_aname.ckpt',
-        'BS Roformer | Vocals by Gabox': 'bs_roformer_vocals_gabox.ckpt',
-        'MelBand Roformer | Vocals by Gabox': 'mel_band_roformer_vocals_gabox.ckpt',
-        'MelBand Roformer | Vocals FV1 by Gabox': 'mel_band_roformer_vocals_fv1_gabox.ckpt',
-        'MelBand Roformer | Vocals FV2 by Gabox': 'mel_band_roformer_vocals_fv2_gabox.ckpt',
-        'MelBand Roformer | Vocals FV3 by Gabox': 'mel_band_roformer_vocals_fv3_gabox.ckpt',
-        'MelBand Roformer | Vocals FV4 by Gabox': 'mel_band_roformer_vocals_fv4_gabox.ckpt',
-        'BS Roformer | Chorus Male-Female by Sucial': 'model_chorus_bs_roformer_ep_267_sdr_24.1275.ckpt',
-        'BS Roformer | Male-Female by aufr33': 'bs_roformer_male_female_by_aufr33_sdr_7.2889.ckpt',
-    },
-    "Instrumentals": {
-        'MelBand Roformer | FVX by Gabox': 'mel_band_roformer_instrumental_fvx_gabox.ckpt',
-        'MelBand Roformer | INSTV8N by Gabox': 'mel_band_roformer_instrumental_instv8n_gabox.ckpt',
-        'MelBand Roformer | INSTV8 by Gabox': 'mel_band_roformer_instrumental_instv8_gabox.ckpt',
-        'MelBand Roformer | INSTV7N by Gabox': 'mel_band_roformer_instrumental_instv7n_gabox.ckpt',
-        'MelBand Roformer | Instrumental Bleedless V3 by Gabox': 'mel_band_roformer_instrumental_bleedless_v3_gabox.ckpt',
-        'MelBand Roformer Kim | Inst V1 (E) Plus by Unwa': 'melband_roformer_inst_v1e_plus.ckpt',
-        'MelBand Roformer Kim | Inst V1 Plus by Unwa': 'melband_roformer_inst_v1_plus.ckpt',
-        'MelBand Roformer Kim | Inst V1 by Unwa': 'melband_roformer_inst_v1.ckpt',
-        'MelBand Roformer Kim | Inst V1 (E) by Unwa': 'melband_roformer_inst_v1e.ckpt',
-        'MelBand Roformer Kim | Inst V2 by Unwa': 'melband_roformer_inst_v2.ckpt',
-        'MelBand Roformer | Instrumental by becruily': 'mel_band_roformer_instrumental_becruily.ckpt',
-        'MelBand Roformer | Instrumental by Gabox': 'mel_band_roformer_instrumental_gabox.ckpt',
-        'MelBand Roformer | Instrumental 2 by Gabox': 'mel_band_roformer_instrumental_2_gabox.ckpt',
-        'MelBand Roformer | Instrumental 3 by Gabox': 'mel_band_roformer_instrumental_3_gabox.ckpt',
-        'MelBand Roformer | Instrumental Bleedless V1 by Gabox': 'mel_band_roformer_instrumental_bleedless_v1_gabox.ckpt',
-        'MelBand Roformer | Instrumental Bleedless V2 by Gabox': 'mel_band_roformer_instrumental_bleedless_v2_gabox.ckpt',
-        'MelBand Roformer | Instrumental Fullness V1 by Gabox': 'mel_band_roformer_instrumental_fullness_v1_gabox.ckpt',
-        'MelBand Roformer | Instrumental Fullness V2 by Gabox': 'mel_band_roformer_instrumental_fullness_v2_gabox.ckpt',
-        'MelBand Roformer | Instrumental Fullness V3 by Gabox': 'mel_band_roformer_instrumental_fullness_v3_gabox.ckpt',
-        'MelBand Roformer | Instrumental Fullness Noisy V4 by Gabox': 'mel_band_roformer_instrumental_fullness_noise_v4_gabox.ckpt',
-        'MelBand Roformer | INSTV5 by Gabox': 'mel_band_roformer_instrumental_instv5_gabox.ckpt',
-        'MelBand Roformer | INSTV5N by Gabox': 'mel_band_roformer_instrumental_instv5n_gabox.ckpt',
-        'MelBand Roformer | INSTV6 by Gabox': 'mel_band_roformer_instrumental_instv6_gabox.ckpt',
-        'MelBand Roformer | INSTV6N by Gabox': 'mel_band_roformer_instrumental_instv6n_gabox.ckpt',
-        'MelBand Roformer | INSTV7 by Gabox': 'mel_band_roformer_instrumental_instv7_gabox.ckpt',
-    },
-    "InstVoc Duality": {
-        'MelBand Roformer Kim | InstVoc Duality V1 by Unwa': 'melband_roformer_instvoc_duality_v1.ckpt',
-        'MelBand Roformer Kim | InstVoc Duality V2 by Unwa': 'melband_roformer_instvox_duality_v2.ckpt',
-    },
-    "De-Reverb": {
-        'BS-Roformer-De-Reverb': 'deverb_bs_roformer_8_384dim_10depth.ckpt',
-        'MelBand Roformer | De-Reverb by anvuew': 'dereverb_mel_band_roformer_anvuew_sdr_19.1729.ckpt',
-        'MelBand Roformer | De-Reverb Less Aggressive by anvuew': 'dereverb_mel_band_roformer_less_aggressive_anvuew_sdr_18.8050.ckpt',
-        'MelBand Roformer | De-Reverb Mono by anvuew': 'dereverb_mel_band_roformer_mono_anvuew.ckpt',
-        'MelBand Roformer | De-Reverb Big by Sucial': 'dereverb_big_mbr_ep_362.ckpt',
-        'MelBand Roformer | De-Reverb Super Big by Sucial': 'dereverb_super_big_mbr_ep_346.ckpt',
-        'MelBand Roformer | De-Reverb-Echo by Sucial': 'dereverb-echo_mel_band_roformer_sdr_10.0169.ckpt',
-        'MelBand Roformer | De-Reverb-Echo V2 by Sucial': 'dereverb-echo_mel_band_roformer_sdr_13.4843_v2.ckpt',
-        'MelBand Roformer | De-Reverb-Echo Fused by Sucial': 'dereverb_echo_mbr_fused.ckpt',
-    },
-    "Denoise": {
-        'Mel-Roformer-Denoise-Aufr33': 'denoise_mel_band_roformer_aufr33_sdr_27.9959.ckpt',
-        'Mel-Roformer-Denoise-Aufr33-Aggr': 'denoise_mel_band_roformer_aufr33_aggr_sdr_27.9768.ckpt',
-        'MelBand Roformer | Denoise-Debleed by Gabox': 'mel_band_roformer_denoise_debleed_gabox.ckpt',
-        'MelBand Roformer | Bleed Suppressor V1 by unwa-97chris': 'mel_band_roformer_bleed_suppressor_v1.ckpt',
-    },
-    "Karaoke": {
-        'Mel-Roformer-Karaoke-Aufr33-Viperx': 'mel_band_roformer_karaoke_aufr33_viperx_sdr_10.1956.ckpt',
-        'MelBand Roformer | Karaoke by Gabox': 'mel_band_roformer_karaoke_gabox.ckpt',
-        'MelBand Roformer | Karaoke by becruily': 'mel_band_roformer_karaoke_becruily.ckpt',
-    },
-    "General Purpose": {
-        'BS-Roformer-Viperx-1297': 'model_bs_roformer_ep_317_sdr_12.9755.ckpt',
-        'BS-Roformer-Viperx-1296': 'model_bs_roformer_ep_368_sdr_12.9628.ckpt',
-        'BS-Roformer-Viperx-1053': 'model_bs_roformer_ep_937_sdr_10.5309.ckpt',
-        'Mel-Roformer-Viperx-1143': 'model_mel_band_roformer_ep_3005_sdr_11.4360.ckpt',
-        'Mel-Roformer-Crowd-Aufr33-Viperx': 'mel_band_roformer_crowd_aufr33_viperx_sdr_8.7144.ckpt',
-        'MelBand Roformer Kim | SYHFT by SYH99999': 'MelBandRoformerSYHFT.ckpt',
-        'MelBand Roformer Kim | SYHFT V2 by SYH99999': 'MelBandRoformerSYHFTV2.ckpt',
-        'MelBand Roformer Kim | SYHFT V2.5 by SYH99999': 'MelBandRoformerSYHFTV2.5.ckpt',
-        'MelBand Roformer Kim | SYHFT V3 by SYH99999': 'MelBandRoformerSYHFTV3Epsilon.ckpt',
-        'MelBand Roformer Kim | Big SYHFT V1 by SYH99999': 'MelBandRoformerBigSYHFTV1.ckpt',
-        'MelBand Roformer | Aspiration by Sucial': 'aspiration_mel_band_roformer_sdr_18.9845.ckpt',
-        'MelBand Roformer | Aspiration Less Aggressive by Sucial': 'aspiration_mel_band_roformer_less_aggr_sdr_18.1201.ckpt',
-    }
-}
+# ROFORMER_MODELS - now using EXTENDED_MODELS from models_config
+ROFORMER_MODELS = get_all_models()
 
 OUTPUT_FORMATS = ['wav', 'flac', 'mp3', 'ogg', 'opus', 'm4a', 'aiff', 'ac3']
 
@@ -509,12 +427,13 @@ def download_from_google_drive(url):
             except Exception as e:
                 logger.warning(f"Failed to delete temporary file {temp_output_path}: {str(e)}")
 
-@spaces.GPU(duration=60)
+@spaces.GPU(duration=300)
 def roformer_separator(audio, model_key, seg_size, override_seg_size, overlap, pitch_shift, model_dir, output_dir, out_format, norm_thresh, amp_thresh, batch_size, exclude_stems="", progress=gr.Progress(track_tqdm=True)):
     if not audio:
         raise ValueError("No audio or video file provided.")
     temp_audio_path = None
     extracted_audio_path = None
+    segment_temp_dir = None
     try:
         file_extension = os.path.splitext(audio)[1].lower().lstrip('.')
         supported_formats = ['wav', 'mp3', 'flac', 'ogg', 'opus', 'm4a', 'aiff', 'ac3', 'mp4', 'mov', 'avi', 'mkv', 'flv', 'wmv', 'webm', 'mpeg', 'mpg', 'ts', 'vob']
@@ -554,29 +473,88 @@ def roformer_separator(audio, model_key, seg_size, override_seg_size, overlap, p
         if os.path.exists(output_dir):
             shutil.rmtree(output_dir)
         os.makedirs(output_dir, exist_ok=True)
-        base_name = os.path.splitext(os.path.basename(audio))[0].replace(' ', '_')  # Boşlukları alt çizgi ile değiştir
-        for category, models in ROFORMER_MODELS.items():
-            if model_key in models:
-                model = models[model_key]
-                break
-        else:
+        base_name = os.path.splitext(os.path.basename(audio))[0].replace(' ', '_')
+
+        # Find model from EXTENDED_MODELS + custom models
+        model = find_model_filename(model_key)
+        if not model:
             raise ValueError(f"Model '{model_key}' not found.")
+
+        # Pre-download model files (checkpoint + config YAML) before loading
+        # This is required for the separator.py bypass to work
+        dl_success, dl_msg = ensure_model_files_downloaded(model, model_dir)
+        if not dl_success:
+            logger.warning(f"Pre-download warning for {model}: {dl_msg}")
+
         logger.info(f"Separating {base_name} with {model_key} on {device}")
-        separator = Separator(
-            log_level=logging.INFO,
-            model_file_dir=model_dir,
-            output_dir=output_dir,
-            output_format=out_format,
-            normalization_threshold=norm_thresh,
-            amplification_threshold=amp_thresh,
-            use_autocast=use_autocast,
-            mdxc_params={"segment_size": seg_size, "override_model_segment_size": override_seg_size, "batch_size": batch_size, "overlap": overlap, "pitch_shift": pitch_shift}
-        )
-        progress(0.2, desc="Loading model...")
-        separator.load_model(model_filename=model)
-        progress(0.7, desc="Separating audio...")
-        separation = separator.separate(audio_to_process)
-        stems = [os.path.join(output_dir, file_name) for file_name in separation]
+
+        # ── Large file segmentation ──
+        audio_duration = get_audio_duration(audio_to_process)
+        was_segmented = False
+        if audio_duration > MAX_UNSPLIT_DURATION:
+            duration_min = audio_duration / 60
+            logger.info(f"⚠️ Large audio detected: {duration_min:.0f} min. Splitting to prevent OOM...")
+            progress(0.05, desc=f"Splitting {duration_min:.0f} min audio into segments...")
+            segment_temp_dir = os.path.join("/tmp", f"sesa_segments_{base_name}")
+            os.makedirs(segment_temp_dir, exist_ok=True)
+            segments = split_audio_segments(audio_to_process, segment_temp_dir, SEGMENT_DURATION)
+            if segments:
+                was_segmented = True
+                logger.info(f"Split into {len(segments)} segments")
+                # Process each segment
+                seg_output_dir = os.path.join("/tmp", f"sesa_seg_output_{base_name}")
+                os.makedirs(seg_output_dir, exist_ok=True)
+                for i, seg_path in enumerate(segments):
+                    progress(0.1 + 0.7 * (i / len(segments)), desc=f"Processing segment {i+1}/{len(segments)}...")
+                    separator = Separator(
+                        log_level=logging.INFO,
+                        model_file_dir=model_dir,
+                        output_dir=seg_output_dir,
+                        output_format=out_format,
+                        normalization_threshold=norm_thresh,
+                        amplification_threshold=amp_thresh,
+                        use_autocast=use_autocast,
+                        mdxc_params={"segment_size": seg_size, "override_model_segment_size": override_seg_size, "batch_size": batch_size, "overlap": overlap, "pitch_shift": pitch_shift}
+                    )
+                    separator.load_model(model_filename=model)
+                    separator.separate(seg_path)
+                    # Free GPU memory between segments
+                    del separator
+                    if torch.cuda.is_available():
+                        torch.cuda.empty_cache()
+                    gc.collect()
+                # Concatenate segment outputs
+                progress(0.85, desc="Concatenating segments...")
+                concatenate_segment_outputs(seg_output_dir, out_format)
+                # Move final concatenated files to output_dir
+                for f in os.listdir(seg_output_dir):
+                    if '_seg' not in f.lower():  # Only move final merged files
+                        shutil.move(os.path.join(seg_output_dir, f), os.path.join(output_dir, f))
+                # Cleanup temp dirs
+                shutil.rmtree(segment_temp_dir, ignore_errors=True)
+                shutil.rmtree(seg_output_dir, ignore_errors=True)
+                segment_temp_dir = None
+
+        if not was_segmented:
+            # Normal processing (no segmentation)
+            separator = Separator(
+                log_level=logging.INFO,
+                model_file_dir=model_dir,
+                output_dir=output_dir,
+                output_format=out_format,
+                normalization_threshold=norm_thresh,
+                amplification_threshold=amp_thresh,
+                use_autocast=use_autocast,
+                mdxc_params={"segment_size": seg_size, "override_model_segment_size": override_seg_size, "batch_size": batch_size, "overlap": overlap, "pitch_shift": pitch_shift}
+            )
+            progress(0.2, desc="Loading model...")
+            separator.load_model(model_filename=model)
+            progress(0.7, desc="Separating audio...")
+            separator.separate(audio_to_process)
+
+        # Collect all output stems
+        output_files = os.listdir(output_dir)
+        stems = [os.path.join(output_dir, f) for f in output_files if os.path.isfile(os.path.join(output_dir, f))]
         file_list = []
         if exclude_stems.strip():
             excluded = [s.strip().lower() for s in exclude_stems.split(',')]
@@ -586,7 +564,7 @@ def roformer_separator(audio, model_key, seg_size, override_seg_size, overlap, p
             stem2 = filtered_stems[1] if len(filtered_stems) > 1 else None
         else:
             file_list = stems
-            stem1 = stems[0]
+            stem1 = stems[0] if stems else None
             stem2 = stems[1] if len(stems) > 1 else None
 
         return stem1, stem2, file_list
@@ -611,7 +589,7 @@ def roformer_separator(audio, model_key, seg_size, override_seg_size, overlap, p
             torch.cuda.empty_cache()
             logger.info("GPU memory cleared")
 
-@spaces.GPU(duration=60)
+@spaces.GPU(duration=300)
 def auto_ensemble_process(audio, model_keys, state, seg_size=64, overlap=0.1, out_format="wav", use_tta="False", model_dir="/tmp/audio-separator-models/", output_dir="output", norm_thresh=0.9, amp_thresh=0.9, batch_size=1, ensemble_method="avg_wave", exclude_stems="", weights_str="", progress=gr.Progress(track_tqdm=True)):
     temp_audio_path = None
     extracted_audio_path = None
@@ -801,6 +779,10 @@ def auto_ensemble_process(audio, model_keys, state, seg_size=64, overlap=0.1, ou
 
                     if model_key not in model_cache:
                         logger.info(f"Loading {model_key} into cache")
+                        # Pre-download model files for bypass
+                        dl_ok, dl_msg = ensure_model_files_downloaded(model, model_dir)
+                        if not dl_ok:
+                            logger.warning(f"Pre-download warning: {dl_msg}")
                         separator = Separator(
                             log_level=logging.INFO,
                             model_file_dir=model_dir,
@@ -895,12 +877,14 @@ def auto_ensemble_process(audio, model_keys, state, seg_size=64, overlap=0.1, ou
             logger.info("GPU memory cleared")
             
 def update_roformer_models(category):
-    choices = list(ROFORMER_MODELS.get(category, {}).keys()) or []
+    all_models = get_all_models()
+    choices = list(all_models.get(category, {}).keys()) or []
     logger.debug(f"Updating roformer models for category {category}: {choices}")
     return gr.update(choices=choices, value=choices[0] if choices else None)
 
 def update_ensemble_models(category):
-    choices = list(ROFORMER_MODELS.get(category, {}).keys()) or []
+    all_models = get_all_models()
+    choices = list(all_models.get(category, {}).keys()) or []
     logger.debug(f"Updating ensemble models for category {category}: {choices}")
     return gr.update(choices=choices, value=[])
 
@@ -908,6 +892,59 @@ def download_audio_wrapper(url, cookie_file):
     file_path, status, audio_data = download_audio(url, cookie_file)
     return file_path, status  # Return file_path instead of audio_data
 
+# ─── Batch Processing ────────────────────────────────────────────────────────
+@spaces.GPU(duration=300)
+def batch_separator(audio_files, model_key, seg_size, override_seg_size, overlap, pitch_shift, model_dir, output_dir, out_format, norm_thresh, amp_thresh, batch_size, exclude_stems="", progress=gr.Progress(track_tqdm=True)):
+    """Process up to 10 audio files sequentially."""
+    if not audio_files:
+        raise ValueError("No audio files provided.")
+    if len(audio_files) > 10:
+        raise ValueError("Maximum 10 files per batch.")
+
+    all_output_files = []
+    status_lines = []
+    for i, audio in enumerate(audio_files):
+        # Handle gr.File objects
+        audio_path = audio.name if hasattr(audio, 'name') else audio
+        base = os.path.splitext(os.path.basename(audio_path))[0]
+        progress((i) / len(audio_files), desc=f"Processing file {i+1}/{len(audio_files)}: {base}")
+        try:
+            stem1, stem2, files = roformer_separator(
+                audio_path, model_key, seg_size, override_seg_size, overlap, pitch_shift,
+                model_dir, output_dir, out_format, norm_thresh, amp_thresh, batch_size,
+                exclude_stems, progress
+            )
+            all_output_files.extend(files)
+            status_lines.append(f"✅ {base}: {len(files)} stems")
+        except Exception as e:
+            status_lines.append(f"❌ {base}: {str(e)[:100]}")
+            logger.error(f"Batch processing error for {base}: {e}")
+
+    status_text = "\n".join(status_lines)
+    return status_text, all_output_files
+
+# ─── Custom Model Management UI handlers ─────────────────────────────────────
+def add_custom_model_handler(name, checkpoint_url, config_url, custom_py_url):
+    success, msg = add_custom_model(name, checkpoint_url, config_url, custom_py_url)
+    # Refresh ROFORMER_MODELS
+    global ROFORMER_MODELS
+    ROFORMER_MODELS = get_all_models()
+    # Get updated custom model list
+    custom_list_data = get_custom_models_list()
+    custom_list = "\n".join([f"• {n}: {u}" for n, u in custom_list_data]) if custom_list_data else "No custom models"
+    # Return updated categories
+    cats = get_categories()
+    return msg, custom_list, gr.update(choices=cats), gr.update(choices=cats)
+
+def delete_custom_model_handler(name):
+    success, msg = delete_custom_model(name)
+    global ROFORMER_MODELS
+    ROFORMER_MODELS = get_all_models()
+    custom_list_data = get_custom_models_list()
+    custom_list = "\n".join([f"• {n}: {u}" for n, u in custom_list_data]) if custom_list_data else "No custom models"
+    cats = get_categories()
+    return msg, custom_list, gr.update(choices=cats), gr.update(choices=cats)
+
 def create_interface():
     with gr.Blocks(title="🎵 SESA Fast Separation 🎵", css=CSS, elem_id="app-container") as app:
         gr.Markdown("<h1 class='header-text'>🎵 SESA Fast Separation 🎵</h1>")
@@ -940,8 +977,8 @@ def create_interface():
                     roformer_download_status = gr.Textbox(label="📢 Download Status", interactive=False)
                     roformer_exclude_stems = gr.Textbox(label="🚫 Exclude Stems", placeholder="e.g., vocals, drums (comma-separated)", interactive=True)
                     with gr.Row():
-                        roformer_category = gr.Dropdown(label="📚 Category", choices=list(ROFORMER_MODELS.keys()), value="General Purpose", interactive=True)
-                        roformer_model = gr.Dropdown(label="🛠️ Model", choices=list(ROFORMER_MODELS["General Purpose"].keys()), interactive=True, allow_custom_value=True)
+                        roformer_category = gr.Dropdown(label="📚 Category", choices=get_categories(), value="Vocals", interactive=True)
+                        roformer_model = gr.Dropdown(label="🛠️ Model", choices=get_model_choices("Vocals"), interactive=True, allow_custom_value=True)
                     with gr.Row():
                         roformer_seg_size = gr.Slider(32, 512, value=64, step=32, label="📏 Segment Size", interactive=True)
                         roformer_overlap = gr.Slider(2, 10, value=8, step=1, label="🔄 Overlap", interactive=True)
@@ -965,8 +1002,8 @@ def create_interface():
                     ensemble_download_status = gr.Textbox(label="📢 Download Status", interactive=False)
                     ensemble_exclude_stems = gr.Textbox(label="🚫 Exclude Stems", placeholder="e.g., vocals, drums (comma-separated)", interactive=True)
                     with gr.Row():
-                        ensemble_category = gr.Dropdown(label="📚 Category", choices=list(ROFORMER_MODELS.keys()), value="Instrumentals", interactive=True)
-                        ensemble_models = gr.Dropdown(label="🛠️ Models (Max 6)", choices=list(ROFORMER_MODELS["Instrumentals"].keys()), multiselect=True, interactive=True, allow_custom_value=True)
+                        ensemble_category = gr.Dropdown(label="📚 Category", choices=get_categories(), value="Instrumentals", interactive=True)
+                        ensemble_models = gr.Dropdown(label="🛠️ Models (Max 6)", choices=get_model_choices("Instrumentals"), multiselect=True, interactive=True, allow_custom_value=True)
                     with gr.Row():
                         ensemble_seg_size = gr.Slider(32, 512, value=64, step=32, label="📏 Segment Size", interactive=True)
                         ensemble_overlap = gr.Slider(2, 10, value=8, step=1, label="🔄 Overlap", interactive=True)
@@ -977,6 +1014,42 @@ def create_interface():
                     ensemble_output = gr.Audio(label="🎶 Ensemble Result", type="filepath", interactive=False)
                     ensemble_status = gr.HTML(label="📢 Status")
                     ensemble_files = gr.File(label="📥 Download Ensemble and Stems", interactive=False)
+            with gr.Tab("📦 Batch Processing"):
+                with gr.Group(elem_classes="dubbing-theme"):
+                    gr.Markdown("### Batch Processing (Max 10 Files)")
+                    gr.Markdown("Upload multiple audio files and process them all with the same model.")
+                    batch_audio = gr.File(label="🎧 Upload Audio Files", file_count="multiple", file_types=['.wav', '.mp3', '.flac', '.ogg', '.opus', '.m4a', '.aiff', '.ac3', '.mp4', '.mov', '.avi', '.mkv'], interactive=True)
+                    with gr.Row():
+                        batch_category = gr.Dropdown(label="📚 Category", choices=get_categories(), value="Vocals", interactive=True)
+                        batch_model = gr.Dropdown(label="🛠️ Model", choices=get_model_choices("Vocals"), interactive=True, allow_custom_value=True)
+                    with gr.Row():
+                        batch_seg_size = gr.Slider(32, 512, value=64, step=32, label="📏 Segment Size", interactive=True)
+                        batch_overlap = gr.Slider(2, 10, value=8, step=1, label="🔄 Overlap", interactive=True)
+                    batch_pitch_shift = gr.Slider(-12, 12, value=0, step=1, label="🎵 Pitch Shift", interactive=True)
+                    batch_override_seg = gr.Dropdown(choices=["True", "False"], value="False", label="🔧 Override Segment Size", interactive=True)
+                    batch_exclude = gr.Textbox(label="🚫 Exclude Stems", placeholder="e.g., vocals, drums (comma-separated)", interactive=True)
+                    batch_button = gr.Button("🚀 Process Batch!", variant="primary")
+                    batch_status = gr.Textbox(label="📢 Batch Status", interactive=False, lines=5)
+                    batch_files = gr.File(label="📥 Download All Stems", interactive=False)
+            with gr.Tab("🔧 Custom Models"):
+                with gr.Group(elem_classes="dubbing-theme"):
+                    gr.Markdown("### Custom Model Management")
+                    gr.Markdown("Add custom models from HuggingFace or other sources by providing download URLs. The model will be automatically downloaded when used.")
+                    with gr.Row():
+                        custom_model_name = gr.Textbox(label="📝 Model Display Name", placeholder="e.g., My Custom Vocal Model", interactive=True)
+                    with gr.Row():
+                        custom_checkpoint_url = gr.Textbox(label="📦 Checkpoint URL (required)", placeholder="https://huggingface.co/.../resolve/main/model.ckpt", interactive=True)
+                    with gr.Row():
+                        custom_config_url = gr.Textbox(label="📄 Config URL (optional)", placeholder="https://huggingface.co/.../resolve/main/config.yaml", interactive=True)
+                    with gr.Row():
+                        custom_py_url = gr.Textbox(label="🐍 Custom .py URL (optional)", placeholder="https://huggingface.co/.../resolve/main/bs_roformer.py", interactive=True)
+                    with gr.Row():
+                        add_model_btn = gr.Button("➕ Add Model", variant="primary")
+                        del_model_name = gr.Textbox(label="🗑️ Model Name to Delete", placeholder="Exact model name", interactive=True)
+                        del_model_btn = gr.Button("🗑️ Delete Model", variant="stop")
+                    custom_model_status = gr.Textbox(label="📢 Status", interactive=False)
+                    custom_model_list = gr.Textbox(label="📋 Custom Models", interactive=False, lines=8,
+                        value="\n".join([f"• {n}: {u}" for n, u in get_custom_models_list()]) or "No custom models")
         gr.HTML("<div class='footer'>Powered by Audio-Separator 🌟🎶 | Made with ❤️</div>")
         roformer_category.change(update_roformer_models, inputs=[roformer_category], outputs=[roformer_model])
         download_roformer.click(
@@ -1009,6 +1082,28 @@ def create_interface():
             ],
             outputs=[ensemble_output, ensemble_status, ensemble_files, ensemble_state]
         )
+        # Batch processing events
+        batch_category.change(update_roformer_models, inputs=[batch_category], outputs=[batch_model])
+        batch_button.click(
+            fn=batch_separator,
+            inputs=[
+                batch_audio, batch_model, batch_seg_size, batch_override_seg, batch_overlap,
+                batch_pitch_shift, model_file_dir, output_dir, output_format,
+                norm_threshold, amp_threshold, batch_size, batch_exclude
+            ],
+            outputs=[batch_status, batch_files]
+        )
+        # Custom model events
+        add_model_btn.click(
+            fn=add_custom_model_handler,
+            inputs=[custom_model_name, custom_checkpoint_url, custom_config_url, custom_py_url],
+            outputs=[custom_model_status, custom_model_list, roformer_category, ensemble_category]
+        )
+        del_model_btn.click(
+            fn=delete_custom_model_handler,
+            inputs=[del_model_name],
+            outputs=[custom_model_status, custom_model_list, roformer_category, ensemble_category]
+        )
     return app
 
 if __name__ == "__main__":