Update app.py

app.py

@@ -1,20 +1,10 @@
 import gradio as gr
-import torch
 import os
-import spaces
+import subprocess
+from typing import Tuple, List, Dict
 from pydub import AudioSegment
-from typing import Tuple, Dict, List
-from demucs.apply import apply_model
-from demucs.separate import load_track
-from demucs.pretrained import get_model
-from demucs.audio import save_audio
 
-device: str = "cuda" if torch.cuda.is_available() else "cpu"
-
-# Define the inference function
-@spaces.GPU
 def inference(audio_file: str, model_name: str, vocals: bool, drums: bool, bass: bool, other: bool, mp3: bool, mp3_bitrate: int) -> Tuple[str, gr.HTML]:
-    separator = get_model(name=model_name)
     log_messages = []
 
     def stream_log(message):
@@ -24,77 +14,60 @@ def inference(audio_file: str, model_name: str, vocals: bool, drums: bool, bass:
     yield None, stream_log("Starting separation process...")
     yield None, stream_log(f"Loading audio file: {audio_file}")
     
-    # Check if audio_file is None
     if audio_file is None:
         yield None, stream_log("Error: No audio file provided")
         raise gr.Error("Please upload an audio file")
 
-    # Load the audio file with the correct samplerate and audio channels
-    try:
-        wav, sr = load_track(audio_file, samplerate=separator.samplerate, audio_channels=2)
-    except Exception as e:
-        yield None, stream_log(f"Error loading audio file: {str(e)}")
-        raise gr.Error(f"Failed to load audio file: {str(e)}")
-
-    # Check the number of channels and adjust if necessary
-    if wav.dim() == 1:
-        wav = wav.unsqueeze(0)  # Add channel dimension if mono
-    if wav.shape[0] == 1:
-        wav = wav.repeat(2, 1)  # If mono, duplicate to stereo
-    elif wav.shape[0] > 2:
-        wav = wav[:2]  # If more than 2 channels, keep only the first two
-
-    wav = wav.to(device)
-    
-    ref = wav.mean(0)
-    wav = (wav - ref.view(1, -1))
-    yield None, stream_log("Audio loaded successfully. Applying model...")
+    output_dir = os.path.join("separated", model_name, os.path.splitext(os.path.basename(audio_file))[0])
+    os.makedirs(output_dir, exist_ok=True)
+
+    # Construct the Demucs command
+    cmd = [
+        "python", "-m", "demucs", 
+        "--out", output_dir,
+        "-n", model_name,
+        audio_file
+    ]
+
+    yield None, stream_log(f"Running Demucs command: {' '.join(cmd)}")
 
-    # Use apply_model as a standalone function
     try:
-        result = apply_model(separator, wav.to(device), device=device)
-        yield None, stream_log(f"Model application result type: {type(result)}")
-        yield None, stream_log(f"Model application result shape: {result.shape if hasattr(result, 'shape') else 'N/A'}")
+        # Run the Demucs command
+        process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True)
         
-        if isinstance(result, tuple) and len(result) == 2:
-            sources, _ = result
-        elif isinstance(result, torch.Tensor):
-            sources = result
-        else:
-            raise ValueError(f"Unexpected result type from apply_model: {type(result)}")
+        # Stream the output
+        for line in process.stdout:
+            yield None, stream_log(line.strip())
         
-        yield None, stream_log(f"Sources shape: {sources.shape}")
-    except ValueError as e:
-        yield None, stream_log(f"Error applying model: {str(e)}")
-        yield None, stream_log(f"Separator sources: {separator.sources}")
-        yield None, stream_log(f"WAV shape: {wav.shape}")
-        yield None, stream_log(f"Separator model: {separator.__class__.__name__}")
-        yield None, stream_log(f"Separator config: {separator.config}")
-        raise gr.Error(f"Failed to apply model: {str(e)}. This might be due to incompatible audio format or model configuration.")
+        # Wait for the process to complete
+        process.wait()
+
+        if process.returncode != 0:
+            error_output = process.stderr.read()
+            yield None, stream_log(f"Error: Demucs command failed with return code {process.returncode}")
+            yield None, stream_log(f"Error output: {error_output}")
+            raise gr.Error(f"Demucs separation failed. Check the logs for details.")
+
     except Exception as e:
-        yield None, stream_log(f"Unexpected error applying model: {str(e)}")
-        raise gr.Error(f"An unexpected error occurred while applying the model: {str(e)}")
+        yield None, stream_log(f"Unexpected error: {str(e)}")
+        raise gr.Error(f"An unexpected error occurred: {str(e)}")
 
-    # Process the sources
-    sources = [source * ref.view(1, -1) + ref.view(1, -1) for source in sources]
-    
-    yield None, stream_log("Model applied. Processing stems...")
+    yield None, stream_log("Separation completed. Processing stems...")
 
-    output_dir: str = os.path.join("separated", model_name, os.path.splitext(os.path.basename(audio_file))[0])
-    os.makedirs(output_dir, exist_ok=True)
     stems: Dict[str, str] = {}
-    for stem, source in zip(separator.sources, sources):
-        stem_path: str = os.path.join(output_dir, f"{stem}.wav")
-        save_audio(source, stem_path, separator.samplerate)
-        stems[stem] = stem_path
-        yield None, stream_log(f"Saved {stem} stem")
+    for stem in ["vocals", "drums", "bass", "other"]:
+        stem_path = os.path.join(output_dir, model_name, f"{stem}.wav")
+        if os.path.exists(stem_path):
+            stems[stem] = stem_path
+            yield None, stream_log(f"Found {stem} stem")
 
-    selected_stems: List[str] = [stems[stem] for stem, include in zip(["vocals", "drums", "bass", "other"], [vocals, drums, bass, other]) if include]
+    selected_stems: List[str] = [stems[stem] for stem in stems if locals()[stem]]
     if not selected_stems:
         raise gr.Error("Please select at least one stem to mix.")
 
     output_file: str = os.path.join(output_dir, "mixed.wav")
     yield None, stream_log("Mixing selected stems...")
+    
     if len(selected_stems) == 1:
         os.rename(selected_stems[0], output_file)
     else: