Update app.py

app.py

@@ -1,14 +1,15 @@
+# ----------- START app.py -----------
 import os
 import uuid
 import tempfile
 import logging
 import asyncio
 from typing import List, Optional, Dict, Any
+import io
+import zipfile
 
 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
@@ -19,20 +20,26 @@ 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
+    # E.g. for Demucs V4 (Hybrid Transformer): from demucs.hdemucs import HDemucs
+    # from demucs.pretrained import hdemucs_mmi
     import soundfile as sf
     import numpy as np
     import librosa # For resampling if needed
+    AI_LIBRARIES_AVAILABLE = True
     print("AI and advanced audio libraries loaded.")
 except ImportError as e:
-    print(f"Error importing AI/Audio libraries: {e}")
+    print(f"Warning: Error importing AI/Audio libraries: {e}")
     print("Ensure torch, transformers, soundfile, librosa are installed.")
     print("AI features will be unavailable.")
+    AI_LIBRARIES_AVAILABLE = False
+    # Define dummy placeholders if needed, or just rely on AI_LIBRARIES_AVAILABLE flag
     torch = None
     pipeline = None
     sf = None
     np = None
     librosa = None
 
+
 # --- Configuration & Setup ---
 TEMP_DIR = tempfile.gettempdir()
 os.makedirs(TEMP_DIR, exist_ok=True)
@@ -43,15 +50,28 @@ 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
+enhancement_models: Dict[str, Any] = {} # Store model/processor or pipeline
+separation_models: Dict[str, Any] = {} # Store model/processor or pipeline
 
 # Target sampling rates for models (check model cards on Hugging Face!)
-ENHANCEMENT_SR = 16000 # Example for speechbrain/sepformer
-DEMUCS_SR = 44100      # Demucs default
+# These MUST match the models being loaded in download_models.py and load_hf_models
+ENHANCEMENT_MODEL_ID = "speechbrain/sepformer-whamr-enhancement"
+ENHANCEMENT_SR = 16000 # Sepformer uses 16kHz
 
-# --- Helper Functions ---
+# Note: facebook/demucs is deprecated in transformers >4.26. Use specific variants.
+# Using facebook/htdemucs_ft for example (requires Demucs v4 style loading)
+# Or find a model suitable for AutoModel if needed.
+SEPARATION_MODEL_ID = "facebook/demucs_quantized" # Example using a quantized version (smaller, faster CPU)
+# SEPARATION_MODEL_ID = "facebook/hdemucs_mmi" # Example for Multi-Media Instructions model (if using demucs lib)
+DEMUCS_SR = 44100      # Demucs default is 44.1kHz
 
+# Define HF_HOME cache directory *within* the container if downloading during build
+HF_CACHE_DIR = os.environ.get("HF_HOME", "/app/hf_cache") # Use HF_HOME from Dockerfile or default
+
+
+# --- Helper Functions (cleanup_file, save_upload_file, load_audio_for_hf, save_hf_audio) ---
+# (Include the helper functions from the previous app.py example here)
+# ...
 def cleanup_file(file_path: str):
     """Safely remove a file."""
     try:
@@ -63,7 +83,6 @@ def cleanup_file(file_path: str):
 
 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}")
@@ -79,26 +98,32 @@ async def save_upload_file(upload_file: UploadFile, prefix: str = "upload_") ->
     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."""
+    if not AI_LIBRARIES_AVAILABLE or sf is None or np is None:
+        raise HTTPException(status_code=501, detail="Audio processing libraries (soundfile, numpy) not available.")
     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")
+        if audio.ndim > 1 and audio.shape[-1] > 1: # Check last dimension for channels
+             if audio.shape[0] == min(audio.shape): # If channels are first dim
+                 audio = audio.T # Transpose to (samples, channels)
+             audio = np.mean(audio, axis=1)
+             logger.info(f"Converted audio to mono, new shape: {audio.shape}")
+        elif audio.ndim > 1: # If shape is like (1, N) or (N, 1)
+             audio = audio.squeeze() # Remove singleton dimension
+             logger.info(f"Squeezed audio to 1D, new shape: {audio.shape}")
+
 
-        # 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)
+            # Ensure audio is contiguous before resampling if necessary
+            if not audio.flags['C_CONTIGUOUS']:
+                 audio = np.ascontiguousarray(audio)
+            audio = librosa.resample(y=audio, orig_sr=orig_sr, target_sr=target_sr)
             logger.info(f"Resampled audio shape: {audio.shape}")
             current_sr = target_sr
         else:
@@ -112,22 +137,34 @@ def load_audio_for_hf(file_path: str, target_sr: Optional[int] = None) -> tuple[
 
 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."""
