app.py

import os
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import io
import json
import math
import base64
import random
import tempfile

import spaces
import torch
import numpy as np
import gradio as gr
from PIL import Image
from diffusers import DiffusionPipeline

from outpaint import prepare_source, composite

# ---------------------------------------------------------------------------
# Model loading (Krea 2 Turbo + registered outpaint LoRA)
# ---------------------------------------------------------------------------
REPO_ID = "yijunwang2/krea2-outpaint"
WEIGHT_NAME = "krea2_outpaint_rank32.safetensors"
BASE_MODEL = "krea/Krea-2-Turbo"

pipe = DiffusionPipeline.from_pretrained(
    BASE_MODEL,
    custom_pipeline=REPO_ID,
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
).to("cuda")
pipe.load_lora_weights(REPO_ID, weight_name=WEIGHT_NAME, adapter_name="outpaint")
pipe.set_adapters(["outpaint"], weights=[1.0])

MAX_CANVAS = 1280          # longest canvas edge we allow (keeps 8-step fast)
MAX_SEED = 2**31 - 1

# Aspect ratio presets shown as visual tiles. (label, w_units, h_units)
ASPECT_RATIOS = [
    ("1:1", 1, 1),
    ("4:3", 4, 3),
    ("3:4", 3, 4),
    ("16:9", 16, 9),
    ("9:16", 9, 16),
]

# Minimum/maximum proportional scale of the source within its full-span box.
MIN_SCALE = 0.2
MAX_SCALE = 1.0

DEFAULT_RATIO = "16:9"


def _round16(v: int) -> int:
    return max(16, int(round(v / 16.0)) * 16)


def _ratio_of(label):
    for lbl, rw, rh in ASPECT_RATIOS:
        if lbl == label:
            return rw, rh
    return 1, 1


def plan_canvas(src_w: int, src_h: int, ratio_w: float, ratio_h: float,
                offset: float, scale: float = 1.0, offset_cross: float = 0.5):
    """Given a source image size, a target canvas aspect ratio, a *continuous*
    placement offset and a proportional ``scale`` factor, compute the output
    canvas (multiple of 16) and the source bounding box.

    ``offset`` is the normalized position of the source along the axis being
    extended, in [0, 1]:
      - horizontal extend: 0 => flush left, 0.5 => centered, 1 => flush right
      - vertical   extend: 0 => flush top,  0.5 => centered, 1 => flush bottom

    ``scale`` (in [MIN_SCALE, MAX_SCALE]) proportionally shrinks the source box
    below the full-span size while keeping its aspect ratio locked. When scale
    < 1 the source no longer spans a full canvas dimension, so the *cross* axis
    also becomes free; ``offset_cross`` positions the source along it (0.5 =>
    centered).

    This single function is the source of truth for BOTH the interactive preview
    and the real generation, so what the user sees is exactly what the model
    receives. At scale == 1 the source spans one FULL canvas dimension:
    - target wider than source  -> source spans full HEIGHT, extend horizontally
    - target taller than source -> source spans full WIDTH,  extend vertically
    - equal ratios              -> source fills the canvas (no outpaint)
    """
    try:
        offset = float(offset)
    except (TypeError, ValueError):
        offset = 0.5
    offset = min(1.0, max(0.0, offset))
    try:
        offset_cross = float(offset_cross)
    except (TypeError, ValueError):
        offset_cross = 0.5
    offset_cross = min(1.0, max(0.0, offset_cross))
    try:
        scale = float(scale)
    except (TypeError, ValueError):
        scale = 1.0
    scale = min(MAX_SCALE, max(MIN_SCALE, scale))

    src_ratio = src_w / src_h
    target_ratio = ratio_w / ratio_h

    if target_ratio >= src_ratio:
        # Canvas is wider (relative to source): full-span keeps source height full.
        canvas_h = _round16(min(src_h, MAX_CANVAS))
        box_h = canvas_h
        box_w = _round16(box_h * src_ratio)
        canvas_w = _round16(canvas_h * target_ratio)
        canvas_w = max(canvas_w, box_w)
        # cap longest edge
        if canvas_w > MAX_CANVAS:
            cap = MAX_CANVAS / canvas_w
            canvas_w = _round16(canvas_w * cap)
            canvas_h = _round16(canvas_h * cap)
            box_h = canvas_h
            box_w = _round16(box_h * src_ratio)
    else:
        # Canvas is taller (relative to source): full-span keeps source width full.
        canvas_w = _round16(min(src_w, MAX_CANVAS))
        box_w = canvas_w
        box_h = _round16(box_w / src_ratio)
        canvas_h = _round16(canvas_w / target_ratio)
        canvas_h = max(canvas_h, box_h)
        if canvas_h > MAX_CANVAS:
            cap = MAX_CANVAS / canvas_h
            canvas_h = _round16(canvas_h * cap)
            canvas_w = _round16(canvas_w * cap)
            box_w = canvas_w
            box_h = _round16(box_w / src_ratio)

    # Apply the proportional resize (aspect ratio locked). The bbox itself does
    # NOT need to be 16-aligned (only the canvas does), so derive box_h from box_w
    # via the true source ratio to keep the aspect ratio tight even when small.
    if scale < 1.0:
        box_w = max(16, int(round(box_w * scale)))
        box_h = max(16, int(round(box_w / src_ratio)))
    box_w = min(box_w, canvas_w)
    box_h = min(box_h, canvas_h)

    # Place along both axes; the extend axis uses ``offset``, the cross axis
    # uses ``offset_cross`` (only free when scale < 1, otherwise it is pinned).
    if target_ratio >= src_ratio:
        x0 = int(round((canvas_w - box_w) * offset))
        y0 = int(round((canvas_h - box_h) * offset_cross))
    else:
        x0 = int(round((canvas_w - box_w) * offset_cross))
        y0 = int(round((canvas_h - box_h) * offset))
    x1, y1 = x0 + box_w, y0 + box_h

    # clamp to canvas & enforce validity
    x0 = max(0, min(x0, canvas_w - 16))
    y0 = max(0, min(y0, canvas_h - 16))
    x1 = max(x0 + 16, min(x1, canvas_w))
    y1 = max(y0 + 16, min(y1, canvas_h))
    return canvas_w, canvas_h, (int(x0), int(y0), int(x1), int(y1))


