feat: per-image moderation check + WebP upload pipeline

- Run the VLM image-context (safety-rating JSON) check on every generated
image inside the same @spaces.GPU window, parse it best-effort, and return
it alongside the image, seed and R2 URL (generate tab + prompt_to_video_assets).
- Encode assets as WebP before R2 upload and return WebP files/presigned URLs;
downstream now only ever sees WebP. All 3 Spaces decode WebP input via PIL.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

app.py

@@ -1,5 +1,6 @@
 import inspect
 import io
+import json
 import os
 import queue
 import random
@@ -151,6 +152,82 @@ vlm_model.to("cuda").eval()
 print("Assistant loaded!")
 
 
+# =============================================================================
+# Image-context (moderation) check
+# =============================================================================
+# Run on every freshly generated image, *inside the same @spaces.GPU window that
+# produced it* — so it shares that GPU allocation rather than paying for a second
+# ZeroGPU acquisition. The VLM fills a fixed safety-rating JSON; the reply is
+# parsed best-effort and returned alongside the image and its R2 URL.
+MODERATION_PROMPT = """fill the following json based on the image, do not output anything else:
+{
+  "rating": "0_SAFE | 1_MILD_SUGGESTIVE | 2_STRONG_SUGGESTIVE | 3_NON_EXPLICIT_NUDITY | 4_EXPLICIT_SEXUAL | 5_PROHIBITED",
+  "confidence": 0.0,
+  "flags": {
+    "minor_or_young_looking": false,
+    "sexualized_minor_risk": false,
+    "explicit_genitals": false,
+    "sexual_act": false,
+    "non_consensual_or_abuse": false,
+    "bestiality": false,
+    "violence_or_gore": false,
+    "anime_or_drawn": false
+  }
+}"""
+
+
+def _parse_moderation_json(text):
+    """Best-effort parse of the moderation JSON the VLM returns.
+
+    The model is asked for JSON only, but a reasoning finetune can still wrap it
+    in <think> tags or stray prose, so we strip those and slice the outermost
+    ``{...}`` before ``json.loads``. Returns a dict with the parsed ``rating`` /
+    ``confidence`` / ``flags`` keys (None when absent) plus ``ok`` and the
+    original ``raw`` text, so nothing is lost even when parsing fails.
+    """
+    raw = (text or "").strip()
+    cleaned = raw.split("</think>")[-1].strip() if "</think>" in raw else raw
+    start, end = cleaned.find("{"), cleaned.rfind("}")
+    parsed = None
+    if start != -1 and end > start:
+        try:
+            parsed = json.loads(cleaned[start:end + 1])
+        except (ValueError, TypeError):
+            parsed = None
+    if isinstance(parsed, dict):
+        return {
+            "ok": True,
+            "rating": parsed.get("rating"),
+            "confidence": parsed.get("confidence"),
+            "flags": parsed.get("flags"),
+            "raw": raw,
+        }
+    return {"ok": False, "rating": None, "confidence": None, "flags": None, "raw": raw}
+
+
+def _moderate_image_inner(image, max_new_tokens=320):
+    """Synchronously run the VLM moderation check on a PIL image (GPU-side)."""
+    messages = [{
+        "role": "user",
+        "content": [
+            {"type": "image", "image": image},
+            {"type": "text", "text": MODERATION_PROMPT},
+        ],
+    }]
+    inputs = vlm_processor.apply_chat_template(
+        messages, tokenize=True, add_generation_prompt=True,
+        return_dict=True, return_tensors="pt", enable_thinking=False,
+    ).to(vlm_model.device)
+    with torch.inference_mode():
+        generated = vlm_model.generate(
+            **inputs, max_new_tokens=int(max_new_tokens), do_sample=False,
+        )
+    trimmed = generated[0][inputs["input_ids"].shape[-1]:]
+    tokenizer = getattr(vlm_processor, "tokenizer", vlm_processor)
+    text = tokenizer.decode(trimmed, skip_special_tokens=True)
+    return _parse_moderation_json(text)
+
+
 # =============================================================================
 # Compel-based prompt weighting helpers (ported from the NoobXL11 reference
 # space so long prompts and (weight:1.2) syntax work for the SDXL model).
