app.py

from __future__ import annotations

import spaces  # imported first, as ZeroGPU requires

import uuid
from pathlib import Path
from tempfile import gettempdir

import gradio as gr
import soundfile as sf

from pyharp import ModelCard, build_endpoint, SampleRate, Channels
from pyharp.tags import Subcategory

from microtcn_runtime import load_model, process_audio

MIN_AUDIO_SECONDS = 0.5
MAX_AUDIO_SECONDS = 30
OUTPUT_ROOT = Path(gettempdir()) / "microtcn_outputs"

# A lossless input keeps its container. A lossy one comes back as WAV, since encoding the result
# lossily again would expose the artifacts of the first encode and add a second set.
LOSSLESS = {".wav", ".flac"}

model_card = ModelCard(
    name="micro-TCN",
    description=(
        "Emulate a Teletronix LA-2A optical compressor/limiter using a "
        "lightweight temporal convolutional network. The model is fixed-rate "
        "and mono, so the result is resampled back to the input's rate and "
        "copied across its channels: the output has the format it came in with, "
        "but no stereo image. Repository: "
        "https://github.com/csteinmetz1/micro-tcn Paper: "
        "https://arxiv.org/abs/2102.06200"
    ),
    author="Christian J. Steinmetz & Joshua D. Reiss",
    tags=[Subcategory.NEURAL_ANALOG_EFFECTS, SampleRate(44100), Channels(1), 'compressor', 'limiter', 'analog-emulation', 'domain:music'],
)


def match_input_format(output_path: str, input_path: str) -> str:
    """
    Rewrite an output file at the sample rate and channel count it came in with.

    The models here are fixed-rate and mono, and hand the writing to vendored
    upstream code, so the format is corrected afterwards rather than threaded
    through inference. Duplicated channels are copies, not a recovered stereo
    image; the model card says so.

    Args:
        output_path (str): File the model wrote. Overwritten in place.
        input_path (str): File the request supplied.

    Returns:
        output_path (str): The same path, for use in a return statement.
    """

    import librosa
    import numpy as np
    import soundfile

    want = soundfile.info(input_path)
    have = soundfile.info(output_path)

    if have.samplerate == want.samplerate and have.channels == want.channels:
        return output_path

    data, rate = soundfile.read(output_path, always_2d=True)

    if rate != want.samplerate:
        data = librosa.resample(
            data.T, orig_sr=rate, target_sr=want.samplerate
        ).T

    if data.shape[1] == 1 and want.channels > 1:
        data = np.tile(data, (1, want.channels))
    elif data.shape[1] > 1 and want.channels == 1:
        data = data.mean(axis=1, keepdims=True)

    soundfile.write(output_path, data, want.samplerate)

    return output_path


@spaces.GPU(duration=120)
def compress(audio_path: str, limit_mode: float, peak_reduction: float):
    return process_audio(
        audio_path=audio_path,
        limit_mode=limit_mode,
        peak_reduction=peak_reduction,
    )


def process_fn(
    audio_path: str | None,
    mode: str,
    peak_reduction: float,
) -> str:
    # Checked before the GPU is reserved, so a request the model cannot take spends none of the
    # caller's ZeroGPU allowance
    if not audio_path:
        raise gr.Error("Please upload an audio file.")

    try:
        duration = sf.info(audio_path).duration
    except Exception as exc:
        raise gr.Error(f"Could not read the audio file: {exc}") from exc

    if duration < MIN_AUDIO_SECONDS:
        raise gr.Error("Audio must be at least 0.5 seconds long.")

    if duration > MAX_AUDIO_SECONDS:
        raise gr.Error(
            f"Audio must be no longer than {MAX_AUDIO_SECONDS} seconds. "
            f"Received {duration:.1f} seconds."
        )

    limit_mode = 1.0 if mode == "Limit" else 0.0

    try:
        audio, sample_rate = compress(audio_path, limit_mode, peak_reduction)
    except Exception as exc:
        raise gr.Error(f"micro-TCN inference failed: {exc}") from exc

    input_ext = Path(audio_path).suffix.lower()
    output_dir = OUTPUT_ROOT / uuid.uuid4().hex
    output_dir.mkdir(parents=True, exist_ok=True)
    output_path = output_dir / f"microtcn_output{input_ext if input_ext in LOSSLESS else '.wav'}"
    sf.write(str(output_path), audio, sample_rate)

    return match_input_format(str(output_path), audio_path)


# pyharp runs process_fn in a worker, which imports this file as __mp_main__ and starts with the
# server, so the model is loaded there, before the first request rather than during it.
if __name__ != "__main__":
    load_model()


if __name__ == "__main__":
    with gr.Blocks(title="micro-TCN LA-2A Emulation") as demo:
        input_components = [
            gr.Audio(type="filepath", label="Input Audio")
            .harp_required(True)
            .set_info("Audio clip to process, up to 30 seconds."),
            gr.Dropdown(
                choices=["Compress", "Limit"],
                value="Compress",
                label="Mode",
            ).set_info(
                "Compress for gentle gain reduction, Limit for harder peak control."
            ),
            gr.Slider(
                minimum=0.0,
                maximum=1.0,
                value=0.42,
                step=0.01,
                label="Peak Reduction",
            ).set_info("Normalized peak-reduction control, from 0 to 1."),
        ]
        output_components = [
            gr.Audio(
                type="filepath",
                label="Processed Audio",
            ).set_info("Audio processed through the LA-2A emulation."),
        ]
        build_endpoint(
            model_card=model_card,
            input_components=input_components,
            output_components=output_components,
            process_fn=process_fn,
        )


    demo.queue(default_concurrency_limit=1).launch(show_error=True, pwa=True)