def _axis_of(src_w, src_h, ratio_w, ratio_h):
    """Which axis is being extended for this ratio vs. source."""
    src_ratio = src_w / src_h
    target_ratio = ratio_w / ratio_h
    if abs(target_ratio - src_ratio) < 1e-3:
        return "none"
    if target_ratio > src_ratio:
        return "horizontal"
    return "vertical"


# ---------------------------------------------------------------------------
# Placement state helpers
# ---------------------------------------------------------------------------
def parse_placement(state):
    """Parse the JSON placement state -> (ratio_label, offset, scale, offset_cross)."""
    ratio, offset, scale, offset_cross = DEFAULT_RATIO, 0.5, 1.0, 0.5
    d = None
    if isinstance(state, dict):
        d = state
    elif isinstance(state, str) and state.strip():
        s = state.strip()
        if s.startswith("{"):
            try:
                d = json.loads(s)
            except Exception:
                d = None
                ratio = s
        else:
            # legacy: a bare ratio label
            ratio = s
    if isinstance(d, dict):
        ratio = d.get("ratio", DEFAULT_RATIO)
        offset = d.get("offset", 0.5)
        scale = d.get("scale", 1.0)
        offset_cross = d.get("offset_cross", 0.5)
    if ratio not in [lbl for lbl, _, _ in ASPECT_RATIOS]:
        ratio = DEFAULT_RATIO
    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
    return ratio, offset, scale, offset_cross


# ---------------------------------------------------------------------------
# Unified interactive canvas + placement component (single custom gr.HTML)
# ---------------------------------------------------------------------------
def _img_data_uri(image: Image.Image, max_edge: int = 480) -> str:
    img = image.convert("RGB")
    if max(img.size) > max_edge:
        scale = max_edge / max(img.size)
        img = img.resize((max(1, round(img.width * scale)), max(1, round(img.height * scale))), Image.LANCZOS)
    buf = io.BytesIO()
    img.save(buf, format="JPEG", quality=88)
    return "data:image/jpeg;base64," + base64.b64encode(buf.getvalue()).decode("ascii")


