Restore FLUX.2 Klein 9B (Fibo evaluation complete)

README.md

@@ -1,6 +1,6 @@
 ---
 title: klein
-emoji: "🧪"
+emoji: "🎨"
 colorFrom: gray
 colorTo: gray
 sdk: gradio

app.py

@@ -1,202 +1,201 @@
-# FIBO-Edit 1.5 turbo — evaluation build.
-#
-# EVALUATION ONLY. The weights are CC BY-NC 4.0 and the gate form states that
-# commercial use needs a separate agreement with Bria. DressFit is a commercial
-# storefront, so nothing here may be pointed at vdf.tecwyze.lk. The purpose is
-# to find out whether that agreement would be worth buying.
-#
-# Why it is interesting for garment fitting specifically:
-#   - 8.3B, so it fits a ZeroGPU slot in bf16 with no FP8 and no torchao pin,
-#     which is the chain that made the 20B Qwen build painful.
-#   - Native multi-reference: image= takes a LIST, which is exactly the
-#     person-plus-garment shape DressFit needs.
-#   - Trained on fully licensed data, which is worth something to a retailer
-#     quite apart from output quality.
-#
-# `import spaces` first and stays first — it patches torch.cuda.*, and once CUDA
-# is initialised in the main process it is too late to intercept.
-import spaces
-
-import base64
-import random
-import re
-import traceback
-from io import BytesIO
-
-import gradio as gr
-import numpy as np
-import torch
-from diffusers import BriaFiboEditPipeline
-
-MODEL = "briaai/Fibo-Edit-1.5-turbo"
-
-dtype = torch.bfloat16
-pipe = BriaFiboEditPipeline.from_pretrained(MODEL, torch_dtype=dtype)
-
-# Eagerly at module scope: ZeroGPU intercepts this, packs the weights to disk
-# and streams them in on the first decorated call.
-pipe = pipe.to("cuda")
-
-MAX_SEED = np.iinfo(np.int32).max
-
-# The card is explicit: distilled from 30 steps to 4, and CFG 1.0 because no
-# classifier-free guidance is needed. Raising either is wasted quota.
-STEPS = 4
-GUIDANCE = 1.0
-
-# PROMPTS ARE JSON, NOT PROSE, AND THIS IS ENFORCED.
-#
-# check_inputs runs is_valid_edit_json on every prompt and raises ValueError on
-# anything without an "edit_instruction" key, so plain prose does not reach the
-# model at all — it dies before the first step. Bria build that JSON with a
-# Gemini-backed VLM ("VGL JSON") which needs a GEMINI_API_KEY we do not have.
-#
-# The validator is far laxer than the schema: any JSON object carrying
-# edit_instruction passes, and the pipeline then json.dumps it and tokenises the
-# result as ordinary text. So a wrapper is enough to make the model reachable,
-# which is what _to_vgl below does. Two consequences worth stating plainly for
-# whoever reads the evaluation:
-#
-#   - The other VGL fields (short_description, objects, lighting, aesthetics …)
-#     need a captioner to fill and are omitted. The model gets less conditioning
-#     than Bria intend it to have.
-#   - The format wants ONE imperative sentence in edit_instruction and forbids
-#     negative constraints elsewhere. DressFit prompts are the opposite: long,
-#     and mostly a list of things not to do. That mismatch is a real handicap,
-#     so the harness also sends a short imperative variant of each case to tell
-#     a weak model apart from a badly-shaped prompt.
-
-
-def _to_vgl(prompt: str) -> str:
-    """Coerce a prompt into the edit JSON the pipeline demands.
-
-    Already-JSON prompts carrying edit_instruction pass through untouched, so a
-    caller that builds proper VGL keeps full control. Prose is wrapped, with
-    "image 1"/"image 2" rewritten to the <image_N> reference tokens the
-    multi-reference format uses.
-    """
-    import json as _json
-    text = (prompt or "").strip()
-    if text.startswith("{"):
-        try:
-            parsed = _json.loads(text)
-            if isinstance(parsed, dict) and "edit_instruction" in parsed:
-                return text
-        except _json.JSONDecodeError:
-            pass
-
-    instruction = re.sub(r"image\s*1", "<image_1>", text, flags=re.I)
-    instruction = re.sub(r"image\s*2", "<image_2>", instruction, flags=re.I)
-    return _json.dumps({"edit_instruction": instruction})
-
-
-def _duration(*a, **k):
-    w = int(k.get("width") or 1024)
-    h = int(k.get("height") or 1024)
-    return min(120, max(30, round((w * h) / 50000)))
-
-
-@spaces.GPU(duration=_duration)
-def infer(image_1, image_2, prompt, seed=0, randomize_seed=True,
-          guidance_scale=GUIDANCE, num_inference_steps=STEPS,
-          width=1024, height=1024,
-          progress=gr.Progress(track_tqdm=True)):
-    """Edit one or more reference images according to a prompt."""
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-
-    images = [im for im in (image_1, image_2) if im is not None]
-    if not images:
-        raise gr.Error("At least one reference image is required.")
-
-    out = pipe(
-        # A LIST is multi-reference, not a batch — the card is explicit that
-        # several references combine into one result.
-        image=images if len(images) > 1 else images[0],
-        prompt=_to_vgl(prompt),
-        num_inference_steps=int(num_inference_steps),
-        guidance_scale=float(guidance_scale),
-        generator=torch.Generator("cuda").manual_seed(int(seed)),
-    ).images[0]
-
-    return out, seed
-
-
-with gr.Blocks(delete_cache=(3600, 86400)) as demo:
-    with gr.Column():
-        gr.HTML(
-            """
-            <div style="text-align:center;padding:4rem 1rem;
-                        font-family:system-ui,sans-serif;color:#999;font-size:.9rem">
-                This Space has no public interface.
-            </div>
-            """
-        )
-        # Hidden, but present: removing the components removes the event
-        # registration that publishes the API with them.
-        with gr.Row(visible=False):
-            image_1 = gr.Image(label="Image 1", type="pil")
-            image_2 = gr.Image(label="Image 2", type="pil")
-            prompt = gr.Text(label="Prompt")
-            seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
-            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-            guidance = gr.Slider(label="Guidance", minimum=0.0, maximum=10.0, step=0.1, value=GUIDANCE)
-            steps = gr.Slider(label="Steps", minimum=1, maximum=40, step=1, value=STEPS)
-            run = gr.Button("Edit")
-            result = gr.Image(label="Result")
-            out_seed = gr.Number(label="Seed used")
-
-    gr.on(
-        triggers=[run.click, prompt.submit],
-        fn=infer,
-        inputs=[image_1, image_2, prompt, seed, randomize_seed, guidance, steps],
-        outputs=[result, out_seed],
-    )
-
-
-# ============================================================
-# API — same positional contract as the other edit Spaces, so the existing
-# client can call this one by changing a URL and nothing else.
-# ============================================================
-def _load(src):
-    """Fetch an https URL or decode a data URL into a PIL image."""
-    import base64 as _b64
-    import urllib.request as _url
-    from PIL import Image as _Image
-    if src.startswith("http://") or src.startswith("https://"):
-        req = _url.Request(src, headers={"User-Agent": "tecwyze/1.0"})
-        with _url.urlopen(req, timeout=30) as resp:
-            raw = resp.read(25 * 1024 * 1024)
-    else:
-        if src.startswith("data:"):
-            src = src.split(",", 1)[1]
-        raw = _b64.b64decode(src)
-    return _Image.open(BytesIO(raw)).convert("RGB")
-
-
-def dressfit_tryon(person: str, garment: str = "", prompt: str = "",
-                   steps: float = STEPS, guidance: float = GUIDANCE,
-                   seed: float = 0, randomize: bool = True) -> dict:
-    """Edit an image. person is the image to edit, garment is an optional second
-    reference, prompt says what to change. Both are https or data URLs."""
-    try:
-        first = _load(person)
-        second = _load(garment) if garment else None
-        w, h = first.size
-        w = max(256, min(2048, round(w / 32) * 32))
-        h = max(256, min(2048, round(h / 32) * 32))
-        out, used = infer(first, second, prompt, int(seed), bool(randomize),
-                          float(guidance), int(steps), w, h)
-        buf = BytesIO()
-        out.save(buf, format="PNG")
-        return {"status": "success", "seed": int(used),
-                "image": "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()}
-    except Exception as exc:
-        return {"status": "error", "message": "".join(traceback.format_exception(exc))[-2500:]}
-
-
-with demo:
-    gr.api(dressfit_tryon, api_name="dressfit_tryon")
-
-if __name__ == "__main__":
-    demo.launch()
+# FLUX.2 Klein 9B — reference-image editing with runtime LoRA.
+#
+# This is the model that answers both open requests at once. Qwen-Image-Edit
+# holds identity but renders glossy; Z-Image looks photographic but its editing
+# pipeline needs siglip weights Tongyi has not released; neither supports LoRA
+# on our stack. Klein is tagged image-to-image, takes one or more reference
+# images, and Flux2KleinPipeline carries load_lora_weights — so one checkpoint
+# covers generation, editing and per-character LoRA.
+#
+# Started on the 4B because it is Apache-2.0 and ungated. Its output was not
+# convincing enough, so this runs the 9B — see the note at MODEL below.
+#
+# `import spaces` first and stays first — it patches torch.cuda.*, and once CUDA
+# is initialised in the main process it is too late to intercept.
+import spaces
+
+import base64
+import random
+import traceback
+from io import BytesIO
+
+import gradio as gr
+import numpy as np
+import torch
+from diffusers import Flux2KleinPipeline
+
+# 9B, not the 4B this started on — the 4B's output was not convincing enough.
+# It is gated behind Black Forest Labs' flux-non-commercial-license, so the
+# account running this Space must accept the terms on the model page first, and
+# the Space needs an HF_TOKEN secret to authenticate the download. Without both,
+# the build dies on a 403 fetching model_index.json.
+MODEL = "black-forest-labs/FLUX.2-klein-9B"
+
+dtype = torch.bfloat16
+pipe = Flux2KleinPipeline.from_pretrained(MODEL, torch_dtype=dtype)
+
+# Eagerly at module scope: ZeroGPU intercepts this, packs the weights to disk
+# and streams them into VRAM on the first decorated call. Loading lazily inside
+# the decorator makes every caller pay the cost instead of nobody.
+pipe = pipe.to("cuda")
+
+MAX_SEED = np.iinfo(np.int32).max
+STEPS = 4          # Klein is distilled; the reference Space defaults to 4
+GUIDANCE = 2.5
+
+# Adapters are cached by name and never re-downloaded. Loading a LoRA per
+# request would re-fetch tens of megabytes every call, and the Space process
+# survives between calls, so the cache is free. Adapters also STACK if you do
+# not manage them — set_adapters replaces the active set, and disable_lora
+# turns them all off, which is what an unadorned request needs.
+LOADED = set()
+
+
+def _apply_loras(specs):
+    """specs is a list of {repo, weight_name, scale}. Empty disables all."""
+    if not specs:
+        pipe.disable_lora()
+        return []
+    names, scales = [], []
+    for i, spec in enumerate(specs):
+        repo = spec.get("repo")
+        if not repo:
+            continue
+        weight_name = spec.get("weight_name") or None
+        name = spec.get("adapter_name") or f"a{abs(hash((repo, weight_name))) % 10**8}"
+        if name not in LOADED:
+            pipe.load_lora_weights(repo, weight_name=weight_name, adapter_name=name)
+            LOADED.add(name)
+        names.append(name)
+        scales.append(float(spec.get("scale", 1.0)))
+    if not names:
+        pipe.disable_lora()
+        return []
+    pipe.set_adapters(names, adapter_weights=scales)
+    return names
+
+
+def _duration(*a, **k):
+    # ZeroGPU bills the duration REQUESTED but ABORTS anything that overruns,
+    # so one flat number either wastes quota on small images or kills big ones.
+    w = int(k.get("width") or 1024)
+    h = int(k.get("height") or 1024)
+    return min(120, max(30, round((w * h) / 40000)))
+
+
+@spaces.GPU(duration=_duration)
+def infer(image_1, image_2, prompt, seed=0, randomize_seed=True,
+          guidance_scale=GUIDANCE, num_inference_steps=STEPS,
+          width=1024, height=1024, loras=None,
+          progress=gr.Progress(track_tqdm=True)):
+    """Edit one or two reference images according to a prompt."""
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+
+    _apply_loras(loras or [])
+
+    images = [im for im in (image_1, image_2) if im is not None]
+    if not images:
+        raise gr.Error("At least one reference image is required.")
+
+    out = pipe(
+        image=images if len(images) > 1 else images[0],
+        prompt=prompt,
+        width=int(width),
+        height=int(height),
+        num_inference_steps=int(num_inference_steps),
+        guidance_scale=float(guidance_scale),
+        generator=torch.Generator("cuda").manual_seed(int(seed)),
+    ).images[0]
+
+    return out, seed
+
+
+with gr.Blocks(delete_cache=(3600, 86400)) as demo:
+    with gr.Column():
+        # No public interface. This is a backend for one application, and a
+        # landing page only attracts traffic that spends the ZeroGPU quota. The
+        # controls below still exist, hidden, because hiding them removes the
+        # interface without removing the event registration that publishes the
+        # API — delete the components and the endpoint goes with them.
+        gr.HTML(
+            """
+            <div style="text-align:center;padding:4rem 1rem;
+                        font-family:system-ui,sans-serif;color:#999;font-size:.9rem">
+                This Space has no public interface.
+            </div>
+            """
+        )
+        with gr.Row(visible=False):
+            image_1 = gr.Image(label="Image 1", type="pil")
+            image_2 = gr.Image(label="Image 2", type="pil")
+            prompt = gr.Text(label="Prompt")
+            seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
+            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+            guidance = gr.Slider(label="Guidance", minimum=0.0, maximum=10.0, step=0.1, value=GUIDANCE)
+            steps = gr.Slider(label="Steps", minimum=1, maximum=50, step=1, value=STEPS)
+            run = gr.Button("Edit")
+            result = gr.Image(label="Result")
+            out_seed = gr.Number(label="Seed used")
+
+    gr.on(
+        triggers=[run.click, prompt.submit],
+        fn=infer,
+        inputs=[image_1, image_2, prompt, seed, randomize_seed, guidance, steps],
+        outputs=[result, out_seed],
+    )
+
+
+# ============================================================
+# API
+# ============================================================
+def _load(src):
+    """Fetch an https URL or decode a data URL into a PIL image."""
+    import base64 as _b64
+    import urllib.request as _url
+    from PIL import Image as _Image
+    if src.startswith("http://") or src.startswith("https://"):
+        req = _url.Request(src, headers={"User-Agent": "tecwyze/1.0"})
+        with _url.urlopen(req, timeout=30) as resp:
+            raw = resp.read(25 * 1024 * 1024)
+    else:
+        if src.startswith("data:"):
+            src = src.split(",", 1)[1]
+        raw = _b64.b64decode(src)
+    return _Image.open(BytesIO(raw)).convert("RGB")
+
+
+# Named dressfit_tryon because the existing client already calls that with this
+# positional order. A tidier name would break every caller for no gain.
+def dressfit_tryon(person: str, garment: str = "", prompt: str = "",
+                   steps: float = STEPS, guidance: float = GUIDANCE,
+                   seed: float = 0, randomize: bool = True,
+                   loras: list | None = None) -> dict:
+    """Edit an image. person is the image to edit, garment is an optional second
+    image, prompt says what to change, loras is an optional list of
+    {repo, weight_name, scale}. Images are https or data URLs."""
+    try:
+        first = _load(person)
+        second = _load(garment) if garment else None
+        # Match the reference's own aspect. The pipeline draws a fresh canvas at
+        # whatever size it is given, so forcing a square distorts a portrait.
+        w, h = first.size
+        w = max(256, min(2048, round(w / 32) * 32))
+        h = max(256, min(2048, round(h / 32) * 32))
+        out, used = infer(first, second, prompt, int(seed), bool(randomize),
+                          float(guidance), int(steps), w, h, loras or [])
+        buf = BytesIO()
+        out.save(buf, format="PNG")
+        return {"status": "success", "seed": int(used),
+                "image": "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()}
+    except Exception as exc:
+        # Return the traceback to the caller, not only the logs — the client is
+        # the only thing watching and each log round trip costs an attempt.
+        return {"status": "error", "message": "".join(traceback.format_exception(exc))[-2500:]}
+
+
+with demo:
+    gr.api(dressfit_tryon, api_name="dressfit_tryon")
+
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -3,3 +3,4 @@ transformers
 accelerate
 safetensors
 sentencepiece
+peft