app.py

try:
    import spaces
except ImportError:
    # local dev without the `spaces` package: keep the @spaces.GPU decorator
    # syntax below working as a no-op, since HF's ZeroGPU runtime statically
    # scans source files for a literal `@spaces.GPU` decorator at startup —
    # any custom-named wrapper around it won't be detected.
    class spaces:
        class GPU:
            def __init__(self, func=None, duration=60):
                self.func = func

            def __call__(self, *args, **kwargs):
                if self.func is not None:
                    return self.func(*args, **kwargs)
                func = args[0]
                return func

import os
import shutil
import tempfile
import threading

import numpy as np
import torch
import torchaudio
import soundfile as sf
import gradio as gr
from pyharp import ModelCard, build_endpoint

# ---- Model card ----

model_card = ModelCard(
    name="MEGAMI",
    description=(
        "Automatic music mixing via a generative model of audio effect embeddings."
        "Upload 2–6 dry (unprocessed) instrument stems and MEGAMI will generate "
        "a professionally mixed stereo output."
    ),
    author="Eloi Moliner, Marco A. Martínez-Ramírez, Junghyun Koo, Wei-Hsiang Liao, Kin Wai Cheuk, Joan Serrà, Vesa Välimäki, Yuki Mitsufuji",
    tags=["mixing", "music-production", "generative", "diffusion"],
)

# ---- Background checkpoint download ----
# Only network I/O runs here — CUDA can only be touched inside @spaces.GPU,
# so model construction is deferred to _get_inference(), called from process_fn.

_inference = None
_inference_lock = threading.Lock()  # guards first-time construction in _get_inference()
_download_error = None
_downloads_ready = threading.Event()

MIN_AUDIO_LEN = 525312   # samples @ 44100 Hz ≈ 11.9 s
SAMPLE_RATE   = 44100


def _download_checkpoints():
    import urllib.request
    from huggingface_hub import hf_hub_download

    os.makedirs("checkpoints", exist_ok=True)

    github_base = "https://github.com/SonyResearch/MEGAMI/releases/download/v0/"
    for fname in ["FxGenerator_public.pt", "FxProcessor_public.pt", "CLAP_DA_public.pt"]:
        dest = os.path.join("checkpoints", fname)
        if not os.path.exists(dest):
            print(f"[download] {fname} …")
            urllib.request.urlretrieve(github_base + fname, dest)
            print(f"[download] {fname} done.")

    # The MEGAMI README cites the original LAION-CLAP file as the source for this checkpoint;
    # no patching script is documented — we download the original and save under the expected filename
    clap_dest = "checkpoints/music_audioset_epoch_15_esc_90.14.patched.pt"
    if not os.path.exists(clap_dest):
        print("[download] music_audioset_epoch_15_esc_90.14.patched.pt …")
        cached = hf_hub_download(
            repo_id="lukewys/laion_clap",
            filename="music_audioset_epoch_15_esc_90.14.pt",
        )
        shutil.copy(cached, clap_dest)
        print("[download] music_audioset_epoch_15_esc_90.14.patched.pt done.")

    fxenc_path = "checkpoints/fxenc_plusplus_default.pt"
    if not os.path.exists(fxenc_path):
        print("[download] fxenc_plusplus_default.pt …")
        cached = hf_hub_download(
            repo_id="yytung/fxencoder-plusplus",
            filename="fxenc_plusplus_default.pt",
        )
        shutil.copy(cached, fxenc_path)
        print("[download] fxenc_plusplus_default.pt done.")

    print("[download] All checkpoints ready.")


def _download_thread_fn():
    global _download_error
    try:
        _download_checkpoints()
        print("[download] Ready for model construction.")
    except Exception as exc:
        import traceback
        _download_error = str(exc)
        print(f"[download] FAILED: {exc}")
        traceback.print_exc()
    finally:
        _downloads_ready.set()


threading.Thread(target=_download_thread_fn, daemon=True).start()


def _get_inference():
    """Build the Inference object on first use (must happen inside @spaces.GPU)."""
    global _inference
    if _inference is not None:
        return _inference
    with _inference_lock:
        if _inference is None:
            import omegaconf
            from inference.inference import Inference

            method_args = omegaconf.OmegaConf.create(
                {
                    "FxGenerator_code": "public",
                    "FxProcessor_code": "public",
                    "T": 30,
                    "cfg_scale": 1.0,
                    "Schurn": 0,
                }
            )
            _inference = Inference(method_args=method_args)
    return _inference


# ---- Process function ----