def build_stage_html(image, state):
    """Render the ONE unified interactive component: ratio presets on top,
    then a live stage that shows the target canvas with the source image placed
    inside it. The source can be dragged freely (or snapped to anchor presets),
    and the aspect ratio is chosen from the same component. All interaction feeds
    the same {ratio, offset} placement used by generation."""
    ratio, offset, scale, offset_cross = parse_placement(state)

    if image is None:
        return (
            "<div style='padding:32px;text-align:center;color:var(--body-text-color-subdued);"
            "border:1px dashed var(--border-color-primary);border-radius:16px;'>"
            "Upload an image to compose your outpaint canvas.</div>"
        )

    rw, rh = _ratio_of(ratio)
    cw, ch, bbox = plan_canvas(image.width, image.height, rw, rh, offset, scale, offset_cross)
    x0, y0, x1, y1 = bbox
    axis = _axis_of(image.width, image.height, rw, rh)

    src_uri = _img_data_uri(image)

    # ---- aspect ratio preset tiles ----
    tiles = []
    for label, trw, trh in ASPECT_RATIOS:
        tcw, tch, tbbox = plan_canvas(image.width, image.height, trw, trh, offset, scale, offset_cross)
        tx0, ty0, tx1, ty1 = tbbox
        maxedge = max(tcw, tch)
        tscale = 60.0 / maxedge
        fw, fh = tcw * tscale, tch * tscale
        bx, by = tx0 * tscale, ty0 * tscale
        bw, bh = (tx1 - tx0) * tscale, (ty1 - ty0) * tscale
        selected = (label == ratio)
        border = "3px solid var(--color-accent)" if selected else "2px solid var(--border-color-primary)"
        bg = "var(--background-fill-secondary)" if selected else "var(--background-fill-primary)"
        tile = f"""
        <div class='ar-tile' data-ratio='{label}' onclick='window._krSelectRatio("{label}")'
             style='cursor:pointer;display:flex;flex-direction:column;align-items:center;gap:4px;
                    padding:8px 4px;border-radius:12px;border:{border};background:{bg};
                    transition:all .12s ease;flex:1 1 0;min-width:0;'>
          <div style='position:relative;width:66px;height:66px;display:flex;align-items:center;justify-content:center;'>
            <div style='position:relative;width:{fw:.1f}px;height:{fh:.1f}px;
                        background:repeating-linear-gradient(45deg,var(--border-color-primary) 0 5px,transparent 5px 10px);
                        border:1px solid var(--border-color-accent);border-radius:3px;'>
              <div style='position:absolute;left:{bx:.1f}px;top:{by:.1f}px;width:{bw:.1f}px;height:{bh:.1f}px;
                          background:linear-gradient(135deg,#3b82f6cc,#6366f1cc);border-radius:2px;
                          box-shadow:0 0 0 1px #ffffff55;'></div>
            </div>
          </div>
          <div style='font-weight:600;font-size:13px;'>{label}</div>
          <div style='font-size:10px;color:var(--body-text-color-subdued);'>{tcw}×{tch}</div>
        </div>"""
        tiles.append(tile)

    ratio_grid = (
        "<div style='display:flex;flex-wrap:nowrap;gap:8px;justify-content:center;"
        "align-items:stretch;margin-top:14px;'>"
        + "".join(tiles)
        + "</div>"
    )

    # ---- the big interactive stage ----
    # scale canvas so its longest edge is ~STAGE px
    STAGE = 520.0
    sc = STAGE / max(cw, ch)
    stage_w, stage_h = cw * sc, ch * sc
    src_left, src_top = x0 * sc, y0 * sc
    src_w_px, src_h_px = (x1 - x0) * sc, (y1 - y0) * sc

    # anchor buttons depend on the extend axis
    if axis == "horizontal":
        anchors = [("Left", 0.0), ("Center", 0.5), ("Right", 1.0)]
        drag_hint = "Drag the image left/right, or tap an anchor. Drag the corner handle to resize."
    elif axis == "vertical":
        anchors = [("Top", 0.0), ("Center", 0.5), ("Bottom", 1.0)]
        drag_hint = "Drag the image up/down, or tap an anchor. Drag the corner handle to resize."
    else:
        anchors = []
        drag_hint = "Drag the corner handle to resize, then drag the image to place it."

    def _anchor_active(val):
        return abs(offset - val) < 0.02
    anchor_btns = "".join(
        f"<button type='button' onclick='window._krSetOffset({val})' "
        f"style='cursor:pointer;padding:5px 14px;border-radius:999px;font-size:12px;font-weight:600;"
        f"border:{'2px solid var(--color-accent)' if _anchor_active(val) else '1px solid var(--border-color-primary)'};"
        f"background:{'var(--background-fill-secondary)' if _anchor_active(val) else 'var(--background-fill-primary)'};"
        f"color:var(--body-text-color);'>{name}</button>"
        for name, val in anchors
    )
    anchor_row = (
        f"<div style='display:flex;gap:8px;justify-content:center;margin-top:12px;flex-wrap:wrap;'>{anchor_btns}</div>"
        if anchors else ""
    )

    # Full-span box size (scale == 1) for this ratio, so JS can map a resize
    # drag back onto the [MIN_SCALE, MAX_SCALE] scale value.
    _fcw, _fch, _fbbox = plan_canvas(image.width, image.height, rw, rh, offset, 1.0, offset_cross)
    full_box_w = _fbbox[2] - _fbbox[0]
    full_box_h = _fbbox[3] - _fbbox[1]

    # The source can always be repositioned/resized now (resize frees both axes).
    draggable = "true"
    stage = f"""
    <div style='display:flex;justify-content:center;'>
      <div id='kr-stage' data-axis='{axis}' data-cw='{cw}' data-ch='{ch}'
           data-boxw='{x1 - x0}' data-boxh='{y1 - y0}' data-draggable='{draggable}'
           data-sc='{sc:.6f}' data-fullw='{full_box_w}' data-fullh='{full_box_h}'
           data-minscale='{MIN_SCALE}' data-maxscale='{MAX_SCALE}' data-scale='{scale:.4f}'
           style='position:relative;width:{stage_w:.1f}px;height:{stage_h:.1f}px;
                  border-radius:10px;overflow:visible;
                  background:repeating-conic-gradient(#e9e9f2 0% 25%, #f6f6fb 0% 50%) 50% / 22px 22px;
                  box-shadow:inset 0 0 0 2px var(--border-color-accent);
                  background-clip:padding-box;'>
        <img id='kr-src' src='{src_uri}' draggable='false'
             style='position:absolute;left:{src_left:.1f}px;top:{src_top:.1f}px;
                    width:{src_w_px:.1f}px;height:{src_h_px:.1f}px;
                    border-radius:2px;box-shadow:0 0 0 2px #6366f1, 0 4px 18px #0004;
                    cursor:grab;user-select:none;touch-action:none;'/>
        <div id='kr-close' title='Remove image' onclick='window._krClearSource()'
             style='position:absolute;left:{src_left - 11:.1f}px;top:{src_top - 11:.1f}px;
                    width:22px;height:22px;border-radius:50%;background:#ef4444;color:#fff;
                    display:flex;align-items:center;justify-content:center;font-size:15px;
                    font-weight:700;line-height:1;cursor:pointer;
                    box-shadow:0 0 0 2px #fff, 0 2px 6px #0006;z-index:6;
                    user-select:none;touch-action:none;'>×</div>
        <div id='kr-resize'
             style='position:absolute;left:{src_left + src_w_px - 9:.1f}px;top:{src_top + src_h_px - 9:.1f}px;
                    width:18px;height:18px;border-radius:50%;background:#6366f1;
                    box-shadow:0 0 0 2px #fff, 0 2px 6px #0006;cursor:nwse-resize;
                    touch-action:none;z-index:5;'></div>
      </div>
    </div>"""

    hint_html = (
        f"<div style='text-align:center;margin-top:10px;font-size:13px;"
        f"color:var(--body-text-color-subdued);'>{drag_hint}</div>"
    )
    info = (
        f"<div style='text-align:center;margin-top:2px;font-size:11px;"
        f"color:var(--body-text-color-subdued);'>Canvas {cw}×{ch} · "
        f"source box ({x0},{y0})–({x1},{y1})</div>"
    )
    legend = (
        "<div style='display:flex;gap:18px;justify-content:center;align-items:center;"
        "margin-top:8px;font-size:12px;color:var(--body-text-color-subdued);flex-wrap:wrap;'>"
        "<span style='display:inline-flex;align-items:center;gap:6px;'>"
        "<span style='width:14px;height:14px;border-radius:3px;background:linear-gradient(135deg,#3b82f6,#6366f1);display:inline-block;'></span>"
        "your image</span>"
        "<span style='display:inline-flex;align-items:center;gap:6px;'>"
        "<span style='width:14px;height:14px;border-radius:3px;background:repeating-conic-gradient(#e9e9f2 0% 25%, #f6f6fb 0% 50%) 50% / 6px 6px;border:1px solid var(--border-color-accent);display:inline-block;'></span>"
        "generated fill</span></div>"
    )

    # New vertical order inside the unified component:
    #   canvas (placement preview) at top, then the aspect-ratio option controls
    #   (preset tiles) below the canvas. Anchors / hint / info / legend stay
    #   attached to the canvas preview.
    return stage + anchor_row + hint_html + info + legend + ratio_grid


