backend_worker.py

#!/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())