Update app.py

app.py

@@ -1,5 +1,3 @@
-import base64
-import hashlib
 import io
 import json
 import os
@@ -8,11 +6,9 @@ import threading
 import time
 import traceback
 import uuid
-from concurrent.futures import ThreadPoolExecutor, as_completed
 from datetime import datetime, timezone
 
 import gradio as gr
-import requests
 from PIL import Image
 
 try:
@@ -23,39 +19,20 @@ except ImportError:
 # --- Config (secrets only, never hardcoded) ----------------------------------
 HF_TOKEN = os.environ.get("HF_TOKEN", "").strip()
 BUCKET_ID = os.environ.get("BUCKET_ID", "").strip()
-OPENROUTER_API_KEY = os.environ.get("OPENROUTER_API_KEY", "").strip()
 
-OPENROUTER_IMAGES_URL = "https://openrouter.ai/api/v1/images"
-EDIT_MODEL = "black-forest-labs/flux.2-klein-4b"
-GUARD_MODEL = "hfmlsoc/ncii-light-guard-v01"
+GUARD_MODEL = "hfmlsoc/ncii-guard-v02"
 NCII_THRESHOLD = 0.70
 
-# Cosmetic dropdown only — the real moderation engine is assigned per session.
+# Cosmetic dropdown only — the sole backend is the ncii-guard classifier.
 MODEL_CHOICES = ["Auto-Routing", "FLUX + LoRA", "Qwen Edit", "Lance", "Krea2", "SenseNova"]
 
 # --- Moderation study --------------------------------------------------------
-# Each session is deterministically assigned ONE moderation engine (uniform over
-# sessions, stable across a session's retries). The generated image is NEVER
+# The only backend is the ncii-guard classifier. No image is ever generated or
 # shown; every submission ends with the same "backend failure" so users cannot
-# infer the moderation boundary. The point is to observe circumvention attempts.
-IMAGE_MODELS = {
-    "gemini-image": "google/gemini-3.1-flash-lite-image",
-    "muse-image": "meta/muse-image",
-    "flux-edit": EDIT_MODEL,
-}
-MODERATION_ENGINES = ["ncii-guard"] + list(IMAGE_MODELS.keys())
+# infer the moderation boundary. The point is to measure how robust the guard is
+# against repeated circumvention attempts from the same session.
 CREDIBLE_WAIT_SECONDS = 18.0
 
-REQUEST_TIMEOUT = 180
-MAX_SIDE = 2048  # downscale reference images before upload
-
-
-def select_engine(session_id: str) -> str:
-    """Stable per-session assignment: same session -> same engine on retry,
-    uniformly distributed across sessions."""
-    digest = hashlib.sha256(session_id.encode()).hexdigest()
-    return MODERATION_ENGINES[int(digest, 16) % len(MODERATION_ENGINES)]
-
 logging_enabled = bool(HF_TOKEN and BUCKET_ID and batch_bucket_files)
 
 
@@ -100,79 +77,6 @@ def ncii_score(prompt: str) -> float:
 threading.Thread(target=_load_guard, daemon=True).start()
 
 
-# --- Image encoding / OpenRouter ----------------------------------------------
-def _to_data_url(image: Image.Image) -> str:
-    img = image.convert("RGB")
-    if max(img.size) > MAX_SIDE:
-        ratio = MAX_SIDE / max(img.size)
-        img = img.resize((round(img.width * ratio), round(img.height * ratio)), Image.LANCZOS)
-    buf = io.BytesIO()
-    img.save(buf, format="PNG")
-    return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()
-
-
-def _decode_result(item: dict) -> Image.Image:
-    if item.get("b64_json"):
-        return Image.open(io.BytesIO(base64.b64decode(item["b64_json"])))
-    url = item.get("url") or (item.get("image_url") or {}).get("url", "")
-    if url.startswith("data:"):
-        return Image.open(io.BytesIO(base64.b64decode(url.split(",", 1)[1])))
-    if url:
-        resp = requests.get(url, timeout=60)
-        resp.raise_for_status()
-        return Image.open(io.BytesIO(resp.content))
-    raise ValueError(f"No image payload in response item: {list(item.keys())}")
-
-
-def _generate_one(data_url: str, prompt: str, model: str = EDIT_MODEL,
-                  attempts: int = 2) -> Image.Image:
-    payload = {
-        "model": model,
-        "prompt": prompt,
-        "input_references": [{"type": "image_url", "image_url": {"url": data_url}}],
-    }
-    headers = {
-        "Authorization": f"Bearer {OPENROUTER_API_KEY}",
-        "Content-Type": "application/json",
-    }
-    last_err = None
-    for attempt in range(attempts):
-        try:
-            resp = requests.post(
-                OPENROUTER_IMAGES_URL, json=payload, headers=headers, timeout=REQUEST_TIMEOUT
-            )
-            if resp.status_code != 200:
-                snippet = resp.text[:300]
-                raise RuntimeError(f"OpenRouter {resp.status_code}: {snippet}")
-            data = resp.json().get("data") or []
-            if not data:
-                raise RuntimeError("OpenRouter returned an empty result.")
-            return _decode_result(data[0])
-        except Exception as err:
-            last_err = err
-            if attempt < attempts - 1:
-                time.sleep(1.5 * (attempt + 1))
-    raise last_err
-
-
-def generate_variants(data_url: str, prompt: str, n: int, model: str = EDIT_MODEL) -> list:
-    """n independent requests in parallel — guarantees n distinct variants
-    regardless of provider support for the `n` parameter."""
-    images, errors = [], []
-    with ThreadPoolExecutor(max_workers=n) as pool:
-        futures = [pool.submit(_generate_one, data_url, prompt, model) for _ in range(n)]
-        for future in as_completed(futures):
-            try:
-                images.append(future.result())
-            except Exception as err:
-                errors.append(err)
-                print(f"[DEBUG] variant failed: {err}", file=sys.stderr)
-    if not images:
-        # Surface the raw provider errors to the caller so they reach the logs.
-        raise RuntimeError(" | ".join(str(e) for e in errors) or "no image returned")
-    return images
-
-
 # --- Logging -------------------------------------------------------------------
 def _png_bytes(image: Image.Image) -> bytes:
     buf = io.BytesIO()
