File size: 3,798 Bytes
1e9fec0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
"""
Stem Separation Module
Uses Demucs for audio source separation
"""

import subprocess
from pathlib import Path
import sys


def separate_stems(input_file, output_dir, two_stem=True):
    """
    Separate audio into stems using Demucs (CPU mode)
    
    Args:
        input_file: Path to input audio file
        output_dir: Directory to save stems
        two_stem: If True, only separate vocals/instrumental (faster)
                  If False, separate into 6 stems (slower)
    
    Returns:
        Path to directory containing stems
    """
    
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)
    
    if two_stem:
        print(f"   Using 2-stem mode (vocal + instrumental)")
        print(f"   Estimated time: 2-4 minutes")
        
        cmd = [
            'demucs',
            '--two-stems=vocals',
            '-o', str(output_path),
            '--mp3',
            '--mp3-bitrate=320',
            '--jobs=0',  # Use all CPU cores
            input_file
        ]
    else:
        print(f"   Using 6-stem mode (vocal, drums, bass, guitar, keys, other)")
        print(f"   Estimated time: 8-12 minutes")
        
        cmd = [
            'demucs',
            '-n', 'htdemucs_6s',
            '-o', str(output_path),
            '--mp3',
            '--mp3-bitrate=320',
            '--jobs=0',
            input_file
        ]
    
    try:
        # Run Demucs with live output
        process = subprocess.Popen(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            universal_newlines=True,
            bufsize=1
        )
        
        # Print progress
        for line in process.stdout:
            line = line.strip()
            if line:
                # Filter for progress info
                if '%' in line or 'Separated' in line or 'Processing' in line:
                    print(f"   {line}")
        
        process.wait()
        
        if process.returncode != 0:
            raise Exception(f"Demucs exited with code {process.returncode}")
        
        # Locate output directory
        # Demucs structure: output_dir/model_name/track_name/stems
        song_stem = Path(input_file).stem
        
        # Try different possible structures
        possible_paths = [
            output_path / 'htdemucs' / song_stem,
            output_path / 'htdemucs_6s' / song_stem,
            output_path / song_stem
        ]
        
        for stems_dir in possible_paths:
            if stems_dir.exists():
                # Verify stems are present
                stem_files = list(stems_dir.glob('*.mp3'))
                if stem_files:
                    return stems_dir
        
        # If we get here, something went wrong
        raise Exception(f"Could not find output stems in {output_path}")
        
    except FileNotFoundError:
        print("\n[FAIL] Error: 'demucs' command not found")
        print("   Make sure Demucs is installed:")
        print("   pip install demucs")
        sys.exit(1)
        
    except Exception as e:
        raise Exception(f"Stem separation failed: {str(e)}")


def list_stems(stems_dir):
    """List all stem files in directory"""
    stems_dir = Path(stems_dir)
    return sorted(stems_dir.glob('*.mp3'))


if __name__ == "__main__":
    # Test module
    import sys
    
    if len(sys.argv) < 2:
        print("Usage: python stems.py <audio_file>")
        sys.exit(1)
    
    test_file = sys.argv[1]
    test_output = "./test_output"
    
    print(f"Testing stem separation on: {test_file}")
    result = separate_stems(test_file, test_output, two_stem=True)
    
    print(f"\nStems created in: {result}")
    print("\nFiles:")
    for stem in list_stems(result):
        print(f"  • {stem.name}")