app.py

import os
import crepe.crepe as crepe
import matplotlib.pyplot as plt
import librosa
import librosa.display
import numpy as np
import torch
import torchaudio
import pretty_midi
import gradio as gr
from scipy.io import wavfile

# Helper functions
def dynamic_range_compression_torch(x, C=1, clip_val=1e-5):
    return torch.log(torch.clamp(x, min=clip_val) * C)

def dynamic_range_decompression_torch(x, C=1):
    return torch.exp(x) / C

def spectral_normalize_torch(magnitudes):
    output = dynamic_range_compression_torch(magnitudes)
    return output

def spectral_de_normalize_torch(magnitudes):
    output = dynamic_range_decompression_torch(magnitudes)
    return output

def get_melspec(y: np.ndarray, sr: int = 16000, n_mels: int = 64,
                window_length: int = 1024, hop_length: int = 160, fmax: int = 8000, fmin: int = 0) -> torch.tensor:
    y = np.pad(y, (int((window_length - hop_length) / 2), int((window_length - hop_length) / 2)), mode="reflect")
    S = librosa.feature.melspectrogram(y=y, sr=sr, n_mels=n_mels,
                                       n_fft=window_length,
                                       hop_length=hop_length,
                                       win_length=window_length, window='hann',
                                       fmin=fmin, fmax=fmax,
                                       center=False, pad_mode="reflect")
    S_norm = spectral_normalize_torch(torch.tensor(S))
    return S_norm

def get_f0_melspec(frequency: np.ndarray, confidence: np.ndarray = None,
                   n_mels: int = 64, fmax: int = 8000, fmin: int = 0) -> torch.tensor:
    n_frames = len(frequency)
    S = np.zeros((n_mels, n_frames))
    mel_bin_frequencies = librosa.core.mel_frequencies(n_mels=n_mels + 1, fmin=fmin, fmax=fmax, htk=False)
    mel_band_indices = [
        np.clip(np.searchsorted(mel_bin_frequencies, mel_f, side='right') - 1, 0, n_mels - 1) if mel_f is not None else None
        for mel_f in frequency
    ]
    for t, mel_band in enumerate(mel_band_indices):
        if mel_band is None:
            continue
        if confidence is not None:
            S[mel_band, t] = confidence[t]
        else:
            S[mel_band, t] = 1.0
    S_norm = spectral_normalize_torch(torch.tensor(S))
    return S_norm

def midi_to_frequency(midi_file, sr=16000, duration=10.24, hop_length=160):
    midi_data = pretty_midi.PrettyMIDI(midi_file)
    max_len = int(duration * sr // hop_length)
    frequency = np.full(max_len, None)
    
    for instrument in midi_data.instruments:
        if instrument.is_drum:
            continue
        for note in instrument.notes:
            start_frame = int(note.start * sr // hop_length)
            end_frame = int(note.end * sr // hop_length)
            frequency[start_frame:end_frame] = note.pitch

    frequency = np.array([pretty_midi.note_number_to_hz(note) if note is not None else None for note in frequency])
    return frequency

def get_item(filepath: str, config: dict, input_type: str = 'midi'):
    max_len = int(config["audio_length"] * config["sr"])
    
    if input_type == 'audio':
        if filepath.endswith('.flac'):
            audio, sr = librosa.load(filepath)
        elif filepath.endswith('.wav'):
            sr, audio = wavfile.read(filepath)
        else:
            raise ValueError('Unsupported file format')
        
        if audio.ndim > 1:
            audio = np.mean(audio, axis=1)
        
        if sr != config["sr"]:
            audio = torchaudio.functional.resample(torch.Tensor(audio), sr, config['sr']).numpy()
        sr = config["sr"]
        hop_length = config["hop_length"]
    
        if len(audio) > max_len:
            audio = audio[:max_len]
        elif len(audio) < max_len:
            pad_length = max_len - len(audio)
            audio = np.pad(audio, (0, pad_length), mode="constant")

        S_waveform = torch.tensor(audio).unsqueeze(0)
        S = get_melspec(audio, sr=config['sr'], n_mels=config['n_mels'], window_length=config['window_length'],
                        hop_length=config['hop_length'], fmax=config['fmax'], fmin=config['fmin'])

        time, frequency, confidence, activation = crepe.predict(audio, sr, viterbi=True)
        if not config['use_confidence']:
            frequency = np.array([f if c >= config['confidence_threshold'] else None for f, c in zip(frequency, confidence)])
        
        if len(frequency) < S.shape[1]:
            if config['use_confidence']:
                assert len(confidence) == len(frequency)
                confidence = np.pad(confidence, (0, S.shape[1] - len(confidence)), mode="constant", constant_values=None)
            frequency = np.pad(frequency, (0, S.shape[1] - len(frequency)), mode="constant", constant_values=None)
        
        elif len(frequency) > S.shape[1]:
            frequency = frequency[:S.shape[1]]
            if config['use_confidence']:
                confidence = confidence[:S.shape[1]]
            
    elif input_type == 'midi':
        confidence = 1.0
        frequency = midi_to_frequency(filepath, sr=config["sr"], duration=config["audio_length"], hop_length=config["hop_length"])
    
    S_f0 = get_f0_melspec(frequency, confidence=confidence if config['use_confidence'] else None,
                          n_mels=config['n_mels'], fmax=config['fmax'], fmin=config['fmin'])
    
    item_dict = {
        'f0_condition': S_f0
    }
    
    return item_dict

def spec_show(S, sr):
    fig, ax = plt.subplots(figsize=(6, 2))
    img = librosa.display.specshow(S, sr=sr, ax=ax)
    buf = BytesIO()
    fig.savefig(buf, format='png')
    buf.seek(0)
    plt.close(fig)
    return buf

config_AudioLDM = {
    "sr": 16000,
    "audio_length": 10.24,
    "n_mels": 64,
    "window_length": 1024,
    "hop_length": 160,
    "fmax": 8000,
    "fmin": 0,
    "use_confidence": False,
    "confidence_threshold": 0.8,
}

cfg = config_AudioLDM

def process_file(file):
    if file.name.endswith('.wav') or file.name.endswith('.flac'):
        item_dict = get_item(file.name, config=cfg, input_type='audio')
    elif file.name.endswith('.mid') or file.name.endswith('.midi'):
        item_dict = get_item(file.name, config=cfg, input_type='midi')
    else:
        return "Unsupported file format"
    
    f0_condition = item_dict['f0_condition']
    return spec_show(f0_condition.numpy(), cfg["sr"])

demo = gr.Interface(fn=process_file, inputs=gr.inputs.File(), outputs="image", title="Pitch ControlNet")

demo.launch()