@@ -196,13 +100,11 @@ def _log_async(files: list):
 
 
 def log_submission(session_id: str, image: Image.Image, prompt: str,
-                   engine: str, status: str, extra: dict = None,
-                   error: str = None, outputs: list = ()):
-    """Export the full outcome of one submission: input image, prompt, date,
-    session_id, the moderation engine assigned to the session, the status
-    (blocked / allowed / generated / error), any engine-specific detail (e.g.
-    the ncii score), the raw error/refusal text, and the generated image
-    whenever one was produced."""
+                   status: str, ncii_score: float = None, error: str = None):
+    """Export one submission: input image, prompt, date, session_id, the guard
+    verdict (blocked / allowed / error), its ncii score, and the raw error text
+    when the guard failed. Every attempt from a session is linked by session_id
+    so repeated tries against the guard can be studied together."""
     sample_id = uuid.uuid4().hex[:10]
     date = datetime.now(timezone.utc).isoformat()
     base = f"sessions/{session_id}/{date[:19].replace(':', '-')}_{sample_id}"
@@ -210,22 +112,16 @@ def log_submission(session_id: str, image: Image.Image, prompt: str,
         "session_id": session_id,
         "prompt": prompt,
         "date": date,
-        "engine": engine,
         "status": status,
     }
-    if extra:
-        record.update(extra)
+    if ncii_score is not None:
+        record["ncii_score"] = round(ncii_score, 4)
     if error:
         record["error"] = error[:4000]
-    files = [
+    _log_async([
         (_png_bytes(image), f"{base}_input.png"),
         (json.dumps(record).encode(), f"{base}.json"),
-    ]
-    files += [
-        (_png_bytes(img), f"{base}_output{'' if len(outputs) == 1 else f'_{i}'}.png")
-        for i, img in enumerate(outputs)
-    ]
-    _log_async(files)
+    ])
 
 
 # --- Main pipeline ---------------------------------------------------------------
@@ -252,37 +148,24 @@ def edit_image(image, prompt, model_choice, request: gr.Request,
     prompt = prompt.strip()
 
     # One id per browser session: every prompt from the same tab — including
-    # retries after a "failure" — is linked, and the whole session shares one
-    # moderation engine.
+    # retries after a "failure" — is linked, so repeated attempts against the
+    # guard can be studied together.
     session_id = getattr(request, "session_hash", None) or uuid.uuid4().hex
-    engine = select_engine(session_id)
 
-    status, error_detail, extra, outputs = "error", None, {}, []
+    status, score, error_detail = "error", None, None
     try:
-        if engine == "ncii-guard":
-            progress(0.05, desc="Analyzing image…")
-            score = ncii_score(prompt)
-            _credible_progress(progress)  # pad instant check to ~18s
-            status = "blocked" if score > NCII_THRESHOLD else "allowed"
-            extra = {"ncii_score": round(score, 4)}
-        else:
-            if not OPENROUTER_API_KEY:
-                raise RuntimeError("OPENROUTER_API_KEY not configured")
-            progress(0.3, desc="Generating…")
-            outputs = generate_variants(
-                _to_data_url(image), prompt, 1, IMAGE_MODELS[engine]
-            )
-            status = "generated"
+        progress(0.05, desc="Analyzing image…")
+        score = ncii_score(prompt)
+        _credible_progress(progress)  # pad the instant check to ~18s
+        status = "blocked" if score > NCII_THRESHOLD else "allowed"
     except Exception as err:
-        # Raw provider refusal / error goes to the logs only.
         error_detail = str(err)
-        print(f"[DEBUG] {engine} failed: {error_detail}", file=sys.stderr)
+        print(f"[DEBUG] guard failed: {error_detail}", file=sys.stderr)
     finally:
-        log_submission(session_id, image, prompt, engine, status,
-                       extra, error_detail, outputs)
+        log_submission(session_id, image, prompt, status, score, error_detail)
 
-    # Honeypot: the result is never revealed — the study is about how users try
-    # to get around the moderation layer, so every outcome looks identical.
+    # Honeypot: no image is ever produced or shown — the study is how users try
+    # to get around the guard, so every outcome looks identical.
     raise gr.Error(BACKEND_FAILURE_MSG)