Spaces:
Running
Running
| import os | |
| import tempfile | |
| import torch | |
| import soundfile as sf | |
| import numpy as np | |
| from pathlib import Path | |
| from pyharp import ModelCard, build_endpoint, SampleRate, Channels | |
| from pyharp.tags import Subcategory | |
| import gradio as gr | |
| from nam.train.lightning_module import LightningModule as _LightningModule | |
| import nam.data as data | |
| os.environ["TOKENIZERS_PARALLELISM"] = "true" | |
| model = None | |
| CKPT_PATH = Path(__file__).parent / "demo_ckpt.ckpt" | |
| MODEL_CONFIG_PATH = Path(__file__).parent / "default_config_files/models/wavenet.json" | |
| # 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"} | |
| def get_model(): | |
| global model | |
| if model is None: | |
| print("Loading PANAMA model...", flush=True) | |
| import json | |
| with open(MODEL_CONFIG_PATH, "r") as fp: | |
| model_config = json.load(fp) | |
| model = _LightningModule.load_from_checkpoint( | |
| CKPT_PATH, | |
| **_LightningModule.parse_config(model_config) | |
| ) | |
| model.eval() | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| model = model.to(device) | |
| print("Model loaded.", flush=True) | |
| return model | |
| model_card = ModelCard( | |
| name="PANAMA", | |
| description=( | |
| "Parametric neural amp modeling. Apply amp simulation with 6 " | |
| "controllable knobs. Amp modeling is mono: a stereo file is downmixed " | |
| "before the amp and the single result is copied back across its " | |
| "channels, so the output has the channel count it came in with but not " | |
| "its stereo image. The input's sample rate is restored. Repository: " | |
| "https://github.com/ETH-DISCO/PANAMA Paper: " | |
| "https://arxiv.org/abs/2509.26564" | |
| ), | |
| author="Florian Grötschla, Longxiang Jiao, Luca A. Lanzendörfer, Roger Wattenhofer", | |
| tags=[Subcategory.NEURAL_ANALOG_EFFECTS, SampleRate(48000), Channels(1), 'amp-modeling', 'guitar'], | |
| ) | |
| def process_fn(input_audio_path: str, knob1: float, knob2: float, knob3: float, knob4: float, knob5: float, knob6: float) -> str: | |
| if not input_audio_path: | |
| raise gr.Error("Provide the guitar or bass recording to play through the amp.") | |
| try: | |
| info = sf.info(input_audio_path) | |
| except Exception as exc: | |
| raise gr.Error(f"Could not read the audio file: {exc}") from exc | |
| if info.frames == 0: | |
| raise gr.Error("The audio file is empty.") | |
| m = get_model() | |
| device = next(m.parameters()).device | |
| # nam's loader rejects anything but mono, and an amp model is mono by | |
| # nature, so the downmix happens here and the channel count is restored | |
| # after inference rather than refusing the file | |
| import librosa | |
| import soundfile | |
| with soundfile.SoundFile(input_audio_path) as handle: | |
| original_sr = handle.samplerate | |
| original_channels = handle.channels | |
| input_audio_np, _ = librosa.load(input_audio_path, sr=None, mono=True) | |
| if original_sr != m.net.sample_rate: | |
| input_audio_np = librosa.resample( | |
| input_audio_np, orig_sr=original_sr, target_sr=m.net.sample_rate | |
| ) | |
| input_audio = torch.tensor(input_audio_np, dtype=torch.float32) | |
| input_audio = input_audio.to(device) | |
| g_tensor = torch.tensor([[knob1, knob2, knob3, knob4, knob5, knob6]], dtype=torch.float32).to(device) | |
| print(f"Processing with g-vector: {g_tensor}", flush=True) | |
| with torch.no_grad(): | |
| output = m.forward(input_audio, g_tensor) | |
| output_np = output.cpu().numpy() | |
| # Back to the rate the session uses, before the channels are rebuilt, so | |
| # the resampler only ever sees one channel | |
| if original_sr != m.net.sample_rate: | |
| output_np = librosa.resample( | |
| output_np, orig_sr=m.net.sample_rate, target_sr=original_sr | |
| ) | |
| # Two identical channels play back the same as a mono file on a stereo | |
| # track. It does not make the result stereo, which the model card says. | |
| if original_channels > 1: | |
| output_np = np.tile(output_np[:, None], (1, original_channels)) | |
| input_ext = Path(input_audio_path).suffix.lower() | |
| output_path = (Path(tempfile.mkdtemp(prefix="panama-")) | |
| / f"panama{input_ext if input_ext in LOSSLESS else '.wav'}") | |
| sf.write(output_path, output_np, original_sr) | |
| return str(output_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__": | |
| get_model() | |
| if __name__ == "__main__": | |
| with gr.Blocks() as demo: | |
| input_components = [ | |
| gr.Audio(type="filepath", label="Input Audio").harp_required(True).set_info(( | |
| "A dry guitar or bass recording. Feeding it an already-amped tone " | |
| "stacks one amp on another." | |
| )), | |
| gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.5, label="Gain", | |
| info="Preamp drive. Raise it for more distortion and compression."), | |
| gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.5, label="Bass", | |
| info="Low-end tone control."), | |
| gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.5, label="Mid", | |
| info=( | |
| "Midrange tone control. Cutting it gives the scooped sound heavy " | |
| "styles use." | |
| )), | |
| gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.5, label="Treble", | |
| info="High-end tone control."), | |
| gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.5, label="Master", | |
| info=( | |
| "Power amp level, which adds its own saturation as it rises rather " | |
| "than only getting louder." | |
| )), | |
| gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.5, label="Presence", | |
| info="Upper-midrange bite, which helps a part cut through a mix."), | |
| ] | |
| output_components = [ | |
| gr.Audio(type="filepath", label="Amped Audio").set_info("Amp-processed audio."), | |
| ] | |
| app = build_endpoint( | |
| model_card=model_card, | |
| input_components=input_components, | |
| output_components=output_components, | |
| process_fn=process_fn, | |
| ) | |
| print("Launching Gradio...", flush=True) | |
| demo.queue().launch(server_name="0.0.0.0", server_port=7860, show_error=True, pwa=True) | |