File size: 5,463 Bytes
8cc1163 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 | from __future__ import annotations
import gradio as gr
try:
import spaces
except ImportError: # 'spaces' is only provided by Hugging Face Spaces
import types as _types
def _gpu(*args, **kwargs):
if len(args) == 1 and callable(args[0]) and not kwargs:
return args[0]
def _decorator(func):
return func
return _decorator
spaces = _types.SimpleNamespace(GPU=_gpu)
import tensorflow as tf
import gin
import ddsp.training.inference
import ddsp.core
import librosa
import soundfile
import numpy as np
import requests
import zipfile
import io
import tempfile
import os
import crepe
# Global variables for model and parameters
MODEL = None
SAMPLE_RATE = None
N_SAMPLES = None
N_FRAMES = None
def _download_and_extract_model(model_name):
model_url = f'https://storage.googleapis.com/ddsp/models/{model_name}/{model_name}.zip'
print(f"Downloading model from: {model_url}")
response = requests.get(model_url)
response.raise_for_status() # Raise an exception for HTTP errors
# Create a temporary directory for the model
temp_dir = tempfile.mkdtemp()
zip_path = os.path.join(temp_dir, f'{model_name}.zip')
with open(zip_path, 'wb') as f:
f.write(response.content)
with zipfile.ZipFile(zip_path, 'r') as zip_ref:
zip_ref.extractall(temp_dir)
# The actual model checkpoint is usually in a subdirectory
# Find the directory that contains the 'operative_config.gin'
model_ckpt_dir = None
for root, dirs, files in os.walk(temp_dir):
if 'operative_config.gin' in files:
model_ckpt_dir = root
break
if not model_ckpt_dir:
raise FileNotFoundError("Could not find operative_config.gin in the downloaded model.")
print(f"Model extracted to: {model_ckpt_dir}")
return model_ckpt_dir
# --- Model Loading ---
MODEL_NAME = 'violin_ddsp_2020_03_03' # Example model from DDSP demos
CKPT_DIR = _download_and_extract_model(MODEL_NAME)
# Initialize the model with fixed parameters for this deployment
LENGTH_SECONDS = 4 # Default length for output audio
REMOVE_REVERB = True # Whether to remove reverb from the output
# Need to unlock gin config before parsing
with gin.unlock_config():
MODEL = ddsp.training.inference.AutoencoderInference(
ckpt=CKPT_DIR,
length_seconds=LENGTH_SECONDS,
remove_reverb=REMOVE_REVERB
)
# Extract model parameters after initialization
SAMPLE_RATE = MODEL.sample_rate
N_SAMPLES = MODEL.n_samples
N_FRAMES = MODEL.n_frames
# Build the network by running a fake batch to initialize weights
MODEL.build_network()
# --- Helper functions for audio processing ---
def _load_audio_for_inference(audio_path, sr, n_samples):
audio, _ = librosa.load(audio_path, sr=sr, mono=True)
# Pad or trim audio to n_samples
if len(audio) < n_samples:
audio = np.pad(audio, (0, n_samples - len(audio)), mode='constant')
else:
audio = audio[:n_samples]
return audio
def _save_audio_for_inference(audio_tensor, sr):
output_path = tempfile.NamedTemporaryFile(suffix=".wav", delete=False).name
audio_np = audio_tensor.numpy() if tf.is_tensor(audio_tensor) else audio_tensor
soundfile.write(output_path, audio_np, sr)
return output_path
@spaces.GPU
def predict(input_audio):
# Load audio
audio_np = _load_audio_for_inference(input_audio, SAMPLE_RATE, N_SAMPLES)
# Extract pitch using crepe
# crepe.predict expects a 1D numpy array
# Calculate step_size in milliseconds for crepe based on model's frame rate
step_size_ms = int(1000 * N_SAMPLES / (SAMPLE_RATE * N_FRAMES))
_, f0_hz, _, _ = crepe.predict(
audio_np, SAMPLE_RATE, viterbi=True, step_size=step_size_ms
)
# Resample f0_hz to N_FRAMES
f0_hz_resampled = ddsp.core.resample(
tf.convert_to_tensor(f0_hz, dtype=tf.float32), N_FRAMES
)
# Compute loudness
loudness_db = ddsp.core.compute_loudness(
tf.expand_dims(tf.convert_to_tensor(audio_np, dtype=tf.float32), axis=0),
SAMPLE_RATE
)
# Prepare features dictionary
audio_features = {
'audio': tf.expand_dims(tf.convert_to_tensor(audio_np, dtype=tf.float32), axis=0),
'f0_hz': tf.expand_dims(f0_hz_resampled, axis=0),
'loudness_db': loudness_db,
}
# Preprocess features using the model's preprocessor
processed_features = MODEL.preprocessor(audio_features)
# Run inference
outputs = MODEL(processed_features, training=False)
# Get synthesized audio from outputs
synthesized_audio = MODEL.get_audio_from_outputs(outputs)
# Save and return the output audio file path
output_audio_path = _save_audio_for_inference(synthesized_audio[0], SAMPLE_RATE)
return output_audio_path
demo = gr.Interface(
fn=predict,
inputs=[
gr.Audio(type="filepath", label="Input Audio"),
],
outputs=[
gr.Audio(type="filepath", label="Output Audio"),
],
title="DDSP Autoencoder Inference",
description="Resynthesize input audio using a pre-trained DDSP autoencoder model. This model takes an audio input, extracts its fundamental frequency (f0) and loudness, and then uses a DDSP autoencoder to synthesize a new audio output, effectively performing a timbre transfer based on the loaded model.",
)
if __name__ == "__main__":
demo.queue().launch(server_name="0.0.0.0", server_port=7860, show_error=True)
|