#!/usr/bin/env python3 """One-shot backend worker for the isolated interpreter. Reads a JSON request {"inputs": {...}} on stdin and prints a JSON response {"ok": bool, "outputs": {...}} on stdout. Media are exchanged by file path. All library stdout noise is redirected to stderr so stdout carries only the JSON protocol. """ from __future__ import annotations import contextlib import json import sys import traceback import os import json as _json import pickle import urllib.parse import urllib.request import tempfile import warnings warnings.filterwarnings("ignore") os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "3") os.environ.setdefault("CUDA_VISIBLE_DEVICES", "") import numpy as np import soundfile as sf import librosa import ddsp import ddsp.training from ddsp.training.postprocessing import detect_notes, fit_quantile_transform import gin import tensorflow.compat.v2 as tf SAMPLE_RATE = 16000 PRETRAINED_MODELS = ["Violin", "Flute", "Flute2", "Trumpet", "Tenor_Saxophone"] _GCS_BUCKET = "ddsp" _GCS_PREFIX = "models/timbre_transfer_colab/2021-07-08" _MODELS_DIR = "/tmp/ddsp_pretrained" os.makedirs(_MODELS_DIR, exist_ok=True) def _ensure_model(model_name): dest = os.path.join(_MODELS_DIR, model_name) gin_file = os.path.join(dest, "operative_config-0.gin") if os.path.exists(gin_file): return dest os.makedirs(dest, exist_ok=True) prefix = f"{_GCS_PREFIX}/solo_{model_name.lower()}_ckpt/" list_url = ( f"https://storage.googleapis.com/storage/v1/b/{_GCS_BUCKET}/o" f"?prefix={urllib.parse.quote(prefix, safe='')}" ) with urllib.request.urlopen(list_url, timeout=60) as response: listing = _json.loads(response.read().decode("utf-8")) items = listing.get("items") or [] if not items: raise RuntimeError(f"Could not list checkpoint files for {model_name!r}.") for item in items: obj = item.get("name") or "" fname = obj.rsplit("/", 1)[-1] if not fname: continue out = os.path.join(dest, fname) if not os.path.exists(out): url = f"https://storage.googleapis.com/{_GCS_BUCKET}/{urllib.parse.quote(obj, safe='')}" urllib.request.urlretrieve(url, out) if not os.path.exists(gin_file): raise RuntimeError(f"Downloaded {model_name!r} but operative_config-0.gin is missing.") return dest def _shift_ld(af, ld_shift=0.0): af["loudness_db"] += ld_shift return af def _shift_f0(af, pitch_shift=0.0): af["f0_hz"] *= 2.0 ** (pitch_shift) af["f0_hz"] = np.clip(af["f0_hz"], 0.0, librosa.midi_to_hz(110.0)) return af def _get_tuning_factor(f0_midi, f0_confidence, mask_on): tuning_factors = np.linspace(-0.5, 0.5, 101) midi_diffs = (f0_midi[mask_on][:, np.newaxis] - tuning_factors[np.newaxis, :]) % 1.0 midi_diffs[midi_diffs > 0.5] -= 1.0 weights = f0_confidence[mask_on][:, np.newaxis] cost_diffs = np.mean(weights * np.abs(midi_diffs), axis=0) f0_at = f0_midi[mask_on][:, np.newaxis] - midi_diffs deltas = (np.diff(f0_at, axis=0) != 0.0).astype(float) cost_deltas = np.mean(weights[:-1] * deltas, axis=0) norm = lambda x: (x - np.mean(x)) / np.std(x) cost = norm(cost_deltas) + norm(cost_diffs) return tuning_factors[np.argmin(cost)] def _auto_tune(f0_midi, tuning_factor, mask_on, amount=0.0): major_scale = np.ravel([np.array([0, 2, 4, 5, 7, 9, 11]) + 12 * i for i in range(10)]) all_scales = np.stack([major_scale + i for i in range(12)]) f0_on = f0_midi[mask_on] f0_diff_tsn = f0_on[:, np.newaxis, np.newaxis] - all_scales[np.newaxis, :, :] f0_diff_ts = np.min(np.abs(f0_diff_tsn), axis=-1) f0_diff_s = np.mean(f0_diff_ts, axis=0) scale_idx = np.argmin(f0_diff_s) f0_diff_tn = f0_midi[:, np.newaxis] - all_scales[scale_idx][np.newaxis, :] note_idx = np.argmin(np.abs(f0_diff_tn), axis=-1) midi_diff = np.take_along_axis(f0_diff_tn, note_idx[:, np.newaxis], axis=-1)[:, 0] return f0_midi - amount * midi_diff def run_timbre_transfer(inputs): audio_path = inputs["audio"] model_name = inputs.get("model_name") or "Violin" threshold = float(inputs.get("threshold", 1.0)) adjust = bool(inputs.get("adjust", True)) quiet = float(inputs.get("quiet", 20.0)) autotune = float(inputs.get("autotune", 0.0)) pitch_shift = float(inputs.get("pitch_shift", 0.0)) loudness_shift = float(inputs.get("loudness_shift", 0.0)) if not audio_path: raise RuntimeError("No input audio provided.") if model_name not in PRETRAINED_MODELS: raise RuntimeError(f"Unknown model {model_name!r}; choose {PRETRAINED_MODELS}.") audio, _ = librosa.load(audio_path, sr=SAMPLE_RATE, mono=True) audio = audio.astype(np.float32)[np.newaxis, :] model_dir = _ensure_model(model_name) gin_file = os.path.join(model_dir, "operative_config-0.gin") dataset_stats = None stats_file = os.path.join(model_dir, "dataset_statistics.pkl") if os.path.exists(stats_file): with open(stats_file, "rb") as fh: dataset_stats = pickle.load(fh) with gin.unlock_config(): gin.parse_config_file(gin_file, skip_unknown=True) ddsp.spectral_ops.reset_crepe() af = ddsp.training.metrics.compute_audio_features(audio) af = {k: (v.numpy() if hasattr(v, "numpy") else v) for k, v in af.items()} af["loudness_db"] = af["loudness_db"].astype(np.float32) ckpt_files = [f for f in os.listdir(model_dir) if "ckpt" in f] if not ckpt_files: raise RuntimeError(f"No checkpoint files in {model_dir}.") ckpt = os.path.join(model_dir, ckpt_files[0].split(".")[0]) time_steps_train = gin.query_parameter("F0LoudnessPreprocessor.time_steps") n_samples_train = gin.query_parameter("Harmonic.n_samples") hop_size = int(n_samples_train / time_steps_train) time_steps = int(audio.shape[1] / hop_size) n_samples = time_steps * hop_size with gin.unlock_config(): gin.parse_config([ f"Harmonic.n_samples = {n_samples}", f"FilteredNoise.n_samples = {n_samples}", f"F0LoudnessPreprocessor.time_steps = {time_steps}", "oscillator_bank.use_angular_cumsum = True", ]) for key in ["f0_hz", "f0_confidence", "loudness_db"]: af[key] = af[key][:time_steps] af["audio"] = af["audio"][:, :n_samples] af_mod = {k: (v.copy() if hasattr(v, "copy") else v) for k, v in af.items()} if adjust and dataset_stats is not None: mask_on, note_on_value = detect_notes(af["loudness_db"], af["f0_confidence"], threshold) if np.any(mask_on): target_mean_pitch = dataset_stats["mean_pitch"] pitch = ddsp.core.hz_to_midi(af["f0_hz"]) mean_pitch = np.mean(pitch[mask_on]) p_diff = target_mean_pitch - mean_pitch p_diff_octave = p_diff / 12.0 round_fn = np.floor if p_diff_octave > 1.5 else np.ceil af_mod = _shift_f0(af_mod, round_fn(p_diff_octave)) _, loudness_norm = fit_quantile_transform( af["loudness_db"], mask_on, inv_quantile=dataset_stats["quantile_transform"] ) mask_off = np.logical_not(mask_on) loudness_norm[mask_off] -= quiet * (1.0 - note_on_value[mask_off][:, np.newaxis]) loudness_norm = np.reshape(loudness_norm, af["loudness_db"].shape) af_mod["loudness_db"] = loudness_norm if autotune: f0_midi = np.array(ddsp.core.hz_to_midi(af_mod["f0_hz"])) tuning_factor = _get_tuning_factor(f0_midi, af_mod["f0_confidence"], mask_on) f0_midi_at = _auto_tune(f0_midi, tuning_factor, mask_on, amount=autotune) af_mod["f0_hz"] = ddsp.core.midi_to_hz(f0_midi_at) af_mod = _shift_ld(af_mod, loudness_shift) af_mod = _shift_f0(af_mod, pitch_shift) model = ddsp.training.models.Autoencoder() model.restore(ckpt) _ = model(af_mod, training=False) outputs_tf = model(af_mod, training=False) audio_gen = np.array(model.get_audio_from_outputs(outputs_tf)) if audio_gen.ndim == 2: audio_gen = audio_gen[0] out_path = os.path.join(tempfile.mkdtemp(), "ddsp_output.wav") sf.write(out_path, audio_gen.astype(np.float32), SAMPLE_RATE) return out_path def _run(inputs): outputs = {"out_audio": run_timbre_transfer(inputs)} return outputs def main(): try: request = json.load(sys.stdin) except Exception as exc: print(json.dumps({"ok": False, "error": f"invalid request: {exc!r}"}), flush=True) return 2 inputs = request.get("inputs") or {} try: with contextlib.redirect_stdout(sys.stderr): outputs = _run(inputs) payload = {"ok": True, "outputs": outputs} except Exception: payload = {"ok": False, "error": traceback.format_exc()[-3000:]} print(json.dumps(payload), flush=True) return 0 if payload["ok"] else 1 if __name__ == "__main__": raise SystemExit(main())