Create app.py

app.py

@@ -0,0 +1,454 @@
+import os
+import uuid
+import tempfile
+import logging
+import asyncio
+from typing import List, Optional, Dict, Any
+
+from fastapi import FastAPI, File, UploadFile, Form, HTTPException, BackgroundTasks, Query
+from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
+import io # For zip file in memory
+import zipfile
+
+# --- Basic Editing Imports ---
+from pydub import AudioSegment
+from pydub.exceptions import CouldntDecodeError
+
+# --- AI & Advanced Audio Imports ---
+try:
+    import torch
+    from transformers import pipeline, AutoModelForSpeechSeq2Seq, AutoProcessor # Using pipeline for simplicity where possible
+    # Specific model imports might be needed depending on the chosen approach
+    import soundfile as sf
+    import numpy as np
+    import librosa # For resampling if needed
+    print("AI and advanced audio libraries loaded.")
+except ImportError as e:
+    print(f"Error importing AI/Audio libraries: {e}")
+    print("Ensure torch, transformers, soundfile, librosa are installed.")
+    print("AI features will be unavailable.")
+    torch = None
+    pipeline = None
+    sf = None
+    np = None
+    librosa = None
+
+# --- Configuration & Setup ---
+TEMP_DIR = tempfile.gettempdir()
+os.makedirs(TEMP_DIR, exist_ok=True)
+
+# Configure logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+# --- Global Variables for Loaded Models ---
+# Use dictionaries to potentially hold multiple models of each type later
+enhancement_pipelines: Dict[str, Any] = {}
+separation_models: Dict[str, Any] = {} # Might store pipeline or model/processor pair
+
+# Target sampling rates for models (check model cards on Hugging Face!)
+ENHANCEMENT_SR = 16000 # Example for speechbrain/sepformer
+DEMUCS_SR = 44100      # Demucs default
+
+# --- Helper Functions ---
+
+def cleanup_file(file_path: str):
+    """Safely remove a file."""
+    try:
+        if file_path and os.path.exists(file_path):
+            os.remove(file_path)
+            logger.info(f"Cleaned up temporary file: {file_path}")
+    except Exception as e:
+        logger.error(f"Error cleaning up file {file_path}: {e}", exc_info=False)
+
+async def save_upload_file(upload_file: UploadFile, prefix: str = "upload_") -> str:
+    """Saves an uploaded file to a temporary location and returns the path."""
+    # Generate a unique temporary file path
+    _, file_extension = os.path.splitext(upload_file.filename)
+    if not file_extension: file_extension = ".wav" # Default if no extension
+    temp_file_path = os.path.join(TEMP_DIR, f"{prefix}{uuid.uuid4().hex}{file_extension}")
+    try:
+        with open(temp_file_path, "wb") as buffer:
+             while content := await upload_file.read(1024 * 1024): buffer.write(content)
+        logger.info(f"Saved uploaded file '{upload_file.filename}' to temp path: {temp_file_path}")
+        return temp_file_path
+    except Exception as e:
+        logger.error(f"Failed to save uploaded file {upload_file.filename}: {e}", exc_info=True)
+        cleanup_file(temp_file_path)
+        raise HTTPException(status_code=500, detail=f"Could not save uploaded file: {upload_file.filename}")
+    finally:
+         await upload_file.close()
+
+# --- Audio Loading/Saving for AI Models ---
+
+def load_audio_for_hf(file_path: str, target_sr: Optional[int] = None) -> tuple[np.ndarray, int]:
+    """Loads audio using soundfile, converts to mono float32, optionally resamples."""
+    try:
+        audio, orig_sr = sf.read(file_path, dtype='float32', always_2d=False)
+        logger.info(f"Loaded audio '{os.path.basename(file_path)}' with SR={orig_sr}, shape={audio.shape}, dtype={audio.dtype}")
+
+        # Convert to mono if stereo
+        if audio.ndim > 1 and audio.shape[1] > 1:
+            # Simple averaging for mono conversion
+            audio = np.mean(audio, axis=1)
+            logger.info("Converted audio to mono")
+
+        # Resample if necessary
+        if target_sr and orig_sr != target_sr:
+            if librosa is None:
+                 raise RuntimeError("Librosa is required for resampling but not installed.")
+            logger.info(f"Resampling from {orig_sr} Hz to {target_sr} Hz...")
+            audio = librosa.resample(audio, orig_sr=orig_sr, target_sr=target_sr)
+            logger.info(f"Resampled audio shape: {audio.shape}")
+            current_sr = target_sr
+        else:
+            current_sr = orig_sr
+
+        return audio, current_sr
+
+    except Exception as e:
+        logger.error(f"Error loading/processing audio file {file_path} for HF: {e}", exc_info=True)
+        raise HTTPException(status_code=415, detail=f"Could not load or process audio file: {os.path.basename(file_path)}. Ensure it's a valid audio format.")
+
+def save_hf_audio(audio_data: np.ndarray, sampling_rate: int, output_format: str = "wav") -> str:
+    """Saves a NumPy audio array to a temporary file."""
+    output_filename = f"ai_output_{uuid.uuid4().hex}.{output_format}"
+    output_path = os.path.join(TEMP_DIR, output_filename)
+    try:
+        logger.info(f"Saving AI processed audio to {output_path} (SR={sampling_rate}, format={output_format})")
+        # Ensure data is float32 for common formats like wav/flac, pydub handles mp3 etc.
+        if audio_data.dtype != np.float32:
+             audio_data = audio_data.astype(np.float32)
+
+        # Use soundfile for lossless formats
+        if output_format.lower() in ['wav', 'flac']:
+             sf.write(output_path, audio_data, sampling_rate, format=output_format.upper())
+        else:
+             # For lossy formats like mp3, use pydub after converting numpy array
+             # Convert numpy array [-1.0, 1.0] float32 to pydub segment
+             # Scale to 16-bit integer range for pydub if needed
+             audio_int16 = (audio_data * 32767).astype(np.int16)
+             segment = AudioSegment(
+                 audio_int16.tobytes(),
+                 frame_rate=sampling_rate,
+                 sample_width=audio_int16.dtype.itemsize,
+                 channels=1 # Assuming mono output from AI models for now
+             )
+             segment.export(output_path, format=output_format)
+
+        return output_path
+    except Exception as e:
+        logger.error(f"Error saving AI processed audio to {output_path}: {e}", exc_info=True)
+        cleanup_file(output_path)
+        raise HTTPException(status_code=500, detail="Failed to save processed audio.")
+
+# --- Synchronous AI Inference Functions (to be run in threads) ---
+
+def _run_enhancement_sync(model_pipeline: Any, audio_data: np.ndarray, sampling_rate: int) -> np.ndarray:
+    """Synchronous wrapper for enhancement model inference."""
+    if not model_pipeline: raise ValueError("Enhancement model not loaded")
+    try:
+        logger.info(f"Running speech enhancement (input shape: {audio_data.shape}, SR: {sampling_rate})...")
+        # Pipeline usage depends heavily on the specific pipeline
+        # Example for a hypothetical 'audio-enhancement' pipeline:
+        result = model_pipeline({"raw": audio_data, "sampling_rate": sampling_rate})
+        enhanced_audio = result["audio"]["array"] # Adjust based on actual pipeline output
+        logger.info(f"Enhancement complete (output shape: {enhanced_audio.shape})")
+        return enhanced_audio
+    except Exception as e:
+        logger.error(f"Error during synchronous enhancement inference: {e}", exc_info=True)
+        raise # Re-raise to be caught by the async wrapper
+
+def _run_separation_sync(model_pipeline: Any, audio_data: np.ndarray, sampling_rate: int) -> Dict[str, np.ndarray]:
+    """Synchronous wrapper for source separation model inference."""
+    if not model_pipeline: raise ValueError("Separation model not loaded")
+    try:
+        logger.info(f"Running source separation (input shape: {audio_data.shape}, SR: {sampling_rate})...")
+        # Usage depends on the separation model/pipeline
+        # Example for a hypothetical 'audio-source-separation' pipeline:
+        # Note: Actual Demucs might need different handling (e.g., direct model call)
+        # result = model_pipeline({"raw": audio_data, "sampling_rate": sampling_rate})
+
+        # Manual example closer to raw Demucs model (if not using pipeline)
+        # Assuming `separation_models['demucs']` holds the loaded Demucs model instance
+        model = separation_models.get('demucs')
+        if not model: raise ValueError("Demucs model not loaded correctly")
+
+        # Demucs expects stereo input in shape (batch, channels, samples)
+        # Convert mono to stereo if needed, add batch dim
+        if audio_data.ndim == 1:
+            audio_data = np.stack([audio_data, audio_data], axis=0) # Create stereo from mono
+        audio_tensor = torch.tensor(audio_data).unsqueeze(0) # Add batch dimension
+
+        # Move to GPU if available and model is on GPU
+        device = next(model.parameters()).device
+        audio_tensor = audio_tensor.to(device)
+
+        with torch.no_grad():
+            sources = model(audio_tensor)[0] # Output shape (stems, channels, samples)
+
+        # Detach, move to CPU, convert to numpy
+        sources_np = sources.detach().cpu().numpy()
+
+        # Convert back to mono for simplicity (average channels)
+        stems = {
+            'drums': np.mean(sources_np[0], axis=0),
+            'bass': np.mean(sources_np[1], axis=0),
+            'other': np.mean(sources_np[2], axis=0),
+            'vocals': np.mean(sources_np[3], axis=0),
+        }
+        # Important: The order (drums, bass, other, vocals) is specific to Demucs v3/v4 default model
+        logger.info(f"Separation complete. Found stems: {list(stems.keys())}")
+        return stems
+
+    except Exception as e:
+        logger.error(f"Error during synchronous separation inference: {e}", exc_info=True)
+        raise
+
+# --- Model Loading Function ---
+def load_hf_models():
+    """Loads Hugging Face models at startup."""
+    global enhancement_pipelines, separation_models
+    if torch is None or pipeline is None:
+        logger.warning("Torch or Transformers not available. Skipping Hugging Face model loading.")
+        return
+
+    # --- Load Enhancement Model ---
+    # Using speechbrain/sepformer-whamr-enhancement via pipeline (check HF for exact pipeline task)
+    # Or load model/processor manually if no direct pipeline exists
+    enhancement_model_id = "speechbrain/sepformer-whamr-enhancement" # Example ID
+    try:
+        logger.info(f"Loading enhancement model: {enhancement_model_id}...")
+        # Use appropriate task, might be 'audio-enhancement', 'audio-classification' with custom logic, or manual loading
+        # If no pipeline, load manually:
+        # enhancement_processor = AutoProcessor.from_pretrained(...)
+        # enhancement_model = AutoModel...from_pretrained(...)
+        # enhancement_pipelines['speechbrain_sepformer'] = {"processor": enhancement_processor, "model": enhancement_model}
+
+        # For now, let's assume a placeholder pipeline exists or skip if complex
+        # enhancement_pipelines['speechbrain_sepformer'] = pipeline("audio-enhancement", model=enhancement_model_id)
+        logger.warning(f"Skipping load for {enhancement_model_id} - requires specific pipeline or manual setup.")
+
+    except Exception as e:
+        logger.error(f"Failed to load enhancement model '{enhancement_model_id}': {e}", exc_info=False)
+
+
+    # --- Load Separation Model (Demucs) ---
+    # Demucs is often used directly, not via a standard HF pipeline task
+    separation_model_id = "facebook/demucs" # Or specific variant like facebook/hybrid_demucs
+    try:
+        logger.info(f"Loading separation model: {separation_model_id}...")
+        # Demucs usually requires loading the model directly
+        # Using AutoModel might work for some variants if configured correctly in HF hub
+        # separation_models['demucs'] = AutoModel.from_pretrained(separation_model_id) # Check if this works
+
+        # More typically, you might need to install the 'demucs' package itself: pip install -U demucs
+        # import demucs.separate
+        # model = demucs.apply.load_model(separation_model_id or demucs.pretrained.DEFAULT_MODEL) # Using demucs package
+        # separation_models['demucs'] = model
+
+        # For now, simulate loading failure as direct AutoModel might not work
+        raise NotImplementedError("Demucs loading typically requires the 'demucs' package or specific manual loading.")
+        logger.info(f"Separation model '{separation_model_id}' loaded.")
+
+    except Exception as e:
+         logger.error(f"Failed to load separation model '{separation_model_id}': {e}", exc_info=False)
+         logger.warning("Note: Demucs loading often requires 'pip install demucs' and specific loading code, not just AutoModel.")
+
+
+# --- FastAPI App and Endpoints ---
+app = FastAPI(
+    title="AI Audio Editor API",
+    description="API for basic audio editing and AI-powered enhancement & separation. Requires FFmpeg and HF model dependencies.",
+    version="2.0.0",
+)
+
+@app.on_event("startup")
+async def startup_event():
+    """Load models when the application starts."""
+    logger.info("Application startup: Loading AI models...")
+    # Running model loading in a separate thread to avoid blocking startup completely
+    # although startup will wait for this thread to finish.
+    # Consider truly background loading if startup time is critical.
+    await asyncio.to_thread(load_hf_models)
+    logger.info("Model loading process finished (check logs for success/failure).")
+
+
+# --- Basic Editing Endpoints (Mostly Unchanged) ---
+
+@app.get("/", tags=["General"])
+def read_root():
+    """Root endpoint providing a welcome message and available features."""
+    features = ["/trim", "/concat", "/volume", "/convert"]
+    ai_features = []
+    if enhancement_pipelines: ai_features.append("/enhance")
+    if separation_models: ai_features.append("/separate")
+
+    return {
+        "message": "Welcome to the AI Audio Editor API.",
+        "basic_features": features,
+        "ai_features": ai_features if ai_features else "None available (models might have failed to load)",
+        "notes": "Requires FFmpeg. AI features require specific models loaded at startup (check logs)."
+        }
+
+# /trim, /concat, /volume, /convert endpoints remain largely the same as before
+# Ensure they use the updated save_upload_file and cleanup logic
+# (Code for these endpoints omitted for brevity - refer to previous example)
+# ... Add /trim, /concat, /volume, /convert endpoints here ...
+
+
+# --- AI Endpoints ---
+
+@app.post("/enhance", tags=["AI Editing"])
+async def enhance_speech(
+    background_tasks: BackgroundTasks,
+    file: UploadFile = File(..., description="Noisy speech audio file to enhance."),
+    model_id: str = Query("speechbrain_sepformer", description="ID of the enhancement model to use (if multiple loaded)."),
+    output_format: str = Form("wav", description="Output format for the enhanced audio (wav, flac recommended).")
+):
+    """Enhances speech audio using a pre-loaded AI model (experimental)."""
+    if torch is None or sf is None or np is None:
+         raise HTTPException(status_code=501, detail="AI processing libraries not available.")
+    if model_id not in enhancement_pipelines:
+         raise HTTPException(status_code=503, detail=f"Enhancement model '{model_id}' is not loaded or available.")
+
+    logger.info(f"Enhance request: file='{file.filename}', model='{model_id}', format='{output_format}'")
+    input_path = await save_upload_file(file, prefix="enhance_in_")
+    background_tasks.add_task(cleanup_file, input_path)
+
+    output_path = None # Define output_path before try block
+    try:
+        # Load audio, ensure correct SR for the model
+        audio_data, current_sr = load_audio_for_hf(input_path, target_sr=ENHANCEMENT_SR)
+
+        # Run inference in a separate thread
+        logger.info("Submitting enhancement task to background thread...")
+        model_pipeline = enhancement_pipelines[model_id] # Get the specific loaded pipeline/model
+        enhanced_audio = await asyncio.to_thread(
+            _run_enhancement_sync, model_pipeline, audio_data, current_sr
+        )
+        logger.info("Enhancement task completed.")
+
+        # Save the result
+        output_path = save_hf_audio(enhanced_audio, current_sr, output_format) # Use current_sr (which is target_sr)
+        background_tasks.add_task(cleanup_file, output_path)
+
+        return FileResponse(
+            path=output_path,
+            media_type=f"audio/{output_format}",
+            filename=f"enhanced_{file.filename}"
+        )
+
+    except Exception as e:
+        logger.error(f"Error during enhancement operation: {e}", exc_info=True)
+        if output_path: cleanup_file(output_path) # Cleanup if error occurred after output started saving
+        if isinstance(e, HTTPException): raise e
+        else: raise HTTPException(status_code=500, detail=f"An unexpected error occurred during enhancement: {str(e)}")
+
+
+@app.post("/separate", tags=["AI Editing"])
+async def separate_sources(
+    background_tasks: BackgroundTasks,
+    file: UploadFile = File(..., description="Music audio file to separate into stems."),
+    model_id: str = Query("demucs", description="ID of the separation model to use."),
+    stems: List[str] = Form(..., description="List of stems to extract (e.g., 'vocals', 'drums', 'bass', 'other')."),
+    output_format: str = Form("wav", description="Output format for the stems (wav, flac recommended).")
+):
+    """Separates music into stems (vocals, drums, bass, other) using Demucs (experimental). Returns a ZIP archive."""
+    if torch is None or sf is None or np is None:
+         raise HTTPException(status_code=501, detail="AI processing libraries not available.")
+    if model_id not in separation_models:
+         raise HTTPException(status_code=503, detail=f"Separation model '{model_id}' is not loaded or available.")
+
+    valid_stems = {'vocals', 'drums', 'bass', 'other'}
+    requested_stems = set(s.lower() for s in stems)
+    if not requested_stems.issubset(valid_stems):
+        raise HTTPException(status_code=422, detail=f"Invalid stem(s) requested. Valid stems are: {', '.join(valid_stems)}")
+
+    logger.info(f"Separate request: file='{file.filename}', model='{model_id}', stems={requested_stems}, format='{output_format}'")
+    input_path = await save_upload_file(file, prefix="separate_in_")
+    background_tasks.add_task(cleanup_file, input_path)
+
+    stem_output_paths: Dict[str, str] = {}
+    zip_buffer = None
+    try:
+        # Load audio, ensure correct SR for the model (Demucs uses 44.1kHz)
+        audio_data, current_sr = load_audio_for_hf(input_path, target_sr=DEMUCS_SR)
+
+        # Run inference in a separate thread
+        logger.info("Submitting separation task to background thread...")
+        model = separation_models[model_id] # Get the specific loaded model
+        all_separated_stems = await asyncio.to_thread(
+            _run_separation_sync, model, audio_data, current_sr
+        )
+        logger.info("Separation task completed.")
+
+        # --- Create ZIP file in memory ---
+        zip_buffer = io.BytesIO()
+        with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zipf:
+            # Save only the requested stems
+            for stem_name in requested_stems:
+                if stem_name in all_separated_stems:
+                    stem_data = all_separated_stems[stem_name]
+                    # Save stem temporarily to disk first (needed for pydub/sf.write)
+                    stem_path = save_hf_audio(stem_data, current_sr, output_format)
+                    stem_output_paths[stem_name] = stem_path # Keep track for cleanup
+                    background_tasks.add_task(cleanup_file, stem_path) # Schedule cleanup
+
+                    # Add the saved stem file to the ZIP archive
+                    archive_name = f"{stem_name}_{os.path.basename(input_path)}.{output_format}"
+                    zipf.write(stem_path, arcname=archive_name)
+                    logger.info(f"Added '{archive_name}' to ZIP.")
+                else:
+                    logger.warning(f"Requested stem '{stem_name}' not found in model output.")
+
+        zip_buffer.seek(0) # Rewind buffer pointer
+
+        # Return the ZIP file
+        zip_filename = f"separated_stems_{os.path.splitext(file.filename)[0]}.zip"
+        return StreamingResponse(
+            zip_buffer,
+            media_type="application/zip",
+            headers={'Content-Disposition': f'attachment; filename="{zip_filename}"'}
+        )
+
+    except Exception as e:
+        logger.error(f"Error during separation operation: {e}", exc_info=True)
+        # Cleanup any stems that were saved before zipping failed
+        for path in stem_output_paths.values():
+            cleanup_file(path)
+        if zip_buffer: zip_buffer.close() # Close memory buffer
+
+        if isinstance(e, HTTPException): raise e
+        else: raise HTTPException(status_code=500, detail=f"An unexpected error occurred during separation: {str(e)}")
+
+
+# --- How to Run ---
+# 1. Ensure FFmpeg is installed and accessible.
+# 2. Save this code as `app.py`.
+# 3. Create `requirements.txt` (as shown above).
+# 4. Install dependencies: `pip install -r requirements.txt` (This can take time!)
+# 5. Run the FastAPI server: `uvicorn app:app --reload --host 0.0.0.0`
+#    (Use --host 0.0.0.0 for external/Docker access. --reload is optional)
+#
+# --- WARNING ---
+# - AI models require SIGNIFICANT RAM and CPU/GPU. Inference can be SLOW.
+# - The first run will download models, which can take a long time and lots of disk space.
+# - Ensure the specific model IDs used are correct and compatible with HF libraries.
+# - Model loading at startup might fail if dependencies are missing or resources are insufficient. Check logs!
+#
+# --- Example Usage (using curl) ---
+#
+# **Enhance:** (Enhance noisy_speech.wav)
+# curl -X POST "http://127.0.0.1:8000/enhance?model_id=speechbrain_sepformer" \
+#      -F "file=@noisy_speech.wav" \
+#      -F "output_format=wav" \
+#      --output enhanced_speech.wav
+#
+# **Separate:** (Separate vocals and drums from music.mp3)
+# curl -X POST "http://127.0.0.1:8000/separate?model_id=demucs" \
+#      -F "[email protected]" \
+#      -F "stems=vocals" \
+#      -F "stems=drums" \
+#      -F "output_format=mp3" \
+#      --output separated_stems.zip