app.py

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 your modules
from engine import compute_mapss_measures
from models import get_model_config, cleanup_all_models
from config import DEFAULT_ALPHA
from utils import clear_gpu_memory

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
    """
    
    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
            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)
            
            # 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=1,  # Limit to 1 GPU for HF Space
                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)
            
            return str(output_file_path), "Processing completed successfully!"
            
        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()

# Create Gradio interface
def create_interface():
    with gr.Blocks(title="MAPSS - Multi-source Audio Perceptual Separation Scores") as demo:
        gr.Markdown("""
        # MAPSS: Multi-source Audio Perceptual Separation Scores
        
        This tool evaluates audio source separation quality using Perceptual Similarity (PS) and Perceptual Matching (PM) metrics.
        
        ## How to use:
        1. **Prepare your audio files**: Create a ZIP file with the following structure:
           ```
           your_mixture.zip
           ├── references/       # Original clean sources
           │   ├── speaker1.wav
           │   ├── speaker2.wav
           │   └── ...
           └── outputs/         # Separated outputs from your algorithm
               ├── separated1.wav
               ├── separated2.wav
               └── ...
           ```
        2. **Upload the ZIP file** using the file uploader below
        3. **Select model and parameters**
        4. **Click "Process"** to run the evaluation
        5. **Download the results** as a ZIP file containing CSV files with PS/PM scores
        
        ## Models available:
        - **raw**: Raw waveform features (no model)
        - **wavlm**: WavLM Large model (best overall performance)
        - **wav2vec2**: Wav2Vec2 Large model  
        - **hubert**: HuBERT Large model
        - **wavlm_base**: WavLM Base model (faster, good performance)
        - **wav2vec2_base**: Wav2Vec2 Base model
        - **hubert_base**: HuBERT Base model
        - **wav2vec2_xlsr**: Wav2Vec2 XLSR-53 model (multilingual)
        - **ast**: Audio Spectrogram Transformer
        """)
        
        with gr.Row():
            with gr.Column():
                file_input = gr.File(
                    label="Upload ZIP file with audio mixtures",
                    file_types=[".zip"],
                    type="filepath"
                )
                
                model_dropdown = gr.Dropdown(
                    choices=["raw", "wavlm", "wav2vec2", "hubert", 
                            "wavlm_base", "wav2vec2_base", "hubert_base",
                            "wav2vec2_xlsr", "ast"],
                    value="wav2vec2_base",
                    label="Select embedding model"
                )
                
                layer_slider = gr.Slider(
                    minimum=0,
                    maximum=24,
                    step=1,
                    value=12,
                    label="Layer (leave at default for automatic selection)"
                )
                
                alpha_slider = gr.Slider(
                    minimum=0.0,
                    maximum=1.0,
                    step=0.1,
                    value=DEFAULT_ALPHA,
                    label="Diffusion maps alpha parameter"
                )
                
                process_btn = gr.Button("Process Audio Files", variant="primary")
            
            with gr.Column():
                output_file = gr.File(
                    label="Download Results (ZIP)",
                    type="filepath"
                )
                status_text = gr.Textbox(
                    label="Status",
                    lines=3,
                    max_lines=10
                )
        
        gr.Markdown("""
        ## Output format:
        The results ZIP will contain:
        - `ps_scores_{model}.csv`: Perceptual Similarity scores for each speaker/source
        - `pm_scores_{model}.csv`: Perceptual Matching scores for each speaker/source  
        - `params.json`: Experiment parameters
        - `manifest_canonical.json`: Processed file manifest
        
        ## Score interpretation:
        - **PS (Perceptual Similarity)**: 0-1 score, higher is better. Measures how well the separated output matches the reference compared to other sources.
        - **PM (Perceptual Matching)**: 0-1 score, higher is better. Measures robustness to audio distortions.
        
        ## Notes:
        - Processing may take several minutes depending on the audio length and model
        - Audio files are automatically resampled to 16kHz
        - The tool automatically matches outputs to references based on correlation
        - For best results, ensure equal number of reference and output files
        
        ## Citation:
        If you use this tool in your research, please cite our paper (details coming soon).
        """)
        
        # 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()