app.py

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 model loading ----

_inference = None
_load_error = None
_model_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 _load_model():
    global _inference, _load_error
    try:
        _download_checkpoints()

        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)
        print("[model] Inference object ready.")
    except Exception as exc:
        import traceback
        _load_error = str(exc)
        print(f"[model] Load FAILED: {exc}")
        traceback.print_exc()
    finally:
        _model_ready.set()


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

# ---- Process function ----

@torch.inference_mode()
def process_fn(track1, track2, track3, track4, track5, track6, steps):
    _model_ready.wait()

    if _load_error:
        raise gr.Error(f"Model failed to load: {_load_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

    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)