audio_processing.py

import os
import numpy as np
from demucs import pretrained
from demucs.apply import apply_model, BagOfModels
import torch
import soundfile as sf
import librosa

def process_in_chunks(model, wav, chunk_size=44100*15, overlap=0.1, progress_callback=None, status_callback=None):
    """Process audio in chunks to save memory"""
    device = "cuda" if torch.cuda.is_available() else "cpu"
    total_length = wav.shape[-1]
    
    # Process first chunk to get output shape
    with torch.no_grad():
        first_chunk = wav[..., :chunk_size]
        first_stems = apply_model(
            model, 
            first_chunk.to(device), 
            shifts=1, 
            split=True, 
            overlap=0.1,
            progress=False
        )
        
    # Initialize output tensor
    n_stems = first_stems.shape[1]  # Usually 4 for [drums, bass, other, vocals]
    final_stems = torch.zeros((1, n_stems, 2, total_length), dtype=torch.float32)
    final_stems[..., :first_stems.shape[-1]] = first_stems
    
    # Process remaining chunks
    overlap_frames = int(chunk_size * overlap)
    current_pos = chunk_size - overlap_frames
    
    for start in range(chunk_size - overlap_frames, total_length, chunk_size - overlap_frames):
        if status_callback:
            progress = (start / total_length) * 100
            status_callback(f"Processing: {progress:.1f}%")
        if progress_callback:
            progress_callback(int((start / total_length) * 50))
            
        end = min(start + chunk_size, total_length)
        if end - start < chunk_size // 2:  # Skip if chunk is too small
            break
            
        # Process chunk
        with torch.no_grad():
            chunk = wav[..., start:end]
            stems = apply_model(
                model,
                chunk.to(device),
                shifts=1,
                split=True,
                overlap=0.1,
                progress=False
            )
            
        # Apply crossfade
        chunk_length = stems.shape[-1]
        fade_in = torch.linspace(0, 1, overlap_frames).view(1, 1, 1, -1)
        fade_out = torch.linspace(1, 0, overlap_frames).view(1, 1, 1, -1)
        
        # Crossfade region
        final_stems[..., current_pos:current_pos + overlap_frames] *= fade_out
        final_stems[..., current_pos:current_pos + overlap_frames] += stems[..., :overlap_frames] * fade_in
        
        # Copy remaining part
        if current_pos + overlap_frames < total_length:
            end_pos = min(current_pos + chunk_length, total_length)
            final_stems[..., current_pos + overlap_frames:end_pos] = stems[..., overlap_frames:end_pos-current_pos]
        
        current_pos += chunk_size - overlap_frames
        torch.cuda.empty_cache() if torch.cuda.is_available() else None
    
    return final_stems

