songprocessor.py

import torch
import torchaudio
from demucs.pretrained import get_model
from demucs.apply import apply_model
from youtubedl import YouTubeAudioDownloader
import soundfile as sf
import numpy as np

# Use MPS (Apple's GPU acceleration) if available, otherwise CPU
device = 'mps' if torch.backends.mps.is_available() else 'cpu'

# Load model once
model = get_model('htdemucs').to(device)
model.eval()

def remove_drums(input_path, output_path):
    # Use soundfile to load mp3
    audio_data, sr = sf.read(input_path)
    # Convert to torch tensor and transpose to (channels, samples)
    wav = torch.from_numpy(audio_data.T).float()

    if wav.shape[0] == 1:
        wav = wav.repeat(2, 1)

    with torch.no_grad():
        sources = apply_model(
            model,
            wav[None].to(device),
            device=device,
            shifts=1,      # Lower for speed
            overlap=0.25
        )

    # Remove drums (index 0)
    no_drums = sources[0, 1:].sum(0).cpu()

    # Save as WAV (torchaudio can save WAV natively)
    output_path = output_path.replace('.mp3', '.wav')
    torchaudio.save(output_path, no_drums, sr)


#returns path of processed file based on youtube search
def download_and_process(search, temp_folder="temp_music", output_folder="output_music", processing_queue=None, position_in_queue=None):
    try:
        downloader = YouTubeAudioDownloader(output_dir=temp_folder)
        audio_path = downloader.download_from_search(search,output_format='mp3')
        print(f"Downloaded to {audio_path}")
        # Change extension to .wav since remove_drums saves as WAV
        output_path = f"{output_folder}/{audio_path.split('/')[-1]}".replace('.mp3', '.wav')
        remove_drums(audio_path, output_path)
        temp_song = None
        if processing_queue is not None and position_in_queue is not None and len(processing_queue) > position_in_queue:
            temp_song = processing_queue.pop(position_in_queue)
            temp_song['location'] = output_path
        return temp_song            #store user and search term
    except Exception as e:
        print(f"Failed to process '{search}': {e}")
        # Remove from queue on failure
        if processing_queue is not None and position_in_queue is not None and len(processing_queue) > position_in_queue:
            processing_queue.pop(position_in_queue)
        return None