app.py

# 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()