Upload 133 files

app.py

@@ -24,9 +24,14 @@ import os
 import re
 import time
 import threading
+import traceback
 from datetime import datetime
 from pathlib import Path
 
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
+)
 logger = logging.getLogger(__name__)
 
 # ── Container environment fixes ──────────────────────────────────────
@@ -90,18 +95,32 @@ except (ImportError, AttributeError):
 def _is_quota_error(exc: BaseException) -> bool:
     """Return True if *exc* is a ZeroGPU quota or session error.
 
-    Matches quota-exceeded errors ("exceeded your GPU quota") and expired
-    proxy tokens ("Expired ZeroGPU proxy token") — both mean the GPU is
-    unavailable and the user should retry later.
+    Matches quota-exceeded errors ("exceeded your GPU quota"), expired
+    proxy tokens ("Expired ZeroGPU proxy token"), and aborted GPU tasks
+    ("GPU task aborted") — all mean the GPU is unavailable and the user
+    should retry later.
     """
     msg = str(exc).lower()
     if "exceeded" in msg and "gpu quota" in msg:
         return True
     if "expired" in msg and "zerogpu" in msg:
         return True
+    if "gpu task aborted" in msg:
+        return True
     return False
 
 
+def _is_zerogpu_abort(exc: BaseException) -> bool:
+    """Return True if *exc* is specifically a ZeroGPU 'GPU task aborted' error.
+
+    This happens when ZeroGPU's internal multiprocessing kills the worker
+    mid-execution — typically because the GPU allocation timed out, a
+    concurrent request conflicted, or ZeroGPU infrastructure had an issue.
+    """
+    msg = str(exc).lower()
+    return "gpu task aborted" in msg
+
+
 def _load_model_to_device(
     pretrained_path: str,
     *,
@@ -209,8 +228,8 @@ def _persist_session_meta(output_dir: str, label: str, meta: dict) -> None:
         p = Path(output_dir) / _SESSION_META_FILE
         data = {"label": label, **meta}
         p.write_text(_json.dumps(data, indent=2))
-    except Exception:
-        pass  # best-effort
+    except Exception as e:
+        logger.debug("Failed to persist session metadata: %s", e)
 
 
 def _recover_sessions_from_disk() -> None:
@@ -1280,6 +1299,8 @@ def benchmark(
             except Exception as e:
                 nonlocal run_error
                 run_error = e
+                logger.error("Benchmark pipeline failed: %s\n%s", e, traceback.format_exc())
+                on_log(f"\n--- TRACEBACK ---\n{traceback.format_exc()}")
 
         worker = threading.Thread(target=run_pipeline, daemon=True)
         worker.start()
@@ -1635,6 +1656,8 @@ def benchmark_multi_model(
             except Exception as e:
                 nonlocal run_error
                 run_error = e
+                logger.error("Tournament pipeline failed: %s\n%s", e, traceback.format_exc())
+                on_log(f"\n--- TRACEBACK ---\n{traceback.format_exc()}")
 
         worker = threading.Thread(target=run_pipeline, daemon=True)
         worker.start()
@@ -1912,9 +1935,10 @@ def obliterate(model_choice: str, method_choice: str,
                     f"Adaptive: using architecture default `{method}` "
                     f"(no telemetry data yet)"
                 )
-        except Exception:
+        except Exception as e:
+            logger.warning("Adaptive method detection failed: %s", e, exc_info=True)
             method = "advanced"
-            _adaptive_info = "Adaptive: fallback to `advanced` (could not detect architecture)"
+            _adaptive_info = f"Adaptive: fallback to `advanced` (detection error: {e})"
 
     # Early validation: gated model access
     from obliteratus.presets import is_gated
@@ -1982,6 +2006,11 @@ def obliterate(model_choice: str, method_choice: str,
 
     def run_pipeline():
         try:
+            _t_pipeline_start = time.time()
+            on_log(f"[timing] Pipeline thread started (ZeroGPU allocation: {300}s)")
+            if _ZEROGPU_AVAILABLE:
+                on_log("[timing] Running on ZeroGPU — GPU may be deallocated if pipeline exceeds duration")
+
             # Load prompts — custom overrides dataset dropdown
             if use_custom:
                 on_log("Using custom user-provided prompts...")
@@ -1993,6 +2022,7 @@ def obliterate(model_choice: str, method_choice: str,
                 on_log(f"Loading dataset: {dataset_key}...")
                 harmful_all, harmless_all = load_dataset_source(dataset_key)
                 on_log(f"Dataset loaded: {len(harmful_all)} harmful, {len(harmless_all)} harmless prompts")
+            on_log(f"[timing] Dataset loaded at +{time.time() - _t_pipeline_start:.1f}s")
 
             # Apply volume cap (-1 = use all)
             if prompt_volume > 0:
@@ -2076,6 +2106,9 @@ def obliterate(model_choice: str, method_choice: str,
                 pipeline.run()
         except Exception as e:
             error_ref[0] = e
+            tb = traceback.format_exc()
+            logger.error("Obliteration pipeline failed: %s\n%s", e, tb)
+            on_log(f"\n--- TRACEBACK ---\n{tb}")
 
     if use_custom:
         source_label = "Custom (user-provided)"
@@ -2097,32 +2130,112 @@ def obliterate(model_choice: str, method_choice: str,
     worker.start()
 
     # Stream log updates while pipeline runs (max 400 hours for large-model Optuna optimization)
+    # Wrapped in try/except to catch ZeroGPU "GPU task aborted" — the abort is thrown
+    # INTO the generator at the yield/sleep points, not into the worker thread.
     _max_pipeline_secs = 400 * 60 * 60
     _pipeline_start = time.time()
     status_msg = "**Obliterating\u2026** (0s)"
-    while worker.is_alive():
-        status_msg = f"**Obliterating\u2026** ({_elapsed()})"
-        if len(log_lines) > last_yielded[0]:
-            last_yielded[0] = len(log_lines)
-            yield status_msg, "\n".join(log_lines), gr.update(), gr.update(), gr.update(), gr.update()
+    try:
+        while worker.is_alive():
+            status_msg = f"**Obliterating\u2026** ({_elapsed()})"
+            if len(log_lines) > last_yielded[0]:
+                last_yielded[0] = len(log_lines)
+                yield status_msg, "\n".join(log_lines), gr.update(), gr.update(), gr.update(), gr.update()
+            else:
+                yield status_msg, "\n".join(log_lines), gr.update(), gr.update(), gr.update(), gr.update()
+            if time.time() - _pipeline_start > _max_pipeline_secs:
+                log_lines.append("\nTIMEOUT: Pipeline exceeded 400-hour limit.")
+                break
+            time.sleep(0.5)
+    except Exception as e:
+        # ZeroGPU can abort the generator mid-yield with "GPU task aborted"
+        # or other errors.  Catch here so we can show a useful message and
+        # reset state instead of leaving status stuck on "obliterating".
+        tb = traceback.format_exc()
+        logger.error("Obliterate generator interrupted: %s\n%s", e, tb)
+        log_lines.append(f"\n--- INTERRUPTED ---")
+        log_lines.append(f"Generator killed after {_elapsed()}: {type(e).__qualname__}: {e}")
+        log_lines.append(f"\nLast pipeline log before abort:")
+        for line in log_lines[-10:]:
+            if line.startswith("[timing]") or line.startswith("  ["):
+                log_lines.append(f"  {line}")
+        with _lock:
+            _state["status"] = "idle"
+            _state["obliterate_started_at"] = None
+            _state["log"] = log_lines
+        err_msg = str(e).strip() or repr(e)
+        if _is_zerogpu_abort(e):
+            hint = (
+                "\n\n**ZeroGPU aborted the GPU task** after " + _elapsed() + ". "
+                "This is a known ZeroGPU issue — common causes:\n"
+                "- **Timeout:** Model loading + probing exceeded the 5-minute GPU allocation\n"
+                "- **Concurrent users:** Another request conflicted with yours\n"
+                "- **ZeroGPU internal error:** Transient infrastructure issue\n\n"
+                "**Try:** Click Obliterate again (often works on retry). "
+                "If it keeps failing, try a smaller model or reduce prompt volume."
+            )
+        elif _is_quota_error(e):
+            hint = "\n\n**ZeroGPU quota exceeded.** Wait a few minutes and retry."
         else:
-            yield status_msg, "\n".join(log_lines), gr.update(), gr.update(), gr.update(), gr.update()
-        if time.time() - _pipeline_start > _max_pipeline_secs:
-            log_lines.append("\nTIMEOUT: Pipeline exceeded 400-hour limit.")
-            break
-        time.sleep(0.5)
+            hint = ""
+        yield (
+            f"**Error:** {type(e).__qualname__}: {err_msg}{hint}",
+            "\n".join(log_lines), get_chat_header(),
+            gr.update(), gr.update(), gr.update(),
+        )
+        return
 
     worker.join(timeout=30)
 
     # Handle error
     if error_ref[0] is not None:
-        err_msg = str(error_ref[0]) or repr(error_ref[0])
-        log_lines.append(f"\nERROR: {err_msg}")
+        err = error_ref[0]
+        err_type = type(err).__qualname__
+        err_str = str(err).strip()
+        if err_str:
+            err_msg = f"{err_type}: {err_str}"
+        else:
+            err_msg = repr(err)
+        # Classify the error for actionable user guidance
+        err_lower = err_msg.lower()
+        if _is_zerogpu_abort(err):
+            err_hint = (
+                "\n\n**ZeroGPU task aborted.** The GPU worker was killed mid-pipeline. "
+                "This is a known ZeroGPU infrastructure issue — common causes:\n"
+                "- **Timeout:** Model loading + probing exceeded the 5-minute GPU allocation\n"
+                "- **Concurrent users:** Another request conflicted with yours\n"
+                "- **ZeroGPU internal error:** Transient infrastructure issue\n\n"
+                "**Try:** Click Obliterate again (often works on retry). "
+                "If it keeps failing, try a smaller model or reduce prompt volume."
+            )
+        elif _is_quota_error(err):
+            err_hint = (
+                "\n\n**ZeroGPU quota exceeded.** Your HuggingFace GPU quota has "
+                "been used up. Wait a few minutes and try again, or run locally."
+            )
+        elif "cuda" in err_lower or "out of memory" in err_lower:
+            err_hint = (
+                "\n\n**GPU out of memory.** Try a smaller model or enable "
+                "quantization (the pipeline auto-detects this for large models)."
+            )
+        elif "meta" in err_lower and "tensor" in err_lower:
+            err_hint = (
+                "\n\n**ZeroGPU device error.** The GPU was deallocated mid-run. "
+                "This is a transient ZeroGPU issue — please retry."
+            )
+        elif "connection" in err_lower or "timeout" in err_lower or "resolve" in err_lower:
+            err_hint = (
+                "\n\n**Network error.** Could not download model weights. "
+                "Check your internet connection and try again."
+            )
+        else:
+            err_hint = ""
+        log_lines.append(f"\nERROR ({err_type}): {err_msg}")
         with _lock:
             _state["status"] = "idle"
             _state["obliterate_started_at"] = None
             _state["log"] = log_lines
-        yield f"**Error:** {err_msg}", "\n".join(log_lines), get_chat_header(), gr.update(), gr.update(), gr.update()
+        yield f"**Error:** {err_msg}{err_hint}", "\n".join(log_lines), get_chat_header(), gr.update(), gr.update(), gr.update()
         return
 
     # Success — keep model in memory for chat.
@@ -2285,7 +2398,8 @@ def obliterate(model_choice: str, method_choice: str,
                     can_generate = True
                     log_lines.append("Reloaded in 4-bit — chat is ready!")
                 except Exception as e:
-                    log_lines.append(f"4-bit reload failed: {e}")
+                    logger.error("4-bit reload failed: %s\n%s", e, traceback.format_exc())
+                    log_lines.append(f"4-bit reload failed ({type(e).__qualname__}): {e}")
                     _clear_gpu()
 
             # -- Attempt 2: CPU offloading (slower but no extra dependencies)
@@ -2327,7 +2441,8 @@ def obliterate(model_choice: str, method_choice: str,
                     can_generate = True
                     log_lines.append("Reloaded with CPU offload — chat is ready (may be slower).")
                 except Exception as e:
-                    log_lines.append(f"CPU offload reload failed: {e}")
+                    logger.error("CPU offload reload failed: %s\n%s", e, traceback.format_exc())
+                    log_lines.append(f"CPU offload reload failed ({type(e).__qualname__}): {e}")
                     log_lines.append("Chat unavailable. Load the saved model on a larger instance.")
                     with _lock:
                         _state["status"] = "idle"
@@ -2374,8 +2489,12 @@ def obliterate(model_choice: str, method_choice: str,
 
     except Exception as e:
         # Ensure status never gets stuck on "obliterating"
-        err_msg = str(e) or repr(e)
+        tb = traceback.format_exc()
+        logger.error("Post-pipeline error: %s\n%s", e, tb)
+        err_type = type(e).__qualname__
+        err_msg = f"{err_type}: {str(e).strip() or repr(e)}"
         log_lines.append(f"\nERROR (post-pipeline): {err_msg}")
+        log_lines.append(f"\n--- TRACEBACK ---\n{tb}")
         with _lock:
             _state["status"] = "idle"
             _state["obliterate_started_at"] = None
@@ -2460,7 +2579,8 @@ def chat_respond(message: str, history: list[dict], system_prompt: str,
                     _needs_reload = True
                 else:
                     model.to(dev.get_device())
-        except Exception:
+        except Exception as e:
+            logger.warning("Model device check failed, triggering reload: %s", e)
             _needs_reload = True
 
     # Reload from saved checkpoint if model is missing or stale
@@ -2502,8 +2622,15 @@ def chat_respond(message: str, history: list[dict], system_prompt: str,
                     _state["model"] = model
                     _state["tokenizer"] = tokenizer
                     _state["status"] = "ready"
-            except Exception:
-                yield "Model failed to reload from checkpoint. Try re-obliterating."
+            except Exception as e:
+                tb = traceback.format_exc()
+                logger.error("Chat model reload failed: %s\n%s", e, tb)
+                err_type = type(e).__qualname__
+                err_str = str(e).strip() or repr(e)
+                yield (
+                    f"Model failed to reload from checkpoint: **{err_type}:** {err_str}\n\n"
+                    "Try re-obliterating the model. If this persists, check the Space logs."
+                )
                 return
         else:
             yield "No model loaded yet. Go to the **Obliterate** tab first and liberate a model."
@@ -2570,6 +2697,7 @@ def chat_respond(message: str, history: list[dict], system_prompt: str,
                 model.generate(**kwargs)
         except Exception as e:
             gen_error[0] = e
+            logger.error("Chat generation failed: %s\n%s", e, traceback.format_exc())
             # Signal the streamer to stop so the main thread doesn't hang
             try:
                 streamer.end()
@@ -2584,8 +2712,9 @@ def chat_respond(message: str, history: list[dict], system_prompt: str,
         for token in streamer:
             partial += token
             yield partial
-    except Exception:
+    except Exception as e:
         # Streamer timeout or broken pipe — yield whatever we have so far
+        logger.warning("Chat streamer interrupted: %s", e)
         if partial:
             yield partial
 

obliteratus/abliterate.py

@@ -936,11 +936,15 @@ class AbliterationPipeline:
         for h in self._steering_hooks:
             h.remove()
         self._steering_hooks.clear()
+        _t0 = time.time()
         self._summon()
+        self.log(f"[timing] SUMMON complete at +{time.time() - _t0:.1f}s")
         self._free_gpu_memory()
         self._probe()
+        self.log(f"[timing] PROBE complete at +{time.time() - _t0:.1f}s")
         self._free_gpu_memory()
         self._distill()
+        self.log(f"[timing] DISTILL complete at +{time.time() - _t0:.1f}s")
         # Free raw per-prompt activations now that means/subspaces are extracted
         self._harmful_acts.clear()
         self._harmless_acts.clear()
@@ -955,10 +959,14 @@ class AbliterationPipeline:
         self._free_gpu_memory()
         self._capture_baseline_kl_logits()
         self._excise()
+        self.log(f"[timing] EXCISE complete at +{time.time() - _t0:.1f}s")
         self._free_gpu_memory()
         self._verify()
+        self.log(f"[timing] VERIFY complete at +{time.time() - _t0:.1f}s")
         self._free_gpu_memory()
-        return self._rebirth()
+        result = self._rebirth()
+        self.log(f"[timing] REBIRTH complete at +{time.time() - _t0:.1f}s — pipeline finished")
+        return result
 
     # ── Stage 1: SUMMON ─────────────────────────────────────────────────
 

obliteratus/tourney.py

@@ -1135,6 +1135,8 @@ class TourneyRunner:
             return True
         if "expired" in msg and "zerogpu" in msg:
             return True
+        if "gpu task aborted" in msg:
+            return True
         return False
 
     def _run_one_method(self, method, harmful, harmless, save_dir, verify_sz, gpu_wrapper):