Upload 133 files

app.py

@@ -2050,6 +2050,20 @@ def _format_multi_model_results(results: list[dict], context: dict | None = None
 # Staged GPU wrapper for obliteration (tourney-style per-stage allocation)
 # ---------------------------------------------------------------------------
 
+def _noop_callback(*args, **kwargs):
+    """Module-level no-op, used as a picklable placeholder for callbacks."""
+    pass
+
+
+def _restore_and_run_stage(pipeline, stage_method_name):
+    """Restore pipeline to GPU and run the named stage method.
+
+    Module-level function so it is picklable for ZeroGPU serialization.
+    """
+    pipeline._restore_to_gpu()
+    getattr(pipeline, stage_method_name)()
+
+
 @spaces.GPU(duration=300)
 def _obliterate_gpu_run(fn, *args, **kwargs):
     """Execute *fn* inside a ZeroGPU GPU allocation.
@@ -2062,6 +2076,25 @@ def _obliterate_gpu_run(fn, *args, **kwargs):
     return fn(*args, **kwargs)
 
 
+def _gpu_run_picklable(pipeline, fn, *args, **kwargs):
+    """Run *fn* via ``_obliterate_gpu_run`` after stripping unpicklable callbacks.
+
+    ZeroGPU pickles arguments to send them to a GPU worker process.  The
+    pipeline's ``_on_stage`` and ``_on_log`` callbacks are local closures
+    that cannot be pickled, so we temporarily replace them with a
+    module-level no-op before the GPU call and restore them afterwards.
+    """
+    saved_on_stage = pipeline._on_stage
+    saved_on_log = pipeline._on_log
+    pipeline._on_stage = _noop_callback
+    pipeline._on_log = _noop_callback
+    try:
+        return _obliterate_gpu_run(fn, *args, **kwargs)
+    finally:
+        pipeline._on_stage = saved_on_stage
+        pipeline._on_log = saved_on_log
+
+
 def obliterate(model_choice: str, method_choice: str,
                prompt_volume_choice: str, dataset_source_choice: str,
                custom_harmful: str, custom_harmless: str,
@@ -2317,43 +2350,31 @@ def obliterate(model_choice: str, method_choice: str,
                 if method == "informed":
                     # Informed pipeline: SUMMON+PROBE | ANALYZE+DISTILL+EXCISE | VERIFY+REBIRTH
                     on_log("\n\u26a1 [staged] GPU Stage 1/3: SUMMON + PROBE")
-                    _obliterate_gpu_run(pipeline.run_stage_summon_probe, time.time())
+                    _gpu_run_picklable(pipeline, pipeline.run_stage_summon_probe, time.time())
                     pipeline._offload_to_cpu()
                     on_log("[staged] GPU released after Stage 1\n")
 
                     on_log("\u26a1 [staged] GPU Stage 2/3: ANALYZE + DISTILL + EXCISE")
-                    def _informed_s2():
-                        pipeline._restore_to_gpu()
-                        pipeline.run_stage_analyze_distill_excise()
-                    _obliterate_gpu_run(_informed_s2)
+                    _gpu_run_picklable(pipeline, _restore_and_run_stage, pipeline, "run_stage_analyze_distill_excise")
                     pipeline._offload_to_cpu()
                     on_log("[staged] GPU released after Stage 2\n")
 
                     on_log("\u26a1 [staged] GPU Stage 3/3: VERIFY + REBIRTH")
-                    def _informed_s3():
-                        pipeline._restore_to_gpu()
-                        pipeline.run_stage_verify_rebirth_informed()
-                    _obliterate_gpu_run(_informed_s3)
+                    _gpu_run_picklable(pipeline, _restore_and_run_stage, pipeline, "run_stage_verify_rebirth_informed")
                 else:
                     # Standard pipeline: SUMMON+PROBE | DISTILL+EXCISE | VERIFY+REBIRTH
                     on_log("\n\u26a1 [staged] GPU Stage 1/3: SUMMON + PROBE")
-                    _obliterate_gpu_run(pipeline.run_stage_summon_probe, time.time())
+                    _gpu_run_picklable(pipeline, pipeline.run_stage_summon_probe, time.time())
                     pipeline._offload_to_cpu()
                     on_log("[staged] GPU released after Stage 1\n")
 
                     on_log("\u26a1 [staged] GPU Stage 2/3: DISTILL + EXCISE")
-                    def _standard_s2():
-                        pipeline._restore_to_gpu()
-                        pipeline.run_stage_distill_excise()
-                    _obliterate_gpu_run(_standard_s2)
+                    _gpu_run_picklable(pipeline, _restore_and_run_stage, pipeline, "run_stage_distill_excise")
                     pipeline._offload_to_cpu()
                     on_log("[staged] GPU released after Stage 2\n")
 
                     on_log("\u26a1 [staged] GPU Stage 3/3: VERIFY + REBIRTH")
-                    def _standard_s3():
-                        pipeline._restore_to_gpu()
-                        pipeline.run_stage_verify_rebirth()
-                    _obliterate_gpu_run(_standard_s3)
+                    _gpu_run_picklable(pipeline, _restore_and_run_stage, pipeline, "run_stage_verify_rebirth")
             else:
                 # ── Local/non-ZeroGPU: single-shot execution ──────────────
                 on_log(f"[timing] Running locally (no GPU time limit)")