"""HuggingFace Spaces (Gradio) demo for Foveated Diffusion.
Mirrors the behavior of `webgui/server.py` from the project repo, adapted to
Gradio + an HF Spaces deployment:
- LoRA checkpoints (no_fov, fov_random, fov_saliency, fov_bbox) are fetched
from `bchao1/foveated_diffusion` via `hf_hub_download` on startup. Same
fuse/unfuse machinery as the Flask server — switching is fast and bounded
in VRAM.
- The painted mask + "preset circle" mask spec are quantized to LR blocks
under the same any-touch-upgrades-the-block rule used in `webgui/`.
- Generation runs under the same fixed settings (1024x1024, 50 steps, soft
foveation blend on, decode_mode=merge, prediction_type=clean).
Layout (mirrors webgui/static/index.html):
Left column = prompt, seed, LoRA, mask mode, draw/preset controls, Generate.
Right column = side-by-side tokenization mask preview + generated image.
"""
from __future__ import annotations
import io
import os
import sys
import threading
import time
from types import SimpleNamespace
# Match webgui/server.py: disable sageattention before importing torch-deps.
sys.modules["sageattention"] = None
# ZeroGPU runtime. MUST be imported before torch so that `spaces` can intercept
# CUDA initialization during module load. On non-Spaces envs `spaces` isn't
# installed, so we install a no-op shim that makes `@spaces.GPU(...)` a
# transparent decorator. This keeps `python app.py` working locally.
try:
import spaces # type: ignore
except ImportError:
class _SpacesShim:
@staticmethod
def GPU(*args, **kwargs):
def decorator(fn):
return fn
# @spaces.GPU (no parens): the first arg is the function itself.
if args and callable(args[0]) and not kwargs:
return args[0]
return decorator
spaces = _SpacesShim() # type: ignore[assignment]
import numpy as np
import torch
import torch.nn.functional as F
import gradio as gr
from huggingface_hub import hf_hub_download
from PIL import Image
# Make the bundled `src/` package importable.
#
# On Spaces, `src/` is vendored into the Space root by `upload_space.py`, so
# it lives next to app.py at `_HERE/src`. For local dev (running app.py from
# the project repo before upload), `src/` only exists at the repo root, so we
# also add `_HERE/../..` to sys.path. The second insert is a no-op on Spaces
# (the path doesn't exist there).
_HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, _HERE)
_REPO_ROOT = os.path.normpath(os.path.join(_HERE, "..", ".."))
if os.path.isdir(os.path.join(_REPO_ROOT, "src")):
sys.path.insert(0, _REPO_ROOT)
from diffsynth.core import load_state_dict # noqa: E402
from src.inference import load_pipeline # noqa: E402
from src.inference.visualize import create_tokenization_mask_vis # noqa: E402
from src.masks import create_foveation_mask_full_res # noqa: E402
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
# Fixed generation settings (per webgui spec).
MODEL_ID = "black-forest-labs/FLUX.2-klein-base-4B"
HEIGHT = 1024
WIDTH = 1024
NUM_INFERENCE_STEPS = 50
GUIDANCE_SCALE = 4.0
DECODE_MODE = "merge"
PREDICTION_TYPE = "clean"
LR_DOWNSAMPLE_FACTOR = 2
SOFT_FOVEATION_BLEND = True
# Masks are tiny (1024x1024 fp32 ~ 4 MB) and the only thing they need to do
# outside a `@spaces.GPU` call is feed the tokenization preview. Building them
# on CPU keeps `refresh_tokenization` (fires on every slider move) off the GPU
# entirely; we move them to `_pipe.device` once, inside `generate`, just before
# the actual forward.
MASK_DEVICE = "cpu"
# LoRA registry. Files come from the bchao1/foveated_diffusion model repo.
LORA_REPO = "bchao1/foveated_diffusion"
LORA_FILES = {
"no_fov": "image/no_fov.safetensors",
"random": "image/fov_random.safetensors",
"saliency": "image/fov_saliency.safetensors",
# "bbox": "image/fov_bbox.safetensors", # disabled in this UI build
}
DEFAULT_LORA = "random"
_pipe = None
_lora_registry: dict = {} # name -> {"path": str, "state_dict": dict or None}
_lora_order: list = []
# LoRA state. The dropdown (`lora_dd`) is the actual source of truth — its
# value is plumbed into every handler (and `generate`) as an explicit Gradio
# input, so it crosses the @spaces.GPU process boundary safely under either
# fork or spawn. The two globals below are kept only for:
# _selected_lora — parent-only book-keeping for `switch_lora`'s rollback on
# an invalid dropdown value (defensive; not load-bearing).
# _current_lora — worker-local. Tracks what's fused in *this* worker's
# model copy so `_apply_selected_lora` can skip work when
# the dropdown choice already matches.
_selected_lora: str = DEFAULT_LORA
_current_lora: str | None = None
_lora_lock = threading.Lock()
_job_lock = threading.Lock() # serializes generation + LoRA switching
# ---------------------------------------------------------------------------
# LoRA fuse / unfuse (ported verbatim from webgui/server.py)
# ---------------------------------------------------------------------------
def _convert_lora_sd(state_dict):
loader = _pipe.lora_loader(torch_dtype=_pipe.torch_dtype, device="cpu")
return loader.convert_state_dict(state_dict)
def _apply_lora_delta(state_dict, sign):
lora_layer_names = {
k[: -len(".lora_B.weight")]
for k in state_dict if k.endswith(".lora_B.weight")
}
updated = 0
for name, module in _pipe.dit.named_modules():
if name not in lora_layer_names:
continue
wb = state_dict[name + ".lora_B.weight"].to(device=_pipe.device, dtype=_pipe.torch_dtype)
wa = state_dict[name + ".lora_A.weight"].to(device=_pipe.device, dtype=_pipe.torch_dtype)
if wb.dim() == 4:
wb = wb.squeeze(3).squeeze(2)
wa = wa.squeeze(3).squeeze(2)
delta = torch.mm(wb, wa).unsqueeze(2).unsqueeze(3)
else:
delta = torch.mm(wb, wa)
base = module.weight.data
base.add_(delta.to(device=base.device, dtype=base.dtype), alpha=float(sign))
updated += 1
return updated
def _switch_lora_unlocked(target_name: str):
global _current_lora
if target_name == _current_lora:
return
if target_name not in _lora_registry:
raise ValueError(f"unknown lora: {target_name}")
if _current_lora is not None:
cur = _lora_registry[_current_lora]
n = _apply_lora_delta(cur["state_dict"], sign=-1)
print(f"[lora] unfused {_current_lora} ({n} tensors)")
nxt = _lora_registry[target_name]
if nxt["state_dict"] is None:
raw = load_state_dict(nxt["path"], torch_dtype=_pipe.torch_dtype)
nxt["state_dict"] = _convert_lora_sd(raw)
n = _apply_lora_delta(nxt["state_dict"], sign=+1)
print(f"[lora] fused {target_name} ({n} tensors)")
_current_lora = target_name
def _fetch_lora_files() -> dict[str, str]:
"""hf_hub_download each LoRA from LORA_REPO; return {name: local_path}."""
out = {}
for name, fn in LORA_FILES.items():
print(f"[lora] fetching {LORA_REPO}/{fn} …")
out[name] = hf_hub_download(repo_id=LORA_REPO, filename=fn)
return out
def _load_pipe():
global _pipe
# On HF Spaces we MUST pull FLUX.2 from HuggingFace Hub (not ModelScope, the
# diffsynth default). The `preload_from_hub` block in README.md caches these
# files during the Space build, so this call is just a fast cache lookup.
# Locally, $DIFFSYNTH_MODEL_BASE_PATH points at a ModelScope checkout, so we
# keep the diffsynth default there.
if "SPACE_ID" in os.environ or os.environ.get("DIFFSYNTH_USE_HF") == "1":
from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig
from src.diffsynth_fov import Flux2FoveatedImagePipeline
print("[pipe] loading FLUX2 from HuggingFace Hub …")
_pipe = Flux2FoveatedImagePipeline.from_pretrained(
torch_dtype=torch.bfloat16,
device="cuda" if torch.cuda.is_available() else "cpu",
model_configs=[
ModelConfig(model_id=MODEL_ID, origin_file_pattern="transformer/*.safetensors", download_source="huggingface"),
ModelConfig(model_id=MODEL_ID, origin_file_pattern="text_encoder/*.safetensors", download_source="huggingface"),
ModelConfig(model_id=MODEL_ID, origin_file_pattern="vae/diffusion_pytorch_model.safetensors", download_source="huggingface"),
],
tokenizer_config=ModelConfig(model_id=MODEL_ID, origin_file_pattern="tokenizer/", download_source="huggingface"),
)
else:
loader_args = SimpleNamespace(
model_id=MODEL_ID,
lora_checkpoint=None,
dit_checkpoint=None,
experiment="ours",
)
_pipe = load_pipeline(loader_args)
print(f"[pipe] loaded on {_pipe.device}")
paths = _fetch_lora_files()
for name in LORA_FILES: # preserve declared order
_lora_registry[name] = {"path": paths[name], "state_dict": None}
_lora_order.append(name)
# Preload + convert all LoRA state dicts on CPU at startup. They stay in
# parent memory and are visible (copy-on-write) to every `@spaces.GPU`
# worker, so a cold worker doesn't pay disk-read + convert cost inside its
# GPU window. ~hundreds of MB per LoRA in CPU RAM — fine.
print("[lora] preloading state dicts on CPU …")
for name in LORA_FILES:
raw = load_state_dict(_lora_registry[name]["path"], torch_dtype=_pipe.torch_dtype)
_lora_registry[name]["state_dict"] = _convert_lora_sd(raw)
print("[lora] preload done")
# NOTE: the actual fuse (`_apply_lora_delta`, which does `torch.mm` on
# CUDA) is deferred to the first `@spaces.GPU` `generate` call. On
# ZeroGPU the GPU is only attached inside @spaces.GPU-decorated
# functions, AND those run in forked workers — so any fuse done outside
# them either fails (no GPU) or doesn't propagate back to the parent.
def _apply_selected_lora(lora_name: str):
"""Reconcile the worker's fused LoRA with `lora_name`.
Must be called from inside an `@spaces.GPU` function. `lora_name` is
passed in (from the dropdown via Gradio inputs) rather than read from a
global — globals don't reliably cross the @spaces.GPU process boundary
on ZeroGPU."""
if _current_lora == lora_name:
return
with _lora_lock:
if _current_lora == lora_name: # double-checked
return
_switch_lora_unlocked(lora_name)
# ---------------------------------------------------------------------------
# Mask construction (ported from webgui/server.py)
# ---------------------------------------------------------------------------
def _quantize_to_lr_blocks(raw_mask: torch.Tensor):
if raw_mask.shape != (HEIGHT, WIDTH):
raw_mask = F.interpolate(
raw_mask.unsqueeze(0).unsqueeze(0), size=(HEIGHT, WIDTH), mode="nearest",
).squeeze(0).squeeze(0)
lr_h = HEIGHT // 16 // LR_DOWNSAMPLE_FACTOR
lr_w = WIDTH // 16 // LR_DOWNSAMPLE_FACTOR
lr_mask = F.adaptive_max_pool2d(
(raw_mask > 0).float().unsqueeze(0).unsqueeze(0), (lr_h, lr_w),
)
token_grid = F.interpolate(
lr_mask, size=(HEIGHT // 16, WIDTH // 16), mode="nearest",
).squeeze(0).squeeze(0)
full_res = F.interpolate(
lr_mask, size=(HEIGHT, WIDTH), mode="nearest",
).squeeze(0).squeeze(0)
return token_grid, full_res
def _drawn_mask_to_tensors(rgba: np.ndarray):
"""ImageEditor output -> (token_grid_mask, full_res_mask). `rgba` is HxWx4 uint8."""
if rgba.ndim == 3 and rgba.shape[2] == 4:
alpha = rgba[..., 3]
elif rgba.ndim == 3:
alpha = rgba.max(axis=2)
else:
alpha = rgba
raw = torch.from_numpy((alpha > 0).astype(np.float32)).to(MASK_DEVICE)
return _quantize_to_lr_blocks(raw)
def _alpha_union_from_layers(layers) -> np.ndarray | None:
"""OR together alpha channels of all layers — gives painted region in RGBA layers
even when the composite is opaque."""
if not layers:
return None
out = None
for layer in layers:
arr = np.asarray(layer)
if arr.ndim != 3 or arr.shape[2] != 4 or arr.size == 0:
continue
a = arr[..., 3]
out = a if out is None else np.maximum(out, a)
return out
def _preset_mask_to_tensors(cx: float, cy: float, r: float):
full_res = create_foveation_mask_full_res(
HEIGHT, WIDTH, (cx, cy), r, "circular", MASK_DEVICE,
)
return _quantize_to_lr_blocks(full_res)
def _empty_mask_tensors():
token = torch.zeros(HEIGHT // 16, WIDTH // 16, device=MASK_DEVICE, dtype=torch.float32)
full = torch.zeros(HEIGHT, WIDTH, device=MASK_DEVICE, dtype=torch.float32)
return token, full
def _full_hr_mask_tensors():
"""All-HR mask — every token is foveal. Used with the no-fov LoRA."""
token = torch.ones(HEIGHT // 16, WIDTH // 16, device=MASK_DEVICE, dtype=torch.float32)
full = torch.ones(HEIGHT, WIDTH, device=MASK_DEVICE, dtype=torch.float32)
return token, full
def _build_mask(lora_name: str, mode: str, editor_value, cx: float, cy: float, r: float):
"""Return (token_grid_mask, full_res_mask, is_empty).
`lora_name` is passed explicitly (from the Gradio dropdown component, not
a global) so the call works across the @spaces.GPU process boundary even
if ZeroGPU uses spawn semantics."""
# The no-fov LoRA was trained without foveation conditioning — it expects
# a full-HR token mask and ignores the circle controls entirely.
if lora_name == "no_fov":
t, f = _full_hr_mask_tensors()
return t, f, False
if mode == "Preset (circle)":
if r <= 0:
return (*_empty_mask_tensors(), True)
t, f = _preset_mask_to_tensors(float(cx), float(cy), float(r))
return t, f, bool(f.sum().item() == 0)
# Drawn. ImageEditor returns {"background": ..., "layers": [...], "composite": ...}
# We MUST use the layers (not the composite): Gradio's composite has alpha=255
# everywhere (opaque background), so `alpha > 0` matches every pixel and the
# whole image registers as HR. The layer RGBA arrays preserve true paint alpha.
if editor_value is None:
return (*_empty_mask_tensors(), True)
if isinstance(editor_value, dict):
layers = editor_value.get("layers") or []
alpha = _alpha_union_from_layers(layers)
if alpha is None:
# No usable layers; fall back to composite alpha (rare).
composite = editor_value.get("composite")
if composite is None:
return (*_empty_mask_tensors(), True)
arr = np.asarray(composite)
if arr.size == 0:
return (*_empty_mask_tensors(), True)
t, f = _drawn_mask_to_tensors(arr)
return t, f, bool(f.sum().item() == 0)
raw = torch.from_numpy((alpha > 0).astype(np.float32)).to(MASK_DEVICE)
t, f = _quantize_to_lr_blocks(raw)
return t, f, bool(f.sum().item() == 0)
# Bare numpy array — treat as composite-like.
arr = np.asarray(editor_value)
if arr.size == 0:
return (*_empty_mask_tensors(), True)
t, f = _drawn_mask_to_tensors(arr)
return t, f, bool(f.sum().item() == 0)
def _tokenization_vis_image(token_mask) -> np.ndarray:
return create_tokenization_mask_vis(token_mask, HEIGHT, WIDTH, lr_factor=LR_DOWNSAMPLE_FACTOR)
# ---------------------------------------------------------------------------
# Gradio handlers
# ---------------------------------------------------------------------------
def refresh_tokenization(lora_name, mode, editor_value, cx, cy, r):
print(f"[tok] enter lora={lora_name!r} mode={mode!r} cx={cx} cy={cy} r={r} editor={type(editor_value).__name__}", flush=True)
if _pipe is None:
print("[tok] pipe is None, returning", flush=True)
return None
try:
t0 = time.time()
token_mask, _, _ = _build_mask(lora_name, mode, editor_value, cx, cy, r)
print(f"[tok] mask built ({time.time() - t0:.3f}s)", flush=True)
t1 = time.time()
vis = _tokenization_vis_image(token_mask)
print(f"[tok] vis built ({time.time() - t1:.3f}s) shape={getattr(vis, 'shape', None)}", flush=True)
return vis
except Exception as e:
import traceback
print(f"[tok] EXCEPTION: {e!r}", flush=True)
traceback.print_exc()
return None
def on_mode_change(lora_name, mode, editor_value, cx, cy, r):
"""Switch active mode — refresh tokenization + swap CSS active/inactive
classes on the two group wrappers.
We deliberately do NOT toggle visibility of the draw/preset groups: hiding
a Group via `gr.update(visible=False)` in Gradio 6.14 unmounts its children,
and re-showing it later does not re-bind their event listeners — slider
.change/.release would silently stop firing. Both groups stay mounted; CSS
dims the inactive one via the `mode-active` / `mode-inactive` classes.
"""
is_draw = (mode == "Draw")
draw_cls = ["mode-group", "mode-active" if is_draw else "mode-inactive"]
preset_cls = ["mode-group", "mode-inactive" if is_draw else "mode-active"]
tok = refresh_tokenization(lora_name, mode, editor_value, cx, cy, r)
return gr.update(elem_classes=draw_cls), gr.update(elem_classes=preset_cls), tok
def switch_lora(name: str, cx: float, cy: float, r: float):
"""Record the user's LoRA selection and refresh dependent UI.
NOT decorated with `@spaces.GPU`: the actual fuse can't happen here
because @spaces.GPU runs in a forked worker and its mutations to the
model + globals don't propagate back to the parent. Instead, we just
update `_selected_lora` here and let the next `generate` worker fuse.
UI updates: dropdown value, description, status, slider interactivity
(no_fov disables circle controls), tokenization preview.
"""
global _selected_lora
if name not in LORA_FILES:
return (
gr.update(value=_selected_lora),
LORA_DESCRIPTIONS.get(_selected_lora, ""),
f"unknown LoRA: {name}",
gr.update(),
gr.update(),
gr.update(),
None,
)
_selected_lora = name
is_no_fov = (name == "no_fov")
# Pass `name` explicitly so the no_fov full-HR short-circuit triggers
# for the preview the user sees right after switching.
token_mask, _, _ = _build_mask(name, "Preset (circle)", None, cx, cy, r)
tok_vis = _tokenization_vis_image(token_mask)
# interactive=False + CSS `pointer-events: none` so the handles can't fire
# `.change`/`.release` even if the user clicks them. (Belt-and-suspenders:
# the `_build_mask` short-circuit above already forces full-HR for no_fov,
# but this prevents the sliders from looking active.)
slider_classes = ["slider-locked"] if is_no_fov else []
def _slider_update():
return gr.update(interactive=not is_no_fov, elem_classes=slider_classes)
return (
gr.update(value=name),
LORA_DESCRIPTIONS.get(name, ""),
f"Selected: {name} (fuses on next generate)",
_slider_update(),
_slider_update(),
_slider_update(),
tok_vis,
)
def toggle_mode(mode):
is_draw = (mode == "Draw")
return gr.update(visible=is_draw), gr.update(visible=not is_draw)
def fill_all_editor():
"""Paint the entire canvas with opaque black so the whole image becomes HR."""
layer = np.zeros((HEIGHT, WIDTH, 4), dtype=np.uint8)
layer[..., 0:3] = 17
layer[..., 3] = 255
return {"background": None, "layers": [layer], "composite": layer}
def clear_editor():
return {"background": None, "layers": [], "composite": None}
def clear_preset():
return 0.0, 0.0, 0.0
@spaces.GPU(duration=120)
def generate(prompt, seed, lora_name, mode, editor_value, cx, cy, r,
progress=gr.Progress(track_tqdm=False)):
print(f"[gen] CLICK REACHED lora={lora_name!r} mode={mode!r} seed={seed}", flush=True)
if _pipe is None:
raise gr.Error("Pipeline not loaded yet — please wait for cold start.")
if lora_name not in LORA_FILES:
raise gr.Error(f"unknown LoRA: {lora_name}")
prompt = (prompt or "").strip()
if not prompt:
raise gr.Error("Prompt is required.")
seed = int(seed or 0)
# `lora_name` came in as a Gradio input (the live dropdown value), so it
# crosses the @spaces.GPU process boundary safely — no global-state
# propagation required.
_apply_selected_lora(lora_name)
progress(0.0, desc="Building mask")
t_mask = time.time()
token_mask, full_res_mask, is_empty = _build_mask(lora_name, mode, editor_value, cx, cy, r)
print(f"[gen] mask built ({time.time() - t_mask:.2f}s) is_empty={is_empty}", flush=True)
if is_empty:
raise gr.Error("Mask is empty — paint a region or set a circle radius > 0.")
tok_vis = _tokenization_vis_image(token_mask)
print("[gen] tokenization vis ready -> yielding to UI", flush=True)
# Yield tok_vis right away so the tokenization panel updates *before* denoise.
# gr.update() keeps out_img unchanged at this stage.
yield tok_vis, gr.update()
print("[gen] waiting for _job_lock", flush=True)
with _job_lock:
print("[gen] _job_lock acquired; calling _pipe", flush=True)
torch.cuda.empty_cache()
# Move CPU-built masks onto the pipeline's device for the forward.
token_mask = token_mask.to(_pipe.device)
full_res_mask = full_res_mask.to(_pipe.device)
def progress_wrap(iterable):
iterable = list(iterable)
total = len(iterable) or NUM_INFERENCE_STEPS
t_prev = time.time()
for i, item in enumerate(iterable):
yield item
t_now = time.time()
print(f"[gen] step {i + 1}/{total} {t_now - t_prev:.3f}s", flush=True)
t_prev = t_now
progress((i + 1) / total, desc=f"Denoising {i + 1}/{total}")
t0 = time.time()
image = _pipe(
prompt=prompt,
height=HEIGHT,
width=WIDTH,
seed=seed,
rand_device="cuda" if torch.cuda.is_available() else "cpu",
num_inference_steps=NUM_INFERENCE_STEPS,
cfg_scale=GUIDANCE_SCALE,
foveation_mask=token_mask,
full_res_foveation_mask=full_res_mask,
decode_mode=DECODE_MODE,
prediction_type=PREDICTION_TYPE,
soft_foveation_blend=SOFT_FOVEATION_BLEND,
lr_downsample_factor=LR_DOWNSAMPLE_FACTOR,
progress_bar_cmd=progress_wrap,
)
print(f"[gen] total pipeline call: {time.time() - t0:.2f}s", flush=True)
yield tok_vis, image
# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------
# Short description shown under the LoRA dropdown.
LORA_DESCRIPTIONS = {
"no_fov": "Finetuned baseline — **no foveation conditioning**. Always generates a uniform full-resolution image (mask controls disabled).",
"random": "Finetuned with **random gaze locations** during training. Most general-purpose foveation LoRA.",
"saliency": "Finetuned with **saliency-driven** foveation masks (model attends to visually salient regions).",
}
DEFAULT_PROMPT = (
"Plitvice Lakes cascade through a series of sixteen terraced turquoise pools "
"connected by hundreds of waterfalls in a forested Croatian valley. The "
"travertine barriers between each pool, the moss and the mineral deposits, "
"and the submerged fallen logs in the crystal-clear water are all equally "
"sharp. The surrounding beech and fir forest is in full autumn colour—gold, "
"russet, and dark green—and the waterfalls range from wide curtains to thin "
"threads, all frozen sharp in the even overcast light. The water's colour "
"shifts from pale aquamarine in the shallows to deep teal in the pools."
)
CSS = """
.tok-img img, .out-img img { background: #f4f4f4 !important; }
/* Belt-and-suspenders lock for cx/cy/r when no_fov is active: gr.update
(interactive=False) sets the disabled attribute, this kills pointer events
so a stray click on the handle can't fire .change/.release. */
.slider-locked, .slider-locked * { pointer-events: none !important; }
.slider-locked { opacity: 0.35 !important; }
"""
def build_ui() -> gr.Blocks:
with gr.Blocks(title="Foveated Diffusion") as demo:
gr.Markdown(
"## Foveated Diffusion: Efficient Spatially Aware Image and Video Generation \n"
"Image generation demo for Foveated Diffusion. The base model is FLUX.2-klein-base-4B. \n"
" \n"
'Project website · '
'Model weights · '
'Paper · '
'Code'
)
with gr.Row():
# ---------------- Left column: controls ----------------
with gr.Column(scale=1, min_width=340):
prompt = gr.Textbox(label="Prompt", lines=6, value=DEFAULT_PROMPT)
seed = gr.Number(label="Seed", value=0, precision=0)
lora_dd = gr.Dropdown(
label="LoRA model",
choices=list(LORA_FILES.keys()),
value=DEFAULT_LORA,
interactive=True,
)
lora_desc = gr.Markdown(
LORA_DESCRIPTIONS[DEFAULT_LORA],
elem_classes=["hint"],
)
lora_status = gr.Markdown(f"Active: {DEFAULT_LORA}")
# --- Mask mode ---
# The "Draw" path (paint a freeform HR region with gr.ImageEditor)
# is implemented end-to-end in this file but intentionally NOT
# surfaced in the UI for this Spaces demo. The radio and the
# ImageEditor are kept alive as hidden components (visible=False)
# so the wiring + handler logic compiles unchanged — re-enable
# them by flipping visible=True if a future demo wants both modes.
mode = gr.Radio(
choices=["Draw", "Preset (circle)"],
value="Preset (circle)",
visible=False,
)
with gr.Group(visible=False) as draw_group:
editor = gr.ImageEditor(
image_mode="RGBA",
type="numpy",
sources=(),
canvas_size=(1024, 1024),
fixed_canvas=True,
layers=False,
transforms=(),
brush=gr.Brush(
default_size=80,
colors=["#111111"],
color_mode="fixed",
default_color="#111111",
),
eraser=gr.Eraser(default_size=80),
height=420,
)
with gr.Row():
clear_btn = gr.Button("Clear", variant="secondary")
fill_btn = gr.Button("Fill all", variant="secondary")
# Preset (circle) is the only mode exposed in the UI.
with gr.Group() as preset_group:
gr.Markdown(
"**Foveal region (circular)** — place a circular HR region "
"on the image. `cx, cy ∈ [-0.5, 0.5]` (0 = image center); "
"`r` is relative to half the image diagonal.",
elem_classes=["hint"],
)
cx = gr.Slider(-0.5, 0.5, value=0.0, step=0.01, label="Center X")
cy = gr.Slider(-0.5, 0.5, value=0.0, step=0.01, label="Center Y")
r = gr.Slider(0.0, 1.0, value=0.3, step=0.01, label="Radius")
preset_clear_btn = gr.Button("Clear", variant="secondary")
generate_btn = gr.Button("Generate", variant="primary")
# ---------------- Right column: stages ----------------
with gr.Column(scale=2):
with gr.Row():
tok_img = gr.Image(
label="Tokenization mask",
height=520,
interactive=False,
elem_classes=["tok-img"],
)
out_img = gr.Image(
label="Generated image",
height=520,
interactive=False,
elem_classes=["out-img"],
)
gr.Markdown(
"**Tokenization mask** — white = HR (foveal) tokens, gray = LR "
"(peripheral). Refreshes as you move the circle sliders. "
"**Generated image** — appears once denoising finishes."
)
# ---------------- Wiring ----------------
# lora_dd is FIRST so every handler that reads `refresh_in` (and
# `generate`) gets the live dropdown value as an explicit argument —
# the only safe way to ferry it into @spaces.GPU workers.
refresh_in = [lora_dd, mode, editor, cx, cy, r]
# Mode change: visibility toggle + tokenization refresh in ONE handler.
# Registering two separate listeners on the same event in Gradio 6.x
# was unreliable; a single combined handler is the safe pattern.
mode.change(
on_mode_change, refresh_in, [draw_group, preset_group, tok_img], queue=False,
)
# LoRA switch: also disables circle sliders and refreshes tok preview
# so no_fov immediately reflects "full-HR, controls disabled".
lora_dd.change(
switch_lora,
[lora_dd, cx, cy, r],
[lora_dd, lora_desc, lora_status, cx, cy, r, tok_img],
)
# Live tokenization preview on edits / slider moves.
# queue=False so brush/slider activity isn't serialized behind a running
# generation in the queue worker. We wire BOTH .release and .change on
# sliders — .release alone can miss events on some Gradio 6.x builds
# (especially for sliders inside a Group that was initially hidden).
editor.apply(refresh_tokenization, refresh_in, tok_img, queue=False)
for trig in (cx.release, cy.release, r.release,
cx.change, cy.change, r.change):
trig(refresh_tokenization, refresh_in, tok_img, queue=False)
# Draw helpers. queue=False — pure value-replacement, no GPU.
clear_btn.click(clear_editor, None, editor, queue=False)
fill_btn.click(fill_all_editor, None, editor, queue=False)
preset_clear_btn.click(clear_preset, None, [cx, cy, r], queue=False)
# Generate.
# show_progress_on=[out_img] keeps the denoise progress bar from
# overlaying the tokenization panel.
generate_btn.click(
generate,
inputs=[prompt, seed, lora_dd, mode, editor, cx, cy, r],
outputs=[tok_img, out_img],
show_progress_on=[out_img],
)
# On page load: render the initial tokenization grid for the default
# circle (cx=0, cy=0, r=0.3). Mode is fixed to "Preset (circle)" — the
# Draw mode is kept in code (visible=False) but not wired into load.
demo.load(refresh_tokenization, refresh_in, tok_img, queue=False)
return demo
if __name__ == "__main__":
print("Loading pipeline and LoRAs (cold start can take several minutes)…")
_load_pipe()
demo = build_ui()
# Spaces sets PORT=7860 by convention.
demo.queue(max_size=8).launch(
server_name=os.environ.get("HOST", "0.0.0.0"),
server_port=int(os.environ.get("PORT", 7860)),
theme=gr.themes.Default(primary_hue="slate"),
css=CSS,
# SHARE=1 -> create a *.gradio.live tunnel. Use this to bypass flaky
# local port-forwarders (e.g. VS Code) that mangle Gradio's SSE/WS.
share=os.environ.get("SHARE", "").lower() in ("1", "true", "yes"),
)