import os
# The Boogu transformer/pipeline select their attention + norm kernels based on
# this env var at construction time, so it must be set before importing torch.
os.environ.setdefault("device", "cuda:0")
# Use the pure-torch RMSNorm path (not the triton fused kernel) so the block
# parameter layout matches the AoTI graph compiled in the companion Space.
import boogu.utils.import_utils as _import_utils
_import_utils._triton_available = False
import base64
import csv
import io
import json
import sys
# Example caching writes the cached output (which embeds the base64 before/after
# data URIs) through the csv module; bump the field limit so large frames don't
# trip "_csv.Error: field larger than field limit".
csv.field_size_limit(sys.maxsize)
import spaces
import torch
import gradio as gr
from PIL import Image
from boogu.pipelines.boogu.pipeline_boogu import BooguImagePipeline
MODEL_ID = "Boogu/Boogu-Image-0.1-Edit"
# Edit-Turbo is a 4-step DMD-distilled variant of the Edit model. Its model_index
# declares the plain BooguImagePipeline and its mllm / vae / processor / scheduler /
# transformer *architecture* are byte-identical to Edit — only the (distilled)
# transformer weights differ. So it runs through the exact same edit/T2I code path,
# just at ~4 steps with CFG disabled (guidance == 1.0).
TURBO_ID = "Boogu/Boogu-Image-0.1-Edit-Turbo"
# One shared AoTI single-stream-block graph. Because the Edit and Edit-Turbo
# transformers have an identical config, the graph compiled from Edit
# (weights are runtime inputs, not baked constants) applies to both pipelines.
AOTI_REPO = "multimodalart/Boogu-Image-0.1-Edit-aoti"
pipe = BooguImagePipeline.from_pretrained(
MODEL_ID,
torch_dtype=torch.bfloat16,
trust_remote_code=True,
)
pipe.to("cuda")
# Edit-Turbo shares the (byte-identical) mllm / vae / processor / scheduler with
# Edit; only the transformer differs. Load just the Turbo transformer and build a
# second pipeline reusing the already-resident components — no duplicate 17.5GB mllm.
turbo_transformer = type(pipe.transformer).from_pretrained(
TURBO_ID,
subfolder="transformer",
torch_dtype=torch.bfloat16,
)
_turbo_components = dict(pipe.components)
_turbo_components["transformer"] = turbo_transformer
turbo_pipe = BooguImagePipeline(**_turbo_components)
turbo_pipe.text_instruction_rewriter = pipe.text_instruction_rewriter
turbo_pipe.instruction_rewriter_processor = pipe.instruction_rewriter_processor
turbo_pipe.to("cuda")
# Swap the repeated single-stream blocks for their AoTI-compiled graph (one shared
# compiled graph, per-block weights streamed in at runtime). The Edit and Edit-Turbo
# transformers share an identical block architecture, so the same package.pt2 patches
# both. Falls back to eager on any error.
try:
from pathlib import Path
from huggingface_hub import snapshot_download
from spaces.zero.torch.aoti import aoti_load_from_module_dir
_block_dir = Path(snapshot_download(AOTI_REPO)) / "BooguImageTransformerBlock"
if (_block_dir / "package.pt2").exists():
for _name, _p in (("Edit", pipe), ("Edit-Turbo", turbo_pipe)):
aoti_load_from_module_dir(_p.transformer.single_stream_layers, _block_dir)
print(f"AoTI: patched {len(_p.transformer.single_stream_layers)} {_name} single-stream blocks")
else:
print("AoTI: package.pt2 not found, running eager")
except Exception as exc: # noqa: BLE001
print(f"AoTI load failed ({exc!r}); running eager")
MAX_SEED = 2**31 - 1
def _data_uri(img):
buf = io.BytesIO()
img.save(buf, format="WEBP", quality=92)
return "data:image/webp;base64," + base64.b64encode(buf.getvalue()).decode()
# Custom before/after comparison built on gr.HTML (gr.ImageSlider is broken with
# gr.Examples caching on this Gradio build and doesn't keep the two sides aligned).
# Markup/CSS mirror Gradio's native ImageSlider: both images fill the same box with
# object-fit:contain so they line up regardless of native size; the edited ("after")
# image is revealed by a clip-path driven by an accent-pill handle on a 1px divider.
# NOTE: Gradio evaluates html_template via `new Function(..., "return `" + tpl + "`")`,
# i.e. it wraps the whole template in backticks. So the template must NOT contain any
# backticks of its own (nested template literals terminate the wrapper and silently
# blank the component) — build the markup with single-quote string concatenation.
# Native-style floating block label (icon + text), mirroring Gradio's block-label.
_BA_LABEL = (
''
)
_BA_DOWNLOAD_ICON = (
''
)
# value arrives as a JSON string (see edit()); parse it defensively. An IIFE keeps
# this a single ${...} expression with no backticks.
BA_HTML = (
"${(function(){\n"
" var d = {};\n"
" try { d = value ? JSON.parse(value) : {}; } catch (e) { d = {}; }\n"
" return (d && d.after)\n"
" ? '