Update app.py

app.py

@@ -1,10 +1,9 @@
-from typing import Optional, Any  # Optional ve Any için gerekli içe aktarma
-
-# Mevcut diğer içe aktarmalar (örneğin, önceki kodunuzdan)
+from typing import Optional, Any
 import os
 import sys
 import torch
 import logging
+import yt_dlp
 from yt_dlp import YoutubeDL
 import gradio as gr
 import argparse
@@ -20,8 +19,8 @@ import spaces
 import gdown
 import scipy.io.wavfile
 from pydub import AudioSegment
-import yt_dlp
 import gc
+import time
 
 # Logging setup (mevcut)
 logging.basicConfig(level=logging.INFO)
@@ -433,12 +432,15 @@ def roformer_separator(audio, model_key, seg_size, override_seg_size, overlap, p
 def auto_ensemble_process(audio, model_keys, seg_size=128, 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
     chunk_paths = []
-    max_retries = 2  # Retry attempts for ZeroGPU session issues
+    max_retries = 2
     try:
         if not audio:
             raise ValueError("No audio file provided.")
         if not model_keys:
             raise ValueError("No models selected.")
+        if len(model_keys) > 2:
+            logger.warning("Limited to 2 models to avoid ZeroGPU timeouts. Using first two: %s", model_keys[:2])
+            model_keys = model_keys[:2]
         if isinstance(audio, tuple):
             sample_rate, data = audio
             temp_audio_path = os.path.join("/tmp", "temp_audio.wav")
@@ -465,8 +467,7 @@ def auto_ensemble_process(audio, model_keys, seg_size=128, overlap=0.1, out_form
             chunks = [audio]
         use_tta = use_tta == "True"
         if os.path.exists(output_dir):
-            shutil.rmtree(outputwatermark = True)
-            shutil.copyfile(audio, os.path.join(output_dir, os.path.basename(audio)))
+            shutil.rmtree(output_dir)
         os.makedirs(output_dir, exist_ok=True)
         base_name = os.path.splitext(os.path.basename(audio))[0]
         logger.info(f"Ensemble for {base_name} with {model_keys} on {device}")
@@ -487,6 +488,10 @@ def auto_ensemble_process(audio, model_keys, seg_size=128, overlap=0.1, out_form
                 while retry_count <= max_retries:
                     try:
                         progress((model_idx + 0.1) / total_models, desc=f"Loading {model_key} for chunk {chunk_idx}")
+                        # Check if model is cached
+                        model_path = os.path.join(model_dir, model)
+                        if not os.path.exists(model_path):
+                            logger.info(f"Model {model} not cached, will download")
                         separator = Separator(
                             log_level=logging.INFO,
                             model_file_dir=model_dir,
@@ -508,20 +513,24 @@ def auto_ensemble_process(audio, model_keys, seg_size=128, overlap=0.1, out_form
                                 model_stems[model_key]["vocals"].append(stem)
                             elif "other" in os.path.basename(stem).lower() or "instrumental" in os.path.basename(stem).lower():
                                 model_stems[model_key]["other"].append(stem)
-                        break  # Success, exit retry loop
+                        break
                     except Exception as e:
                         retry_count += 1
                         logger.error(f"Error processing {model_key} chunk {chunk_idx}, attempt {retry_count}/{max_retries}: {e}")
+                        if "ZeroGPU" in str(e) or "aborted" in str(e).lower():
+                            logger.error("ZeroGPU task aborted, attempting recovery")
                         if retry_count > max_retries:
                             logger.error(f"Max retries reached for {model_key} chunk {chunk_idx}, skipping")
                             break
-                        time.sleep(2)  # Wait before retrying
+                        time.sleep(1)  # Reduced delay to minimize overhead
                     finally:
                         separator = None
                         gc.collect()
                         if torch.cuda.is_available():
                             torch.cuda.empty_cache()
                             logger.info(f"Cleared CUDA cache after {model_key} chunk {chunk_idx}")
+                # Yield control to ZeroGPU scheduler
+                time.sleep(0.1)
         progress(0.8, desc="Combining stems...")
         for model_key, stems_dict in model_stems.items():
             for stem_type in ["vocals", "other"]:
@@ -541,7 +550,7 @@ def auto_ensemble_process(audio, model_keys, seg_size=128, overlap=0.1, out_form
                         logger.error(f"Error combining {stem_type} for {model_key}: {e}")
         all_stems = [stem for stem in all_stems if os.path.exists(stem)]
         if not all_stems:
-            raise ValueError("No valid stems found for ensemble.")
+            raise ValueError("No valid stems found for ensemble. Try uploading a local WAV file.")
         weights = [float(w.strip()) for w in weights_str.split(',')] if weights_str.strip() else [1.0] * len(all_stems)
         if len(weights) != len(all_stems):
             weights = [1.0] * len(all_stems)
@@ -565,7 +574,11 @@ def auto_ensemble_process(audio, model_keys, seg_size=128, overlap=0.1, out_form
             raise RuntimeError(f"Ensemble processing error: {e}")
     except Exception as e:
         logger.error(f"Ensemble error: {e}")
-        raise RuntimeError(f"Ensemble error: {e}")
+        if "ZeroGPU" in str(e) or "aborted" in str(e).lower():
+            error_msg = "ZeroGPU task aborted. Try using fewer models (max 2), lowering segment size, or uploading a local WAV file."
+        else:
+            error_msg = f"Ensemble error: {e}"
+        raise RuntimeError(error_msg)
     finally:
         for path in chunk_paths + ([temp_audio_path] if temp_audio_path and os.path.exists(temp_audio_path) else []):
             try:
@@ -599,6 +612,7 @@ def create_interface():
         gr.Markdown("<h1 class='header-text'>🎵 SESA Fast Separation 🎵</h1>")
         gr.Markdown("**Note**: If YouTube downloads fail, upload an audio file directly or use a valid cookies file. [Cookie Instructions](https://github.com/yt-dlp/yt-dlp/wiki/Extractors#exporting-youtube-cookies)")
         gr.Markdown("**Warning**: Audio files longer than 15 minutes are automatically split into chunks, which may require more time and resources.")
+        gr.Markdown("**ZeroGPU Notice**: Use up to 2 models for ensemble to avoid timeouts. For large tasks, upload a local WAV file.")
         with gr.Tabs():
             with gr.Tab("⚙️ Settings"):
                 with gr.Group(elem_classes="dubbing-theme"):
@@ -608,7 +622,7 @@ def create_interface():
                     output_format = gr.Dropdown(value="wav", choices=OUTPUT_FORMATS, label="🎶 Output Format", interactive=True)
                     norm_threshold = gr.Slider(0.1, 1.0, value=0.9, step=0.1, label="🔊 Normalization Threshold", interactive=True)
                     amp_threshold = gr.Slider(0.1, 1.0, value=0.3, step=0.1, label="📈 Amplification Threshold", interactive=True)
-                    batch_size = gr.Slider(1, 16, value=1, step=1, label="⚡ Batch Size", interactive=True)
+                    batch_size = gr.Slider(1, 8, value=1, step=1, label="⚡ Batch Size", interactive=True)
             with gr.Tab("🎤 Roformer"):
                 with gr.Group(elem_classes="dubbing-theme"):
                     gr.Markdown("### Audio Separation")
@@ -623,7 +637,7 @@ def create_interface():
                         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)
                     with gr.Row():
-                        roformer_seg_size = gr.Slider(32, 4000, value=256, step=32, label="📏 Segment Size", interactive=True)
+                        roformer_seg_size = gr.Slider(32, 512, value=128, step=32, label="📏 Segment Size", interactive=True)
                         roformer_overlap = gr.Slider(2, 10, value=8, step=1, label="🔄 Overlap", interactive=True)
                     with gr.Row():
                         roformer_pitch_shift = gr.Slider(-12, 12, value=0, step=1, label="🎵 Pitch Shift", interactive=True)
@@ -635,7 +649,7 @@ def create_interface():
             with gr.Tab("🎚️ Auto Ensemble"):
                 with gr.Group(elem_classes="dubbing-theme"):
                     gr.Markdown("### Ensemble Processing")
-                    gr.Markdown("Note: If weights are not specified, equal weights (1.0) are applied to all models.")
+                    gr.Markdown("Note: If weights are not specified, equal weights (1.0) are applied to all models. Max 2 models recommended.")
                     with gr.Row():
                         ensemble_audio = gr.Audio(label="🎧 Upload Audio", type="filepath", interactive=True)
                         url_ensemble = gr.Textbox(label="🔗 Or Paste URL", placeholder="YouTube or audio URL", interactive=True)
@@ -645,9 +659,9 @@ def create_interface():
                     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", choices=list(ROFORMER_MODELS["Instrumentals"].keys()), multiselect=True, interactive=True, allow_custom_value=True)
+                        ensemble_models = gr.Dropdown(label="🛠️ Models (Max 2)", choices=list(ROFORMER_MODELS["Instrumentals"].keys()), multiselect=True, interactive=True, allow_custom_value=True)
                     with gr.Row():
-                        ensemble_seg_size = gr.Slider(32, 4000, value=256, step=32, label="📏 Segment Size", interactive=True)
+                        ensemble_seg_size = gr.Slider(32, 512, value=128, step=32, label="📏 Segment Size", interactive=True)
                         ensemble_overlap = gr.Slider(2, 10, value=8, step=1, label="🔄 Overlap", interactive=True)
                         ensemble_use_tta = gr.Dropdown(choices=["True", "False"], value="False", label="🔍 Use TTA", interactive=True)
                     ensemble_method = gr.Dropdown(label="⚙️ Ensemble Method", choices=['avg_wave', 'median_wave', 'max_wave', 'min_wave', 'avg_fft', 'median_fft', 'max_fft', 'min_fft'], value='avg_wave', interactive=True)