Stillfront-comfyui-sf-nodes/py/wavespeed_api/outpaint.py

# ABOUTME: Geometry for keep-the-original outpainting: canvas, mask, paste-back, Bria Expand placement
# ABOUTME: Pure Pillow + numpy, ported from the Art Pipeline Hub resolution-variant-generator (tested there)

"""Geometry for "keep the original" reframing. Pure Pillow and numpy.

Nothing here talks to a network or imports the runner library, on purpose:
every function takes pictures and numbers and returns pictures and numbers, so
the tests can check the part that decides whether the artist's original
survives (where it sits, what the mask covers, what the paste-back may touch)
without any model in the loop.

THE IDEA. A whole-image edit model (GPT Image, Nano Banana) redraws the
picture and moves the subject. A fill model paints only what it is told to
paint. So:

  1. build a canvas at the TARGET shape that holds the source at its own pixels
     (geometry, canvas_and_mask);
  2. mask only the new strips, reaching a little into the original so the model
     blends the seam (OVERLAP);
  3. after the model answers, paste the original back over the result with a
     soft edge on the grown sides only (paste_back);
  4. scale the whole canvas to the target together, so old and new areas share
     one level of detail.

Ported from the 2026-11-09 bake-off prototype; the prototype's detail-matching
blur is gone because the winning results did not need it.
"""

from __future__ import annotations

import math

import numpy as np
from PIL import Image, ImageFilter

# Same-shape tolerance: a target within 0.5% of the source's aspect ratio is a
# plain resize, not a generation.
SAME_SHAPE_TOLERANCE = 0.005

MODEL_MIN_EDGE = 512
MODEL_MAX_EDGE = 1536
MODEL_SHORT_MIN = 256      # flux-fill-dev refuses anything under 256 on a side
MODEL_STEP = 16


def same_shape(sw: int, sh: int, tw: int, th: int, tolerance: float = SAME_SHAPE_TOLERANCE) -> bool:
    return abs(math.log(tw / th) - math.log(sw / sh)) <= math.log(1 + tolerance)


def overlap_feather(canvas_longest: int) -> tuple[int, int]:
    """OVERLAP 24 / FEATHER 16 at ~800 px, scaled linearly, clamped to 8..64."""
    scale = canvas_longest / 800
    return (
        int(min(64, max(8, round(24 * scale)))),
        int(min(64, max(8, round(16 * scale)))),
    )


def geometry(sw: int, sh: int, tw: int, th: int) -> tuple[int, int, int, int]:
    """Smallest canvas at the target's shape that holds the source at its own
    pixels. Returns (canvas_w, canvas_h, offset_x, offset_y), the source
    centred on the axis that grew."""
    if tw / th >= sw / sh:
        cw, ch = round(sh * tw / th), sh
    else:
        cw, ch = sw, round(sw * th / tw)
    cw, ch = max(cw, sw), max(ch, sh)
    return cw, ch, (cw - sw) // 2, (ch - sh) // 2


def fit_source_to_target(src: Image.Image, largest_tw: int, largest_th: int) -> Image.Image:
    """The downscale rule: if the canvas would be bigger than the largest
    target in the group needs, scale the source down first so the model works
    at the size that will be delivered. Never scales UP — the original's own
    pixels are the point."""
    cw, ch, _, _ = geometry(src.width, src.height, largest_tw, largest_th)
    scale = min(1.0, largest_tw / cw, largest_th / ch)
    if scale >= 1.0:
        return src
    size = (max(1, round(src.width * scale)), max(1, round(src.height * scale)))
    return src.resize(size, Image.LANCZOS)


def _grown(cw: int, ch: int, ox: int, oy: int, sw: int, sh: int) -> tuple[bool, bool, bool, bool]:
    """Which sides of the source have new canvas beside them: left, top, right, bottom."""
    return ox > 0, oy > 0, cw - ox - sw > 0, ch - oy - sh > 0


