Add app.py for the user audio input and mel-spec output

app.py

@@ -1,7 +1,168 @@
+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
 
-def greet(name):
-    return "Hello " + name + "!!"
+# 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 = gr.Interface(fn=greet, inputs="text", outputs="text")
 demo.launch()