Update app.py

app.py

@@ -1,397 +1,215 @@
 import gradio as gr
+import torch
+import torchaudio
+import numpy as np
+from pathlib import Path
 import tempfile
 import os
-import subprocess
-import shutil
-from pathlib import Path
-import logging
-from typing import List, Tuple, Dict
-import json
 
-# Configure logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
+# Check if CUDA is available
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Using device: {device}")
 
-class StemSeparator:
-    """Modern stem separation with multiple model support"""
-    
-    def __init__(self):
-        self.supported_models = {
-            "htdemucs": {
-                "command": "demucs",
-                "stems": 4,
-                "description": "HTDemucs - High quality 4-stem separation"
-            },
-            "htdemucs_ft": {
-                "command": "demucs",
-                "model_name": "htdemucs_ft",
-                "stems": 4,
-                "description": "HTDemucs Fine-tuned - Enhanced 4-stem separation"
-            },
-            "htdemucs_6s": {
-                "command": "demucs",
-                "model_name": "htdemucs_6s", 
-                "stems": 6,
-                "description": "HTDemucs 6-stem - Bass, Drums, Vocals, Other, Guitar, Piano"
-            },
-            "mdx": {
-                "command": "demucs",
-                "model_name": "mdx",
-                "stems": 4,
-                "description": "MDX - Optimized for vocal separation"
-            },
-            "mdx_extra": {
-                "command": "demucs",
-                "model_name": "mdx_extra",
-                "stems": 4,
-                "description": "MDX Extra - Enhanced vocal separation"
-            },
-            "spleeter_4stems": {
-                "command": "spleeter",
-                "model_name": "spleeter:4stems-waveform",
-                "stems": 4,
-                "description": "Spleeter 4-stem - Vocals, Bass, Drums, Other"
-            },
-            "spleeter_5stems": {
-                "command": "spleeter", 
-                "model_name": "spleeter:5stems-waveform",
-                "stems": 5,
-                "description": "Spleeter 5-stem - Vocals, Bass, Drums, Piano, Other"
-            }
-        }
-    
-    def check_dependencies(self) -> Dict[str, bool]:
-        """Check if required tools are installed"""
-        dependencies = {}
-        
-        # Check demucs
-        try:
-            result = subprocess.run(["python", "-m", "demucs", "--help"], 
-                                  capture_output=True, text=True, timeout=10)
-            dependencies["demucs"] = result.returncode == 0
-        except (subprocess.TimeoutExpired, FileNotFoundError):
-            dependencies["demucs"] = False
-            
-        # Check spleeter
-        try:
-            result = subprocess.run(["spleeter", "--help"], 
-                                  capture_output=True, text=True, timeout=10)
-            dependencies["spleeter"] = result.returncode == 0
-        except (subprocess.TimeoutExpired, FileNotFoundError):
-            dependencies["spleeter"] = False
-            
-        return dependencies
+def separate_stems(audio_file, model_name="htdemucs"):
+    """
+    Separate audio stems using Demucs
+    """
+    if audio_file is None:
+        return None, None, None, None, "❌ Please upload an audio file"
     
-    def separate_audio(self, audio_file: str, model_choice: str) -> Tuple[List[str], str]:
-        """Separate audio into stems using the selected model"""
-        if not audio_file:
-            return [], "❌ No audio file provided"
-            
-        if model_choice not in self.supported_models:
-            return [], f"❌ Unsupported model: {model_choice}"
-            
-        model_config = self.supported_models[model_choice]
-        
-        try:
-            with tempfile.TemporaryDirectory() as temp_dir:
-                temp_path = Path(temp_dir)
-                
-                # Copy input file to temp directory with proper extension
-                input_file = Path(audio_file)
-                temp_input = temp_path / f"input{input_file.suffix}"
-                shutil.copy2(audio_file, temp_input)
-                
-                logger.info(f"Processing {temp_input} with {model_choice}")
-                
-                # Build command based on model type
-                if model_config["command"] == "demucs":
-                    command = self._build_demucs_command(temp_input, temp_path, model_config)
-                elif model_config["command"] == "spleeter":
-                    command = self._build_spleeter_command(temp_input, temp_path, model_config)
-                else:
-                    return [], f"❌ Unknown command type: {model_config['command']}"
-                
-                # Execute separation
-                logger.info(f"Running command: {' '.join(command)}")
-                result = subprocess.run(
-                    command, 
-                    capture_output=True, 
-                    text=True, 
-                    timeout=300,  # 5 minute timeout
-                    cwd=temp_dir
-                )
-                
-                if result.returncode != 0:
-                    error_msg = f"❌ Separation failed: {result.stderr}"
-                    logger.error(error_msg)
-                    return [], error_msg
+    try:
+        # Import demucs modules
+        from demucs.pretrained import get_model
+        from demucs.apply import apply_model
+        from demucs.audio import save_audio
+        
+        # Load the model
+        model = get_model(model_name)
+        model.to(device)
+        model.eval()
+        
+        # Load audio
+        wav, sr = torchaudio.load(audio_file)
+        
+        # Ensure stereo
+        if wav.shape[0] == 1:
+            wav = wav.repeat(2, 1)
+        elif wav.shape[0] > 2:
+            wav = wav[:2]
+        
+        # Resample if necessary
+        if sr != model.samplerate:
+            resampler = torchaudio.transforms.Resample(sr, model.samplerate)
+            wav = resampler(wav)
+            sr = model.samplerate
+        
+        # Move to device
+        wav = wav.to(device)
+        
+        # Apply the model
+        with torch.no_grad():
+            sources = apply_model(model, wav.unsqueeze(0))
+        
+        # Get source names
+        source_names = model.sources
+        
+        # Save separated sources
+        output_files = {}
+        with tempfile.TemporaryDirectory() as temp_dir:
+            for i, source in enumerate(source_names):
+                output_path = os.path.join(temp_dir, f"{source}.wav")
+                save_audio(sources[0, i], output_path, sr)
                 
