ASesYusuf1/SESA_Audio_Separation

1import torch2import torchaudio3import os4import gc5 6def frequency_blend_phases(phase1, phase2, freq_bins, low_cutoff=500, high_cutoff=5000, base_factor=0.25, scale_factor=1.85):7    if phase1.shape != phase2.shape:8        raise ValueError("phase1 and phase2 must have the same shape.")9    if len(freq_bins) != phase1.shape[0]:10        raise ValueError("freq_bins must have the same length as the number of frequency bins in phase1 and phase2.")11    if low_cutoff >= high_cutoff:12        raise ValueError("low_cutoff must be less than high_cutoff.")13 14    blended_phase = torch.zeros_like(phase1)15    blend_factors = torch.zeros_like(freq_bins)16 17    blend_factors[freq_bins < low_cutoff] = base_factor18    blend_factors[freq_bins > high_cutoff] = base_factor + scale_factor19 20    in_range_mask = (freq_bins >= low_cutoff) & (freq_bins <= high_cutoff)21    blend_factors[in_range_mask] = base_factor + scale_factor * (22        (freq_bins[in_range_mask] - low_cutoff) / (high_cutoff - low_cutoff)23    )24 25    for i in range(phase1.shape[0]):26        blended_phase[i, :] = (1 - blend_factors[i]) * phase1[i, :] + blend_factors[i] * phase2[i, :]27 28    blended_phase = torch.remainder(blended_phase + torch.pi, 2 * torch.pi) - torch.pi29 30    return blended_phase31 32def transfer_magnitude_phase(source_file, target_file, output_folder, transfer_magnitude=False, transfer_phase=True, 33                              low_cutoff=500, high_cutoff=9000, scale_factor=1.4, output_format='flac'):34    target_name, target_ext = os.path.splitext(os.path.basename(target_file))35    36    target_name = target_name.replace("_other", "").replace("_vocals", "").replace("_instrumental", "")37    target_name = target_name.replace("_Other", "").replace("_Vocals", "").replace("_Instrumental", "").strip()38    39    ext = '.flac' if output_format == 'flac' else '.wav'40    output_file = os.path.join(output_folder, f"{target_name} (Fixed Instrumental){ext}")41 42    print(f"Phase Fixing: {os.path.basename(target_file)}...")43    source_waveform, source_sr = torchaudio.load(source_file)44    target_waveform, target_sr = torchaudio.load(target_file)45 46    if source_sr != target_sr:47        raise ValueError("Sample rates of source and target audio files must match.")48 49    n_fft = 204850    hop_length = 51251    window = torch.hann_window(n_fft)52 53    source_stfts = torch.stft(source_waveform, n_fft=n_fft, hop_length=hop_length, window=window, return_complex=True, pad_mode="reflect")54    target_stfts = torch.stft(target_waveform, n_fft=n_fft, hop_length=hop_length, window=window, return_complex=True, pad_mode="reflect")55 56    freqs = torch.linspace(0, source_sr // 2, steps=n_fft // 2 + 1)57 58    modified_stfts = []59    for source_stft, target_stft in zip(source_stfts, target_stfts):60        source_mag, source_phs = torch.abs(source_stft), torch.angle(source_stft)61        target_mag, target_phs = torch.abs(target_stft), torch.angle(target_stft)62 63        modified_stft = target_stft.clone()64        if transfer_magnitude:65            modified_stft = source_mag * torch.exp(1j * torch.angle(modified_stft))66 67        if transfer_phase:68            blended_phase = frequency_blend_phases(target_phs, source_phs, freqs, low_cutoff, high_cutoff, scale_factor=scale_factor)69            modified_stft = torch.abs(modified_stft) * torch.exp(1j * blended_phase)70 71        modified_stfts.append(modified_stft)72 73    modified_waveform = torch.istft(74        torch.stack(modified_stfts),75        n_fft=n_fft,76        hop_length=hop_length,77        window=window,78        length=source_waveform.size(1)79    )80 81    if output_format == 'flac':82        torchaudio.save(output_file, modified_waveform, target_sr, format="flac", bits_per_sample=16)83    else:84        torchaudio.save(output_file, modified_waveform, target_sr)85    86    print(f"Saved: {output_file}")87    return output_file88 89def process_phase_fix(source_file, target_file, output_folder, low_cutoff=500, high_cutoff=9000, 90                      scale_factor=1.4, output_format='flac'):91    os.makedirs(output_folder, exist_ok=True)92    93    try:94        output_file = transfer_magnitude_phase(95            source_file=source_file,96            target_file=target_file,97            output_folder=output_folder,98            transfer_magnitude=False,99            transfer_phase=True,100            low_cutoff=low_cutoff,101            high_cutoff=high_cutoff,102            scale_factor=scale_factor,103            output_format=output_format104        )105        gc.collect()106        torch.cuda.empty_cache() if torch.cuda.is_available() else None107        return output_file, "Phase fix completed successfully!"108    except Exception as e:109        return None, f"Error during phase fix: {str(e)}"110 111SOURCE_MODELS = [112    'VOCALS-MelBand-Roformer (by Becruily)',113    'VOCALS-Mel-Roformer big beta 4 (by unwa)',114    'VOCALS-Melband-Roformer BigBeta5e (by unwa)',115    'VOCALS-big_beta6 (by Unwa)',116    'VOCALS-big_beta6X (by Unwa)',117    'VOCALS-MelBand-Roformer (by KimberleyJSN)',118    'VOCALS-MelBand-Roformer Kim FT (by Unwa)',119    'VOCALS-MelBand-Roformer Kim FT 2 (by Unwa)',120    'VOCALS-MelBand-Roformer Kim FT 2 Blendless (by unwa)',121    'VOCALS-Mel-Roformer FT 3 Preview (by unwa)',122    'VOCALS-BS-Roformer_1296 (by viperx)',123    'VOCALS-BS-Roformer_1297 (by viperx)',124    'VOCALS-BS-RoformerLargev1 (by unwa)',125    'bs_roformer_revive (by unwa)'126]127 128TARGET_MODELS = [129    'INST-MelBand-Roformer (by Becruily)',130    'INST-Mel-Roformer v1 (by unwa)',131    'INST-Mel-Roformer v2 (by unwa)',132    'inst_v1e (by unwa)',133    'INST-Mel-Roformer v1e+ (by unwa)',134    'Inst_GaboxV7 (by Gabox)',135    'INST-VOC-Mel-Roformer a.k.a. duality (by unwa)',136    'INST-VOC-Mel-Roformer a.k.a. duality v2 (by unwa)',137    'inst_gabox (by Gabox)',138    'inst_gaboxFlowersV10 (by Gabox)'139]140