def canvas_and_mask(
    src: Image.Image, cw: int, ch: int, ox: int, oy: int
) -> tuple[Image.Image, Image.Image]:
    """The canvas the model is given and the mask over it (255 = paint here).

    The new area is prefilled with the source stretched and blurred so the
    model starts from the right colours. The mask is 255 over everything new
    and 0 over the original, EXCEPT an OVERLAP band reaching into the original
    on the sides that grew — a side that did not grow keeps its 0 right to the
    edge."""
    sw, sh = src.size
    overlap, _ = overlap_feather(max(cw, ch))
    overlap = max(1, min(overlap, sw // 4, sh // 4))
    blur = min(80, max(8, round(40 * max(cw, ch) / 800)))
    canvas = src.resize((cw, ch), Image.BICUBIC).filter(ImageFilter.GaussianBlur(blur))
    canvas.paste(src, (ox, oy))

    left, top, right, bottom = _grown(cw, ch, ox, oy, sw, sh)
    mask = Image.new("L", (cw, ch), 255)
    inner = (
        ox + (overlap if left else 0),
        oy + (overlap if top else 0),
        ox + sw - (overlap if right else 0),
        oy + sh - (overlap if bottom else 0),
    )
    mask.paste(0, inner)
    return canvas, mask


def model_size(cw: int, ch: int) -> tuple[int, int]:
    """Both sides multiples of 16; longest edge clamp(canvas longest, 512,
    1536), raised only as far as needed to keep the short side at 256+ (a
    very wide canvas would otherwise fall under the model's minimum)."""
    longest = min(MODEL_MAX_EDGE, max(MODEL_MIN_EDGE, max(cw, ch)))
    shortest = min(cw, ch)
    longest = min(MODEL_MAX_EDGE, max(longest, math.ceil(MODEL_SHORT_MIN * max(cw, ch) / shortest)))
    scale = longest / max(cw, ch)
    step = MODEL_STEP
    return (
        max(MODEL_SHORT_MIN, round(cw * scale / step) * step),
        max(MODEL_SHORT_MIN, round(ch * scale / step) * step),
    )


def paste_back(gen: Image.Image, src: Image.Image, ox: int, oy: int) -> Image.Image:
    """The original's own pixels win everywhere except a soft band at the seam
    on the sides that grew. On a side that did not grow the original is kept
    to the very edge."""
    sw, sh = src.size
    _, feather = overlap_feather(max(gen.width, gen.height))
    feather = max(1, min(feather, sw // 4, sh // 4))
    left, top, right, bottom = _grown(gen.width, gen.height, ox, oy, sw, sh)

    # Built on a padded sheet and cropped, so the blur cannot fade the edges
    # of a side that did not grow: there the 255 runs on into the padding.
    pad = feather * 4
    sheet = Image.new("L", (sw + 2 * pad, sh + 2 * pad), 0)
    sheet.paste(255, (
        pad + (feather if left else -pad),
        pad + (feather if top else -pad),
        pad + sw - (feather if right else -pad),
        pad + sh - (feather if bottom else -pad),
    ))
    alpha = sheet.filter(ImageFilter.GaussianBlur(feather / 2)).crop((pad, pad, pad + sw, pad + sh))
    out = gen.copy()
    out.paste(src, (ox, oy), alpha)
    return out


def new_boxes(cw: int, ch: int, ox: int, oy: int, sw: int, sh: int) -> list[tuple[int, int, int, int]]:
    """The strips of new canvas, one box per side that grew."""
    left, top, right, bottom = _grown(cw, ch, ox, oy, sw, sh)
    boxes = []
    if left:
        boxes.append((0, 0, ox, ch))
    if right:
        boxes.append((ox + sw, 0, cw, ch))
    if top:
        boxes.append((0, 0, cw, oy))
    if bottom:
        boxes.append((0, oy + sh, cw, ch))
    return boxes


# ---- Bria Expand: it takes only a ratio from a short list ----

def _ratio(label: str) -> float:
    left, _, right = str(label).partition(":")
    return float(left) / float(right)


def pick_expand_ratio(
    sw: int, sh: int, tw: int, th: int, enum_labels: list[str] | tuple[str, ...]
) -> str | None:
    """The enum ratio whose expanded canvas still contains the target's
    minimal canvas, nearest to the target. None when no ratio can.

    The expanded canvas holds the source at its own scale, centred. A wide
    target (wider than the source) is contained by a canvas at least as wide;
    a narrow one by a canvas at least as narrow; a canvas on the other side of
    the source's own shape never contains it. So a 16:9 source cannot be made
    wider than 16:9 here."""
    target = tw / th
    source = sw / sh
    slack = 1 + SAME_SHAPE_TOLERANCE
    best: tuple[float, str] | None = None
    for label in enum_labels:
        r = _ratio(label)
        if target >= source:
            ok = r >= target / slack and r >= source / slack
            distance = r - target
        else:
            ok = r <= target * slack and r <= source * slack
            distance = target - r
        if ok and (best is None or abs(distance) < best[0]):
            best = (abs(distance), label)
    return best[1] if best else None


def locate_original(
    out: Image.Image,
    src: Image.Image,
    max_mean_diff: float = 20.0,
) -> tuple[float, int, int] | None:
    """Where did the model put the original? Returns (scale, x, y) with the
    source scaled by `scale` landing at (x, y) in `out`, or None when nothing
    in `out` looks like the source — a refusal beats a guess, because a guess
    would paste the original over the wrong place.

    The scale starts from the axis that filled (the one that did not grow) and
    is tried at 0.97-1.03 of that. A coarse search on a small grayscale copy
    finds the neighbourhood; a finer one at full resolution settles the pixel."""
    sw, sh = src.size
    base = min(out.width / sw, out.height / sh)
    candidates = [round(base * f, 4) for f in np.linspace(0.97, 1.03, 7)]

    out_gray = np.asarray(out.convert("L"), dtype=np.float32)
    src_img = src.convert("L")

    coarse_f = min(1.0, 128 / max(out.size))
    out_small = np.asarray(out.convert("L").resize(
        (max(1, round(out.width * coarse_f)), max(1, round(out.height * coarse_f))), Image.BILINEAR
    ), dtype=np.float32)

    per_scale: list[tuple[float, float, int, int]] = []   # coarse diff, scale, x, y (full res)
    for scale in candidates:
        w, h = round(sw * scale), round(sh * scale)
        if w < 8 or h < 8 or w > out.width or h > out.height:
            continue
        sw_s, sh_s = max(1, round(w * coarse_f)), max(1, round(h * coarse_f))
        if sw_s > out_small.shape[1] or sh_s > out_small.shape[0]:
            continue
        template = np.asarray(src_img.resize((sw_s, sh_s), Image.BILINEAR), dtype=np.float32)
        here: tuple[float, int, int] | None = None
        for y in range(out_small.shape[0] - sh_s + 1):
            for x in range(out_small.shape[1] - sw_s + 1):
                diff = float(np.abs(out_small[y:y + sh_s, x:x + sw_s] - template).mean())
                if here is None or diff < here[0]:
                    here = (diff, x, y)
        if here is not None:
            per_scale.append((here[0], float(scale), round(here[1] / coarse_f), round(here[2] / coarse_f)))
    if not per_scale:
        return None

    # The coarse pass cannot tell 0.99 from 1.00 on smooth pictures, so the
    # few best scales are all settled at full resolution and the best of those wins.
    reach = int(math.ceil(1 / coarse_f)) + 4
    best: tuple[float, float, int, int] | None = None   # diff, scale, x, y
    for _, scale, cx, cy in sorted(per_scale)[:3]:
        w, h = round(sw * scale), round(sh * scale)
        template = np.asarray(src_img.resize((w, h), Image.LANCZOS), dtype=np.float32)[::4, ::4]
        for y in range(max(0, cy - reach), min(out.height - h, cy + reach) + 1):
            for x in range(max(0, cx - reach), min(out.width - w, cx + reach) + 1):
                window = out_gray[y:y + h:4, x:x + w:4][:template.shape[0], :template.shape[1]]
                diff = float(np.abs(window - template).mean())
                if best is None or diff < best[0]:
                    best = (diff, scale, x, y)
    if best is None or best[0] > max_mean_diff:
        return None
    return best[1], best[2], best[3]


def crop_box_around_original(
    image_size: tuple[int, int], box: tuple[int, int, int, int], tw: int, th: int
) -> tuple[int, int, int, int]:
    """The smallest window at the target's ratio that contains `box`, centred
    on it and slid (never shrunk) to stay inside the image. Only generated
    margin is ever cut away."""
    iw, ih = image_size
    x0, y0, x1, y1 = box
    bw, bh = x1 - x0, y1 - y0
    target = tw / th
    if target >= bw / bh:
        w, h = math.ceil(bh * target), bh
    else:
        w, h = bw, math.ceil(bw / target)
    w, h = min(w, iw), min(h, ih)
    cx, cy = (x0 + x1) / 2, (y0 + y1) / 2
    left = int(round(min(max(cx - w / 2, 0), iw - w)))
    top = int(round(min(max(cy - h / 2, 0), ih - h)))
    return left, top, left + w, top + h


def crop_around_original(
    img: Image.Image, box: tuple[int, int, int, int], tw: int, th: int
) -> Image.Image:
    return img.crop(crop_box_around_original(img.size, box, tw, th))