-                # Collect output stems
-                stems = self._collect_stems(temp_path, model_choice)
-                
-                if not stems:
-                    return [], "❌ No stems were generated"
-                
-                success_msg = f"✅ Successfully separated into {len(stems)} stems"
-                logger.info(success_msg)
-                return stems, success_msg
-                
-        except subprocess.TimeoutExpired:
-            return [], "❌ Process timed out - file may be too large"
-        except Exception as e:
-            error_msg = f"❌ Error during separation: {str(e)}"
-            logger.error(error_msg)
-            return [], error_msg
-    
-    def _build_demucs_command(self, input_file: Path, output_dir: Path, model_config: Dict) -> List[str]:
-        """Build demucs command"""
-        command = ["python", "-m", "demucs"]
+                # Read the saved file for Gradio
+                output_files[source] = output_path
         
-        if "model_name" in model_config:
-            command.extend(["-n", model_config["model_name"]])
+        # Return the separated stems (assuming 4 stems: drums, bass, other, vocals)
+        stems = [None] * 4
+        status_msg = f"✅ Successfully separated into {len(source_names)} stems"
         
-        command.extend([
-            "-o", str(output_dir),
-            "--filename", "{track}/{stem}.{ext}",  # Organized output structure
-            str(input_file)
-        ])
+        for i, source in enumerate(source_names[:4]):  # Limit to 4 for UI
+            if source in output_files:
+                stems[i] = output_files[source]
         
-        return command
-    
-    def _build_spleeter_command(self, input_file: Path, output_dir: Path, model_config: Dict) -> List[str]:
-        """Build spleeter command"""
-        model_name = model_config.get("model_name", "spleeter:4stems-waveform")
-        
-        command = [
-            "spleeter", "separate",
-            "-p", model_name,
-            "-o", str(output_dir),
-            "--filename_format", "{instrument}.{codec}",
-            str(input_file)
-        ]
-        
-        return command
-    
-    def _collect_stems(self, output_dir: Path, model_choice: str) -> List[str]:
-        """Collect generated stem files"""
-        stems = []
+        return tuple(stems + [status_msg])
         
-        # Search for audio files in output directory
-        for audio_file in output_dir.rglob("*.wav"):
-            if audio_file.is_file() and audio_file.stat().st_size > 0:
-                # Copy to a permanent location that Gradio can access
-                permanent_path = self._copy_to_permanent_location(audio_file)
-                if permanent_path:
-                    stems.append(permanent_path)
-        
-        # Also check for other common audio formats
-        for ext in ["*.mp3", "*.flac", "*.m4a"]:
-            for audio_file in output_dir.rglob(ext):
-                if audio_file.is_file() and audio_file.stat().st_size > 0:
-                    permanent_path = self._copy_to_permanent_location(audio_file)
-                    if permanent_path:
-                        stems.append(permanent_path)
-        
-        return sorted(stems)
-    
-    def _copy_to_permanent_location(self, temp_file: Path) -> str:
-        """Copy temporary file to permanent location for Gradio"""
-        try:
-            # Create output directory if it doesn't exist
-            output_dir = Path("./separated_stems")
-            output_dir.mkdir(exist_ok=True)
-            
-            # Generate unique filename
-            import time
-            timestamp = int(time.time() * 1000)
-            permanent_file = output_dir / f"{temp_file.stem}_{timestamp}{temp_file.suffix}"
-            
-            shutil.copy2(temp_file, permanent_file)
-            return str(permanent_file)
-        except Exception as e:
-            logger.error(f"Failed to copy {temp_file}: {e}")
-            return None
-
-# Initialize separator
-separator = StemSeparator()
+    except Exception as e:
+        error_msg = f"❌ Error during separation: {str(e)}"
+        return None, None, None, None, error_msg
 
