fix(zerogpu): build Separator inside @spaces.GPU to dodge RLock pickling error

Bug: ZeroGPU throws _pickle.PicklingError: cannot pickle '_thread.RLock'
when @separate runs on HuggingFace ZeroGPU.

Root cause: The @spaces.GPU decorator ships the function's arguments to
a worker process via multiprocessing.Queue (which uses pickle). The
previous code passed a live Separator instance as the first argument:

@spaces.GPU
def separate_on_gpu(separator: Separator, input_file: str): ...
^^^^^^^^^^^^^^^^^^^^
not picklable — holds a Logger (RLock)
and possibly torch CUDA state

Fix: Move Separator construction AND model loading INSIDE the
@spaces.GPU-decorated function. Only picklable primitives
(str, int, float, bool, dict) cross the queue boundary.

Before: After:
--------- --------
separator = build_separator(...) output = _run_separation_on_gpu(
separator.load_model(...) architecture=..., # str
output = separate_on_gpu( model_file_dir=..., # str
separator, # <-- crashes out_dir=..., # str
input_file out_format=..., # str
) norm_thresh=..., # float
...
model_filename=..., # str
input_file=..., # str
)
# Inside the worker:
# separator = build_separator(...)
# separator.load_model(...)
# return separator.separate(input_file)

Additional safety:
- Added _sanitize_for_pickle() that recursively strips non-picklable
items from arch_params (defensive — UI should already only pass
primitives, but a stray torch tensor would re-introduce the bug)
- Updated module docstring documenting the ZeroGPU pickling rule so
future contributors don't re-introduce the bug
- The retry logic now re-creates the Separator on each attempt so a
half-broken state from a failed librosa call doesn't poison the retry

Verified locally:
- pickle.dumps(all_args) succeeds (363 bytes) — would have raised
PicklingError before this fix
- CPU end-to-end: 2 stems from 2s test tone in 16s
- All 4 API endpoints still work

main/separator.py

@@ -3,8 +3,14 @@
 Thin wrapper around `audio_separator.separator.Separator` that:
   * Maps UI-friendly architecture names to the right params dict
     (`mdxc_params`, `mdx_params`, `vr_params`, `demucs_params`).
-  * Wraps the actual `separator.separate()` call in `@spaces.GPU` so it can
-    run on HuggingFace ZeroGPU hardware.
+  * Wraps the actual Separator construction + model loading + `separate()`
+    call in `@spaces.GPU` so it can run on HuggingFace ZeroGPU hardware.
+    CRITICAL: the `@spaces.GPU` decorator ships arguments to a worker
+    process via pickle, so the decorated function must take ONLY picklable
+    primitives (strings, dicts of primitives, numbers). A live `Separator`
+    instance holds a logger / torch objects containing `_thread.RLock`
+    which cannot be pickled — hence the construction must happen INSIDE
+    the decorator, not outside.
   * Parses the produced filenames back into stem-label → absolute-path pairs.
   * Applies known-good defaults for output bitrate and retries on the
     sporadic `librosa.ParameterError` that affects ~1 in 5 runs
@@ -74,6 +80,11 @@ def stem_label_from_filename(fname: str) -> str:
 # --------------------------------------------------------------------------- #
 #  Separator construction                                                     #
 # --------------------------------------------------------------------------- #
