Audio Transformation

# Demucs: State-of-the-art music source separation
# Separates audio into drums, bass, vocals, and other

import torch
from demucs import pretrained
from demucs.apply import apply_model
from demucs.audio import AudioFile, save_audio

# Load the model (htdemucs is the 4-stem hybrid model)
# Other options: htdemucs_6s (6 stems), htdemucs_ft (fine-tuned)
model = pretrained.get_model('htdemucs')
model.cpu()  # Use .cuda() for GPU acceleration

# Load audio file
# Demucs expects stereo 44.1kHz
audio_file = AudioFile("song.mp3")
waveform = audio_file.read(
    seek_time=0,       # Start position (seconds)
    duration=None,     # Duration (None = full file)
    streams=0          # 0 = first audio stream
)

# waveform shape: [channels, samples]
# For stereo 44.1kHz: [2, 44100 * duration]

# Apply separation with overlap for quality
sources = apply_model(
    model,
    waveform[None],    # Add batch dimension: [1, 2, samples]
    split=True,        # Split into chunks for memory
    overlap=0.25,      # 25% overlap between chunks
    progress=True      # Show progress bar
)[0]  # Remove batch dimension

# sources shape: [4, 2, samples]
# Index 0=drums, 1=bass, 2=other, 3=vocals

# Save separated stems
source_names = ['drums', 'bass', 'other', 'vocals']
for idx, name in enumerate(source_names):
    save_audio(
        sources[idx],
        f"output/{name}.wav",
        samplerate=model.samplerate
    )
    print(f"Saved {name}.wav")

# For just vocals (common use case):
vocals = sources[3]  # Index 3 is vocals
instrumental = sources[0] + sources[1] + sources[2]  # Everything else
save_audio(instrumental, "instrumental.wav", samplerate=model.samplerate)