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)