Process encrypted bucket jobs and deliver results from the Space

app.py

@@ -21,6 +21,7 @@ from ignored_preview import save_ignored_preview
 from clean import load_clean, clean_image
 from illustrious import load_illustrious, prepare_illustrious, infer_illustrious, save_illustrious
 from watermark_remover import WatermarkRemover
+from bucket_jobs import run_bucket_job
 
 INPUT_IMAGE_HEIGHT = 600
 TEXT_WORKER = [sys.executable, "-m", "text_detection"]
@@ -267,6 +268,10 @@ def remove_watermark(file_obj, strength, steps, model, ignore_eyes=False, ignore
     return make_terminal_html(cmd_html + f"\n\nDone!{preserved}{cleaned}", cursor=False), str(output_path), str(preview_path)
 
 
+def process_bucket_job(payload):
+    return run_bucket_job(payload, remove_watermark)
+
+
 def check_ai_metadata(file_obj):
     if file_obj is None:
         raise gr.Error("Please upload an image first.")
@@ -327,6 +332,19 @@ MODELS = load_models()
 
 with gr.Blocks(title="noai-watermark") as demo:
 
+    bucket_payload = gr.JSON(visible=False)
+    bucket_result = gr.JSON(visible=False)
+    bucket_submit = gr.Button(visible=False)
+    bucket_submit.click(
+        fn=process_bucket_job,
+        inputs=[bucket_payload],
+        outputs=[bucket_result],
+        api_name="process_bucket_job",
+        api_visibility="undocumented",
+        concurrency_id="watermark",
+        concurrency_limit=1,
+    )
+
     gr.HTML("""
         <h1 style="text-align:center;">🛡️ noai-watermark</h1>
         <p style="text-align:center; color:#888;">
@@ -394,6 +412,7 @@ with gr.Blocks(title="noai-watermark") as demo:
                 fn=remove_watermark,
                 inputs=[upload_wm, strength, steps, model, ignore_eyes, ignore_hands, ignore_feet, clean_after, ignore_text],
                 outputs=[wm_terminal, output_image, ignored_image],
+                concurrency_id="watermark",
                 concurrency_limit=1,
             )
 

bucket_jobs.py

@@ -0,0 +1,211 @@
+import hashlib
+import hmac
+import json
+import math
+import re
+import sqlite3
+import time
+from contextlib import closing
+from html.parser import HTMLParser
+from pathlib import Path
+from secrets import token_bytes
+from tempfile import TemporaryDirectory
+from urllib.parse import urlencode
+
+import httpx
+import xxhash
+from cryptography.hazmat.primitives.ciphers.aead import AESGCM
+from huggingface_hub import batch_bucket_files
+
+API_URL = "https://hailxd.vercel.app"
+DISCORD_API = "https://discord.com/api/v9"
+DEV_CHANNEL = "1352311755682091111"
+BOT_USER_AGENT = f"DiscordBot ({API_URL}, 1.0.0)"
+JOB_DB = (Path("/data") if Path("/data").is_dir() else Path.home() / ".cache/no-watermark") / "bucket-jobs.sqlite3"
+MAX_BYTES = 128 * 1024 * 1024
+NONCE_BYTES = 12
+TAG_BYTES = 16
+RETRY_DELAYS = (1, 2, 4, 8, 16)
+MODELS = {"CtrlRegen", "CtrlRegen Illustrious (SDXL)", "Default"}
+IMAGE_TYPES = {"image/png": ".png", "image/jpeg": ".jpg", "image/webp": ".webp"}
+CREATE_SQL = (
+    "CREATE TABLE IF NOT EXISTS jobs (id TEXT PRIMARY KEY, fingerprint TEXT NOT NULL, "
+    "output_hash TEXT NOT NULL DEFAULT '', discord_id TEXT NOT NULL DEFAULT '', error TEXT NOT NULL DEFAULT '')"
+)
+
+
+class JobError(ValueError):
+    pass
+
+
+class PlainText(HTMLParser):
+    def __init__(self):
+        super().__init__()
+        self.parts = []
+
+    def handle_data(self, data):
+        self.parts.append(data)
+
+
+def _validate(payload):
+    if not isinstance(payload, dict):
+        raise JobError("Invalid job")
+    job = dict(payload)
+    for name in ("id", "input_hash", "input_url", "filename", "content_type", "bucket", "suffix", "hf_token", "discord_token", "channel"):
+        if not isinstance(job.get(name), str) or not job[name]:
+            raise JobError("Invalid job settings")
+    if not re.fullmatch(r"[a-f0-9-]{36}", job["id"]):
+        raise JobError("Invalid job ID")
+    if not re.fullmatch(r"[A-Za-z0-9_-]{8}", job["suffix"]):
+        raise JobError("Invalid encryption suffix")
+    if not re.fullmatch(r"[a-f0-9]{32}" + re.escape(job["suffix"]), job["input_hash"]):
+        raise JobError("Invalid encrypted image hash")
+    if not re.fullmatch(r"[A-Za-z0-9_-]+/[A-Za-z0-9_.-]+", job["bucket"]):
+        raise JobError("Invalid Hugging Face bucket")
+    expected_url = f"https://huggingface.co/buckets/{job['bucket']}/resolve/{job['input_hash']}"
+    if job["input_url"] != expected_url:
+        raise JobError("Invalid input bucket URL")
+    if job["channel"] != DEV_CHANNEL:
+        raise JobError("Delivery must use the Discord Dev channel")
+    if job["content_type"] not in IMAGE_TYPES or len(job["filename"]) > 240 or any(ord(char) < 32 for char in job["filename"]):
+        raise JobError("Invalid image metadata")
+    key = job.get("key")
+    if not isinstance(key, list) or len(key) != 16 or any(type(value) is not int or not 0 <= value <= 255 for value in key):
+        raise JobError("Invalid encryption key")
+    settings = job.get("options")
+    if not isinstance(settings, list) or len(settings) != 8:
+        raise JobError("Invalid image settings")
+    strength, steps, model, *flags = settings
+    if type(strength) not in (int, float) or not math.isfinite(strength) or not 0.01 <= strength <= 1:
+        raise JobError("Invalid strength")
+    if type(steps) is not int or not 10 <= steps <= 100 or not isinstance(model, str) or model not in MODELS:
+        raise JobError("Invalid model or steps")
+    if any(type(flag) is not bool for flag in flags):
+        raise JobError("Invalid image options")
+    return job
+
+
+def _record(job):
+    fingerprint = hashlib.sha256(json.dumps(
+        {name: value for name, value in job.items() if name not in {"hf_token", "discord_token"}},
+        sort_keys=True, separators=(",", ":"),
+    ).encode()).hexdigest()
+    JOB_DB.parent.mkdir(parents=True, exist_ok=True)
+    with closing(sqlite3.connect(JOB_DB)) as db, db:
+        db.execute(CREATE_SQL)
+        db.execute("INSERT INTO jobs (id, fingerprint) VALUES (?, ?) ON CONFLICT(id) DO NOTHING", (job["id"], fingerprint))
+        row = db.execute("SELECT fingerprint, output_hash, discord_id, error FROM jobs WHERE id = ?", (job["id"],)).fetchone()
+    if row[0] != fingerprint:
+        raise JobError("Job ID already belongs to another image or settings")
+    return {"output_hash": row[1], "discord_id": row[2], "error": row[3]}
+
+
+def _save(job_id, **values):
+    if not values or not set(values) <= {"output_hash", "discord_id", "error"}:
+        raise JobError("Invalid job state")
+    with closing(sqlite3.connect(JOB_DB)) as db, db:
+        db.execute("UPDATE jobs SET " + ", ".join(f"{name} = ?" for name in values) + " WHERE id = ?", (*values.values(), job_id))
+
+
+def _retry(action):
+    for attempt in range(len(RETRY_DELAYS) + 1):
+        try:
+            return action()
+        except Exception:
+            if attempt == len(RETRY_DELAYS):
+                raise
+            time.sleep(RETRY_DELAYS[attempt])
+
+
+def _download(job):
+    with httpx.stream("GET", job["input_url"], headers={"Authorization": f"Bearer {job['hf_token']}"}, follow_redirects=True, timeout=60) as response:
+        if not response.is_success:
+            raise JobError(f"Input download failed ({response.status_code})")
+        data = bytearray()
+        for chunk in response.iter_bytes():
+            data.extend(chunk)
+            if len(data) > MAX_BYTES + NONCE_BYTES + TAG_BYTES:
+                raise JobError("Input image exceeds 128 MB")
+    try:
+        image = AESGCM(bytes(job["key"])).decrypt(bytes(data[:NONCE_BYTES]), bytes(data[NONCE_BYTES:]), None)
+    except Exception as exc:
+        raise JobError("Unable to decrypt the input image") from exc
+    if xxhash.xxh3_128_hexdigest(image) + job["suffix"] != job["input_hash"]:
+        raise JobError("Input image hash mismatch")
+    return image
+
+
+def _upload(job, content):
+    file_hash = xxhash.xxh3_128_hexdigest(content) + job["suffix"]
+    nonce = token_bytes(NONCE_BYTES)
+    encrypted = nonce + AESGCM(bytes(job["key"])).encrypt(nonce, content, None)
+    _retry(lambda: batch_bucket_files(job["bucket"], add=[(encrypted, file_hash)], token=job["hf_token"]))
+    return file_hash
+
+
+def _notify(job, record):
+    if record["output_hash"]:
+        file_hash = record["output_hash"]
+        signature = hmac.new(bytes(job["key"]), f"no-watermark:{file_hash}".encode(), hashlib.sha256).hexdigest()
+        url = API_URL + "/no-watermark/?" + urlencode({"file": file_hash, "signature": signature})
+        content = f"No Watermark — {job['filename']}\n{url}"
+    else:
+        content = f"No Watermark failed — {job['filename']}\n{record['error']}"
+    payload = {
+        "content": content,
+        "allowed_mentions": {"parse": []},
+        "nonce": hashlib.blake2s(job["id"].encode(), digest_size=12).hexdigest(),
+        "enforce_nonce": True,
+    }
+
+    def send():
+        response = httpx.post(
+            f"{DISCORD_API}/channels/{job['channel']}/messages",
+            headers={"Authorization": f"Bot {job['discord_token']}", "User-Agent": BOT_USER_AGENT},
+            json=payload, timeout=30,
+        )
+        if response.status_code == 429:
+            time.sleep(min(float(response.json().get("retry_after", 1)), 60))
+        if not response.is_success:
+            raise JobError(f"Discord delivery failed ({response.status_code})")
+        return response.json()["id"]
+
+    return _retry(send)
+
+
+def run_bucket_job(payload, remove_watermark):
+    try:
+        job = _validate(payload)
+        record = _record(job)
+    except Exception:
+        return {"status": "rejected", "error": "Invalid job or unavailable job storage"}
+    if record["discord_id"]:
+        return {"id": job["id"], "status": "failed" if record["error"] else "completed"}
+    if not record["output_hash"] and not record["error"]:
+        try:
+            content = _retry(lambda: _download(job))
+            with TemporaryDirectory(prefix="bucket-input-") as directory:
+                input_path = Path(directory) / ("image" + IMAGE_TYPES[job["content_type"]])
+                input_path.write_bytes(content)
+                terminal, output, _ = remove_watermark(str(input_path), *job["options"])
+                if not output:
+                    parser = PlainText()
+                    parser.feed(str(terminal))
+                    raise JobError("".join(parser.parts).strip() or "Image processing failed")
+                output_path = Path(output)
+                if output_path.stat().st_size > MAX_BYTES:
+                    raise JobError("Completed image exceeds 128 MB")
+                record["output_hash"] = _upload(job, output_path.read_bytes())
+                _save(job["id"], output_hash=record["output_hash"])
+        except Exception as exc:
+            error = str(exc) if isinstance(exc, JobError) else "Image processing or encrypted upload failed"
+            for secret in (job["hf_token"], job["discord_token"]):
+                error = error.replace(secret, "[redacted]")
+            record["error"] = error[:1000]
+            _save(job["id"], error=record["error"])
+    try:
+        discord_id = _notify(job, record)
+        _save(job["id"], discord_id=discord_id)
+    except Exception:
+        return {"id": job["id"], "status": "delivery_failed"}
+    return {"id": job["id"], "status": "failed" if record["error"] else "completed"}

requirements.txt

@@ -8,3 +8,7 @@ ultralytics
 spandrel
 paddleocr>=3.0,<4
 paddlepaddle==3.2.2
+cryptography
+httpx
+huggingface_hub>=1.6.0
+xxhash