Fix ZeroGPU pickling error by extracting plain values before GPU boundary

_vamp previously received the raw Gradio data dict (with unpicklable
component-object keys), causing PicklingError on lambdas. Now _vamp
accepts explicit plain-value params, and a wrapper extracts values from
the Gradio dict before the call crosses the ZeroGPU process boundary.

Made-with: Cursor

vampnet/app.py

@@ -99,111 +99,134 @@ def load_example_audio():
 
 
 @gpu
-def _vamp(data, return_mask=False):
+def _vamp(
+    _input_audio,
+    _num_steps,
+    _masktemp,
+    _sampletemp,
+    _top_p,
+    _prefix_s,
+    _suffix_s,
+    _rand_mask_intensity,
+    _periodic_p,
+    _periodic_w,
+    _n_conditioning_codebooks,
+    _dropout,
+    _use_coarse2fine,
+    _stretch_factor,
+    _onset_mask_width,
+    _typical_filtering,
+    _typical_mass,
+    _typical_min_tokens,
+    _beat_mask_width,
+    _beat_mask_downbeats,
+    _seed,
+    _n_mask_codebooks,
+    _pitch_shift_amt,
+    _sample_cutoff,
+    return_mask=False,
+):
     interface.to("cuda")
 
     out_dir = OUT_DIR / str(uuid.uuid4())
     out_dir.mkdir()
-    sig = at.AudioSignal(data[input_audio])
+    sig = at.AudioSignal(_input_audio)
     sig = interface.preprocess(sig)
 
     loudness = sig.loudness()
     print(f"input loudness is {loudness}")
 
-    if data[pitch_shift_amt] != 0:
-        sig = shift_pitch(sig, data[pitch_shift_amt])
+    if _pitch_shift_amt != 0:
+        sig = shift_pitch(sig, _pitch_shift_amt)
 
     z = interface.encode(sig)
 
-    ncc = data[n_conditioning_codebooks]
+    ncc = _n_conditioning_codebooks
 
     # build the mask
-    mask = pmask.linear_random(z, data[rand_mask_intensity])
+    mask = pmask.linear_random(z, _rand_mask_intensity)
     mask = pmask.mask_and(
         mask, pmask.inpaint(
             z,
-            interface.s2t(data[prefix_s]),
-            interface.s2t(data[suffix_s])
+            interface.s2t(_prefix_s),
+            interface.s2t(_suffix_s)
         )
     )
     mask = pmask.mask_and(
         mask, pmask.periodic_mask(
             z,
-            data[periodic_p],
-            data[periodic_w],
+            _periodic_p,
+            _periodic_w,
             random_roll=True
         )
     )
-    if data[onset_mask_width] > 0:
+    if _onset_mask_width > 0:
         mask = pmask.mask_or(
-            mask, pmask.onset_mask(sig, z, interface, width=data[onset_mask_width])
+            mask, pmask.onset_mask(sig, z, interface, width=_onset_mask_width)
         )
-    if data[beat_mask_width] > 0:
+    if _beat_mask_width > 0:
         beat_mask = interface.make_beat_mask(
             sig,
-            after_beat_s=(data[beat_mask_width]/1000), 
-            mask_upbeats=not data[beat_mask_downbeats],
+            after_beat_s=(_beat_mask_width/1000), 
+            mask_upbeats=not _beat_mask_downbeats,
         )
         mask = pmask.mask_and(mask, beat_mask)
 
     # these should be the last two mask ops
-    mask = pmask.dropout(mask, data[dropout])
+    mask = pmask.dropout(mask, _dropout)
     mask = pmask.codebook_unmask(mask, ncc)
-    mask = pmask.codebook_mask(mask, int(data[n_mask_codebooks]))
-
-
-
-    print(f"dropout {data[dropout]}")
-    print(f"masktemp {data[masktemp]}")
-    print(f"sampletemp {data[sampletemp]}")
-    print(f"top_p {data[top_p]}")
-    print(f"prefix_s {data[prefix_s]}")
-    print(f"suffix_s {data[suffix_s]}")
-    print(f"rand_mask_intensity {data[rand_mask_intensity]}")
-    print(f"num_steps {data[num_steps]}")
-    print(f"periodic_p {data[periodic_p]}")
-    print(f"periodic_w {data[periodic_w]}")
-    print(f"n_conditioning_codebooks {data[n_conditioning_codebooks]}")
-    print(f"use_coarse2fine {data[use_coarse2fine]}")
-    print(f"onset_mask_width {data[onset_mask_width]}")
-    print(f"beat_mask_width {data[beat_mask_width]}")
-    print(f"beat_mask_downbeats {data[beat_mask_downbeats]}")
-    print(f"stretch_factor {data[stretch_factor]}")
-    print(f"seed {data[seed]}")
-    print(f"pitch_shift_amt {data[pitch_shift_amt]}")
-    print(f"sample_cutoff {data[sample_cutoff]}")
-    
-    
-    _top_p = data[top_p] if data[top_p] > 0 else None
+    mask = pmask.codebook_mask(mask, int(_n_mask_codebooks))
+
+    print(f"dropout {_dropout}")
+    print(f"masktemp {_masktemp}")
+    print(f"sampletemp {_sampletemp}")
+    print(f"top_p {_top_p}")
+    print(f"prefix_s {_prefix_s}")
+    print(f"suffix_s {_suffix_s}")
+    print(f"rand_mask_intensity {_rand_mask_intensity}")
+    print(f"num_steps {_num_steps}")
+    print(f"periodic_p {_periodic_p}")
+    print(f"periodic_w {_periodic_w}")
+    print(f"n_conditioning_codebooks {_n_conditioning_codebooks}")
+    print(f"use_coarse2fine {_use_coarse2fine}")
+    print(f"onset_mask_width {_onset_mask_width}")
+    print(f"beat_mask_width {_beat_mask_width}")
+    print(f"beat_mask_downbeats {_beat_mask_downbeats}")
+    print(f"stretch_factor {_stretch_factor}")
+    print(f"seed {_seed}")
+    print(f"pitch_shift_amt {_pitch_shift_amt}")
+    print(f"sample_cutoff {_sample_cutoff}")
+
+    _top_p_val = _top_p if _top_p > 0 else None
     # save the mask as a txt file
     np.savetxt(out_dir / "mask.txt", mask[:,0,:].long().cpu().numpy())
 