# ---------------------------------------------------------------------------
# Inference
# ---------------------------------------------------------------------------
def _estimate_duration(image, prompt, state, steps, seed, randomize, *a, **k):
    try:
        s = int(steps)
    except Exception:
        s = 8
    # Observed ~8-11s GPU compute for 8 steps; add margin for cold weight
    # streaming (35 GB) on the first worker call.
    return int(30 + s * 4)


@spaces.GPU(duration=_estimate_duration)
def outpaint(image, prompt, state, steps=8, seed=42, randomize=True,
             progress=gr.Progress(track_tqdm=True)):
    """Outpaint an image into a larger canvas of the chosen aspect ratio.

    Args:
        image: the source RGB image to extend.
        prompt: text describing the full desired output scene.
        state: JSON placement string {"ratio": "16:9", "offset": 0.5} produced by
            the unified interactive canvas (ratio + free/anchored placement).
        steps: number of denoising steps (8 recommended for Krea 2 Turbo).
        seed: RNG seed.
        randomize: if True, pick a random seed each run.

    Returns:
        The outpainted image, and the seed actually used.
    """
    if image is None:
        raise gr.Error("Please upload a source image first.")
    if not prompt or not prompt.strip():
        raise gr.Error("Please enter a prompt describing the full output image.")

    if randomize:
        seed = random.randint(0, MAX_SEED)
    seed = int(seed)

    img = image.convert("RGB")
    ratio, offset, scale, offset_cross = parse_placement(state)
    rw, rh = _ratio_of(ratio)
    # Exact same offset/scale -> bbox mapping as the interactive preview.
    cw, ch, bbox = plan_canvas(img.width, img.height, rw, rh, offset, scale, offset_cross)

    prepared = prepare_source(img, (cw, ch), bbox)

    generator = torch.Generator(device="cuda").manual_seed(seed)
    generated = pipe(
        prompt=prompt.strip(),
        image=prepared.condition,
        width=cw,
        height=ch,
        num_inference_steps=int(steps),
        guidance_scale=0.0,
        generator=generator,
        reference_max_pixels=384 * 384,
        reference_placements=[{"bbox_normalized": prepared.bbox_normalized}],
        encode_reference_in_prompt=False,
        kv_cache=True,
    ).images[0]

    result = composite(generated, prepared)
    return result, seed


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------
CSS = """
#col-container { max-width: 1180px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
#kr-stage-wrap { min-height: 40px; }
/* Keep the placement-state textbox in the DOM (so JS can find its textarea) but hidden. */
#ratio_state { position: absolute !important; width: 0; height: 0; overflow: hidden; padding: 0; margin: 0; border: 0; opacity: 0; pointer-events: none; }
/* Hidden trigger button for the on-canvas X (close) button. */
#kr-clear-src { position: absolute !important; width: 0; height: 0; overflow: hidden; padding: 0; margin: 0; border: 0; opacity: 0; pointer-events: none; }
"""