+    if not AI_LIBRARIES_AVAILABLE or sf is None or np is None:
+        raise HTTPException(status_code=501, detail="Audio processing libraries (soundfile, numpy) not available.")
+
     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.
+        logger.info(f"Saving AI processed audio to {output_path} (SR={sampling_rate}, format={output_format}, shape={audio_data.shape})")
+
+        # Ensure data is float32 for common formats like wav/flac
         if audio_data.dtype != np.float32:
+             logger.warning(f"Audio data has dtype {audio_data.dtype}, converting to float32.")
              audio_data = audio_data.astype(np.float32)
 
+        # Clip data to avoid issues with some formats/players if values go beyond [-1, 1]
+        audio_data = np.clip(audio_data, -1.0, 1.0)
+
         # 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
+             logger.debug("Using pydub for lossy format export...")
+             # Scale float32 [-1, 1] to int16 for pydub
              audio_int16 = (audio_data * 32767).astype(np.int16)
+             if audio_int16.ndim > 1: # Should be mono by now, but double check
+                 logger.warning("Audio data still has multiple dimensions before pydub export, attempting mean.")
+                 audio_int16 = np.mean(audio_int16, axis=1).astype(np.int16)
+
              segment = AudioSegment(
                  audio_int16.tobytes(),
                  frame_rate=sampling_rate,
@@ -142,118 +179,191 @@ def save_hf_audio(audio_data: np.ndarray, sampling_rate: int, output_format: str
         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 AI Inference Functions (_run_enhancement_sync, _run_separation_sync) ---
+# (Include the sync functions from the previous app.py example here)
+# Make sure they handle potential model loading issues gracefully
+# ...
+def _run_enhancement_sync(model_key: str, 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")
+    if not AI_LIBRARIES_AVAILABLE or model_key not in enhancement_models:
+        raise ValueError(f"Enhancement model '{model_key}' not available or AI libraries missing.")
+
+    model_info = enhancement_models[model_key]
+    # Adapt based on whether model_info holds a pipeline or model/processor
+    # This example assumes a pipeline-like object is stored
+    enhancer = model_info # Assuming pipeline
+    if not enhancer: raise ValueError(f"Enhancement pipeline '{model_key}' is None.")
+
     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})
+        logger.info(f"Running speech enhancement with '{model_key}' (input shape: {audio_data.shape}, SR: {sampling_rate})...")
+        # Usage depends heavily on the specific model/pipeline interface
+        # For SpeechBrain models often used *without* HF pipeline:
+        # Example: enhanced_wav = enhancer.enhance_batch(torch.tensor(audio_data).unsqueeze(0), lengths=torch.tensor([audio_data.shape[0]]))
+        # enhanced_audio = enhanced_wav.squeeze(0).cpu().numpy()
+
+        # If using a generic HF pipeline:
+        result = enhancer({"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)
+        logger.error(f"Error during synchronous enhancement inference with '{model_key}': {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]:
+
+def _run_separation_sync(model_key: str, 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")
+    if not AI_LIBRARIES_AVAILABLE or model_key not in separation_models:
+        raise ValueError(f"Separation model '{model_key}' not available or AI libraries missing.")
+
+    model_info = separation_models[model_key]
+    model = model_info # Assuming direct model object is stored for Demucs
+
+    if not model: raise ValueError(f"Separation model '{model_key}' is None.")
+
     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
+        logger.info(f"Running source separation with '{model_key}' (input shape: {audio_data.shape}, SR: {sampling_rate})...")
+
+        # Prepare input tensor for Demucs-like models
+        # Expects (batch, channels, samples), float32
         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
+            # Need stereo for standard Demucs
+            logger.debug("Separation input is mono, duplicating to create stereo.")
+            audio_data = np.stack([audio_data, audio_data], axis=0) # (2, samples)
+        if audio_data.shape[0] != 2:
+             # If it's somehow (samples, 2), transpose
+             if audio_data.shape[1] == 2: audio_data = audio_data.T
+             else: raise ValueError(f"Unexpected input audio shape for separation: {audio_data.shape}")
+
+        audio_tensor = torch.tensor(audio_data, dtype=torch.float32).unsqueeze(0) # (1, 2, samples)
 
-        # Move to GPU if available and model is on GPU
+        # Move to model's device (CPU or GPU)
         device = next(model.parameters()).device
+        logger.debug(f"Moving separation tensor to device: {device}")
         audio_tensor = audio_tensor.to(device)
 
+        # Perform inference
         with torch.no_grad():
-            sources = model(audio_tensor)[0] # Output shape (stems, channels, samples)
+             logger.debug("Starting model inference for separation...")
+             # Output shape depends on model, e.g., (batch, stems, channels, samples)
+             sources = model(audio_tensor)[0] # Remove batch dim
+             logger.debug(f"Model inference complete, sources shape: {sources.shape}")
 
         # 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
+        sources_np = sources.detach().cpu().numpy() # (stems, channels, samples)
+
+        # Define stem order based on the *specific* Demucs model used
+        # This order is for default Demucs v3/v4 (facebook/demucs, facebook/htdemucs_ft, etc.)
+        stem_names = ['drums', 'bass', 'other', 'vocals']
+        if sources_np.shape[0] != len(stem_names):
+             logger.warning(f"Model output {sources_np.shape[0]} stems, expected {len(stem_names)}. Stem names might be incorrect.")
+             # Fallback names if shape mismatch
+             stem_names = [f"stem_{i+1}" for i in range(sources_np.shape[0])]
+
+        stems = {}
+        for i, name in enumerate(stem_names):
+             # Average channels to get mono stem
+             mono_stem = np.mean(sources_np[i], axis=0)
+             stems[name] = mono_stem
+             logger.debug(f"Extracted stem '{name}', shape: {mono_stem.shape}")
+
         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)
+        logger.error(f"Error during synchronous separation inference with '{model_key}': {e}", exc_info=True)
         raise
 
 # --- Model Loading Function ---
+# (Include the load_hf_models function from the previous app.py example here)
+# Make sure it uses the correct model IDs and potentially adjusts loading logic
+# if using libraries like `demucs` directly.
+# ...
 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.")
+    if not AI_LIBRARIES_AVAILABLE:
+        logger.warning("Skipping Hugging Face model loading as libraries are missing.")
         return
 
+    global enhancement_models, separation_models
+
     # --- 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
+    enhancement_key = "speechbrain_enhancer" # Internal key
     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.")
+        logger.info(f"Attempting to load enhancement model: {ENHANCEMENT_MODEL_ID}...")
+        # SpeechBrain models often require specific loading from their toolkit or HF spaces
+        # This might involve cloning a repo or using specific classes.
+        # Using HF pipeline if available, otherwise manual load.
+        # Example using pipeline (might not work for all speechbrain models):
+        # enhancement_models[enhancement_key] = pipeline(
+        #     "audio-enhancement", # Or appropriate task
+        #     model=ENHANCEMENT_MODEL_ID,
+        #     cache_dir=HF_CACHE_DIR,
+        #     device=0 if torch.cuda.is_available() else -1 # Use GPU if possible
+        # )
+        # Manual load might be needed:
+        # from speechbrain.pretrained import SepformerEnhancement
+        # enhancer = SepformerEnhancement.from_hparams(
+        #      source=ENHANCEMENT_MODEL_ID,
+        #      savedir=os.path.join(HF_CACHE_DIR, "speechbrain", ENHANCEMENT_MODEL_ID.split('/')[-1]),
+        #      run_opts={"device": "cuda" if torch.cuda.is_available() else "cpu"}
+        # )
+        # enhancement_models[enhancement_key] = enhancer
+        logger.warning(f"Actual loading for {ENHANCEMENT_MODEL_ID} skipped - requires SpeechBrain toolkit or specific pipeline setup.")
+        # To make the endpoint testable without full setup:
+        # enhancement_models[enhancement_key] = None # Or a dummy function
 
     except Exception as e:
-        logger.error(f"Failed to load enhancement model '{enhancement_model_id}': {e}", exc_info=False)
+        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
+    separation_key = "demucs_separator" # Internal key
     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
+        logger.info(f"Attempting to load separation model: {SEPARATION_MODEL_ID}...")
+        # Loading Demucs models can be complex.
+        # Option 1: Use AutoModel if the HF Hub version supports it (less common for Demucs)
+        # Option 2: Use the `demucs` library (recommended if installed: pip install -U demucs)
+        # Option 3: Find a Transformers-compatible version if available.
+
+        # Example using AutoModel (Try this first, might work for some quantized/HF versions)
+        try:
+             # Determine device
+             device = "cuda" if torch.cuda.is_available() else "cpu"
+             logger.info(f"Loading Demucs on device: {device}")
+             # Check if AutoModelForSpeechSeq2Seq is appropriate, might need a different AutoModel class
+             separation_models[separation_key] = AutoModelForSpeechSeq2Seq.from_pretrained(
+                 SEPARATION_MODEL_ID,
+                 cache_dir=HF_CACHE_DIR
+                 # Add trust_remote_code=True if needed for custom model code on HF hub
+             ).to(device)
+
+             # Check if the loaded model has an 'eval' method (common for PyTorch models)
+             if hasattr(separation_models[separation_key], 'eval'):
+                 separation_models[separation_key].eval() # Set to evaluation mode
+
+             logger.info(f"Successfully loaded separation model '{SEPARATION_MODEL_ID}' using AutoModel.")
+
+        except Exception as auto_model_err:
+            logger.warning(f"Failed to load '{SEPARATION_MODEL_ID}' using AutoModel: {auto_model_err}. Consider installing 'demucs' library.")
+            separation_models[separation_key] = None # Ensure it's None if loading failed
+
+            # Example using `demucs` library (if installed)
+            # try:
+            #     import demucs.separate
+            #     model = demucs.apply.load_model(pretrained_model_path_or_url) # Needs adjustment
+            #     separation_models[separation_key] = model
+            #     logger.info(f"Successfully loaded separation model using 'demucs' library.")
+            # except ImportError:
+            #      logger.error("Cannot load Demucs: 'demucs' library not found. Please run 'pip install -U demucs'.")
+            # except Exception as demucs_lib_err:
+            #      logger.error(f"Error loading model using 'demucs' library: {demucs_lib_err}")
 
-        # 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.")
+         logger.error(f"General error loading separation model '{SEPARATION_MODEL_ID}': {e}", exc_info=False)
+         if separation_key in separation_models: separation_models[separation_key] = None
 
 
 # --- FastAPI App and Endpoints ---
@@ -266,82 +376,267 @@ app = FastAPI(
 @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.
+    logger.info("Application startup: Loading AI models (this may take time)...")
     await asyncio.to_thread(load_hf_models)
-    logger.info("Model loading process finished (check logs for success/failure).")
-
+    logger.info("Model loading process finished.")
 
-# --- Basic Editing Endpoints (Mostly Unchanged) ---
 
+# --- API Endpoints ---
+# (Include / , /trim, /concat, /volume, /convert endpoints here - same as previous version)
+# ...
 @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")
+    # Check if models were successfully loaded (i.e., not None)
+    if any(model is not None for model in enhancement_models.values()): ai_features.append("/enhance")
+    if any(model is not None for model in separation_models.values()): 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)",
+        "ai_features": ai_features if ai_features else "None loaded (check logs)",
         "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 ...
+@app.post("/trim", tags=["Basic Editing"])
+async def trim_audio(
+    background_tasks: BackgroundTasks,
+    file: UploadFile = File(..., description="Audio file to trim."),
+    start_ms: int = Form(..., description="Start time in milliseconds."),
+    end_ms: int = Form(..., description="End time in milliseconds.")
+):
+    """Trims an audio file to the specified start and end times (in milliseconds)."""
+    if start_ms < 0 or end_ms <= start_ms:
+        raise HTTPException(status_code=422, detail="Invalid start/end times. Ensure start_ms >= 0 and end_ms > start_ms.")
+
+    logger.info(f"Trim request: file='{file.filename}', start={start_ms}ms, end={end_ms}ms")
+    input_path = None
+    output_path = None
+    try:
+        input_path = await save_upload_file(file, prefix="trim_in_")
+        background_tasks.add_task(cleanup_file, input_path) # Schedule input cleanup
+
+        # Use Pydub for basic trim
+        audio = AudioSegment.from_file(input_path)
+        trimmed_audio = audio[start_ms:end_ms]
+        logger.info(f"Audio trimmed to {len(trimmed_audio)}ms")
+
+        original_format = os.path.splitext(file.filename)[1][1:].lower() or "mp3"
+        if not original_format or original_format == "tmp": original_format = "mp3"
+        output_filename = f"trimmed_{uuid.uuid4().hex}.{original_format}"
+        output_path = os.path.join(TEMP_DIR, output_filename)
+
+        trimmed_audio.export(output_path, format=original_format)
+        background_tasks.add_task(cleanup_file, output_path) # Schedule output cleanup
+
+        return FileResponse(
+            path=output_path,
+            media_type=f"audio/{original_format}", # Attempt correct media type
+            filename=f"trimmed_{file.filename}"
+        )
+    except CouldntDecodeError:
+         logger.warning(f"pydub failed to decode: {file.filename}")
+         raise HTTPException(status_code=415, detail="Unsupported audio format or corrupted file.")
+    except Exception as e:
+        logger.error(f"Error during trim operation: {e}", exc_info=True)
+        if output_path: cleanup_file(output_path) # Immediate cleanup on error
+        if input_path: cleanup_file(input_path)   # Immediate cleanup on error
+        if isinstance(e, HTTPException): raise e
+        else: raise HTTPException(status_code=500, detail=f"An unexpected error occurred during trimming: {str(e)}")
+
+
+@app.post("/concat", tags=["Basic Editing"])
+async def concatenate_audio(
+    background_tasks: BackgroundTasks,
+    files: List[UploadFile] = File(..., description="Two or more audio files to join in order."),
+    output_format: str = Form("mp3", description="Desired output format (e.g., 'mp3', 'wav', 'ogg').")
+):
+    """Concatenates two or more audio files sequentially."""
+    if len(files) < 2:
+        raise HTTPException(status_code=422, detail="Please upload at least two files to concatenate.")
+
+    logger.info(f"Concatenate request: {len(files)} files, output_format='{output_format}'")
+    input_paths = []
+    loaded_audios = []
+    output_path = None
+
+    try:
+        combined_audio = AudioSegment.empty()
+        first_filename_base = "combined"
 
+        for i, file in enumerate(files):
+            input_path = await save_upload_file(file, prefix=f"concat_{i}_")
+            input_paths.append(input_path)
+            background_tasks.add_task(cleanup_file, input_path)
+            audio = AudioSegment.from_file(input_path)
+            combined_audio += audio
+            if i == 0: first_filename_base = os.path.splitext(file.filename)[0]
+            logger.info(f"Appended '{file.filename}', current total duration: {len(combined_audio)}ms")
 
-# --- AI Endpoints ---
+        if len(combined_audio) == 0:
+             raise HTTPException(status_code=500, detail="No audio data after attempting concatenation.")
 
+        output_filename_final = f"concat_{first_filename_base}_and_{len(files)-1}_others.{output_format}"
+        output_path = os.path.join(TEMP_DIR, f"concat_out_{uuid.uuid4().hex}.{output_format}")
+
+        combined_audio.export(output_path, format=output_format)
+        background_tasks.add_task(cleanup_file, output_path) # Schedule output cleanup
+
+        return FileResponse(
+            path=output_path,
+            media_type=f"audio/{output_format}",
+            filename=output_filename_final
+        )
+    except CouldntDecodeError as e:
+         logger.warning(f"pydub failed to decode one of the concat files: {e}")
+         raise HTTPException(status_code=415, detail=f"Unsupported format or corrupted file among inputs: {e}")
+    except Exception as e:
+        logger.error(f"Error during concat operation: {e}", exc_info=True)
+        if output_path: cleanup_file(output_path)
+        for p in input_paths: cleanup_file(p)
+        if isinstance(e, HTTPException): raise e
+        else: raise HTTPException(status_code=500, detail=f"An unexpected error occurred during concatenation: {str(e)}")
+
+@app.post("/volume", tags=["Basic Editing"])
+async def change_volume(
+    background_tasks: BackgroundTasks,
+    file: UploadFile = File(..., description="Audio file to adjust volume for."),
+    change_db: float = Form(..., description="Volume change in decibels (dB). Positive values increase volume, negative values decrease.")
+):
+    """Adjusts the volume of an audio file by a specified decibel amount."""
+    logger.info(f"Volume request: file='{file.filename}', change_db={change_db}dB")
+    input_path = None
+    output_path = None
+    try:
+        input_path = await save_upload_file(file, prefix="volume_in_")
+        background_tasks.add_task(cleanup_file, input_path)
+
+        audio = AudioSegment.from_file(input_path)
+        adjusted_audio = audio + change_db
+        logger.info(f"Volume adjusted by {change_db}dB.")
+
+        original_format = os.path.splitext(file.filename)[1][1:].lower() or "mp3"
+        if not original_format or original_format == "tmp": original_format = "mp3"
+        output_filename_final = f"volume_{change_db}dB_{file.filename}"
+        output_path = os.path.join(TEMP_DIR, f"volume_out_{uuid.uuid4().hex}.{original_format}")
+
+        adjusted_audio.export(output_path, format=original_format)
+        background_tasks.add_task(cleanup_file, output_path)
+
+        return FileResponse(
+            path=output_path,
+            media_type=f"audio/{original_format}",
+            filename=output_filename_final
+        )
+    except CouldntDecodeError:
+         logger.warning(f"pydub failed to decode: {file.filename}")
+         raise HTTPException(status_code=415, detail="Unsupported audio format or corrupted file.")
+    except Exception as e:
+        logger.error(f"Error during volume operation: {e}", exc_info=True)
+        if output_path: cleanup_file(output_path)
+        if input_path: cleanup_file(input_path)
+        if isinstance(e, HTTPException): raise e
+        else: raise HTTPException(status_code=500, detail=f"An unexpected error occurred during volume adjustment: {str(e)}")
+
+@app.post("/convert", tags=["Basic Editing"])
+async def convert_format(
+    background_tasks: BackgroundTasks,
+    file: UploadFile = File(..., description="Audio file to convert."),
+    output_format: str = Form(..., description="Target audio format (e.g., 'mp3', 'wav', 'ogg', 'flac').")
+):
+    """Converts an audio file to a different format."""
+    allowed_formats = {'mp3', 'wav', 'ogg', 'flac', 'aac', 'm4a'}
+    if output_format.lower() not in allowed_formats:
+         raise HTTPException(status_code=422, detail=f"Invalid output format. Allowed: {', '.join(allowed_formats)}")
+
+    logger.info(f"Convert request: file='{file.filename}', output_format='{output_format}'")
+    input_path = None
+    output_path = None
+    try:
+        input_path = await save_upload_file(file, prefix="convert_in_")
+        background_tasks.add_task(cleanup_file, input_path)
+
+        audio = AudioSegment.from_file(input_path)
+
+        output_format_lower = output_format.lower()
+        filename_base = os.path.splitext(file.filename)[0]
+        output_filename_final = f"{filename_base}_converted.{output_format_lower}"
+        output_path = os.path.join(TEMP_DIR, f"convert_out_{uuid.uuid4().hex}.{output_format_lower}")
+
+        audio.export(output_path, format=output_format_lower)
+        background_tasks.add_task(cleanup_file, output_path)
+
+        return FileResponse(
+            path=output_path,
+            media_type=f"audio/{output_format_lower}",
+            filename=output_filename_final
+        )
+    except CouldntDecodeError:
+         logger.warning(f"pydub failed to decode: {file.filename}")
+         raise HTTPException(status_code=415, detail="Unsupported audio format or corrupted file.")
+    except Exception as e:
+        logger.error(f"Error during convert operation: {e}", exc_info=True)
+        if output_path: cleanup_file(output_path)
+        if input_path: cleanup_file(input_path)
+        if isinstance(e, HTTPException): raise e
+        else: raise HTTPException(status_code=500, detail=f"An unexpected error occurred during format conversion: {str(e)}")
+
+
+# (Include /enhance and /separate AI endpoints here - same as previous version)
+# ...
 @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)."),
+    model_key: str = Query("speechbrain_enhancer", description="Internal key of the enhancement model to use."),
     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:
+    if not AI_LIBRARIES_AVAILABLE:
          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.")
+    if model_key not in enhancement_models or enhancement_models[model_key] is None:
+         logger.error(f"Enhancement model key '{model_key}' requested but model not loaded.")
+         raise HTTPException(status_code=503, detail=f"Enhancement model '{model_key}' is not loaded or available. Check server logs.")
 
-    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
+    logger.info(f"Enhance request: file='{file.filename}', model_key='{model_key}', format='{output_format}'")
+    input_path = None
+    output_path = None
     try:
+        input_path = await save_upload_file(file, prefix="enhance_in_")
+        background_tasks.add_task(cleanup_file, input_path)
+
         # Load audio, ensure correct SR for the model
+        logger.debug(f"Loading audio for enhancement, target SR: {ENHANCEMENT_SR}")
         audio_data, current_sr = load_audio_for_hf(input_path, target_sr=ENHANCEMENT_SR)
+        if current_sr != ENHANCEMENT_SR: # Should have been resampled, but double check
+             logger.warning(f"Audio SR after loading is {current_sr}, expected {ENHANCEMENT_SR}. Check resampling.")
+             # Depending on model strictness, could raise error or proceed cautiously.
+             # raise HTTPException(status_code=500, detail="Audio resampling failed.")
 
         # 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
+            _run_enhancement_sync, model_key, audio_data, current_sr # Pass key, data, and ACTUAL sr used
         )
         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)
+        output_path = save_hf_audio(enhanced_audio, ENHANCEMENT_SR, output_format) # Save with model's target SR
         background_tasks.add_task(cleanup_file, output_path)
 
+        output_filename_final = f"enhanced_{os.path.splitext(file.filename)[0]}.{output_format}"
         return FileResponse(
             path=output_path,
             media_type=f"audio/{output_format}",
-            filename=f"enhanced_{file.filename}"
+            filename=output_filename_final
         )
 
     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 output_path: cleanup_file(output_path)
+        if input_path: cleanup_file(input_path)
         if isinstance(e, HTTPException): raise e
         else: raise HTTPException(status_code=500, detail=f"An unexpected error occurred during enhancement: {str(e)}")
 
@@ -350,105 +645,96 @@ async def enhance_speech(
 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."),
+    model_key: str = Query("demucs_separator", description="Internal key 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:
+    """Separates music into stems (vocals, drums, bass, other) using a pre-loaded AI model (experimental). Returns a ZIP archive."""
+    if not AI_LIBRARIES_AVAILABLE:
          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.")
+    if model_key not in separation_models or separation_models[model_key] is None:
+         logger.error(f"Separation model key '{model_key}' requested but model not loaded.")
+         raise HTTPException(status_code=503, detail=f"Separation model '{model_key}' is not loaded or available. Check server logs.")
 
-    valid_stems = {'vocals', 'drums', 'bass', 'other'}
+    valid_stems = {'vocals', 'drums', 'bass', 'other'} # Based on typical Demucs output
     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)
+        # Allow if all stems are requested even if validation set is smaller? Or just error.
+        raise HTTPException(status_code=422, detail=f"Invalid stem(s) requested. Valid stems are generally: {', '.join(valid_stems)}")
 
+    logger.info(f"Separate request: file='{file.filename}', model_key='{model_key}', stems={requested_stems}, format='{output_format}'")
+    input_path = None
     stem_output_paths: Dict[str, str] = {}
-    zip_buffer = None
+    zip_buffer = io.BytesIO() # Use BytesIO for in-memory ZIP
+
     try:
-        # Load audio, ensure correct SR for the model (Demucs uses 44.1kHz)
+        input_path = await save_upload_file(file, prefix="separate_in_")
+        background_tasks.add_task(cleanup_file, input_path) # Schedule input cleanup
+
+        # Load audio, ensure correct SR for the model
+        logger.debug(f"Loading audio for separation, target SR: {DEMUCS_SR}")
         audio_data, current_sr = load_audio_for_hf(input_path, target_sr=DEMUCS_SR)
+        if current_sr != DEMUCS_SR:
+             logger.warning(f"Audio SR after loading is {current_sr}, expected {DEMUCS_SR}. Check resampling.")
+             # raise HTTPException(status_code=500, detail="Audio resampling failed.")
 
         # 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
+            _run_separation_sync, model_key, audio_data, current_sr # Pass key, data, actual SR
         )
         logger.info("Separation task completed.")
 
         # --- Create ZIP file in memory ---
-        zip_buffer = io.BytesIO()
+        zip_filename_base = f"separated_{os.path.splitext(file.filename)[0]}"
         with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zipf:
-            # Save only the requested stems
+            logger.info(f"Creating ZIP archive in memory...")
+            found_stems_count = 0
             for stem_name in requested_stems:
                 if stem_name in all_separated_stems:
                     stem_data = all_separated_stems[stem_name]
+                    if stem_data is None or stem_data.size == 0:
+                        logger.warning(f"Stem '{stem_name}' data is empty, skipping.")
+                        continue
+
                     # Save stem temporarily to disk first (needed for pydub/sf.write)
-                    stem_path = save_hf_audio(stem_data, current_sr, output_format)
+                    logger.debug(f"Saving temporary stem file for '{stem_name}'...")
+                    stem_path = save_hf_audio(stem_data, DEMUCS_SR, output_format) # Save with model's target SR
                     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}"
