Update app.py

app.py

@@ -1,154 +1,499 @@
+import base64
+import io
+import json
+import os
+import sys
+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 numpy as np
-import random
-
-# import spaces #[uncomment to use ZeroGPU]
-from diffusers import DiffusionPipeline
-import torch
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model_repo_id = "stabilityai/sdxl-turbo"  # Replace to the model you would like to use
-
-if torch.cuda.is_available():
-    torch_dtype = torch.float16
-else:
-    torch_dtype = torch.float32
-
-pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
-pipe = pipe.to(device)
-
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1024
-
-
-# @spaces.GPU #[uncomment to use ZeroGPU]
-def infer(
-    prompt,
-    negative_prompt,
-    seed,
-    randomize_seed,
-    width,
-    height,
-    guidance_scale,
-    num_inference_steps,
-    progress=gr.Progress(track_tqdm=True),
-):
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-
-    generator = torch.Generator().manual_seed(seed)
-
-    image = pipe(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
-        width=width,
-        height=height,
-        generator=generator,
-    ).images[0]
-
-    return image, seed
-
-
-examples = [
-    "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
-    "An astronaut riding a green horse",
-    "A delicious ceviche cheesecake slice",
-]
-
-css = """
-#col-container {
-    margin: 0 auto;
-    max-width: 640px;
+import requests
+from PIL import Image
+
+try:
+    from huggingface_hub import batch_bucket_files
+except ImportError:
+    batch_bucket_files = None
+
+# --- 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-guard-v02"
+NCII_THRESHOLD = 0.70
+
+MODEL_CHOICES = ["Auto-Routing", "FLUX + LoRA", "Qwen Edit", "Lance", "Krea2", "SenseNova"]
+# Every choice currently routes to FLUX.2 Klein; swap slugs here to wire real backends.
+MODEL_ROUTES = {choice: EDIT_MODEL for choice in MODEL_CHOICES}
+
+REQUEST_TIMEOUT = 180
+MAX_SIDE = 2048  # downscale reference images before upload
+
+logging_enabled = bool(HF_TOKEN and BUCKET_ID and batch_bucket_files)
+
+
+# --- Prompt guard -------------------------------------------------------------
+_guard = None
+_guard_lock = threading.Lock()
+_guard_failed = False
+
+
+def _load_guard():
+    """Load the classifier once; safe to call from several threads."""
+    global _guard, _guard_failed
+    if _guard is not None:
+        return _guard
+    with _guard_lock:
+        if _guard is None and not _guard_failed:
+            try:
+                from transformers import pipeline
+                _guard = pipeline("text-classification", model=GUARD_MODEL)
+            except Exception:
+                _guard_failed = True
+                print(f"[DEBUG] guard load failed:\n{traceback.format_exc()}", file=sys.stderr)
+    return _guard
+
+
+def ncii_score(prompt: str) -> float:
+    """Return the ncii probability for a prompt. Fails closed on any error."""
+    clf = _load_guard()
+    if clf is None:
+        # Safety system unavailable -> refuse to generate rather than bypass it.
+        raise gr.Error("Safety system is warming up — please try again in a minute.")
+    try:
+        result = clf(prompt[:2048])[0]
+        label = str(result.get("label", "")).strip().lower()
+        score = float(result.get("score", 1.0))
+    except Exception:
+        print(f"[DEBUG] guard inference failed:\n{traceback.format_exc()}", file=sys.stderr)
+        raise gr.Error("Safety check failed — please try again.")
+    return score if label == "ncii" else 1.0 - score
+
+
+# Warm the model in the background so the first user doesn't pay the load time.
+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:
+        raise gr.Error("Generation failed — the model may be busy, please retry.")
+    if errors:
+        gr.Warning(f"{len(images)}/{n} variants completed.")
+    return images
+
+
+# --- Logging -------------------------------------------------------------------
+def log_submission(session_id: str, image: Image.Image, prompt: str):
+    """Fire-and-forget export of {image, prompt, date, session_id} to the bucket."""
+    if not logging_enabled:
+        return
+
+    def _worker():
+        try:
+            sample_id = uuid.uuid4().hex[:10]
+            date = datetime.now(timezone.utc).isoformat()
+            buf = io.BytesIO()
+            image.convert("RGB").save(buf, format="PNG")
+            record = {"session_id": session_id, "prompt": prompt, "date": date}
+            base = f"sessions/{session_id}/{date[:19].replace(':', '-')}_{sample_id}"
+            batch_bucket_files(
+                BUCKET_ID,
+                add=[
+                    (buf.getvalue(), f"{base}.png"),
+                    (json.dumps(record).encode(), f"{base}.json"),
+                ],
+                token=HF_TOKEN,
+            )
+        except Exception:
+            # logging must never take the app down — degrade to no logging
+            print(f"[DEBUG] log_submission failed:\n{traceback.format_exc()}", file=sys.stderr)
+
+    threading.Thread(target=_worker, daemon=True).start()
+
+
+# --- Main pipeline ---------------------------------------------------------------
+def edit_image(image, prompt, model_choice, request: gr.Request,
+               progress=gr.Progress()):
+    if image is None:
+        raise gr.Error("Please upload an image first.")
+    if not prompt or not prompt.strip():
+        raise gr.Error("Please describe the change you would like to make.")
+    prompt = prompt.strip()
+
+    # One id per browser session: every prompt sent from the same tab —
+    # including blocked ones and retries — is linked together.
+    session_id = getattr(request, "session_hash", None) or uuid.uuid4().hex
+    log_submission(session_id, image, prompt)
+
+    progress(0.1, desc="Checking prompt…")
+    if ncii_score(prompt) > NCII_THRESHOLD:
+        raise gr.Error("🚫 Prompt forbidden")
+
+    if not OPENROUTER_API_KEY:
+        raise gr.Error("Server is not configured yet (missing API key).")
+
+    model = MODEL_ROUTES.get(model_choice, EDIT_MODEL)
+    progress(0.3, desc="Generating…")
+    images = generate_variants(_to_data_url(image), prompt, 1, model)
+
+    return images, gr.update(visible=True)
+
+
+# --- Styling: editorial brief — cream paper, grid, ink & orange -----------------
+LAB_CSS = """
+@import url('https://fonts.googleapis.com/css2?family=Archivo+Black&family=JetBrains+Mono:wght@400;500;700&display=swap');
+
+:root {
+    --paper: #f4f0e6;
+    --panel: #fbf9f2;
+    --ink: #16130e;
+    --muted: #8f8a7d;
+    --accent: #ee4f1e;
+    --grid: rgba(22, 19, 14, 0.06);
+}
+
+body, .app, gradio-app {
+    background: var(--paper) !important;
+}
+
+.gradio-container {
+    background-color: var(--paper) !important;
+    background-image:
+        linear-gradient(var(--grid) 1px, transparent 1px),
+        linear-gradient(90deg, var(--grid) 1px, transparent 1px);
+    background-size: 44px 44px;
+    font-family: 'JetBrains Mono', monospace !important;
+    color: var(--ink) !important;
+    max-width: 1180px !important;
+    margin: 0 auto !important;
+}
+
+/* ---------- header frame ---------- */
+#brief-frame {
+    position: relative;
+    border: 2px solid var(--ink);
+    background: var(--paper);
+    padding: 1.1rem 1.6rem 0.4rem;
+    margin: 1.6rem 0 1.8rem;
+}
+#brief-frame .tick {
+    position: absolute;
+    background: var(--ink);
+}
+#brief-frame .tick.t1 { top: -12px; left: 18%;  width: 2px; height: 24px; }
+#brief-frame .tick.t2 { top: -12px; right: 8%;  width: 2px; height: 24px; }
+#brief-frame .tick.t3 { bottom: -12px; left: 40%; width: 2px; height: 24px; }
+#brief-frame .tick.t4 { top: 30%; left: -12px;  width: 24px; height: 2px; }
+#brief-frame .tick.t5 { top: 62%; right: -12px; width: 24px; height: 2px; }
+
+.brief-kicker {
+    display: flex;
+    justify-content: space-between;
+    gap: 1rem;
+    color: var(--ink);
+    font-size: 0.72rem;
+    font-weight: 700;
+    letter-spacing: 4px;
+    text-transform: uppercase;
+    padding-bottom: 0.9rem;
+}
+.brief-kicker span { color: var(--ink) !important; }
+.brief-kicker span.dim { color: var(--muted) !important; font-weight: 500; }
+
+.brief-headline {
+    font-family: 'Archivo Black', 'JetBrains Mono', sans-serif;
+    font-size: clamp(2.1rem, 5.2vw, 3.6rem);
+    line-height: 1.04;
+    letter-spacing: 1px;
+    text-transform: uppercase;
+    margin: 1.4rem 0 1rem;
+    color: var(--ink);
+}
+.brief-headline .accent { color: var(--accent); }
+
+.brief-sub {
+    font-size: 0.8rem;
+    letter-spacing: 3.5px;
+    text-transform: uppercase;
+    color: var(--muted);
+    margin: 0 0 1.6rem;
+}
+
+/* floating component labels (e.g. on the image inputs) */
+.block label.float, .block .label {
+    background: var(--ink) !important;
+    color: var(--paper) !important;
+    border-radius: 0 !important;
+}
+
+/* ---------- panels & fields ---------- */
+.gr-panel, .block, .form {
+    background: var(--panel) !important;
+    border: 2px solid var(--ink) !important;
+    border-radius: 0 !important;
+    box-shadow: none !important;
+}
+
+textarea, input, select {
+    background: var(--panel) !important;
+    color: var(--ink) !important;
+    font-family: 'JetBrains Mono', monospace !important;
+    border: 2px solid var(--ink) !important;
+    border-radius: 0 !important;
+}
+textarea:focus, input:focus { border-color: var(--accent) !important; }
+
+label, label span, .gr-check-radio span, span[data-testid="block-info"] {
+    color: var(--ink) !important;
+    font-family: 'JetBrains Mono', monospace !important;
+    font-size: 0.72rem !important;
+    font-weight: 700 !important;
+    text-transform: uppercase;
+    letter-spacing: 2px;
+}
+
+button {
+    font-family: 'JetBrains Mono', monospace !important;
+    font-weight: 700 !important;
+    text-transform: uppercase;
+    letter-spacing: 2.5px;
+    border-radius: 0 !important;
+    transition: all 0.15s ease-in-out;
+}
+
+#submit-btn {
+    background: var(--ink) !important;
+    color: var(--paper) !important;
+    border: 2px solid var(--ink) !important;
+    padding: 0.9rem !important;
+    font-size: 0.9rem !important;
+}
+#submit-btn:hover {
+    background: var(--accent) !important;
+    border-color: var(--accent) !important;
+    color: #fff !important;
+}
+
+/* ---------- like callout ---------- */
+#like-callout {
+    border: 2px solid var(--accent) !important;
+    background: var(--panel) !important;
+    text-align: center;
+    padding: 0.8rem;
+    font-size: 0.78rem;
+    letter-spacing: 2.5px;
+    text-transform: uppercase;
+    color: var(--ink);
+}
+#like-callout .accent { color: var(--accent); font-weight: 700; }
+
+footer { visibility: hidden; }
+
+/* ---------- privacy ---------- */
+#privacy-footer {
+    text-align: center;
+    font-size: 0.7rem;
+    color: var(--muted);
+    margin-top: 2rem;
+    letter-spacing: 1.5px;
+    text-transform: uppercase;
+}
+#privacy-footer a { color: var(--muted); text-decoration: underline; cursor: pointer; }
+
+#privacy-modal {
+    display: none;
+    position: fixed;
+    top: 0; left: 0; width: 100%; height: 100%;
+    background: rgba(22, 19, 14, 0.6);
+    z-index: 9999;
+    align-items: center;
+    justify-content: center;
+}
+#privacy-modal.open { display: flex; }
+#privacy-modal-box {
+    background: var(--paper);
+    border: 2px solid var(--ink);
+    padding: 2rem;
+    max-width: 480px;
+    font-family: 'JetBrains Mono', monospace;
+    font-size: 0.8rem;
+    line-height: 1.55;
+}
+#privacy-modal-box button {
+    margin-top: 1rem;
+    background: var(--ink);
+    color: var(--paper);
+    border: 2px solid var(--ink);
+    padding: 0.5rem 1.4rem;
 }
 """
 
-with gr.Blocks(css=css) as demo:
-    with gr.Column(elem_id="col-container"):
-        gr.Markdown(" # Text-to-Image Gradio Template")
-
-        with gr.Row():
-            prompt = gr.Text(
-                label="Prompt",
-                show_label=False,
-                max_lines=1,
-                placeholder="Enter your prompt",
-                container=False,
-            )
+HEADER_HTML = """
+<div id="brief-frame">
+    <span class="tick t1"></span><span class="tick t2"></span><span class="tick t3"></span>
+    <span class="tick t4"></span><span class="tick t5"></span>
+    <div class="brief-kicker">
+        <span>WanGen &middot; Image Editing Studio</span>
+        <span class="dim">Open Models &middot; Free For All</span>
+    </div>
+    <h1 class="brief-headline">Describe It. <span class="accent">Done.</span></h1>
+    <p class="brief-sub">Multi-model editing &mdash; for the beauty of open source.</p>
+</div>
+"""
+
+PRIVACY_HTML = """
+<p id='privacy-footer'>
+    <a onclick="document.getElementById('privacy-modal').classList.add('open')">Privacy Policy</a>
+</p>
 