@@ -380,8 +457,11 @@ def generate_image(
 
     This is a generator (so ZeroGPU streams its yields back over SSE). It yields
     ``("progress", step, total)`` tuples during sampling and a final
-    ``("image", image, seed)`` tuple. The sampler runs in a worker thread feeding
-    a queue so the main thread can yield progress as each diffusion step lands.
+    ``("image", image, seed, moderation)`` tuple. The sampler runs in a worker
+    thread feeding a queue so the main thread can yield progress as each diffusion
+    step lands. Once the image is ready, the VLM image-context (moderation) check
+    runs in this same GPU window and its best-effort-parsed result rides along on
+    the final tuple.
     """
     _gpu_start = time.time()
     total_steps = int(num_inference_steps)
@@ -422,7 +502,15 @@ def generate_image(
         )
     if "error" in result:
         raise result["error"]
-    yield ("image", result["image"], result["seed"])
+    image = result["image"]
+    # Same GPU window: run the image-context (moderation) check on the freshly
+    # generated image before the GPU allocation is released.
+    try:
+        moderation = _moderate_image_inner(image)
+    except Exception as exc:  # noqa: BLE001 - never let moderation break generation
+        moderation = {"ok": False, "rating": None, "confidence": None,
+                      "flags": None, "raw": "", "error": f"{type(exc).__name__}: {exc}"}
+    yield ("image", image, result["seed"], moderation)
 
 
 def _generate_image_inner(
@@ -605,11 +693,13 @@ def generate_and_upload(
     Generator yielding ``(image, seed, r2_status, progress)``. Streams per-step
     sampling progress (image still None) and finishes with the real image, seed
     and R2 status once the upload completes. The image is always the original
-    HF-generated asset; ``r2_status`` reports the uploaded filekey on success or
-    the error on failure. The caller's unique id (``uid`` cookie) is recorded in
-    the uploaded object's metadata.
+    HF-generated asset; ``r2_status`` reports the uploaded filekey, bucket and a
+    presigned ``r2_url`` on success (or the error on failure) plus the VLM
+    ``moderation`` result. The asset is encoded as **WebP** before upload (and
+    the returned file is WebP too) so downstream consumers only ever see WebP.
+    The caller's unique id (``uid`` cookie) is recorded in the object's metadata.
     """
-    image, used = None, None
+    image, used, moderation = None, None, None
     for ev in generate_image(
         model_name, prompt, negative_prompt, use_negative_prompt,
         height, width, num_inference_steps, guidance_scale, seed, randomize_seed,
@@ -623,12 +713,12 @@ def generate_and_upload(
                 "image", frac, step, total, f"Sampling {step}/{total}"
             )
         else:
-            _, image, used = ev
+            _, image, used, moderation = ev
 
     yield None, used, None, _progress("image", 0.97, label="Uploading")
     uid = r2_uploader.uid_from_request(request)
     buf = io.BytesIO()
-    image.save(buf, format="PNG")
+    image.save(buf, format="WEBP", quality=95, method=6)
     params = {
         "model": model_name,
         "prompt": prompt,
@@ -645,14 +735,19 @@ def generate_and_upload(
         prompt=prompt,
         params=params,
         data=buf.getvalue(),
-        ext=".png",
-        content_type="image/png",
+        ext=".webp",
+        content_type="image/webp",
         uid=uid,
     )
     if result.get("ok"):
-        status = {"r2_filekey": result["filekey"], "r2_bucket": result["bucket"]}
+        status = {
+            "r2_filekey": result["filekey"],
+            "r2_bucket": result["bucket"],
+            "r2_url": r2_uploader.presign_get_url(result["filekey"], result["bucket"]),
+        }
     else:
         status = {"r2_error": result.get("error", "unknown error")}
+    status["moderation"] = moderation
     yield image, used, status, _progress("done", 1.0, label="Done")
 
 
@@ -727,9 +822,11 @@ def prompt_to_video_assets(
     """Text -> first-frame image (uploaded to R2) + video generation prompt.
 
     Streams progress and returns a dict with ``video_prompt``, ``first_frame_url``
-    (a presigned GET, usable directly), ``r2_filekey``/``r2_bucket`` (for callers
-    that resolve their own public URL), ``first_frame_prompt`` (the intermediate
-    image prompt) and ``seed``. Has no UI (registered via ``gr.api``).
+    (a presigned GET to the **WebP** first frame, usable directly),
+    ``r2_filekey``/``r2_bucket`` (for callers that resolve their own public URL),
+    ``first_frame_prompt`` (the intermediate image prompt), ``moderation`` (the
+    VLM image-context check on the rendered frame) and ``seed``. Has no UI
+    (registered via ``gr.api``).
     """
     state = {
         "stage": "frame_prompt", "p": 0.0, "step": 0, "total": 0, "label": "",
@@ -738,6 +835,7 @@ def prompt_to_video_assets(
         "r2_filekey": None,
         "r2_bucket": None,
         "video_prompt": None,
+        "moderation": None,
         "seed": None,
         "done": False,
         "error": None,
@@ -778,7 +876,7 @@ def prompt_to_video_assets(
 
     # --- Stage 2: image model renders the first frame -------------------------
     base += _P2V_W_FRAME_PROMPT
-    image, used_seed = None, seed
+    image, used_seed, moderation = None, seed, None
     use_negative_prompt = (model_name == MODEL_NOOBXL)
     for ev in generate_image(
         model_name, frame_prompt, NOOBXL_NEGATIVE, use_negative_prompt,
@@ -790,15 +888,16 @@ def prompt_to_video_assets(
                         frac=step / max(total, 1), step=step, total=total,
                         label=f"Rendering first frame {step}/{total}")
         else:
-            _, image, used_seed = ev
+            _, image, used_seed, moderation = ev
     state["seed"] = int(used_seed)
+    state["moderation"] = moderation
 
-    # --- Stage 3: upload the first frame to R2 --------------------------------
+    # --- Stage 3: upload the first frame to R2 (WebP) -------------------------
     base += _P2V_W_IMAGE
     yield frame("upload", base, _P2V_W_UPLOAD, frac=0.1, label="Uploading first frame")
     uid = r2_uploader.uid_from_request(request)
     buf = io.BytesIO()
-    image.save(buf, format="PNG")
+    image.save(buf, format="WEBP", quality=95, method=6)
     params = {
         "model": model_name,
         "prompt": frame_prompt,
@@ -814,7 +913,7 @@ def prompt_to_video_assets(
     }
     up = r2_uploader.upload_asset(
         namespace=R2_NAMESPACE, prompt=frame_prompt, params=params,
-        data=buf.getvalue(), ext=".png", content_type="image/png", uid=uid,
+        data=buf.getvalue(), ext=".webp", content_type="image/webp", uid=uid,
     )
     if up.get("ok"):
         state["r2_filekey"] = up["filekey"]
@@ -1015,7 +1114,7 @@ with gr.Blocks(fill_height=True) as demo:
                     output_image = gr.Image(
                         label="Generated Image",
                         type="pil",
-                        format="png",
+                        format="webp",
                         show_label=False,
                         height=600,
                         buttons=["download", "share"],