Upload app.py

app.py

@@ -15,12 +15,11 @@ from models import get_model_config, cleanup_all_models
 from config import DEFAULT_ALPHA
 from utils import clear_gpu_memory
 
-@spaces.GPU(duration=300)
-def process_audio_files(zip_file, model_name, layer, alpha):
-    """Process uploaded ZIP file containing audio mixtures."""
+def process_audio_files_cpu(zip_file, model_name, layer, alpha):
+    """Process uploaded ZIP file containing audio mixtures - CPU part."""
     
     if zip_file is None:
-        return None, "Please upload a ZIP file"
+        return None, "Please upload a ZIP file", None
     
     try:
         # Use a fixed extraction path
@@ -56,18 +55,18 @@ def process_audio_files(zip_file, model_name, layer, alpha):
                                 outs_dir = subitem
         
         if refs_dir is None or outs_dir is None:
-            return None, "Could not find 'references' and 'outputs' directories in the ZIP file"
+            return None, "Could not find 'references' and 'outputs' directories in the ZIP file", None
         
         # Get audio files
         ref_files = sorted([f for f in refs_dir.glob("*.wav")])
         out_files = sorted([f for f in outs_dir.glob("*.wav")])
         
         if len(ref_files) == 0:
-            return None, "No reference WAV files found"
+            return None, "No reference WAV files found", None
         if len(out_files) == 0:
-            return None, "No output WAV files found"
+            return None, "No output WAV files found", None
         if len(ref_files) != len(out_files):
-            return None, f"Number of reference files ({len(ref_files)}) must match number of output files ({len(out_files)}). Files must be in the same order."
+            return None, f"Number of reference files ({len(ref_files)}) must match number of output files ({len(out_files)}). Files must be in the same order.", None
         
         # Create manifest
         manifest = [{
@@ -81,7 +80,7 @@ def process_audio_files(zip_file, model_name, layer, alpha):
         # Validate model
         allowed_models = set(get_model_config(0).keys())
         if model_name not in allowed_models:
-            return None, f"Invalid model. Allowed: {', '.join(sorted(allowed_models))}"
+            return None, f"Invalid model. Allowed: {', '.join(sorted(allowed_models))}", None
         
         # Set layer
         if model_name == "raw":
@@ -94,17 +93,33 @@ def process_audio_files(zip_file, model_name, layer, alpha):
             }
             layer_final = layer if layer is not None else model_defaults.get(model_name, 12)
         
-        # Check GPU availability - use all available GPUs on the space
-        max_gpus = torch.cuda.device_count() if torch.cuda.is_available() else 0
+        # Return preprocessed data for GPU processing
+        return manifest, layer_final, alpha
+        
+    except Exception as e:
+        error_msg = f"Error in preprocessing: {str(e)}\n{traceback.format_exc()}"
+        return None, error_msg, None
+
+@spaces.GPU(duration=300)
+def process_audio_files_gpu(manifest, model_name, layer_final, alpha):
+    """GPU processing part - only called when GPU is allocated."""
+    
+    if manifest is None:
+        return None, "Invalid input data"
+    
+    try:
+        # Force single GPU mode in Spaces environment
+        # The spaces decorator handles GPU allocation
+        max_gpus = 1 if torch.cuda.is_available() else 0
         
-        # Run experiment
+        # Run experiment with forced single GPU
         results_dir = compute_mapss_measures(
             models=[model_name],
             mixtures=manifest,
             layer=layer_final,
             alpha=alpha,
             verbose=True,
-            max_gpus=max_gpus,
+            max_gpus=max_gpus,  # Force single GPU
             add_ci=False  # Disable CI for faster processing in demo
         )
         
@@ -128,13 +143,30 @@ def process_audio_files(zip_file, model_name, layer, alpha):
             return None, f"Processing completed but no output files were generated. Check if embeddings were computed."
             
     except Exception as e:
-        error_msg = f"Error: {str(e)}\n{traceback.format_exc()}"
+        error_msg = f"Error in GPU processing: {str(e)}\n{traceback.format_exc()}"
         return None, error_msg
     
     finally:
+        # Cleanup is critical in Spaces environment
         cleanup_all_models()
         clear_gpu_memory()
         gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+def process_audio_files(zip_file, model_name, layer, alpha):
+    """Main processing function that combines CPU and GPU parts."""
+    
+    # First, do CPU preprocessing
+    manifest, layer_final, alpha_processed = process_audio_files_cpu(
+        zip_file, model_name, layer, alpha
+    )
+    
+    if manifest is None:
+        return None, layer_final  # layer_final contains error message in this case
+    
+    # Then do GPU processing
+    return process_audio_files_gpu(manifest, model_name, layer_final, alpha_processed)
 
 def create_interface():
     with gr.Blocks(title="MAPSS - Multi-source Audio Perceptual Separation Scores") as demo:
@@ -214,6 +246,7 @@ def create_interface():
         - PS/PM scores are only computed between active speakers
         - Processing time scales with number of sources and audio length
         - GPU acceleration is automatically used when available
+        - **Note**: This Hugging Face Space runs with a single GPU allocation
         
         ## Citation