-def get_available_models() -> List[Tuple[str, str]]:
-    """Get list of available models based on installed dependencies"""
-    deps = separator.check_dependencies()
-    available_models = []
-    
-    for model_id, config in separator.supported_models.items():
-        if config["command"] in deps and deps[config["command"]]:
-            label = f"{model_id} ({config['stems']} stems) - {config['description']}"
-            available_models.append((label, model_id))
+def create_hf_interface():
+    """Create Hugging Face Spaces compatible interface"""
     
-    if not available_models:
-        available_models = [("No models available - install demucs or spleeter", "none")]
-    
-    return available_models
-
-def separate_stems_ui(audio_file: str, model_choice: str) -> Tuple[List[str], str]:
-    """UI wrapper for stem separation"""
-    if model_choice == "none":
-        return [], "❌ Please install demucs and/or spleeter first"
-    
-    stems, message = separator.separate_audio(audio_file, model_choice)
-    return stems, message
-
-def create_audio_gallery(stems: List[str]) -> List[gr.Audio]:
-    """Create audio components for each stem"""
-    if not stems:
-        return []
-    
-    audio_components = []
-    for i, stem_path in enumerate(stems):
-        stem_name = Path(stem_path).stem
-        audio_comp = gr.Audio(
-            value=stem_path,
-            label=f"Stem {i+1}: {stem_name}",
-            interactive=False,
-            show_download_button=True
-        )
-        audio_components.append(audio_comp)
-    
-    return audio_components
-
-# Create Gradio interface
-def create_interface():
     with gr.Blocks(
-        title="🎵 Advanced Music Stem Separator",
+        title="🎵 Music Stem Separator",
         theme=gr.themes.Soft(),
-        css="""
-        .audio-container { margin: 10px 0; }
-        .status-success { color: #22c55e; font-weight: bold; }
-        .status-error { color: #ef4444; font-weight: bold; }
-        """
-    ) as demo:
+    ) as interface:
         
         gr.Markdown("""
-        # 🎵 Advanced Music Stem Separator
-        
-        Separate music into individual stems (vocals, instruments, etc.) using state-of-the-art AI models.
-        Supports up to 6 stems depending on the model chosen.
+        # 🎵 Music Stem Separator
         
-        **Supported Models:**
-        - **Demucs Models**: HTDemucs, HTDemucs-FT, HTDemucs-6s, MDX, MDX-Extra
-        - **Spleeter Models**: 4-stem and 5-stem separation
+        Separate music into individual stems using **Meta's Demucs** model.
+        Upload an audio file and get separated tracks for **drums**, **bass**, **other instruments**, and **vocals**.
         
-        **Requirements**: Install `demucs` and/or `spleeter` packages
+        ⚡ **Powered by Demucs** - State-of-the-art source separation
         """)
         
         with gr.Row():
-            with gr.Column(scale=2):
+            with gr.Column():
+                # Input
                 audio_input = gr.Audio(
                     type="filepath",
-                    label="🎼 Upload Audio File",
-                    info="Supported formats: WAV, MP3, FLAC, M4A"
+                    label="🎼 Upload Music File",
+                    info="Supported: MP3, WAV, FLAC (max 10MB recommended)"
                 )
                 
