feat: upload generated images to R2 with prompt/params + requester uid in metadata

Add r2_uploader helper that pushes each generated image to the
free-generated-assets R2 bucket (key yyyymmddhhmmss-s01-<rand>.png) with the
prompt, request params and the caller's uid cookie recorded as object metadata.
The generate handler now returns the uploaded filekey alongside the original
image, or an error plus the original image on failure.

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

app.py

@@ -1,3 +1,4 @@
+import io
 import random
 import re
 
@@ -8,6 +9,13 @@ import gradio as gr
 from diffusers import DiffusionPipeline, AutoPipelineForText2Image, AutoencoderKL
 from compel import Compel, ReturnedEmbeddingsType
 
+import r2_uploader
+
+# Per-Space namespace embedded in every uploaded object key. Deliberately opaque
+# (not the readable Space name) but stable so the owner can tell assets apart.
+# Legend: s01=ImageStudio, s02=LTX2.3-Studio, s03=wan2-2-fp8da-aoti-preview-2.
+R2_NAMESPACE = "s01"
+
 # =============================================================================
 # Model registry
 # =============================================================================
@@ -318,6 +326,63 @@ def generate_image(
     return image, seed
 
 
+def generate_and_upload(
+    model_name,
+    prompt,
+    negative_prompt,
+    use_negative_prompt,
+    height,
+    width,
+    num_inference_steps,
+    guidance_scale,
+    seed,
+    randomize_seed,
+    request: gr.Request = None,
+    progress=gr.Progress(track_tqdm=True),
+):
+    """Generate, then upload the result to R2 outside the GPU window.
+
+    Returns ``(image, seed, r2_status)``. 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.
+    """
+    image, used = generate_image(
+        model_name, prompt, negative_prompt, use_negative_prompt,
+        height, width, num_inference_steps, guidance_scale, seed, randomize_seed,
+        progress=progress,
+    )
+
+    uid = r2_uploader.uid_from_request(request)
+    buf = io.BytesIO()
+    image.save(buf, format="PNG")
+    params = {
+        "model": model_name,
+        "prompt": prompt,
+        "negative_prompt": (negative_prompt if use_negative_prompt else ""),
+        "height": int(height),
+        "width": int(width),
+        "num_inference_steps": int(num_inference_steps),
+        "guidance_scale": float(guidance_scale),
+        "seed": int(used),
+        "uid": uid,
+    }
+    result = r2_uploader.upload_asset(
+        namespace=R2_NAMESPACE,
+        prompt=prompt,
+        params=params,
+        data=buf.getvalue(),
+        ext=".png",
+        content_type="image/png",
+        uid=uid,
+    )
+    if result.get("ok"):
+        status = {"r2_filekey": result["filekey"], "r2_bucket": result["bucket"]}
+    else:
+        status = {"r2_error": result.get("error", "unknown error")}
+    return image, used, status
+
+
 # Recommended defaults per model: (steps, guidance, height, width)
 MODEL_DEFAULTS = {
     MODEL_ZIMAGE: dict(steps=9, guidance=0.0, height=1024, width=1024),
@@ -495,6 +560,7 @@ with gr.Blocks(fill_height=True) as demo:
             used_seed = gr.Number(
                 label="🎲 Seed Used", interactive=False, container=True,
             )
+            r2_status = gr.JSON(label="☁️ R2 Upload")
 
     gr.Markdown(
         """
@@ -522,11 +588,15 @@ with gr.Blocks(fill_height=True) as demo:
         height, width, num_inference_steps, guidance_scale, seed, randomize_seed,
     ]
 
+    # api_name pinned so the generator's "/generate_image" endpoint keeps
+    # resolving even though the click now runs the upload wrapper.
     generate_btn.click(
-        fn=generate_image, inputs=gen_inputs, outputs=[output_image, used_seed],
+        fn=generate_and_upload, inputs=gen_inputs,
+        outputs=[output_image, used_seed, r2_status], api_name="generate_image",
     )
     prompt.submit(
-        fn=generate_image, inputs=gen_inputs, outputs=[output_image, used_seed],
+        fn=generate_and_upload, inputs=gen_inputs,
+        outputs=[output_image, used_seed, r2_status], show_api=False,
     )
 
 if __name__ == "__main__":

r2_uploader.py

@@ -0,0 +1,185 @@
+"""Cloudflare R2 asset uploader.
+
+Reads R2 credentials from the ``SCONFIG`` Space secret (a JSON blob) and uploads
+generated assets to a preset bucket. Object keys follow
+``<yyyymmddhhmmss>-<namespace>-<rand>.<ext>``; the prompt and request params are
+attached as JSON object metadata so the owner can trace any asset back to the
+request that produced it.
+
+The ``SCONFIG`` secret is JSON shaped like::
+
+    {
+      "account_id": "<cloudflare account id>",
+      "akey":       "<R2 access key id>",
+      "skey":       "<R2 secret access key>",
+      "token":      "<cloudflare api token value>",   # optional, not used for S3
+      "bucket":     "free-generated-assets"            # optional, this is the default
+    }
+
+``upload_asset`` never raises: callers always also return the original
+HF-generated asset, so an R2 failure degrades to "asset shown, upload errored"
+rather than breaking generation.
+"""
+from __future__ import annotations
+
+import datetime
+import json
+import mimetypes
+import os
+import threading
+import uuid
+
+CONFIG_ENV = "SCONFIG"
+DEFAULT_BUCKET = "free-generated-assets"
+_MAX_META_BYTES = 2000  # R2/S3 cap user metadata; stay well under header limits.
+
+_ensured_buckets: set[str] = set()
+_ensure_lock = threading.Lock()
+
+
+def _load_cfg() -> dict:
+    raw = (os.environ.get(CONFIG_ENV) or "").strip()
+    if not raw:
+        raise RuntimeError(f"{CONFIG_ENV} secret is not set")
+    return json.loads(raw)
+
+
+def uid_from_request(request, cookie_name: str = "uid") -> str:
+    """Best-effort read of the caller's unique id from a gr.Request cookie.
+
+    The web app forwards the per-browser anonymous id and the generator sets it
+    as the ``uid`` cookie on the Space call; we record it in object metadata so
+    every asset traces back to the requester. Returns "" if absent.
+    """
+    if request is None:
+        return ""
+    try:
+        cookies = getattr(request, "cookies", None)
+        if isinstance(cookies, dict) and cookies.get(cookie_name):
+            return str(cookies[cookie_name])
+        headers = getattr(request, "headers", None)
+        raw = ""
+        if headers is not None:
+            raw = (headers.get("cookie") if hasattr(headers, "get") else "") or ""
+        for part in raw.split(";"):
+            k, _, v = part.strip().partition("=")
+            if k == cookie_name and v:
+                from urllib.parse import unquote
+
+                return unquote(v)
+    except Exception:  # noqa: BLE001 - identity is best-effort, never fatal
+        return ""
+    return ""
+
+
+def _client(cfg: dict):
+    import boto3
+    from botocore.config import Config
+
+    return boto3.client(
+        "s3",
+        endpoint_url=f"https://{cfg['account_id']}.r2.cloudflarestorage.com",
+        aws_access_key_id=cfg["akey"],
+        aws_secret_access_key=cfg["skey"],
+        region_name="auto",
+        config=Config(signature_version="s3v4", retries={"max_attempts": 3}),
+    )
+
+
+def _now_stamp() -> str:
+    return datetime.datetime.now(datetime.timezone.utc).strftime("%Y%m%d%H%M%S")
+
+
+def _meta_blob(prompt, params, uid: str = "") -> str:
+    """Compact JSON of uid + prompt + params, truncated to fit R2 metadata limits."""
+    blob = json.dumps({"uid": uid, "prompt": prompt, "params": params}, ensure_ascii=True, default=str)
+    if len(blob.encode("utf-8")) <= _MAX_META_BYTES:
+        return blob
+    keep = max(0, _MAX_META_BYTES - len(json.dumps({"uid": uid, "prompt": "", "params": params}, default=str)) - 16)
+    short = {"uid": uid, "prompt": str(prompt)[:keep] + "...[truncated]", "params": params}
+    return json.dumps(short, ensure_ascii=True, default=str)[:_MAX_META_BYTES]
+
+
+def _ensure_bucket(s3, bucket: str) -> None:
+    """Create the bucket on first use if it does not already exist (best effort)."""
+    if bucket in _ensured_buckets:
+        return
+    with _ensure_lock:
+        if bucket in _ensured_buckets:
+            return
+        from botocore.exceptions import ClientError
+
+        try:
+            s3.head_bucket(Bucket=bucket)
+        except ClientError as exc:
+            code = str(exc.response.get("Error", {}).get("Code", ""))
+            if code in ("404", "NoSuchBucket", "NotFound"):
+                # Best effort: object-scoped tokens can't create buckets. If this
+                # fails, let the subsequent put_object surface the real error
+                # (NoSuchBucket) rather than masking it with an AccessDenied here.
+                try:
+                    s3.create_bucket(Bucket=bucket)
+                except ClientError:
+                    pass
+            # Other codes (e.g. 403 with object-only tokens) are ignored; the
+            # subsequent put_object is the real test.
+        _ensured_buckets.add(bucket)
+
+
+def upload_asset(
+    *,
+    namespace: str,
+    prompt,
+    params: dict,
+    data: bytes | None = None,
+    path: str | None = None,
+    ext: str | None = None,
+    content_type: str | None = None,
+    uid: str = "",
+) -> dict:
+    """Upload one asset to R2.
+
+    Provide either ``data`` (raw bytes) or ``path`` (a local file). ``uid`` is
+    the caller's unique id (see :func:`uid_from_request`) and is recorded both as
+    a dedicated object-metadata field and inside the generation JSON. Returns
+    ``{"ok": True, "filekey": ..., "bucket": ...}`` on success or
+    ``{"ok": False, "error": ...}`` on failure. Never raises.
+    """
+    try:
+        if data is None and path is None:
+            raise ValueError("upload_asset requires either data or path")
+
+        cfg = _load_cfg()
+        bucket = cfg.get("bucket") or DEFAULT_BUCKET
+
+        if ext is None and path is not None:
+            ext = os.path.splitext(path)[1]
+        ext = ext or ""
+        if ext and not ext.startswith("."):
+            ext = "." + ext
+
+        if content_type is None:
+            guessed = mimetypes.guess_type("x" + ext)[0] if ext else None
+            content_type = guessed or "application/octet-stream"
+
+        if data is None:
+            with open(path, "rb") as fh:
+                data = fh.read()
+
+        filekey = f"{_now_stamp()}-{namespace}-{uuid.uuid4().hex[:8]}{ext}"
+
+        s3 = _client(cfg)
+        _ensure_bucket(s3, bucket)
+        metadata = {"generation": _meta_blob(prompt, params, uid)}
+        if uid:
+            metadata["uid"] = uid
+        s3.put_object(
+            Bucket=bucket,
+            Key=filekey,
+            Body=data,
+            ContentType=content_type,
+            Metadata=metadata,
+        )
+        return {"ok": True, "filekey": filekey, "bucket": bucket}
+    except Exception as exc:  # noqa: BLE001 - report any failure to the caller
+        return {"ok": False, "error": f"{type(exc).__name__}: {exc}"}

requirements.txt

@@ -5,3 +5,4 @@ accelerate
 compel
 sentencepiece
 numpy
+boto3