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)