-                model_dropdown = gr.Dropdown(
-                    choices=get_available_models(),
-                    value=get_available_models()[0][1] if get_available_models() else "none",
-                    label="🧠 Separation Model",
-                    info="Choose the AI model for stem separation"
+                model_choice = gr.Dropdown(
+                    choices=[
+                        ("HTDemucs (4 stems)", "htdemucs"),
+                        ("HTDemucs FT (4 stems)", "htdemucs_ft"), 
+                        ("MDX Extra (4 stems)", "mdx_extra")
+                    ],
+                    value="htdemucs",
+                    label="🤖 Model",
+                    info="Choose separation model"
                 )
                 
                 separate_btn = gr.Button(
-                    "🎛️ Separate Stems",
+                    "🎛️ Separate Stems", 
                     variant="primary",
                     size="lg"
                 )
             
-            with gr.Column(scale=1):
+            with gr.Column():
                 gr.Markdown("""
-                ### ℹ️ Model Info
-                - **4-stem**: Vocals, Bass, Drums, Other
-                - **5-stem**: + Piano
-                - **6-stem**: + Guitar
+                ### ℹ️ About Stem Separation
                 
-                ### 💡 Tips
+                **What you'll get:**
+                - 🥁 **Drums**: Percussion and rhythm
+                - 🎸 **Bass**: Bass lines and low frequencies  
+                - 🎹 **Other**: Instruments, synths, effects
+                - 🎤 **Vocals**: Lead and backing vocals
+                
+                **Tips:**
                 - Higher quality input = better separation
-                - Processing time varies by model and file length
-                - Results will appear below after processing
+                - Processing takes 1-3 minutes depending on length
+                - Works best with modern pop/rock music
                 """)
         
-        # Status display
-        status_display = gr.Textbox(
+        # Status
+        status_output = gr.Textbox(
             label="Status",
             interactive=False,
-            visible=True
+            show_label=True
         )
         
-        # Dynamic audio outputs
-        stems_state = gr.State([])
-        audio_outputs = gr.Column(visible=False)
+        # Output stems
+        gr.Markdown("### 🎶 Separated Stems")
         
-        def process_and_display(audio_file, model_choice):
-            if not audio_file:
-                return [], "❌ Please upload an audio file", gr.Column(visible=False)
-            
-            # Process the audio
-            stems, message = separate_stems_ui(audio_file, model_choice)
-            
-            # Create audio components
-            if stems:
-                with gr.Column() as output_col:
-                    gr.Markdown(f"### 🎶 Separated Stems ({len(stems)} files)")
-                    for i, stem_path in enumerate(stems):
-                        stem_name = Path(stem_path).stem.replace("_", " ").title()
-                        gr.Audio(
-                            value=stem_path,
-                            label=f"🎵 {stem_name}",
-                            show_download_button=True,
-                            interactive=False
-                        )
-                return stems, message, gr.Column(visible=True)
-            else:
-                return [], message, gr.Column(visible=False)
+        with gr.Row():
+            drums_output = gr.Audio(
+                label="🥁 Drums",
+                interactive=False,
+                show_download_button=True
+            )
+            bass_output = gr.Audio(
+                label="🎸 Bass", 
+                interactive=False,
+                show_download_button=True
+            )
+        
+        with gr.Row():
+            other_output = gr.Audio(
+                label="🎹 Other",
+                interactive=False, 
+                show_download_button=True
+            )
+            vocals_output = gr.Audio(
+                label="🎤 Vocals",
+                interactive=False,
+                show_download_button=True
+            )
         
+        # Connect the interface
         separate_btn.click(
-            fn=process_and_display,
-            inputs=[audio_input, model_dropdown],
-            outputs=[stems_state, status_display, audio_outputs],
+            fn=separate_stems,
+            inputs=[audio_input, model_choice],
+            outputs=[
+                drums_output,
+                bass_output, 
+                other_output,
+                vocals_output,
+                status_output
+            ],
             show_progress=True
         )
         
-        # Dependency check display
-        with gr.Accordion("🔧 System Status", open=False):
-            def check_system():
-                deps = separator.check_dependencies()
-                status_text = "**Dependency Status:**\n"
-                for tool, available in deps.items():
-                    status = "✅ Available" if available else "❌ Not installed"
-                    status_text += f"- {tool}: {status}\n"
-                
-                if not any(deps.values()):
-                    status_text += "\n**Installation Instructions:**\n"
-                    status_text += "```bash\n"
-                    status_text += "# Install Demucs (recommended)\n"
-                    status_text += "pip install demucs\n\n"
-                    status_text += "# Install Spleeter (alternative)\n"
-                    status_text += "pip install spleeter tensorflow\n"
-                    status_text += "```"
-                
-                return status_text
-            
-            system_status = gr.Markdown(value=check_system())
-            
-            gr.Button("🔄 Refresh Status").click(
-                fn=check_system,
-                outputs=system_status
-            )
+        # Examples
+        gr.Markdown("### 🎵 Try with example audio")
+        gr.Examples(
+            examples=[
+                ["example1.wav", "htdemucs"],
+                ["example2.mp3", "htdemucs"],
+            ],
+            inputs=[audio_input, model_choice],
+            outputs=[drums_output, bass_output, other_output, vocals_output, status_output],
+            fn=separate_stems,
+            cache_examples=False
+        )
+        
+        gr.Markdown("""
+        ---
+        **Note**: This space uses Meta's Demucs for stem separation. Processing time depends on audio length and available compute resources.
+        
+        **Limitations**: 
+        - Max file size: ~50MB
+        - Processing time: 1-5 minutes
+        - Works best with clear, well-produced music
+        """)
     
-    return demo
+    return interface
 
-# Launch the interface
+# Create and launch the interface
 if __name__ == "__main__":
-    demo = create_interface()
-    demo.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        share=False,
-        show_error=True,
-        debug=True
-    )
+    demo = create_hf_interface()
+    demo.launch()