Upload app.py

app.py

@@ -1,15 +1,13 @@
 import gradio as gr
 import zipfile
-import tempfile
 import shutil
 from pathlib import Path
-import pandas as pd
 import json
 import os
 import traceback
 import gc
 import torch
-import spaces  # Import spaces for ZeroGPU
+import spaces
 
 # Import your modules
 from engine import compute_mapss_measures
@@ -17,146 +15,125 @@ from models import get_model_config, cleanup_all_models
 from config import DEFAULT_ALPHA
 from utils import clear_gpu_memory
 
-@spaces.GPU(duration=300)  # Request GPU for up to 5 minutes
+@spaces.GPU(duration=300)
 def process_audio_files(zip_file, model_name, layer, alpha):
-    """
-    Process uploaded ZIP file containing audio mixtures.
-    
-    Expected ZIP structure:
-    - references/: Contains N reference audio files
-    - outputs/: Contains N output audio files
-    """
+    """Process uploaded ZIP file containing audio mixtures."""
     
     if zip_file is None:
         return None, "Please upload a ZIP file"
     
-    # Create temporary directory for processing
-    with tempfile.TemporaryDirectory() as temp_dir:
-        temp_path = Path(temp_dir)
-        
-        try:
-            # Extract ZIP file
-            extract_path = temp_path / "extracted"
-            extract_path.mkdir(exist_ok=True)
-            
-            with zipfile.ZipFile(zip_file.name, 'r') as zip_ref:
-                zip_ref.extractall(extract_path)
-            
-            # Find references and outputs directories
-            refs_dir = None
-            outs_dir = None
-            
-            # Check for standard structure
+    try:
+        # Use a fixed extraction path
+        extract_path = Path("/tmp/mapss_extract")
+        if extract_path.exists():
+            shutil.rmtree(extract_path)
+        extract_path.mkdir(parents=True)
+        
+        # Extract ZIP
+        with zipfile.ZipFile(zip_file.name, 'r') as zip_ref:
+            zip_ref.extractall(extract_path)
+        
+        # Find references and outputs directories
+        refs_dir = None
+        outs_dir = None
+        
+        for item in extract_path.iterdir():
+            if item.is_dir():
+                if item.name.lower() in ['references', 'refs', 'reference']:
+                    refs_dir = item
+                elif item.name.lower() in ['outputs', 'outs', 'output', 'separated']:
+                    outs_dir = item
+        
+        # Check one level deeper if not found
+        if refs_dir is None or outs_dir is None:
             for item in extract_path.iterdir():
                 if item.is_dir():