+# NOTE: `build_separator` is now only called from INSIDE the `@spaces.GPU`
+# boundary (in `_run_separation_on_gpu` below). It must NOT be called from
+# the main process, because the returned `Separator` instance cannot be
+# pickled across the ZeroGPU worker queue (it holds a `_thread.RLock`
+# via its logger and possibly torch CUDA state).
 
 def build_separator(
     *,
@@ -87,7 +98,8 @@ def build_separator(
     arch_params: dict[str, Any],
 ) -> "Separator":
     """Instantiate an `audio_separator.Separator` pre-configured for the
-    requested architecture.
+    requested architecture. Must be called from inside the `@spaces.GPU`
+    boundary — see module docstring.
     """
     if not HAS_SEPARATOR:
         raise RuntimeError(
@@ -150,23 +162,62 @@ def build_separator(
 
 
 # --------------------------------------------------------------------------- #
-#  GPU-wrapped separation                                                     #
+#  GPU-wrapped separation                                                      #
 # --------------------------------------------------------------------------- #
+# CRITICAL ZeroGPU pickling rule:
+#   The `@spaces.GPU` decorator ships the function's arguments to a worker
+#   process via `multiprocessing.Queue`, which uses pickle. The arguments
+#   here MUST all be picklable primitives:
+#     - str, int, float, bool
+#     - dict / list / tuple of the above
+#   A live `Separator` instance is NOT picklable — it holds a Python
+#   `logging.Logger` (which transitively owns a `_thread.RLock`) and may
+#   hold torch CUDA tensors / ONNX runtime sessions. Passing one across
+#   the queue raises `_pickle.PicklingError: cannot pickle '_thread.RLock'
+#   object`.
+#
+#   Therefore: build the Separator INSIDE this function, not outside.
 
 @spaces.GPU
-def separate_on_gpu(separator: "Separator", input_file: str) -> list[str]:
-    """Run the actual separation. Wrapped in `@spaces.GPU` so HuggingFace
-    ZeroGPU allocates a GPU for this call. When running locally (no `spaces`
-    module available), the decorator becomes a no-op and execution falls
-    through to CPU.
+def _run_separation_on_gpu(
+    *,
+    architecture: str,
+    model_file_dir: str,
+    out_dir: str,
+    out_format: str,
+    norm_thresh: float,
+    amp_thresh: float,
+    batch_size: int,
+    arch_params: dict[str, Any],
+    model_filename: str,
+    input_file: str,
+) -> list[str]:
+    """Build the Separator, load the model, and run separation — all inside
+    the `@spaces.GPU` boundary so ZeroGPU allocates a GPU for the whole
+    process. Returns the list of output file paths produced by
+    `Separator.separate()`.
 
     Includes a one-shot retry on the sporadic `librosa.ParameterError`
     ("Audio buffer is not finite") that affects ~1 in 5 runs on certain
-    inputs (upstream issue #57).
+    inputs (upstream issue #57). The retry re-creates the Separator from
+    scratch to dodge any half-broken state.
     """
     last_exc: Exception | None = None
     for attempt in range(_MAX_RETRIES + 1):
         try:
+            # Build a fresh Separator on every attempt — if librosa left
+            # the previous one in a bad state, we want a clean slate.
+            separator = build_separator(
+                architecture=architecture,
+                model_file_dir=model_file_dir,
+                out_dir=out_dir,
+                out_format=out_format,
+                norm_thresh=norm_thresh,
+                amp_thresh=amp_thresh,
+                batch_size=batch_size,
+                arch_params=arch_params,
+            )
+            separator.load_model(model_filename=model_filename)
             return separator.separate(input_file)
         except Exception as exc:
             last_exc = exc
@@ -204,8 +255,13 @@ def separate_core(
 ) -> dict[str, str]:
     """Run separation and return ``{stem_label: absolute_path}``.
 
-    Returns an empty dict on missing input. Raises on actual errors so the UI
-    layer can surface them via `gr.Warning`.
+    Returns an empty dict on missing input. Raises on actual errors so the
+    UI layer can surface them via `gr.Warning`.
+
+    This function runs in the main process. It only prepares picklable
+    arguments and then calls `_run_separation_on_gpu(...)` — the
+    `@spaces.GPU`-decorated worker function — which builds the Separator
+    and runs inference inside the ZeroGPU-allocated worker.
     """
     if not input_file:
         return {}
@@ -221,8 +277,27 @@ def separate_core(
         model_display_name,
     )
 
+    # Prepare the per-input output directory from the main process — this
+    # is just filesystem work, no GPU needed, and the directory needs to
+    # exist before the worker writes stems into it.
     out_dir = prepare_output_dir(input_file, output_dir)
-    separator = build_separator(
+
+    # Translate the UI-friendly model name into the filename the library
+    # expects. This is a pure dict lookup — safe to do in the main process.
+    model_filename = resolve_model_filename(architecture, model_display_name)
+
+    # Make sure arch_params only contains picklable primitives before
+    # shipping it across the ZeroGPU queue. (Defensive — the UI layer
+    # should already only pass primitives, but a stray torch tensor here
+    # would re-introduce the same PicklingError.)
+    arch_params_clean = _sanitize_for_pickle(arch_params)
+
+    progress(0.10, desc="Loading model (first use downloads it)…")
+    progress(0.30, desc="Separating audio…")
+
+    # All heavy work (Separator construction, model load, inference)
+    # happens inside the @spaces.GPU boundary.
+    output_files = _run_separation_on_gpu(
         architecture=architecture,
         model_file_dir=model_file_dir,
         out_dir=out_dir,
@@ -230,16 +305,11 @@ def separate_core(
         norm_thresh=norm_thresh,
         amp_thresh=amp_thresh,
         batch_size=batch_size,
-        arch_params=arch_params,
+        arch_params=arch_params_clean,
+        model_filename=model_filename,
+        input_file=input_file,
     )
 
-    progress(0.10, desc="Loading model (first use downloads it)…")
-    model_filename = resolve_model_filename(architecture, model_display_name)
-    separator.load_model(model_filename=model_filename)
-
-    progress(0.30, desc="Separating audio…")
-    output_files = separate_on_gpu(separator, input_file)
-
     progress(0.95, desc="Finalizing stems…")
     stems: dict[str, str] = {}
     for rel_path in output_files:
@@ -258,3 +328,23 @@ def separate_core(
 
     log.info("Stems produced: %s", ", ".join(stems.keys()) or "(none)")
     return stems
+
+
+def _sanitize_for_pickle(obj: Any) -> Any:
+    """Recursively strip non-picklable items from a dict/list structure.
+
+    `@spaces.GPU` ships args across a multiprocessing queue, which uses
+    pickle. This helper ensures we never accidentally pass a torch tensor,
+    numpy array, or other unpicklable object through. Numbers, strings,
+    bools, lists, and dicts of primitives pass through unchanged.
+    """
+    if isinstance(obj, dict):
+        return {str(k): _sanitize_for_pickle(v) for k, v in obj.items()}
+    if isinstance(obj, (list, tuple)):
+        return [_sanitize_for_pickle(v) for v in obj]
+    if isinstance(obj, (str, int, float, bool)) or obj is None:
+        return obj
+    # Anything else (torch.Tensor, numpy array, custom object) — convert
+    # to its string representation so we don't crash the queue.
+    log.warning("Stripping non-picklable %s from arch_params: %r", type(obj).__name__, obj)
+    return str(obj)