from __future__ import annotations
import base64
import io
import json
import os
import random
import time
from html import escape
from pathlib import Path
import gradio as gr
import numpy as np
from PIL import Image
try:
import spaces
except ImportError:
class _LocalSpaces:
@staticmethod
def GPU(*_args, **_kwargs):
return lambda function: function
spaces = _LocalSpaces()
import torch
from anypaint import prepare_anypaint
from quantized_runtime import build_quantized_pipeline, install_fixed_lora
REPO_ID = "yijunwang2/krea2-anypaint"
WEIGHT_NAME = "krea2_anypaint_rank32.safetensors"
MAX_CANVAS_EDGE = 1280
MAX_SEED = 2**31 - 1
MIN_SCALE = 0.2
MAX_SCALE = 1.0
DEFAULT_RATIO = "16:9"
ASPECT_RATIOS = (
("1:1", 1, 1),
("4:3", 4, 3),
("3:4", 3, 4),
("16:9", 16, 9),
("9:16", 9, 16),
)
ASSETS = Path(__file__).resolve().parent / "assets"
if os.environ.get("ANYPAINT_UI_ONLY") == "1":
pipe = None
else:
pipe = build_quantized_pipeline(REPO_ID)
pipe._anypaint_lora_hooks = install_fixed_lora(pipe.transformer, REPO_ID, WEIGHT_NAME)
print(f"[AnyPaint runtime] {json.dumps(pipe._anypaint_runtime, sort_keys=True)}", flush=True)
def _round16(value: float) -> int:
return max(16, int(round(value / 16.0)) * 16)
def _ratio_of(label: str) -> tuple[int, int]:
return next(((rw, rh) for name, rw, rh in ASPECT_RATIOS if name == label), (1, 1))
def _default_state(ratio: str = DEFAULT_RATIO) -> dict:
return {
"ratio": ratio,
"offset": 0.5,
"offset_cross": 0.5,
"scale": 1.0,
"tool": "move",
"brush_size": 48,
"mask_data": "",
}
def parse_state(value) -> dict:
state = _default_state()
if isinstance(value, dict):
incoming = value
elif isinstance(value, str) and value.strip().startswith("{"):
try:
incoming = json.loads(value)
except json.JSONDecodeError:
incoming = {}
else:
incoming = {}
state.update({key: incoming[key] for key in state if key in incoming})
if state["ratio"] not in {item[0] for item in ASPECT_RATIOS}:
state["ratio"] = DEFAULT_RATIO
for key in ("offset", "offset_cross"):
try:
state[key] = min(1.0, max(0.0, float(state[key])))
except (TypeError, ValueError):
state[key] = 0.5
try:
state["scale"] = min(MAX_SCALE, max(MIN_SCALE, float(state["scale"])))
except (TypeError, ValueError):
state["scale"] = 1.0
if state["tool"] not in {"move", "paint", "erase"}:
state["tool"] = "move"
try:
state["brush_size"] = min(256, max(4, int(state["brush_size"])))
except (TypeError, ValueError):
state["brush_size"] = 48
state["mask_data"] = str(state.get("mask_data") or "")
return state
def plan_canvas(
source_width: int,
source_height: int,
ratio_width: int,
ratio_height: int,
offset: float,
scale: float,
offset_cross: float,
) -> tuple[int, int, tuple[int, int, int, int]]:
try:
offset = min(1.0, max(0.0, float(offset)))
except (TypeError, ValueError):
offset = 0.5
try:
offset_cross = min(1.0, max(0.0, float(offset_cross)))
except (TypeError, ValueError):
offset_cross = 0.5
try:
scale = min(MAX_SCALE, max(MIN_SCALE, float(scale)))
except (TypeError, ValueError):
scale = 1.0
source_ratio = source_width / source_height
target_ratio = ratio_width / ratio_height
if target_ratio >= source_ratio:
canvas_height = _round16(min(source_height, MAX_CANVAS_EDGE))
full_height = canvas_height
full_width = _round16(full_height * source_ratio)
canvas_width = max(_round16(canvas_height * target_ratio), full_width)
if canvas_width > MAX_CANVAS_EDGE:
factor = MAX_CANVAS_EDGE / canvas_width
canvas_width = _round16(canvas_width * factor)
canvas_height = _round16(canvas_height * factor)
full_height = canvas_height
full_width = _round16(full_height * source_ratio)
else:
canvas_width = _round16(min(source_width, MAX_CANVAS_EDGE))
full_width = canvas_width
full_height = _round16(full_width / source_ratio)
canvas_height = max(_round16(canvas_width / target_ratio), full_height)
if canvas_height > MAX_CANVAS_EDGE:
factor = MAX_CANVAS_EDGE / canvas_height
canvas_height = _round16(canvas_height * factor)
canvas_width = _round16(canvas_width * factor)
full_width = canvas_width
full_height = _round16(full_width / source_ratio)
box_width = full_width
box_height = full_height
if scale < 1.0:
box_width = max(16, round(box_width * scale))
box_height = max(16, round(box_width / source_ratio))
box_width = min(box_width, canvas_width)
box_height = min(box_height, canvas_height)
if target_ratio >= source_ratio:
x0 = round((canvas_width - box_width) * offset)
y0 = round((canvas_height - box_height) * offset_cross)
else:
x0 = round((canvas_width - box_width) * offset_cross)
y0 = round((canvas_height - box_height) * offset)
x1, y1 = x0 + box_width, y0 + box_height
x0 = max(0, min(x0, canvas_width - 16))
y0 = max(0, min(y0, canvas_height - 16))
x1 = max(x0 + 16, min(x1, canvas_width))
y1 = max(y0 + 16, min(y1, canvas_height))
return canvas_width, canvas_height, (int(x0), int(y0), int(x1), int(y1))
def _axis_of(source_width: int, source_height: int, ratio_width: int, ratio_height: int) -> str:
source_ratio = source_width / source_height
target_ratio = ratio_width / ratio_height
if abs(target_ratio - source_ratio) < 1e-3:
return "none"
return "horizontal" if target_ratio > source_ratio else "vertical"
def _image_data_uri(image: Image.Image, max_edge: int = 720) -> str:
image = image.convert("RGB")
if max(image.size) > max_edge:
scale = max_edge / max(image.size)
image = image.resize(
(max(1, round(image.width * scale)), max(1, round(image.height * scale))),
Image.Resampling.LANCZOS,
)
with io.BytesIO() as buffer:
image.save(buffer, format="JPEG", quality=90)
return "data:image/jpeg;base64," + base64.b64encode(buffer.getvalue()).decode("ascii")
def _mask_data_uri(path: Path) -> str:
with Image.open(path) as image, io.BytesIO() as buffer:
image.convert("L").save(buffer, format="PNG")
return "data:image/png;base64," + base64.b64encode(buffer.getvalue()).decode("ascii")
def _space_examples() -> tuple[list[list], dict[str, str]]:
manifest_path = ASSETS / "manifest.json"
if not manifest_path.is_file():
return [], {}
records = json.loads(manifest_path.read_text(encoding="utf-8"))
examples = []
states = {}
for record in records:
source_path = ASSETS / record["source"]
mask_path = ASSETS / record["mask"]
if not source_path.is_file() or not mask_path.is_file():
continue
canvas_width, canvas_height = record["canvas_size"]
target_ratio = canvas_width / canvas_height
ratio = min(ASPECT_RATIOS, key=lambda item: abs(item[1] / item[2] - target_ratio))[0]
state = _default_state(ratio)
state["mask_data"] = _mask_data_uri(mask_path)
states[record["id"]] = json.dumps(state)
examples.append([str(source_path), record["prompt"], record["id"]])
return examples, states
SPACE_EXAMPLES, EXAMPLE_STATES = _space_examples()
def _mask_from_data(value: str, size: tuple[int, int]) -> Image.Image:
if not value.startswith("data:image/"):
return Image.new("L", size, 0)
try:
payload = base64.b64decode(value.split(",", 1)[1])
mask = Image.open(io.BytesIO(payload)).convert("L")
except Exception:
return Image.new("L", size, 0)
if mask.size != size:
mask = mask.resize(size, Image.Resampling.NEAREST)
return mask.point(lambda pixel: 255 if pixel > 32 else 0)
def _stage_html(source: Image.Image | None, state_value: str) -> str:
state = parse_state(state_value)
if source is None:
return (
""
"Upload an image to start painting or extending the canvas.
"
)
source = source.convert("RGB")
rw, rh = _ratio_of(state["ratio"])
canvas_width, canvas_height, bbox = plan_canvas(
source.width,
source.height,
rw,
rh,
state["offset"],
state["scale"],
state["offset_cross"],
)
x0, y0, x1, y1 = bbox
axis = _axis_of(source.width, source.height, rw, rh)
source_uri = _image_data_uri(source)
mask_uri = escape(state["mask_data"], quote=True)
ratio_tiles = []
for label, tile_rw, tile_rh in ASPECT_RATIOS:
tile_width, tile_height, tile_bbox = plan_canvas(
source.width,
source.height,
tile_rw,
tile_rh,
state["offset"],
state["scale"],
state["offset_cross"],
)
tile_scale = 58 / max(tile_width, tile_height)
tx0, ty0, tx1, ty1 = tile_bbox
selected = label == state["ratio"]
ratio_tiles.append(
f""
"
"
f"
{label}
"
f"
{tile_width}x{tile_height}
"
"
" + "".join(ratio_tiles) + "
"
)
stage_size = 520.0
display_scale = stage_size / max(canvas_width, canvas_height)
display_width = canvas_width * display_scale
display_height = canvas_height * display_scale
source_left = x0 * display_scale
source_top = y0 * display_scale
source_width = (x1 - x0) * display_scale
source_height = (y1 - y0) * display_scale
if axis == "horizontal":
anchors = (("Left", 0.0), ("Center", 0.5), ("Right", 1.0))
elif axis == "vertical":
anchors = (("Top", 0.0), ("Center", 0.5), ("Bottom", 1.0))
else:
anchors = ()
anchor_buttons = "".join(
f"" for name, value in anchors
)
anchor_row = (
f"{anchor_buttons}
"
if anchors else ""
)
_, _, full_bbox = plan_canvas(
source.width, source.height, rw, rh, state["offset"], 1.0, state["offset_cross"]
)
full_width = full_bbox[2] - full_bbox[0]
full_height = full_bbox[3] - full_bbox[1]
def tool_button(label: str, value: str) -> str:
active = state["tool"] == value
return (
f""
)
tools = (
""
+ tool_button("Move", "move") + tool_button("Paint", "paint") + tool_button("Erase", "erase")
+ ""
+ "
"
)
return f"""
{tools}
×
{anchor_row}
Canvas {canvas_width}x{canvas_height} ยท source ({x0},{y0})-({x1},{y1})
source
generated
{ratio_grid}