-    _seed = data[seed] if data[seed] > 0 else None
+    _seed_val = _seed if _seed > 0 else None
     zv, mask_z = interface.coarse_vamp(
         z, 
         mask=mask,
-        sampling_steps=data[num_steps],
-        mask_temperature=data[masktemp]*10,
-        sampling_temperature=data[sampletemp],
+        sampling_steps=_num_steps,
+        mask_temperature=_masktemp*10,
+        sampling_temperature=_sampletemp,
         return_mask=True, 
-        typical_filtering=data[typical_filtering], 
-        typical_mass=data[typical_mass], 
-        typical_min_tokens=data[typical_min_tokens], 
-        top_p=_top_p,
+        typical_filtering=_typical_filtering, 
+        typical_mass=_typical_mass, 
+        typical_min_tokens=_typical_min_tokens, 
+        top_p=_top_p_val,
         gen_fn=interface.coarse.generate,
-        seed=_seed,
-        sample_cutoff=data[sample_cutoff],
+        seed=_seed_val,
+        sample_cutoff=_sample_cutoff,
     )
 
-    if use_coarse2fine: 
+    if _use_coarse2fine: 
         zv = interface.coarse_to_fine(
             zv, 
-            mask_temperature=data[masktemp]*10, 
-            sampling_temperature=data[sampletemp],
+            mask_temperature=_masktemp*10, 
+            sampling_temperature=_sampletemp,
             mask=mask,
-            sampling_steps=data[num_steps],
-            sample_cutoff=data[sample_cutoff], 
-            seed=_seed,
+            sampling_steps=_num_steps,
+            sample_cutoff=_sample_cutoff, 
+            seed=_seed_val,
         )
 
     sig = interface.to_signal(zv).cpu()
@@ -222,11 +245,41 @@ def _vamp(data, return_mask=False):
     else:
         return sig.path_to_file
 
+def _extract_and_call_vamp(data, return_mask):
+    """Extract plain values from Gradio data dict so only picklable args cross the ZeroGPU boundary."""
+    return _vamp(
+        _input_audio=data[input_audio],
+        _num_steps=data[num_steps],
+        _masktemp=data[masktemp],
+        _sampletemp=data[sampletemp],
+        _top_p=data[top_p],
+        _prefix_s=data[prefix_s],
+        _suffix_s=data[suffix_s],
+        _rand_mask_intensity=data[rand_mask_intensity],
+        _periodic_p=data[periodic_p],
+        _periodic_w=data[periodic_w],
+        _n_conditioning_codebooks=data[n_conditioning_codebooks],
+        _dropout=data[dropout],
+        _use_coarse2fine=data[use_coarse2fine],
+        _stretch_factor=data[stretch_factor],
+        _onset_mask_width=data[onset_mask_width],
+        _typical_filtering=data[typical_filtering],
+        _typical_mass=data[typical_mass],
+        _typical_min_tokens=data[typical_min_tokens],
+        _beat_mask_width=data[beat_mask_width],
+        _beat_mask_downbeats=data[beat_mask_downbeats],
+        _seed=data[seed],
+        _n_mask_codebooks=data[n_mask_codebooks],
+        _pitch_shift_amt=data[pitch_shift_amt],
+        _sample_cutoff=data[sample_cutoff],
+        return_mask=return_mask,
+    )
+
 def vamp(data):
-    return _vamp(data, return_mask=True)
+    return _extract_and_call_vamp(data, return_mask=True)
 
 def api_vamp(data):
-    return _vamp(data, return_mask=False)
+    return _extract_and_call_vamp(data, return_mask=False)
         
 def save_vamp(data):
     out_dir = OUT_DIR / "saved" / str(uuid.uuid4())