-            run_button = gr.Button("Run", scale=0, variant="primary")
+<div id='privacy-modal'>
+    <div id='privacy-modal-box'>
+        <strong>Privacy Policy</strong><br><br>
+        Please do not upload personal information, or images you do not have the right to use.<br><br>
+        No personal data beyond what you explicitly submitted is collected. Only the data required for the system to function &mdash; your prompt and the submitted image &mdash; is processed, for AI research purposes.<br><br>
+        <button onclick="document.getElementById('privacy-modal').classList.remove('open')">Close</button>
+    </div>
+</div>
+"""
 
-        result = gr.Image(label="Result", show_label=False)
+# Pin the light theme — the design is paper-based and must not invert.
+FORCE_LIGHT_HEAD = """
+<script>
+(function () {
+    const url = new URL(window.location);
+    if (url.searchParams.get('__theme') !== 'light') {
+        url.searchParams.set('__theme', 'light');
+        window.location.replace(url.href);
+    }
+})();
+</script>
+"""
 
-        with gr.Accordion("Advanced Settings", open=False):
-            negative_prompt = gr.Text(
-                label="Negative prompt",
-                max_lines=1,
-                placeholder="Enter a negative prompt",
-                visible=False,
+# Gradio 6 moved css/head from the Blocks constructor to launch().
+GRADIO_MAJOR = int(gr.__version__.split(".")[0])
+_style_kwargs = {"css": LAB_CSS, "head": FORCE_LIGHT_HEAD}
+_blocks_kwargs = {} if GRADIO_MAJOR >= 6 else dict(_style_kwargs)
+_launch_kwargs = dict(_style_kwargs) if GRADIO_MAJOR >= 6 else {}
+
+with gr.Blocks(title="Describe It. Done.", **_blocks_kwargs) as demo:
+    gr.HTML(HEADER_HTML)
+
+    with gr.Row(equal_height=False):
+        with gr.Column(scale=5):
+            image_in = gr.Image(
+                type="pil",
+                label="Input Image — drop, paste or upload",
+                sources=["upload", "clipboard"],
+                height=320,
+            )
+            prompt_in = gr.Textbox(
+                label="What changes would you like to make ?",
+                placeholder="e.g. change the background to a forest at dusk",
+                lines=3,
             )
+            model_in = gr.Dropdown(
+                choices=MODEL_CHOICES,
+                value="Auto-Routing",
+                label="Model",
+            )
+            submit_btn = gr.Button("Submit", elem_id="submit-btn")
 
-            seed = gr.Slider(
-                label="Seed",
-                minimum=0,
-                maximum=MAX_SEED,
-                step=1,
-                value=0,
+        with gr.Column(scale=5):
+            gallery_out = gr.Gallery(
+                label="Output",
+                columns=1,
+                height=460,
+                object_fit="contain",
+            )
+            like_callout = gr.HTML(
+                "<div id='like-callout'>Enjoying the results ? "
+                "<span class='accent'>&hearts; Like this Space</span> 🙂</div>",
+                visible=False,
             )
 
-            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-
-            with gr.Row():
-                width = gr.Slider(
-                    label="Width",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,  # Replace with defaults that work for your model
-                )
-
-                height = gr.Slider(
-                    label="Height",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,  # Replace with defaults that work for your model
-                )
-
-            with gr.Row():
-                guidance_scale = gr.Slider(
-                    label="Guidance scale",
-                    minimum=0.0,
-                    maximum=10.0,
-                    step=0.1,
-                    value=0.0,  # Replace with defaults that work for your model
-                )
-
-                num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
-                    minimum=1,
-                    maximum=50,
-                    step=1,
-                    value=2,  # Replace with defaults that work for your model
-                )
-
-        gr.Examples(examples=examples, inputs=[prompt])
-    gr.on(
-        triggers=[run_button.click, prompt.submit],
-        fn=infer,
-        inputs=[
-            prompt,
-            negative_prompt,
-            seed,
-            randomize_seed,
-            width,
-            height,
-            guidance_scale,
-            num_inference_steps,
-        ],
-        outputs=[result, seed],
+    submit_btn.click(
+        fn=edit_image,
+        inputs=[image_in, prompt_in, model_in],
+        outputs=[gallery_out, like_callout],
+        show_progress="full",
     )
 
+    gr.HTML(PRIVACY_HTML)
+
+demo.queue(max_size=20, default_concurrency_limit=4)
+
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(share=False, **_launch_kwargs)