+                    archive_name = f"{stem_name}.{output_format}" # Simple name inside zip
                     zipf.write(stem_path, arcname=archive_name)
                     logger.info(f"Added '{archive_name}' to ZIP.")
+                    found_stems_count += 1
                 else:
-                    logger.warning(f"Requested stem '{stem_name}' not found in model output.")
+                    logger.warning(f"Requested stem '{stem_name}' not found in model output keys: {list(all_separated_stems.keys())}")
+
+            if found_stems_count == 0:
+                raise HTTPException(status_code=404, detail="None of the requested stems were found or generated successfully.")
 
         zip_buffer.seek(0) # Rewind buffer pointer
 
-        # Return the ZIP file
-        zip_filename = f"separated_stems_{os.path.splitext(file.filename)[0]}.zip"
+        # Return the ZIP file via StreamingResponse
+        zip_filename_download = f"{zip_filename_base}.zip"
+        logger.info(f"Sending ZIP file '{zip_filename_download}'")
         return StreamingResponse(
-            zip_buffer,
+            zip_buffer, # Pass the BytesIO buffer directly
             media_type="application/zip",
-            headers={'Content-Disposition': f'attachment; filename="{zip_filename}"'}
+            headers={'Content-Disposition': f'attachment; filename="{zip_filename_download}"'}
         )
 
     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
+        # Cleanup temporary stem files if they exist
+        for path in stem_output_paths.values(): cleanup_file(path)
+        # Close buffer just in case (though StreamingResponse should handle it)
+        # if zip_buffer and not zip_buffer.closed: zip_buffer.close()
+        if input_path: cleanup_file(input_path)
 
         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
+# ----------- END app.py -----------