def separate_stems(input_file, output_dir, progress_callback=None, status_callback=None):
    """
    Separates audio into stems using the advanced HTQ model of Demucs.

    Args:
        input_file (str): Path to the input audio file.
        output_dir (str): Directory to save the separated stems.
    """
    try:
        if status_callback:
            status_callback("Initializing...")
        
        # Check GPU availability and provide detailed feedback
        cuda_available = torch.cuda.is_available()
        if cuda_available:
            gpu_count = torch.cuda.device_count()
            gpu_name = torch.cuda.get_device_name(0)
            print(f"CUDA GPU detected: {gpu_name} (Device {0})")
            if status_callback:
                status_callback(f"GPU detected: {gpu_name}")
        else:
            print("CUDA GPU not available - using CPU (this will be slower)")
            print("To enable GPU acceleration, install CUDA-enabled PyTorch:")
            print("pip install torch==2.1.0+cu118 torchvision==0.16.0+cu118 torchaudio==2.1.0+cu118 --index-url https://download.pytorch.org/whl/cu118")
            if status_callback:
                status_callback("Using CPU - GPU not detected")
        
        # Ensure the output directory exists
        os.makedirs(output_dir, exist_ok=True)
        print(f"Output directory created: {output_dir}")

        # Load the pre-trained HTDemucs model with optimized settings
        model = pretrained.get_model("htdemucs_ft")  # Using fine-tuned model
        if cuda_available:
            model.cuda()
            print(f"Model loaded on GPU: {torch.cuda.get_device_name(0)}")
        else:
            print("Model loaded on CPU")
        print(f"Loaded model: htdemucs_ft, Model type: {type(model)}")

        if status_callback:
            status_callback("Loading audio file...")
        # Load and prepare audio with optimal settings
        print(f"Loading audio file: {input_file}")
        wav, sr = librosa.load(input_file, sr=44100, mono=False, duration=None)
        
        # Enhanced preprocessing
        if wav.ndim == 1:
            wav = np.expand_dims(wav, axis=0)
        
        # Improved normalization with headroom
        wav = torch.tensor(wav, dtype=torch.float32)
        if wav.dim() == 2:
            wav = wav.unsqueeze(0)
        
        # Dynamic range compression for better vocal presence
        wav = wav / (wav.abs().max() * 1.1)  # Leave headroom
        
        print(f"Audio file loaded: {input_file}, Sample rate: {sr}")
        print(f"Audio tensor shape: {wav.shape}")

        # Move to GPU if available
        if torch.cuda.is_available():
            wav = wav.cuda()

        if status_callback:
            status_callback("Performing stem separation...")
        # Perform stem separation with optimized parameters
        print(f"Performing stem separation")
        sources = process_in_chunks(model, wav, 
                                  progress_callback=progress_callback,
                                  status_callback=status_callback)
        print(f"Sources shape: {sources.shape}")  # Add debug print
        stems = torch.split(sources.squeeze(0), 1, dim=0)
        print(f"Stem separation completed")

        if status_callback:
            status_callback("Saving stems...")

        # Extract base name from input file
        base_name = os.path.splitext(os.path.basename(input_file))[0].replace(" ", "_")

        # Save stems with enhanced quality
        stem_names = ["drums", "bass", "other", "vocals"]
        base_name = os.path.splitext(os.path.basename(input_file))[0]
        for i, stem in enumerate(stems):
            if status_callback:
                status_callback(f"Saving {stem_names[i]}...")
            if progress_callback:
                progress_callback(50 + int((i / len(stem_names)) * 50))  # Last 50% of progress
                
            output_path = os.path.join(output_dir, f"{base_name}_{stem_names[i]}.wav")
            stem = stem.squeeze(0).cpu().numpy()
            
            # Ensure proper shape and apply final processing
            if stem.shape[0] == 2:
                stem = stem.T
            
            # Apply subtle enhancement for vocals
            if stem_names[i] == "vocals":
                stem = enhance_vocals(stem, sr)
                
            print(f"Saving stem: {stem_names[i]} to {output_path}, shape: {stem.shape}")
            sf.write(output_path, stem, sr, subtype='FLOAT')
            print(f"Stem saved: {output_path}")

        if status_callback:
            status_callback("All stems saved successfully!")
        if progress_callback:
            progress_callback(100)

        print(f"All stems saved in: {output_dir}")

    except Exception as e:
        if status_callback:
            status_callback(f"Error: {str(e)}")
        print(f"Error occurred: {str(e)}")
        raise

def segment_size(audio_length):
    """Calculate optimal segment size based on audio length"""
    base_segment = 44100 * 10  # 10 seconds base
    return min(base_segment, audio_length // 4)

def enhance_vocals(vocals, sr):
    """Apply subtle enhancements to vocal stem"""
    # Gentle normalization to preserve dynamics
    max_val = np.abs(vocals).max()
    if max_val > 0:
        vocals = vocals / max_val * 0.95
    
    # Preserve stereo field
    if vocals.ndim == 2:
        vocals = vocals.copy()  # Prevent modifying original array
        # Apply mild stereo enhancement
        mid = np.mean(vocals, axis=1, keepdims=True)
        side = vocals - mid
        vocals = mid + side * 1.2  # Enhance stereo width by 20%
    
    return vocals