# JS that wires the unified interactive canvas (ratio tiles + anchor buttons +
# free drag of the source) to a hidden Gradio textbox holding a JSON placement.
# We must (1) locate the textarea even inside a <gradio-app> shadow root, and
# (2) set its value via the native setter so Gradio's (Svelte) controlled input
# fires .change(); (3) attach fresh drag handlers every time the HTML re-renders.
SELECT_JS = """
() => {
  const findBox = () => {
    let box = document.querySelector('#ratio_state textarea');
    if (box) return box;
    const app = document.querySelector('gradio-app');
    if (app && app.shadowRoot) {
      box = app.shadowRoot.querySelector('#ratio_state textarea');
      if (box) return box;
    }
    return null;
  };
  const readState = () => {
    const box = findBox();
    let st = { ratio: "16:9", offset: 0.5, scale: 1.0, offset_cross: 0.5 };
    if (box && box.value) {
      try {
        const v = box.value.trim();
        if (v.startsWith('{')) st = Object.assign(st, JSON.parse(v));
        else st.ratio = v;
      } catch (e) {}
    }
    if (typeof st.offset !== 'number' || isNaN(st.offset)) st.offset = 0.5;
    if (typeof st.offset_cross !== 'number' || isNaN(st.offset_cross)) st.offset_cross = 0.5;
    if (typeof st.scale !== 'number' || isNaN(st.scale)) st.scale = 1.0;
    if (!st.ratio) st.ratio = "16:9";
    return st;
  };
  const writeState = (st) => {
    const box = findBox();
    if (!box) return;
    const setter = Object.getOwnPropertyDescriptor(
      window.HTMLTextAreaElement.prototype, 'value'
    ).set;
    setter.call(box, JSON.stringify(st));
    box.dispatchEvent(new Event('input', { bubbles: true }));
    box.dispatchEvent(new Event('change', { bubbles: true }));
  };
  window._krSelectRatio = (label) => {
    const st = readState();
    st.ratio = label;
    writeState(st);
  };
  window._krSetOffset = (off) => {
    const st = readState();
    st.offset = Math.min(1, Math.max(0, off));
    writeState(st);
  };
  window._krSetPlacement = (off, offCross, scale) => {
    const st = readState();
    if (typeof off === 'number' && !isNaN(off)) st.offset = Math.min(1, Math.max(0, off));
    if (typeof offCross === 'number' && !isNaN(offCross)) st.offset_cross = Math.min(1, Math.max(0, offCross));
    if (typeof scale === 'number' && !isNaN(scale)) st.scale = scale;
    writeState(st);
  };
  // X close button on the source image: reuse the exact same clear/reset logic
  // as clearing the Source Image component, by clicking a hidden Gradio button.
  window._krClearSource = () => {
    const app = document.querySelector('gradio-app');
    const root = (app && app.shadowRoot) ? app.shadowRoot : document;
    const btn = root.getElementById
      ? root.getElementById('kr-clear-src')
      : root.querySelector('#kr-clear-src');
    if (btn) {
      const real = btn.querySelector('button') || btn;
      real.click();
    }
  };

  // Free-drag: update the image position live, commit offset on release.
  const rootOf = () => {
    const app = document.querySelector('gradio-app');
    return (app && app.shadowRoot) ? app.shadowRoot : document;
  };
  const attachDrag = () => {
    const root = rootOf();
    const stage = root.getElementById ? root.getElementById('kr-stage') : root.querySelector('#kr-stage');
    if (!stage || stage.dataset.krBound === '1') return;
    if (stage.dataset.draggable !== 'true') return;
    stage.dataset.krBound = '1';
    const img = stage.querySelector('#kr-src');
    const handle = stage.querySelector('#kr-resize');
    const closeBtn = stage.querySelector('#kr-close');
    if (!img) return;
    const axis = stage.dataset.axis;
    const sc = parseFloat(stage.dataset.sc) || 1;              // display px per canvas px
    const fullW = (parseFloat(stage.dataset.fullw) || 0) * sc;  // full-span box in display px
    const fullH = (parseFloat(stage.dataset.fullh) || 0) * sc;
    const minScale = parseFloat(stage.dataset.minscale) || 0.2;
    const maxScale = parseFloat(stage.dataset.maxscale) || 1.0;

    const freeX = () => stage.clientWidth - img.offsetWidth;
    const freeY = () => stage.clientHeight - img.offsetHeight;

    const syncHandle = () => {
      if (handle) {
        handle.style.left = (img.offsetLeft + img.offsetWidth - 9) + 'px';
        handle.style.top = (img.offsetTop + img.offsetHeight - 9) + 'px';
      }
      // Keep the close (X) button pinned to the image's top-left corner,
      // updated live client-side just like the resize handle above.
      if (closeBtn) {
        closeBtn.style.left = (img.offsetLeft - 11) + 'px';
        closeBtn.style.top = (img.offsetTop - 11) + 'px';
      }
    };
    const commit = () => {
      const fx = freeX(), fy = freeY();
      const offX = fx > 0 ? img.offsetLeft / fx : 0.5;
      const offY = fy > 0 ? img.offsetTop / fy : 0.5;
      // map current box size back to a scale value along the full-span axis
      let scale = 1.0;
      if (fullW > 0 && fullH > 0) {
        const s = (axis === 'vertical') ? (img.offsetWidth / fullW) : (img.offsetHeight / fullH);
        scale = Math.min(maxScale, Math.max(minScale, s || 1.0));
      }
      // which axis is the extend (offset) axis vs. the cross axis
      let off, offCross;
      if (axis === 'vertical') { off = offY; offCross = offX; }
      else { off = offX; offCross = offY; }   // horizontal & none
      window._krSetPlacement(off, offCross, scale);
    };

    // ---- reposition (drag the image on both axes) ----
    let dragging = false, sX = 0, sY = 0, sL = 0, sT = 0;
    const onDown = (e) => {
      dragging = true;
      img.style.cursor = 'grabbing';
      const p = (e.touches ? e.touches[0] : e);
      sX = p.clientX; sY = p.clientY; sL = img.offsetLeft; sT = img.offsetTop;
      e.preventDefault(); e.stopPropagation();
    };
    const onMove = (e) => {
      if (!dragging) return;
      const p = (e.touches ? e.touches[0] : e);
      let nl = Math.min(Math.max(sL + (p.clientX - sX), 0), Math.max(0, freeX()));
      let nt = Math.min(Math.max(sT + (p.clientY - sY), 0), Math.max(0, freeY()));
      img.style.left = nl + 'px';
      img.style.top = nt + 'px';
      syncHandle();
      e.preventDefault();
    };
    const onUp = () => {
      if (!dragging) return;
      dragging = false;
      img.style.cursor = 'grab';
      commit();
    };
    img.addEventListener('mousedown', onDown);
    img.addEventListener('touchstart', onDown, { passive: false });

    // ---- resize (drag the corner handle, aspect ratio locked) ----
    let resizing = false, rX = 0, rW0 = 0, rH0 = 0, rL0 = 0, rT0 = 0, rRatio = 1;
    const onRDown = (e) => {
      resizing = true;
      const p = (e.touches ? e.touches[0] : e);
      rX = p.clientX;
      rW0 = img.offsetWidth; rH0 = img.offsetHeight;
      rL0 = img.offsetLeft; rT0 = img.offsetTop;
      rRatio = rH0 / rW0;
      e.preventDefault(); e.stopPropagation();
    };
    const onRMove = (e) => {
      if (!resizing) return;
      const p = (e.touches ? e.touches[0] : e);
      let nw = rW0 + (p.clientX - rX);
      // clamp width to [minScale*full, full] and keep inside stage from top-left anchor
      const minW = (fullW > 0 ? fullW : rW0) * minScale;
      const maxW = Math.min((fullW > 0 ? fullW : stage.clientWidth) * maxScale,
                            stage.clientWidth - rL0);
      nw = Math.min(Math.max(nw, minW), Math.max(minW, maxW));
      let nh = nw * rRatio;
      if (rT0 + nh > stage.clientHeight) {
        nh = stage.clientHeight - rT0;
        nw = nh / rRatio;
      }
      img.style.width = nw + 'px';
      img.style.height = nh + 'px';
      syncHandle();
      e.preventDefault();
    };
    const onRUp = () => {
      if (!resizing) return;
      resizing = false;
      commit();
    };
    if (handle) {
      handle.addEventListener('mousedown', onRDown);
      handle.addEventListener('touchstart', onRDown, { passive: false });
    }

    window.addEventListener('mousemove', (e) => { onMove(e); onRMove(e); });
    window.addEventListener('touchmove', (e) => { onMove(e); onRMove(e); }, { passive: false });
    window.addEventListener('mouseup', () => { onUp(); onRUp(); });
    window.addEventListener('touchend', () => { onUp(); onRUp(); });
  };
  // Re-attach whenever Gradio re-renders the HTML component.
  const mo = new MutationObserver(() => attachDrag());
  mo.observe(document.body, { childList: true, subtree: true });
  const app = document.querySelector('gradio-app');
  if (app && app.shadowRoot) mo.observe(app.shadowRoot, { childList: true, subtree: true });
  attachDrag();
  setInterval(attachDrag, 800);
}
"""