def _gpu_duration(track1, track2, track3, track4, track5, track6, steps):
    # Rough estimate: dominated by diffusion steps (T); pad generously for I/O and effect application.
    # First call also has to build the model (move checkpoints onto the GPU), so budget extra time.
    build_time = 0 if _inference is not None else 90
    return min(build_time + int(steps) * 3 + 60, 300)


@spaces.GPU(duration=_gpu_duration)
@torch.inference_mode()
def process_fn(track1, track2, track3, track4, track5, track6, steps):
    _downloads_ready.wait()

    if _download_error:
        raise gr.Error(f"Checkpoint download failed: {_download_error}")

    track_paths = [t for t in [track1, track2, track3, track4, track5, track6] if t is not None]
    if len(track_paths) < 2:
        raise gr.Error("Please upload at least 2 tracks.")

    from utils.feature_extractors.dsp_features import compute_log_rms_gated_max

    inference = _get_inference()

    dry_tracks   = []
    dry_segments = []

    for path in track_paths:
        audio, file_sr = sf.read(path, dtype="float32")
        if audio.ndim == 1:
            audio = np.stack([audio, audio], axis=-1)
        elif audio.shape[1] == 1:
            audio = np.concatenate([audio, audio], axis=-1)

        audio_tensor = torch.from_numpy(audio).T  # (channels, samples)
        if file_sr != SAMPLE_RATE:
            audio_tensor = torchaudio.functional.resample(audio_tensor, file_sr, SAMPLE_RATE)

        x_mono = audio_tensor.mean(dim=0, keepdim=True)  # (1, samples)

        if x_mono.shape[-1] < MIN_AUDIO_LEN:
            needed_seconds = MIN_AUDIO_LEN / SAMPLE_RATE
            got_seconds    = x_mono.shape[-1] / SAMPLE_RATE
            raise gr.Error(
                f"Each track must be at least {needed_seconds:.1f} s long "
                f"(got {got_seconds:.1f} s)."
            )

        x_mono = x_mono.to(inference.device)
        segment = inference.select_high_energy_segment(x_mono)

        dry_tracks.append(x_mono)
        dry_segments.append(segment)

    # Uploaded stems can differ slightly in length; pad with trailing silence so
    # they stack into one tensor without clipping the longer tracks.
    max_len = max(t.shape[-1] for t in dry_tracks)
    dry_tracks = [
        torch.nn.functional.pad(t, (0, max_len - t.shape[-1])) for t in dry_tracks
    ]

    x_dry          = torch.stack(dry_tracks)    # (N, 1, L)
    x_dry_segments = torch.stack(dry_segments)  # (N, 1, MIN_AUDIO_LEN)

    rms    = compute_log_rms_gated_max(x_dry_segments)
    silent = (rms < -60).squeeze()
    if silent.ndim == 0:
        silent = silent.unsqueeze(0)
    if silent.all():
        raise gr.Error("All uploaded tracks appear to be silent.")
    if silent.any():
        x_dry          = x_dry[~silent]
        x_dry_segments = x_dry_segments[~silent]

    inference.method_args.T = int(steps)  # slider returns float, but T is used as a loop bound so casting to int

    embedding_preds = inference.generate_Fx(x_dry_segments, num_samples=1)
    fx_embeddings   = inference.embedding_post_processing(embedding_preds)
    fx_embedding    = fx_embeddings[0]

    y_final = inference.apply_effects(x_dry.clone(), fx_embedding)  # (N, 2, L)
    mix     = y_final.sum(dim=0)                             # (2, L)

    peak = torch.max(torch.abs(mix)).clamp(min=1e-8)  # avoiding div by 0 for silent output
    mix  = mix / peak

    output_file = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
    sf.write(
        output_file.name,
        mix.cpu().clamp(-1, 1).numpy().T,
        SAMPLE_RATE,
        subtype="PCM_16",
    )
    return output_file.name


# ---- Gradio interface ----

with gr.Blocks() as demo:
    inputs = [
        gr.Audio(type="filepath", label="Track 1 (required)").harp_required(True),
        gr.Audio(type="filepath", label="Track 2 (required)").harp_required(True),
        gr.Audio(type="filepath", label="Track 3"),
        gr.Audio(type="filepath", label="Track 4"),
        gr.Audio(type="filepath", label="Track 5"),
        gr.Audio(type="filepath", label="Track 6"),
        gr.Slider(
            minimum=10,
            maximum=100,
            step=10,
            value=30,
            label="Diffusion Steps (higher = better quality, slower)",
        ),
    ]

    outputs = [
        gr.Audio(type="filepath", label="MEGAMI Mix"),
    ]

    build_endpoint(
        model_card=model_card,
        input_components=inputs,
        output_components=outputs,
        process_fn=process_fn,
    )

if __name__ == "__main__":
    demo.queue().launch(pwa=True)