-                    if item.name.lower() in ['references', 'refs', 'reference']:
-                        refs_dir = item
-                    elif item.name.lower() in ['outputs', 'outs', 'output', 'separated']:
-                        outs_dir = item
-            
-            # If not found at root, check one level deeper
-            if refs_dir is None or outs_dir is None:
-                for item in extract_path.iterdir():
-                    if item.is_dir():
-                        for subitem in item.iterdir():
-                            if subitem.is_dir():
-                                if subitem.name.lower() in ['references', 'refs', 'reference']:
-                                    refs_dir = subitem
-                                elif subitem.name.lower() in ['outputs', 'outs', 'output', 'separated']:
-                                    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"
-            
-            # 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"
-            if len(out_files) == 0:
-                return None, "No output WAV files found"
-            
-            # Create manifest
-            manifest = [{
-                "mixture_id": "uploaded_mixture",
-                "references": [str(f) for f in ref_files],
-                "systems": {
-                    "uploaded_system": [str(f) for f in out_files]
-                }
-            }]
-            
-            # Validate model and layer
-            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))}"
-            
-            # Set default layer if needed
-            if model_name == "raw":
-                layer_final = 0
-            else:
-                model_defaults = {
-                    "wavlm": 24, "wav2vec2": 24, "hubert": 24,
-                    "wavlm_base": 12, "wav2vec2_base": 12, "hubert_base": 12,
-                    "wav2vec2_xlsr": 24, "ast": 12
-                }
-                layer_final = layer if layer is not None else model_defaults.get(model_name, 12)
-            
-            # Check if GPU is available
-            max_gpus = 1 if torch.cuda.is_available() else 0
-            
-            # Run experiment with compute_mapss_measures
-            results_dir = compute_mapss_measures(
-                models=[model_name],
-                mixtures=manifest,
-                layer=layer_final,
-                alpha=alpha,
-                verbose=True,
-                max_gpus=max_gpus,  # Use GPU if available, otherwise CPU
-                add_ci=False  # Disable CI for faster processing
-            )
-            
-            # Create output ZIP with results
-            output_zip = temp_path / "results.zip"
-            
-            with zipfile.ZipFile(output_zip, 'w') as zipf:
-                # Add all CSV files from results
-                results_path = Path(results_dir)
-                for csv_file in results_path.rglob("*.csv"):
-                    arcname = str(csv_file.relative_to(results_path.parent))
-                    zipf.write(csv_file, arcname)
-                
-                # Add params.json
-                params_file = results_path / "params.json"
-                if params_file.exists():
-                    zipf.write(params_file, str(params_file.relative_to(results_path.parent)))
-                
-                # Add manifest
-                manifest_file = results_path / "manifest_canonical.json"
-                if manifest_file.exists():
-                    zipf.write(manifest_file, str(manifest_file.relative_to(results_path.parent)))
-            
-            # Read the ZIP file to return
-            with open(output_zip, 'rb') as f:
-                output_data = f.read()
-            
-            # Create a proper file object for Gradio
-            output_file_path = temp_path / "download_results.zip"
-            with open(output_file_path, 'wb') as f:
-                f.write(output_data)
+                    for subitem in item.iterdir():
+                        if subitem.is_dir():
+                            if subitem.name.lower() in ['references', 'refs', 'reference']:
+                                refs_dir = subitem
+                            elif subitem.name.lower() in ['outputs', 'outs', 'output', 'separated']:
+                                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"
+        
+        # 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"
+        if len(out_files) == 0:
+            return None, "No output WAV files found"
+        
+        # Create manifest
+        manifest = [{
+            "mixture_id": "uploaded_mixture",
+            "references": [str(f) for f in ref_files],
+            "systems": {
+                "uploaded_system": [str(f) for f in out_files]
+            }
+        }]
+        
+        # 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))}"
+        
+        # Set layer
+        if model_name == "raw":
+            layer_final = 0
+        else:
+            model_defaults = {
+                "wavlm": 24, "wav2vec2": 24, "hubert": 24,
+                "wavlm_base": 12, "wav2vec2_base": 12, "hubert_base": 12,
+                "wav2vec2_xlsr": 24, "ast": 12
+            }
+            layer_final = layer if layer is not None else model_defaults.get(model_name, 12)
+        
+        # Check GPU availability
+        max_gpus = 1 if torch.cuda.is_available() else 0
+        
+        # Run experiment
+        results_dir = compute_mapss_measures(
+            models=[model_name],
+            mixtures=manifest,
+            layer=layer_final,
+            alpha=alpha,
+            verbose=True,
+            max_gpus=max_gpus,
+            add_ci=False
+        )
+        
+        # Create output ZIP at a fixed location
+        output_zip = Path("/tmp/mapss_results.zip")
+        
+        with zipfile.ZipFile(output_zip, 'w') as zipf:
+            results_path = Path(results_dir)
+            files_added = 0
             
-            return str(output_file_path), "Processing completed successfully!"
+            # Add all files from results
+            for file_path in results_path.rglob("*"):
+                if file_path.is_file():
+                    arcname = str(file_path.relative_to(results_path.parent))
+                    zipf.write(file_path, arcname)
+                    files_added += 1
+        
+        if output_zip.exists() and files_added > 0:
+            return str(output_zip), f"Processing completed! Created ZIP with {files_added} files."
+        else:
+            return None, f"Processing completed but no output files were generated. Check if embeddings were computed."
             
-        except Exception as e:
-            error_msg = f"Error processing files: {str(e)}\n{traceback.format_exc()}"
-            return None, error_msg
-        finally:
-            # Ensure cleanup happens
-            cleanup_all_models()
-            clear_gpu_memory()
-            gc.collect()
+    except Exception as e:
+        error_msg = f"Error: {str(e)}\n{traceback.format_exc()}"
+        return None, error_msg
+    
+    finally:
+        cleanup_all_models()
+        clear_gpu_memory()
+        gc.collect()
 
-# Create Gradio interface
 def create_interface():
     with gr.Blocks(title="MAPSS - Multi-source Audio Perceptual Separation Scores") as demo:
         gr.Markdown("""
@@ -278,7 +255,6 @@ def create_interface():
                     label="Diffusion maps alpha parameter"
                 )
                 
-                # Define model layer configurations
                 def update_layer_slider(model_name):
                     """Update layer slider based on selected model"""
                     model_configs = {
@@ -303,7 +279,6 @@ def create_interface():
                         interactive=config["interactive"]
                     )
                 
-                # Connect the model dropdown to update the layer slider
                 model_dropdown.change(
                     fn=update_layer_slider,
                     inputs=[model_dropdown],
@@ -323,24 +298,14 @@ def create_interface():
                     max_lines=10
                 )
         
-        # Set up the processing
         process_btn.click(
             fn=process_audio_files,
             inputs=[file_input, model_dropdown, layer_slider, alpha_slider],
             outputs=[output_file, status_text]
         )
-        
-        # Add examples if you want
-        gr.Examples(
-            examples=[
-                # You can add example ZIP files here if you have them
-            ],
-            inputs=[file_input]
-        )
     
     return demo
 
-# Create and launch the app
 if __name__ == "__main__":
     demo = create_interface()
     demo.launch()