with gr.Blocks(theme=gr.themes.Citrus(), css=CSS, title="Krea 2 Outpaint") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            "# 🖼️ Krea 2 Outpaint\n"
            "Extend an image into a larger canvas while keeping the original pixels. "
            "Powered by [`yijunwang2/krea2-outpaint`](https://huggingface.co/yijunwang2/krea2-outpaint) "
            "(a rank-32 LoRA on [`krea/Krea-2-Turbo`](https://huggingface.co/krea/Krea-2-Turbo), 8-step distilled)."
        )

        ratio_state = gr.Textbox(
            value=json.dumps({"ratio": DEFAULT_RATIO, "offset": 0.5, "scale": 1.0, "offset_cross": 0.5}),
            label="", show_label=False, container=False,
            elem_id="ratio_state",
        )

        with gr.Row():
            with gr.Column(scale=1):
                prompt = gr.Textbox(
                    label="Prompt",
                    placeholder="Describe the FULL output scene, e.g. 'a wide sunlit living room, photorealistic'",
                    lines=2,
                )
                input_image = gr.Image(type="pil", label="Source image", height=320)
                stage = gr.HTML(
                    build_stage_html(None, json.dumps({"ratio": DEFAULT_RATIO, "offset": 0.5, "scale": 1.0, "offset_cross": 0.5})),
                    elem_id="kr-stage-wrap",
                )
                # Hidden button clicked by the on-canvas X button; reuses the
                # same clear/reset logic as clearing the Source Image component.
                clear_src_btn = gr.Button("clear source", elem_id="kr-clear-src")
                run = gr.Button("Outpaint", variant="primary")

            with gr.Column(scale=1):
                result = gr.Image(label="Outpainted result", height=380, interactive=False)

        with gr.Accordion("Advanced settings", open=False):
            steps = gr.Slider(4, 16, value=8, step=1, label="Steps (8 recommended)")
            with gr.Row():
                seed = gr.Number(value=42, precision=0, label="Seed")
                randomize = gr.Checkbox(value=True, label="Randomize seed")

        # --- keep the unified stage in sync with image + placement state ---
        def refresh(image, state):
            return build_stage_html(image, state)

        def reset_on_upload(image, state):
            """When a new image arrives, keep the ratio but recenter placement and
            reset scale so the stage always renders a valid composition. Hide the
            Source Image input so the unified placement canvas takes over."""
            ratio, _, _, _ = parse_placement(state)
            new_state = json.dumps({"ratio": ratio, "offset": 0.5, "scale": 1.0, "offset_cross": 0.5})
            # Hide the Source Image input once an image is present; if the value
            # was cleared (image is None) keep it visible so the user can re-upload.
            return build_stage_html(image, new_state), new_state, gr.update(visible=image is None)

        def clear_source(state):
            """When the X (clear) button is clicked: clear the uploaded image,
            reset the canvas back to its empty state, and make the Source Image
            input reappear so the user can upload a new image."""
            ratio, _, _, _ = parse_placement(state)
            new_state = json.dumps({"ratio": ratio, "offset": 0.5, "scale": 1.0, "offset_cross": 0.5})
            return build_stage_html(None, new_state), new_state, gr.update(visible=True)

        def clear_source_from_canvas(state):
            """Same end-state as clearing the Source Image component, triggered
            by the on-canvas X button: clear the image value AND make the
            component visible again while resetting the canvas to empty."""
            ratio, _, _, _ = parse_placement(state)
            new_state = json.dumps({"ratio": ratio, "offset": 0.5, "scale": 1.0, "offset_cross": 0.5})
            return build_stage_html(None, new_state), new_state, gr.update(value=None, visible=True)

        # events
        input_image.change(
            fn=reset_on_upload, inputs=[input_image, ratio_state],
            outputs=[stage, ratio_state, input_image],
        )
        input_image.clear(
            fn=clear_source, inputs=[ratio_state],
            outputs=[stage, ratio_state, input_image],
        )
        # On-canvas X button: same clear/reset end-state as clearing the
        # Source Image component (also clears the underlying image value).
        clear_src_btn.click(
            fn=clear_source_from_canvas, inputs=[ratio_state],
            outputs=[stage, ratio_state, input_image],
        )
        # This event fires continuously during drag/resize of the custom canvas;
        # hide Gradio's loading/progress indicator so it isn't distracting.
        ratio_state.change(
            fn=refresh, inputs=[input_image, ratio_state],
            outputs=[stage], show_progress="hidden",
        )

        run.click(
            fn=outpaint,
            inputs=[input_image, prompt, ratio_state, steps, seed, randomize],
            outputs=[result, seed],
            api_name="outpaint",
        )

        gr.Examples(
            examples=[
                ["assets/photo_source.webp", "a wide photorealistic modern living room interior with large windows and warm daylight", json.dumps({"ratio": "16:9", "offset": 0.5})],
                ["assets/watercolor_source.webp", "a serene watercolor village landscape with rolling hills and a soft sky", json.dumps({"ratio": "3:4", "offset": 0.0})],
                ["assets/3d_source.webp", "a bustling stylized 3D night market street stretching into the distance, neon signs", json.dumps({"ratio": "16:9", "offset": 0.5})],
            ],
            inputs=[input_image, prompt, ratio_state],
            outputs=[result, seed],
            fn=outpaint,
            cache_examples=False,
            run_on_click=True,
        )

    demo.load(fn=None, js=SELECT_JS)

if __name__ == "__main__":
    demo.launch(mcp_server=True)