Build the default model in the worker, clean up, and declare the weights' license

Moves the model card to module scope, builds and places the default model in the worker and fetches the others there, removes debug prints, dead code, and the nested and unused GPU functions, and declares CC-BY-NC-SA-4.0, the license the authors state for the weights.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

README.md

@@ -8,7 +8,7 @@ sdk_version: 6.24.0
 python_version: 3.11
 app_file: app.py
 pinned: false
-license: mit
+license: cc-by-nc-sa-4.0
 short_description: Remixes and extends audio by masked generation
 tags:
   - category:generation
@@ -22,237 +22,13 @@ tags:
 
 # VampNet
 
-# Table of contents
-
-- [setting up](#setting-up)
-- [programmatic usage](#programmatic-usage)
-- [launching the web app](#launching-the-web-app)
-- [training / fine-tuning](#training--fine-tuning)
-  - [training a model](#training-a-model)
-  - [debugging training](#debugging-training)
-  - [fine-tuning](#fine-tuning)
-- [exporting your model](#exporting-your-model)
-- [unloop](#unloop)
-- [token telephone](#token-telephone)
-- [a note on argbind](#a-note-on-argbind)
-- [take a look at the pretrained models](#take-a-look-at-the-pretrained-models)
-- [licensing for pretrained models](#licensing-for-pretrained-models)
-
-## setting up
-
-python 3.9-3.11 works well. (for example, using conda)
-```bash
-conda create -n vampnet python=3.9
-conda activate vampnet
-```
-
-install VampNet
-
-```bash
-git clone https://github.com/hugofloresgarcia/vampnet.git
-pip install -e ./vampnet
-```
-
-## programmatic usage
-
-quick start!
-```python
-import random
-import vampnet
-import audiotools as at
-
-# load the default vampnet model
-interface = vampnet.interface.Interface.default()
-
-# list available finetuned models
-finetuned_model_choices = interface.available_models()
-print(f"available finetuned models: {finetuned_model_choices}")
-
-# pick a random finetuned model
-model_choice = random.choice(finetuned_model_choices)
-print(f"choosing model: {model_choice}")
-
-# load a finetuned model
-interface.load_finetuned(model_choice)
-
-# load an example audio file
-signal = at.AudioSignal("assets/example.wav")
-
-# get the tokens for the audio
-codes = interface.encode(signal)
-
-# build a mask for the audio
-mask = interface.build_mask(
-    codes, signal,
-    periodic_prompt=7, 
-    upper_codebook_mask=3,
-)
-
-# generate the output tokens
-output_tokens = interface.vamp(
-    codes, mask, return_mask=False,
-    temperature=1.0, 
-    typical_filtering=True, 
-)
-
-# convert them to a signal
-output_signal = interface.decode(output_tokens)
-
-# save the output signal
-output_signal.write("scratch/output.wav")
-```
-
-
-# Launching the Web app
-You can launch a gradio UI to play with vampnet. 
-
-```bash
-python app.py 
-```
-
-# Training / Fine-tuning 
-
-## Training a model
-
-To train a model, run the following script: 
-
-```bash
-python scripts/exp/train.py --args.load conf/vampnet.yml --save_path /path/to/checkpoints
-```
-
-for multi-gpu training, use torchrun:
-
-```bash
-torchrun --nproc_per_node gpu scripts/exp/train.py --args.load conf/vampnet.yml --save_path path/to/ckpt
-```
-
-You can edit `conf/vampnet.yml` to change the dataset paths or any training hyperparameters. 
-
-For coarse2fine models, you can use `conf/c2f.yml` as a starting configuration. 
-
-See `python scripts/exp/train.py -h` for a list of options.
-
-## Debugging training
-
-To debug training, it's easier to debug with 1 gpu and 0 workers
-
-```bash
-CUDA_VISIBLE_DEVICES=0 python -m pdb scripts/exp/train.py --args.load conf/vampnet.yml --save_path /path/to/checkpoints --num_workers 0
-```
-
-# Fine-tuning
-
-To fine-tune a model, use the script in `scripts/exp/fine_tune.py` 
-
-for an audio folder
-```bash
-python scripts/exp/fine_tune.py /path/to/audio/folder <fine_tune_name>
-```
-
-for multiple files
-```bash
-python scripts/exp/fine_tune.py "/path/to/audio1.mp3 /path/to/audio2/ /path/to/audio3.wav" <fine_tune_name>
-```
-
-This creates configuration files for a fine tuning train job. The save_paths will be set to `runs/<fine_tune_name>/coarse` and `runs/<fine_tune_name>/c2f`. 
-
-launch the coarse job: 
-```bash
-python scripts/exp/train.py --args.load conf/generated/<fine_tune_name>/coarse.yml 
-```
-
-this will save the coarse model to `runs/<fine_tune_name>/coarse/ckpt/best/`.
-
-launch the c2f job: 
-```bash
-python  scripts/exp/train.py --args.load conf/generated/<fine_tune_name>/c2f.yml 
-```
-
-# Resuming a Training/Finetuning Job from checkpoint. 
-
-To resume from checkpoint, use the `--resume` flag and the `--save_path` to point to the checkpoint you want to resume from.
-```bash
-python scripts/exp/train.py --args.load conf/generated/steve/coarse.yml --save_path runs/steve/coarse --resume
-```
-
-# Exporting your model
-
-Once your model has been fine-tuned, you can export it to a HuggingFace model. 
-
-In order to use your model in `app.py`, you will need to export it to HuggingFace.
-
-**NOTE**: In order to export, you will need a [huggingface account](https://huggingface.co/).
-
-Now, log in to huggingface using the command line:
-```bash
-huggingface-cli login
-```
-
-replace the contents of the file named `./DEFAULT_HF_MODEL_REPO` with your `<HUGGINGFACE_USERNAME>/vampnet`. A model repo will be automatically created for you with `export.py`. The default is `hugggof/vampnet`. 
-
-for example, if my username is `hugggof`, I would run the following command:`
-```bash
-echo 'hugggof/vampnet' > ./DEFAULT_HF_MODEL_REPO
-```
-
-Now, run the following command to export your model (replace `<your_finetuned_model_name>` with the name of your model):
-
-```bash
-python scripts/exp/export.py --name <your_finetuned_model_name> --model latest
-```
-
-Once that's done, your model should appear on the list of available models in the gradio interface.
-Simply run `python app.py` and select your model from the dropdown list.
-
-
-# Unloop
-
-Make sure you have Max installed on your laptop!
-
-**NOTE**: To run unloop (with a GPU-powered server), you will need to install the vampnet repo in both your local machine and your GPU server.
-
-## start a vampnet gradio server
-
-First, **on your GPU server**, run the gradio server:
-```bash
-python app.py --args.load conf/interface.yml --Interface.device cuda
-```
-This will run a vampnet gradio API on your GPU server. Copy the address. It will be something like `https://127.0.0.1:7860/`. 
-
-**IMPORTANT** Make sure that this gradio port (by default `7860`) is forwarded to your local machine, where you have Max installed. 
-
-## start the unloop gradio client
-Now, **on your local machine**, run the unloop gradio client.
-```
-cd unloop
-pip install -r requirements.txt
-python client.py --vampnet_url https://127.0.0.1:7860/ # replace with your gradio server address
-```
-This will start a gradio client that connects to the gradio server running on your GPU server.
-
-## start the unloop Max patch
-Now, open the unloop Max patch. It's located at `unloop/max/unloop.maxpat`.
-
-In the tape controls, check the heartbeat (`<3`) to make sure the connection to the local gradio client is working. 
-
-have fun!
-
-# Token Telephone
-
-Instructions forthcoming, but the sauce is in `token_telephone/tt.py`
-
-## A note on argbind
-This repository relies on [argbind](https://github.com/pseeth/argbind) to manage CLIs and config files. 
-Config files are stored in the `conf/` folder. 
-
-### Take a look at the pretrained models
-All the pretrained models (trained by hugo) are stored here: https://huggingface.co/hugggof/vampnet 
-
-### Licensing for Pretrained Models: 
-The weights for the models are licensed [`CC BY-NC-SA 4.0`](https://creativecommons.org/licenses/by-nc-sa/4.0/deed.ml). Likewise, any VampNet models fine-tuned on the pretrained models are also licensed [`CC BY-NC-SA 4.0`](https://creativecommons.org/licenses/by-nc-sa/4.0/deed.ml).
-
-Download the pretrained models from [this link](https://zenodo.org/record/8136629). Then, extract the models to the `models/` folder. 
-
-
+Remixes and extends audio by masked acoustic token modeling, with
+[VampNet](https://github.com/hugofloresgarcia/vampnet) and its fine-tuned
+models.
 
+The code is MIT. The pretrained weights are
+[CC BY-NC-SA 4.0](https://creativecommons.org/licenses/by-nc-sa/4.0/), as are
+models fine-tuned from them, so results cannot be used commercially and must be
+shared alike.
 
+A HARP deployment. See https://github.com/TEAMuP-dev/HARP

app.py

@@ -1,4 +1,4 @@
-import spaces
+import spaces  # imported first, as ZeroGPU requires
 
 try:  # torch>=2.6 flipped torch.load(weights_only) to True; legacy ckpts need False
     import torch as _torch
@@ -17,165 +17,80 @@ try:  # torch>=2.6 flipped torch.load(weights_only) to True; legacy ckpts need F
 except Exception:  # torch not installed / unexpected API -- nothing to patch
     pass
 
-from pathlib import Path
-import yaml
 import time
-import uuid
 
 import numpy as np
 import soundfile as sf
 import audiotools as at
-import argbind
-import shutil
 import torch
-from datetime import datetime
-from pyharp.core import build_endpoint, ModelCard  
-from pyharp.media.audio import save_audio  
+from pyharp.core import build_endpoint, ModelCard
+from pyharp.media.audio import save_audio
 from pyharp.tags import Subcategory, SampleRate, Channels
 
 import gradio as gr
-from vampnet.interface import Interface, signal_concat
+from vampnet.interface import Interface
 from vampnet import mask as pmask
 
-device="cpu"
-print(f"using device {device}\n"*10)
-
-interface = Interface.default()
+model_card = ModelCard(
+    name="VampNet",
+    description=(
+        "Generates new audio by remixing and extending your input, filling in "
+        "the blanks to create a coherent continuation in the same style. "
+        "Available models include cat10, fine-tuned on expressive cat sounds; "
+        "lazaro-ros-sep, percussion extracted from the Lázaro Ros discography, "
+        "rich in Afro-Cuban rhythm; and supermetroid, trained on Super Metroid "
+        "(SNES) gameplay audio for retro, game-style soundscapes. Generation "
+        "runs in mono at the model's own rate; the result is resampled back to "
+        "the input's rate and copied across its channels, so the output has the "
+        "format it came in with but no stereo image. The code is MIT, but the "
+        "pretrained weights are CC-BY-NC-SA-4.0, so results cannot be used "
+        "commercially and must be shared alike. Repository: "
+        "https://github.com/hugofloresgarcia/vampnet Paper: "
+        "https://arxiv.org/abs/2307.04686"
+    ),
+    author="Hugo Flores Garcia, Prem Seetharaman, Rithesh Kumar, Bryan Pardo",
+    tags=[Subcategory.INFILLING, SampleRate(44100), Channels(1), 'sound-effects']
+)
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 init_model_choice = open("DEFAULT_MODEL").read().strip()
+MODEL_CHOICES = Interface.available_models()
 
-# load the init model
-interface.load_finetuned(init_model_choice)
-
-# The other fine-tuned models the interface offers are fetched in the worker, which runs
-# process_fn, rather than inside the GPU call the first time a request picks one
+# pyharp runs process_fn in a worker, which imports this file as __mp_main__ and starts with the
+# server, so the default model is built and placed there alone. Each fine-tuned model is 2.4 GB,
+# too much to hold all of them, so the others are only fetched here, and built when a request
+# picks one.
 if __name__ != "__main__":
     from vampnet import download_finetuned
-    for _name in interface.available_models():
+
+    interface = Interface.default()
+    interface.load_finetuned(init_model_choice)
+    interface.to(DEVICE)
+    for _name in MODEL_CHOICES:
         if _name != "default":
             download_finetuned(_name)
-print(f"Loaded initial model: {init_model_choice}")
-print(f"Coarse model path: {interface.coarse_path}")
-print(f"C2F model path: {interface.c2f_path}\n")
-
-    
-def to_output(sig):
-    return sig.sample_rate, sig.cpu().detach().numpy()[0][0]
-
-MAX_DURATION_S = 10
-def load_audio(file):
-    print(file)
-    if isinstance(file, str):
-        filepath = file
-    elif isinstance(file, tuple):
-        # not a file
-        sr, samples = file
-        samples = samples / np.iinfo(samples.dtype).max
-        return sr, samples
-    else:
-        filepath = file.name
-    sig = at.AudioSignal.salient_excerpt(
-        filepath, duration=MAX_DURATION_S
-    )
-    sig = at.AudioSignal(filepath)
-    return to_output(sig)
-
+    print(f"Loaded initial model: {init_model_choice}")
 
 
 from torch_pitch_shift import pitch_shift, get_fast_shifts
 def shift_pitch(signal, interval: int):
     signal.samples = pitch_shift(
-        signal.samples, 
-        shift=interval, 
+        signal.samples,
+        shift=interval,
         sample_rate=signal.sample_rate
     )
     return signal
 
 
-def onsets(sig: at.AudioSignal, hop_length: int):
-    assert sig.batch_size == 1, "batch size must be 1"
-    assert sig.num_channels == 1, "mono signals only"
-    import librosa
-    onset_frame_idxs = librosa.onset.onset_detect(
-        y=sig.samples[0][0].detach().cpu().numpy(), sr=sig.sample_rate, 
-        hop_length=hop_length,
-        backtrack=True,
-    )
-    return onset_frame_idxs
-
-
-@spaces.GPU
-def new_vampnet_mask(self, 
-    codes, 
-    onset_idxs, 
-    width: int = 5, 
-    periodic_prompt=2, 
-    upper_codebook_mask=1,
-    drop_amt: float = 0.1
-):
-    from vampnet.newmask import mask_and, mask_or, onset_mask, periodic_mask, drop_ones, codebook_mask
-    mask =  mask_and(
-        periodic_mask(codes, periodic_prompt, 1, random_roll=False),
-        mask_or( # this re-masks the onsets, according to a periodic schedule
-            onset_mask(onset_idxs, codes, width=width),
-            periodic_mask(codes, periodic_prompt, 1, random_roll=False),
-        )
-    ).int()
-    # make sure the onset idxs themselves are unmasked
-    # mask = 1 - mask
-    mask[:, :, onset_idxs] = 0
-    mask = mask.cpu() # debug
-    mask = 1-drop_ones(1-mask, drop_amt)
-    mask = codebook_mask(mask, upper_codebook_mask)
-
-    
-    mask = mask.to(self.device)
-    return mask[:, :, :]
-
-
-
-@spaces.GPU
 def _vamp_internal(
-        seed, input_audio, model_choice, 
-        pitch_shift_amt, periodic_p, 
-        n_mask_codebooks, onset_mask_width, 
-        dropout, sampletemp, typical_filtering, 
-        typical_mass, typical_min_tokens, top_p, 
-        sample_cutoff, stretch_factor, sampling_steps, beat_mask_ms, num_feedback_steps, api=False, harp=False
+        seed, input_audio, model_choice,
+        pitch_shift_amt, periodic_p,
+        n_mask_codebooks, onset_mask_width,
+        dropout, sampletemp, typical_filtering,
+        typical_mass, typical_min_tokens, top_p,
+        sample_cutoff, stretch_factor, sampling_steps, beat_mask_ms, num_feedback_steps
     ):
-    if torch.cuda.is_available():
-        device = "cuda"
-    elif torch.backends.mps.is_available():
-        device = "mps"
-    else:
-        device = "cpu"
-
-
-    print("args!")
-    print(f"seed: {seed}")
-    print(f"input_audio: {input_audio}")
-    print(f"model_choice: {model_choice}")
-    print(f"pitch_shift_amt: {pitch_shift_amt}")
-    print(f"periodic_p: {periodic_p}")
-    print(f"n_mask_codebooks: {n_mask_codebooks}")
-    print(f"onset_mask_width: {onset_mask_width}")
-    print(f"dropout: {dropout}")
-    print(f"sampletemp: {sampletemp}")
-    print(f"typical_filtering: {typical_filtering}")
-    print(f"typical_mass: {typical_mass}")
-    print(f"typical_min_tokens: {typical_min_tokens}")
-    print(f"top_p: {top_p}")
-    print(f"sample_cutoff: {sample_cutoff}")
-    print(f"stretch_factor: {stretch_factor}")
-    print(f"sampling_steps: {sampling_steps}")
-    print(f"api: {api}")
-    print(f"beat_mask_ms: {beat_mask_ms}")
-    print(f"using device {interface.device}")
-    print(f"num feedback steps: {num_feedback_steps}")
-
-
     t0 = time.time()
-    interface.to(device)
-    print(f"using device {interface.device}")
     _seed = seed if seed > 0 else None
     if _seed is None:
         _seed = int(torch.randint(0, 2**32, (1,)).item())
@@ -185,7 +100,7 @@ def _vamp_internal(
         raise gr.Error("no input audio received!")
     sr, input_audio = input_audio
     input_audio = input_audio / np.iinfo(input_audio.dtype).max
-    
+
     sig = at.AudioSignal(input_audio, sr).to_mono()
 
     loudness = sig.loudness()
@@ -199,31 +114,15 @@ def _vamp_internal(
 
     codes = interface.encode(sig)
 
-    # mask = new_vampnet_mask(
-    #     interface, 
-    #     codes, 
-    #     onset_idxs=onsets(sig, hop_length=interface.codec.hop_length),
-    #     width=onset_mask_width,
-    #     periodic_prompt=periodic_p,
-    #     upper_codebook_mask=n_mask_codebooks,
-    #     drop_amt=dropout
-    # ).long()
-
-    
     mask = interface.build_mask(
         codes,
-        sig=sig, 
+        sig=sig,
         periodic_prompt=periodic_p,
         onset_mask_width=onset_mask_width,
         _dropout=dropout,
         upper_codebook_mask=n_mask_codebooks,
     )
     if beat_mask_ms > 0:
-        # bm = pmask.mask_or(
-        #     pmask.periodic_mask(
-        #         codes, periodic_p, random_roll=False
-        #     ),
-        # )
         mask = pmask.mask_and(
             mask, interface.make_beat_mask(
                 sig, after_beat_s=beat_mask_ms/1000.,
@@ -248,9 +147,9 @@ def _vamp_internal(
         time_stretch_factor=stretch_factor,
         return_mask=True,
         temperature=sampletemp,
-        typical_filtering=typical_filtering, 
-        typical_mass=typical_mass, 
-        typical_min_tokens=typical_min_tokens, 
+        typical_filtering=typical_filtering,
+        typical_mass=typical_mass,
+        typical_min_tokens=typical_min_tokens,
         top_p=top_p,
         seed=_seed,
         sample_cutoff=sample_cutoff,
@@ -260,55 +159,7 @@ def _vamp_internal(
     sig = interface.decode(codes)
     sig = sig.normalize(loudness)
 
-    import matplotlib.pyplot as plt
-    plt.clf()
-    # plt.imshow(mask_z[0].cpu().numpy(), aspect='auto
-
-    if harp: 
-        return sig
-
-    return to_output(sig[0]), to_output(sig[1])
-
-@spaces.GPU
-def vamp(input_audio, 
-        sampletemp,
-        top_p,
-        periodic_p, 
-        dropout,
-        stretch_factor, 
-        onset_mask_width, 
-        typical_filtering,
-        typical_mass,
-        typical_min_tokens,
-        seed, 
-        model_choice,
-        n_mask_codebooks,
-        pitch_shift_amt, 
-        sample_cutoff, 
-        sampling_steps, 
-        beat_mask_ms,
-        num_feedback_steps):
-    return _vamp_internal(
-        seed=seed,
-        input_audio=input_audio,
-        model_choice=model_choice,
-        pitch_shift_amt=pitch_shift_amt,
-        periodic_p=periodic_p,
-        n_mask_codebooks=n_mask_codebooks,
-        onset_mask_width=onset_mask_width,
-        dropout=dropout,
-        sampletemp=sampletemp,
-        typical_filtering=typical_filtering,
-        typical_mass=typical_mass,
-        typical_min_tokens=typical_min_tokens,
-        top_p=top_p,
-        sample_cutoff=sample_cutoff,
-        stretch_factor=stretch_factor,
-        sampling_steps=sampling_steps,
-        beat_mask_ms=beat_mask_ms,
-        num_feedback_steps=num_feedback_steps,
-        api=False,
-    )
+    return sig
 
 
 # Generation covers the whole input, stretched, in 10 s chunks, so the cost grows with the
@@ -330,6 +181,8 @@ def process_fn(input_audio, sampletemp, periodic_p, dropout, n_mask_codebooks, m
                stretch_factor):
     # Checked before the GPU is reserved, so a request the model cannot take spends none of the
     # caller's ZeroGPU allowance
+    if model_choice not in MODEL_CHOICES:
+        raise gr.Error(f"Unknown model: {model_choice}")
     if not input_audio:
         raise gr.Error("Provide the audio to rework.")
     try:
@@ -367,7 +220,7 @@ def harp_vamp(input_audio, sampletemp, periodic_p, dropout, n_mask_codebooks, mo
     input_audio = (sig.sample_rate, input_audio)
 
     sig =  _vamp_internal(
-        seed=0, 
+        seed=0,
         input_audio=input_audio,
         model_choice=model_choice,
         pitch_shift_amt=0,
@@ -385,8 +238,6 @@ def harp_vamp(input_audio, sampletemp, periodic_p, dropout, n_mask_codebooks, mo
         sampling_steps=36,
         beat_mask_ms=int(0),
         num_feedback_steps=1,
-        api=False, 
-        harp=True,
     )
 
     sig = sig.detach().cpu()
@@ -403,26 +254,6 @@ def harp_vamp(input_audio, sampletemp, periodic_p, dropout, n_mask_codebooks, mo
 
 if __name__ == "__main__":
     with gr.Blocks() as demo:
-         #NEW: HARP endpoint (new PyHARP API)
-        harp_model_card = ModelCard(
-            name="VampNet",
-            description=(
-                "Generates new audio by remixing and extending your input, filling in "
-                "the blanks to create a coherent continuation in the same style. "
-                "Available models include: cat10: fine-tuned on expressive cat sounds; "
-                "lazaro-ros-sep: percussion extracted from the Lázaro Ros discography, "
-                "rich in Afro-Cuban rhythm supermetroid: trained on Super Metroid "
-                "(SNES) gameplay audio for retro, game-style soundscapes. Generation "
-                "runs in mono at the model's own rate; the result is resampled back to "
-                "the input's rate and copied across its channels, so the output has the "
-                "format it came in with but no stereo image. Repository: "
-                "https://github.com/hugofloresgarcia/vampnet Paper: "
-                "https://arxiv.org/abs/2307.04686"
-            ),
-            author="Hugo Flores Garcia, Prem Seetharaman, Rithesh Kumar, Bryan Pardo",
-            tags=[Subcategory.INFILLING, SampleRate(44100), Channels(1), 'sound-effects']
-        )
-
         harp_input_components = [
             gr.Audio(type="filepath", label="Input Audio").harp_required(True).set_info((
                 "The audio to rework. The model masks out parts of this recording "
@@ -436,7 +267,7 @@ if __name__ == "__main__":
                       info="How much of the input the model changes. 0 = close copy, higher = greater remix."),
             gr.Slider(label="Timbral Reference", value=3, minimum=1, maximum=14, step=1,
                       info="Controls the richness of timbre retained from the source audio. Lower = simpler, higher = richer timbre."),
-            gr.Dropdown(label="Model Choice", choices=list(interface.available_models()), value=init_model_choice, 
+            gr.Dropdown(label="Model Choice", choices=list(MODEL_CHOICES), value=init_model_choice,
                         info="Select which fine-tuned VampNet model to use."),
             gr.Slider(label="Time Stretch Factor", minimum=1, maximum=8, step=1, value=1,
                       info="Scales the output's length relative to the input. 1 = the same length, 2 = twice as long."),
@@ -449,15 +280,11 @@ if __name__ == "__main__":
         ]
 
         build_endpoint(
-            model_card=harp_model_card,
+            model_card=model_card,
             input_components=harp_input_components,
             output_components=harp_output_components,
             process_fn=process_fn,
         )
 
-    try:
-        demo.queue()
-        demo.launch(pwa=True, show_error=True)
-    except KeyboardInterrupt:
-        shutil.rmtree("gradio-outputs", ignore_errors=True)
-        raise
+    demo.queue()
+    demo.launch(pwa=True, show_error=True)