Blame for engine.py

# coding=utf-8
"""Inference engine for OpenHTR-Gen.

Conditioning: a glyph strip is stacked on top of the real crop, the target box is
zeroed in the condition image, and Flux Fill regenerates the canvas. The strip is then
dropped and the edited region is cropped out.

Invariants (must match the training pipeline):
  * canvas is size x (size + strip_height_for(size)), default 512 x 608;
  * the target box must be zeroed in the condition image;
  * call the pipeline with `masked_image_latents=` so it skips `image * (1 - mask)`,
    which requires diffusers 0.32.x (0.33+ rejects `image=None`);
  * only T5 is used; CLIP is not loaded.
"""
from __future__ import annotations

import math
import os
import random
from contextlib import contextmanager

import cv2
import numpy as np
import torch
from PIL import Image, ImageDraw, ImageFont

from diffusers import FluxFillPipeline, FluxTransformer2DModel

HERE = os.path.dirname(os.path.abspath(__file__))

DEFAULT_BASE_MODEL = os.environ.get("BASE_MODEL", "black-forest-labs/FLUX.1-Fill-dev")
DEFAULT_FONT_PATH = os.path.join(HERE, "Arial-Unicode-Regular.ttf")
LORA_FILENAME = os.environ.get("LORA_FILENAME", "pytorch_lora_weights.safetensors")
# Fallback host for the style adapter when no local copy is shipped with the Space.
DEFAULT_LORA_REPO = os.environ.get("DEFAULT_LORA_REPO", "OpenHTR/OpenHTR-Gen-LoRA")

# Must match the training resolution.
DEFAULT_SIZE = 512

_PIPE = None

# =============================================================================
# Prompt
# =============================================================================
CONCAT_PROMPT_TEMPLATE = (
    "The pair of images highlights some white words on a black background, "
    "as well as their style on a real-world scene image. "
    "[IMAGE1] is a template image rendering the text, with the words {words}; "
    "[IMAGE2] shows the text content {words} naturally and correspondingly integrated into the image."
)

# CLIP is not loaded, but pooled_projections is required; training used zeros.
POOLED_PROJECTION_DIM = 768


def generate_prompt_concat(words):
    if isinstance(words, str):
        words = [words]
    words_str = ', '.join(f"'{w}'" for w in words)
    return CONCAT_PROMPT_TEMPLATE.format(words=words_str)


def zero_pooled_projections(batch_size, device, dtype):
    return torch.zeros(batch_size, POOLED_PROJECTION_DIM, device=device, dtype=dtype)


# =============================================================================
# Geometry
# =============================================================================
def polygon_bbox(polygon):
    xs = [p[0] for p in polygon]
    ys = [p[1] for p in polygon]
    return min(xs), min(ys), max(xs), max(ys)


def _window_range(img_w, img_h, ux0, uy0, ux1, uy1, L):
    """Valid top-left range so that [sx0, sx0+L) x [sy0, sy0+L) covers the box."""
    sx0_min = max(0, int(math.ceil(ux1 - L)))
    sx0_max = min(int(math.floor(ux0)), int(img_w - L))
    sy0_min = max(0, int(math.ceil(uy1 - L)))
    sy0_max = min(int(math.floor(uy0)), int(img_h - L))
    return sx0_min, sx0_max, sy0_min, sy0_max


def compute_crop_box(img_w, img_h, target_polygon, other_polygons,
                     target_size=512, rng=None, deterministic=False):
    """Fixed-size square crop window. Returns (sx0, sy0, L, target_size) or None."""
    if rng is None:
        rng = random

    tx0, ty0, tx1, ty1 = polygon_bbox(target_polygon)
    L_max = min(img_w, img_h)
    if L_max <= 1:
        return None
    tw, th = tx1 - tx0, ty1 - ty0
    if max(tw, th) > L_max:
        return None

    def _pick(L, ux0, uy0, ux1, uy1):
        sx0_min, sx0_max, sy0_min, sy0_max = _window_range(img_w, img_h, ux0, uy0, ux1, uy1, L)
        if sx0_max < sx0_min or sy0_max < sy0_min:
            return None
        if deterministic:
            return (sx0_min + sx0_max) // 2, (sy0_min + sy0_max) // 2
        return rng.randint(sx0_min, sx0_max), rng.randint(sy0_min, sy0_max)

    # Prefer a window that also contains a second annotation box.
    best_L_min, best_union = None, None
    for op in other_polygons:
        ox0, oy0, ox1, oy1 = polygon_bbox(op)
        ux0, uy0 = min(tx0, ox0), min(ty0, oy0)
        ux1, uy1 = max(tx1, ox1), max(ty1, oy1)
        L_cand = max(ux1 - ux0, uy1 - uy0)
        if L_cand > L_max:
            continue
        if best_L_min is None or L_cand < best_L_min:
            best_L_min = L_cand
            best_union = (ux0, uy0, ux1, uy1)

    if best_union is not None:
        L_min = max(best_L_min, max(tw, th), 1.0)
        if L_min <= L_max:
            L = min(int(target_size), int(L_max)) if L_min <= target_size \
                else min(int(math.ceil(L_min)), int(L_max))
            got = _pick(L, *best_union)
            if got is None:
                L = int(math.ceil(L_min))
                got = _pick(L, *best_union)
            if got is not None:
                return got[0], got[1], L, target_size

    # Fallback: cover the target box only.
    L = min(int(target_size), int(L_max))
    if max(tw, th) > L:
        L = min(int(math.ceil(max(tw, th))), int(L_max))
    got = _pick(L, tx0, ty0, tx1, ty1)
    if got is None:
        return None
    return got[0], got[1], L, target_size


def _pad_to_fit(real_pil, target_polygon, other_polygons, needed_L):
    """Replicate-pad the page when the box is larger than its short side."""
    W, H = real_pil.size
    if needed_L <= min(W, H):
        return real_pil, target_polygon, other_polygons, 0

    pad = int(math.ceil((needed_L - min(W, H)) / 2.0)) + 1
    arr = cv2.copyMakeBorder(np.array(real_pil), pad, pad, pad, pad, cv2.BORDER_REPLICATE)
    shifted_t = [[px + pad, py + pad] for px, py in target_polygon]
    shifted_o = [[[px + pad, py + pad] for px, py in op] for op in other_polygons]
    return Image.fromarray(arr), shifted_t, shifted_o, pad


def loosen_polygon_width(polygon, other_polygons, canvas_w, max_scale=2.0, rng=None):
    """Random horizontal widening of the mask (training only)."""
    if rng is None:
        rng = random

    x0, y0, x1, y1 = polygon_bbox(polygon)
    width = max(x1 - x0, 1.0)
    cx = (x0 + x1) / 2.0

    extra = width * (rng.uniform(1.0, max_scale) - 1.0)
    left_ratio = rng.random()
    desired_left = extra * left_ratio
    desired_right = extra * (1.0 - left_ratio)

    left_limit, right_limit = 0.0, float(canvas_w)
    for op in other_polygons:
        ox0, oy0, ox1, oy1 = polygon_bbox(op)
        if oy1 <= y0 or oy0 >= y1:
            continue
        if ox1 <= x0:
            left_limit = max(left_limit, ox1)
        elif ox0 >= x1:
            right_limit = min(right_limit, ox0)

    left_add = min(desired_left, max(0.0, x0 - left_limit))
    right_add = min(desired_right, max(0.0, right_limit - x1))

    out = []
    for px, py in polygon:
        out.append([px - left_add, py] if px <= cx else [px + right_add, py])
    return out


def polygon_to_mask(polygon, canvas_w, canvas_h):
    """255 = region to rewrite."""
    m = np.zeros((int(canvas_h), int(canvas_w)), dtype=np.uint8)
    cv2.fillPoly(m, [np.array(polygon, dtype=np.int32)], 255)
    return m


# =============================================================================
# Condition construction (glyph strip on top, real crop below)
# =============================================================================
def strip_height_for(size):
    """Strip height: size // 6 rounded up to a multiple of 16 (VAE x8 + 2x2 patchify)."""
    base = size // 6
    return int(math.ceil(base / 16.0)) * 16


def concat_canvas_size(size):
    return size, size + strip_height_for(size)


def crop_scene_part(image, size):
    """Drop the glyph strip, keeping the bottom size x size scene."""
    w, h = image.size
    return image.crop((0, h - size, size, h))


def render_glyph_strip(font, text, width, height, margin=0.90, min_font_size=8, max_font_size=None):
    """Render white text on black, centered and scaled to fit the strip."""
    canvas = Image.new("L", (int(width), int(height)), 0)
    if not text or not text.strip():
        return canvas

    g_size = 50
    try:
        probe = font.font_variant(size=g_size)
    except Exception:
        probe = font
    left, top, right, bottom = probe.getbbox(text)
    tw = max(right - left, 1)
    th = max(bottom - top, 1)

    ratio = min(width * margin / tw, height * margin / th)
    size = int(g_size * ratio)
    if max_font_size is not None:
        size = min(size, int(max_font_size))
    size = max(size, min_font_size)

    try:
        final_font = font.font_variant(size=size)
    except Exception:
        final_font = font

    ImageDraw.Draw(canvas).text(
        (width / 2.0, height / 2.0), text, font=final_font, fill=255, anchor="mm",
    )
    return canvas


def build_sample_concat(real_pil, text, target_polygon, other_polygons, font,
                        size=512, loose_mask=False, mask_aug=False,
                        rng=None, deterministic=True, allow_pad=True):
    """Build the condition tensors for one edit. Returns a dict or None."""
    if rng is None:
        rng = random

    if allow_pad:
        tx0, ty0, tx1, ty1 = polygon_bbox(target_polygon)
        needed = max(int(size), int(math.ceil(max(tx1 - tx0, ty1 - ty0))))
        real_pil, target_polygon, other_polygons, _ = _pad_to_fit(
            real_pil, target_polygon, other_polygons, needed
        )

    img_w, img_h = real_pil.size
    crop = compute_crop_box(img_w, img_h, target_polygon, other_polygons,
                            target_size=size, rng=rng, deterministic=deterministic)
    if crop is None:
        return None
    sx0, sy0, L, target_size = crop

    real = real_pil.crop((sx0, sy0, sx0 + L, sy0 + L))
    polygon = [[px - sx0, py - sy0] for px, py in target_polygon]
    others = [[[px - sx0, py - sy0] for px, py in op] for op in other_polygons]
    if L != target_size:
        scale = target_size / L
        real = real.resize((target_size, target_size), Image.BICUBIC)
        polygon = [[px * scale, py * scale] for px, py in polygon]
        others = [[[px * scale, py * scale] for px, py in op] for op in others]

    w = h = target_size
    if loose_mask and not deterministic:
        polygon = loosen_polygon_width(polygon, others, canvas_w=w, rng=rng)

    mask_np = polygon_to_mask(polygon, w, h)
    if mask_aug:
        r = rng.random()
        k = np.ones((3, 3), dtype=np.uint8)
        if r < 0.7:
            pass
        elif r < 0.8:
            mask_np = cv2.dilate(mask_np, k, iterations=1)
        elif r < 0.9:
            mask_np = cv2.erode(mask_np, k, iterations=1)
        elif r < 0.95:
            mask_np = cv2.dilate(mask_np, k, iterations=2)
        else:
            mask_np = cv2.erode(mask_np, k, iterations=2)

    strip_h = strip_height_for(size)
    render_img = render_glyph_strip(font, text, w, strip_h)
    combined_image_np = np.vstack((
        np.repeat(np.array(render_img)[..., None], 3, axis=2),
        np.array(real),
    ))
    combined_mask_np = np.vstack((
        np.zeros((strip_h, w), dtype=np.uint8),
        mask_np,
    ))

    # Zero the target box in the condition image; otherwise the answer leaks in.
    m = (combined_mask_np > 127).astype(np.uint8)[..., None]
    combined_cond_image_np = (combined_image_np * (1 - m)).astype(np.uint8)

    return {
        "combined_image": Image.fromarray(combined_image_np),
        "combined_cond_image": Image.fromarray(combined_cond_image_np),
        "combined_mask": Image.fromarray(combined_mask_np),
        "polygon": polygon,
        "strip_height": strip_h,
        "prompt": generate_prompt_concat([text]),
        "text": text,
        "crop_scale": target_size / L,
    }


# =============================================================================
# 320-channel condition latents (64 image + 256 mask)
# =============================================================================
def _pack_latents(x, batch_size, num_channels, height, width):
    """Same layout as FluxFillPipeline._pack_latents."""
    x = x.view(batch_size, num_channels, height // 2, 2, width // 2, 2)
    x = x.permute(0, 2, 4, 1, 3, 5)
    return x.reshape(batch_size, (height // 2) * (width // 2), num_channels * 4)


def build_masked_image_latents(vae, cond_image, mask, dtype, vae_scale_factor=8, generator=None):
    """Encode the condition image and pack it together with the mask."""
    if mask.ndim == 3:
        mask = mask.unsqueeze(1)
    mask = (mask > 0.5).to(dtype)

    cond_image = cond_image.to(device=vae.device, dtype=vae.dtype)
    posterior = vae.encode(cond_image).latent_dist
    latents = posterior.sample(generator=generator) if generator is not None else posterior.sample()
    latents = (latents - vae.config.shift_factor) * vae.config.scaling_factor
    latents = latents.to(dtype)

    bsz, ch, lh, lw = latents.shape
    packed_image = _pack_latents(latents, bsz, ch, lh, lw)

    m = mask.to(device=latents.device, dtype=dtype)[:, 0]
    m = m.view(bsz, lh, vae_scale_factor, lw, vae_scale_factor)
    m = m.permute(0, 2, 4, 1, 3).reshape(bsz, vae_scale_factor ** 2, lh, lw)
    packed_mask = _pack_latents(m, bsz, vae_scale_factor ** 2, lh, lw)

    return torch.cat([packed_image, packed_mask], dim=-1)


# =============================================================================
# Model loading
# =============================================================================
@contextmanager
def _tolerate_redundant_to_on_quantized_models():
    """Work around accelerate calling .to() on bitsandbytes models."""
    try:
        from transformers.modeling_utils import PreTrainedModel
    except Exception:  # pragma: no cover
        yield
        return

    orig_to = PreTrainedModel.to

    def patched_to(self, *args, **kwargs):
        if getattr(self, "is_loaded_in_8bit", False) or getattr(self, "is_loaded_in_4bit", False):
            return self
        return orig_to(self, *args, **kwargs)

    PreTrainedModel.to = patched_to
    try:
        yield
    finally:
        PreTrainedModel.to = orig_to


def _as_lora_file(path):
    if not path:
        return None
    if os.path.isdir(path):
        cand = os.path.join(path, LORA_FILENAME)
        return cand if os.path.isfile(cand) else None
    return path if os.path.isfile(path) else None


def _download_lora(repo_id):
    try:
        from huggingface_hub import hf_hub_download
        return hf_hub_download(
            repo_id, LORA_FILENAME,
            token=os.environ.get("HF_TOKEN") or None,
        )
    except Exception as e:  # pragma: no cover
        print(f"[lora] could not fetch {repo_id}/{LORA_FILENAME}: {e}")
        return None


def resolve_lora_path():
    """Locate the style adapter weights. Priority:
        1) LORA_PATH (local file/dir, or a HF repo id)
        2) LORA_REPO_ID (HF repo id)
        3) a local lora_checkpoint_v2 folder shipped with this app
        4) DEFAULT_LORA_REPO (OpenHTR/OpenHTR-Gen-LoRA unless overridden)
    Returns None when nothing is reachable (the demo still runs on the base model).
    """
    env_path = os.environ.get("LORA_PATH", "").strip()
    if env_path:
        hit = _as_lora_file(env_path) or _download_lora(env_path)
        if hit:
            return hit

    repo_id = os.environ.get("LORA_REPO_ID", "").strip()
    if repo_id:
        hit = _download_lora(repo_id)
        if hit:
            return hit

    candidates = [
        os.path.join(HERE, "lora_checkpoint_v2"),
        os.path.join(HERE, "models", "lora_checkpoint_v2"),
        os.path.join(os.path.dirname(HERE), "models", "lora_checkpoint_v2"),
        os.path.join(os.path.dirname(HERE), "lora_checkpoint_v2"),
    ]
    for cand in candidates:
        hit = _as_lora_file(cand)
        if hit:
            return hit

    if DEFAULT_LORA_REPO:
        return _download_lora(DEFAULT_LORA_REPO)
    return None


def load_pipeline(base_model=DEFAULT_BASE_MODEL, lora_path=None,
                  quantize_8bit=False, device="cuda", zero_gpu=False):
    """Load the base model and style adapter once per process.

    On ZeroGPU, call this at module scope, after importing ``spaces``. Its
    CUDA emulation registers the weights for GPU worker reuse. Avoid the
    Accelerate ``device_map`` loading route in that environment.
    """
    global _PIPE
    if _PIPE is not None:
        return _PIPE

    if quantize_8bit:
        compute_dtype = torch.float16
    elif zero_gpu:
        # ZeroGPU exposes CUDA through an emulation layer before GPU allocation.
        compute_dtype = torch.bfloat16
    else:
        compute_dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16

    transformer_kwargs = dict(torch_dtype=compute_dtype)
    if quantize_8bit:
        from diffusers import BitsAndBytesConfig as DiffusersBnbConfig
        transformer_kwargs["quantization_config"] = DiffusersBnbConfig(load_in_8bit=True)
    elif not zero_gpu and torch.device(device).type == "cuda" and os.environ.get("HW_DEVICE_MAP", "1") != "0":
        # Direct-to-GPU loading is a useful local option, but Accelerate's
        # hooks collide with DiffusionPipeline.to() in diffusers 0.32.2.
        # ZeroGPU must use regular loading + module-scope CUDA registration.
        dev = torch.device(device)
        if dev.index is None:
            dev = torch.device("cuda", torch.cuda.current_device())
        transformer_kwargs["device_map"] = {"": str(dev)}

    print(f"[load] base={base_model} dtype={compute_dtype} 8bit={quantize_8bit}")
    with _tolerate_redundant_to_on_quantized_models():
        try:
            transformer = FluxTransformer2DModel.from_pretrained(
                base_model, subfolder="transformer", **transformer_kwargs
            )
        except (ValueError, TypeError) as e:
            if "device_map" not in transformer_kwargs:
                raise
            print(f"[load] device_map failed ({e}), falling back to plain loading")
            transformer_kwargs.pop("device_map")
            torch.cuda.empty_cache()
            transformer = FluxTransformer2DModel.from_pretrained(
                base_model, subfolder="transformer", **transformer_kwargs
            )

        if lora_path:
            print(f"[load] adapter: {lora_path}")
            state_dict, network_alphas = FluxFillPipeline.lora_state_dict(
                pretrained_model_name_or_path_or_dict=lora_path, return_alphas=True
            )
            if not all("lora" in k or "dora_scale" in k for k in state_dict.keys()):
                raise ValueError(f"invalid adapter weights: {lora_path}")
            FluxFillPipeline.load_lora_into_transformer(
                state_dict=state_dict, network_alphas=network_alphas, transformer=transformer,
            )
        else:
            print("[load] no adapter found, running with the base model only")

        # CLIP is not needed: text_encoder/tokenizer are skipped, T5 (text_encoder_2) stays.
        pipe = FluxFillPipeline.from_pretrained(
            base_model,
            transformer=transformer,
            text_encoder=None,
            tokenizer=None,
            torch_dtype=compute_dtype,
        )
        pipe.vae.enable_tiling()
        pipe.vae.enable_slicing()

        if "device_map" in transformer_kwargs or quantize_8bit:
            # The transformer is already dispatched (or quantized). Moving the
            # whole pipeline would either conflict with Accelerate hooks or
            # attempt an unsupported .to() on a bitsandbytes module.
            # Only the remaining PyTorch modules need to be moved.
            pipe.text_encoder_2.to(device)
            pipe.vae.to(device)
        else:
            # ZeroGPU intercepts this call at module scope and registers the
            # weights for later GPU worker allocation. It also works locally.
            pipe.to(device)

        print(
            f"[load] devices: transformer={pipe.transformer.device}, "
            f"T5={pipe.text_encoder_2.device}, VAE={pipe.vae.device}",
            flush=True,
        )

    pipe.set_progress_bar_config(disable=True)
    pipe._hw_compute_dtype = compute_dtype
    _PIPE = pipe
    return _PIPE


def pipeline_dtype(pipe):
    return getattr(pipe, "_hw_compute_dtype", torch.bfloat16)


def pipeline_device(pipe):
    try:
        return pipe._execution_device
    except Exception:
        return pipe.transformer.device


# =============================================================================
# Inference
# =============================================================================
def _encode_t5(pipe, prompt, max_sequence_length=512):
    dtype = pipeline_dtype(pipe)
    inputs = pipe.tokenizer_2(
        [prompt] if isinstance(prompt, str) else prompt,
        padding="max_length", max_length=max_sequence_length, truncation=True,
        return_length=False, return_overflowing_tokens=False, return_tensors="pt",
    )
    device = pipe.text_encoder_2.device
    with torch.no_grad():
        return pipe.text_encoder_2(inputs.input_ids.to(device))[0].to(dtype)


@torch.no_grad()
def run_one_inference(pipe, cond_image, mask, prompt, num_steps, guidance_scale, seed,
                      max_sequence_length=512):
    """Run Flux Fill on one condition canvas; returns a PIL image of the same size."""
    assert cond_image.size == mask.size, f"cond_image{cond_image.size} vs mask{mask.size}"
    width, height = cond_image.size

    dtype = pipeline_dtype(pipe)
    device = pipeline_device(pipe)

    cond_t = torch.from_numpy((np.array(cond_image, dtype=np.float32) / 127.5) - 1.0)
    cond_t = cond_t.permute(2, 0, 1).unsqueeze(0).to(device=device, dtype=dtype)
    mask_t = torch.from_numpy((np.array(mask) > 127).astype(np.float32))
    mask_t = mask_t.unsqueeze(0).to(device=device, dtype=dtype)

    masked_image_latents = build_masked_image_latents(
        pipe.vae, cond_t, mask_t, dtype, vae_scale_factor=pipe.vae_scale_factor
    ).to(device)

    prompt_embeds = _encode_t5(pipe, prompt, max_sequence_length).to(device)
    pooled_prompt_embeds = zero_pooled_projections(prompt_embeds.shape[0], device, dtype)

    generator = torch.Generator(device=device).manual_seed(int(seed))
    return pipe(
        height=height,
        width=width,
        masked_image_latents=masked_image_latents,   # skips the internal image*(1-mask)
        prompt_embeds=prompt_embeds,
        pooled_prompt_embeds=pooled_prompt_embeds,
        num_inference_steps=num_steps,
        guidance_scale=guidance_scale,
        generator=generator,
        max_sequence_length=max_sequence_length,
    ).images[0]


def run_editing_inference(pipe, real_pil, text, target_polygon, other_polygons, font,
                          num_steps, guidance_scale, seed, size=512,
                          deterministic=True, allow_pad=True, max_sequence_length=512):
    """Build the condition, run inference, drop the glyph strip.

    Returns (scene image, polygon in scene coordinates) or (None, None).
    """
    sample = build_sample_concat(
        real_pil, text, target_polygon, other_polygons, font,
        size=size, loose_mask=False, mask_aug=False,
        rng=None, deterministic=deterministic, allow_pad=allow_pad,
    )
    if sample is None:
        return None, None

    gen = run_one_inference(
        pipe, sample["combined_cond_image"], sample["combined_mask"], sample["prompt"],
        num_steps=num_steps, guidance_scale=guidance_scale, seed=seed,
        max_sequence_length=max_sequence_length,
    )
    return crop_scene_part(gen, size), sample["polygon"]


def crop_with_padding(img, polygon, pad_ratio=0.0, min_pad=0):
    """Crop the polygon bounding box with optional padding."""
    x0, y0, x1, y1 = polygon_bbox(polygon)
    w = max(x1 - x0, 1.0)
    h = max(y1 - y0, 1.0)
    pad_x = max(w * pad_ratio, min_pad)
    pad_y = max(h * pad_ratio, min_pad)
    img_w, img_h = img.size
    cx0 = max(0, int(round(x0 - pad_x)))
    cy0 = max(0, int(round(y0 - pad_y)))
    cx1 = min(img_w, int(round(x1 + pad_x)))
    cy1 = min(img_h, int(round(y1 + pad_y)))
    if cx1 <= cx0:
        cx1 = min(img_w, cx0 + 1)
    if cy1 <= cy0:
        cy1 = min(img_h, cy0 + 1)
    return img.crop((cx0, cy0, cx1, cy1)), (cx0, cy0, cx1, cy1)


# =============================================================================
# Demo API: page + box + text -> crop (+ page preview)
# =============================================================================
def box_to_polygon(box):
    x0, y0, x1, y1 = [float(v) for v in box]
    if x1 < x0:
        x0, x1 = x1, x0
    if y1 < y0:
        y0, y1 = y1, y0
    return [[x0, y0], [x1, y0], [x1, y1], [x0, y1]]


def _locate_window(real_pil, target_polygon, size):
    """Recompute the crop window (same deterministic path as build_sample_concat)."""
    tx0, ty0, tx1, ty1 = polygon_bbox(target_polygon)
    needed = max(int(size), int(math.ceil(max(tx1 - tx0, ty1 - ty0))))
    padded, poly_p, _, pad = _pad_to_fit(real_pil, target_polygon, [], needed)
    crop = compute_crop_box(padded.size[0], padded.size[1], poly_p, [],
                            target_size=size, deterministic=True)
    if crop is None:
        return None
    sx0, sy0, L, _ = crop
    return padded, poly_p, pad, sx0, sy0, L


def paste_back(real_pil, gen_scene, target_polygon, size):
    """Paste the edited box back into the full page (preview only)."""
    loc = _locate_window(real_pil, target_polygon, size)
    if loc is None:
        return real_pil
    padded, poly_p, pad, sx0, sy0, L = loc
    scale = size / float(L)

    local = [[(px - sx0) * scale, (py - sy0) * scale] for px, py in poly_p]
    bx0, by0, bx1, by1 = polygon_bbox(local)
    gx0 = max(0, int(math.floor(bx0)))
    gy0 = max(0, int(math.floor(by0)))
    gx1 = min(gen_scene.width, int(math.ceil(bx1)))
    gy1 = min(gen_scene.height, int(math.ceil(by1)))
    if gx1 <= gx0 or gy1 <= gy0:
        return real_pil

    patch = gen_scene.crop((gx0, gy0, gx1, gy1))
    ox0 = int(round(sx0 + gx0 / scale))
    oy0 = int(round(sy0 + gy0 / scale))
    ox1 = int(round(sx0 + gx1 / scale))
    oy1 = int(round(sy0 + gy1 / scale))

    Wp, Hp = padded.size
    ox0, oy0 = max(0, min(ox0, Wp - 1)), max(0, min(oy0, Hp - 1))
    ox1, oy1 = max(ox0 + 1, min(ox1, Wp)), max(oy0 + 1, min(oy1, Hp))

    patch = patch.resize((ox1 - ox0, oy1 - oy0), Image.BICUBIC)
    out = padded.copy()
    out.paste(patch, (ox0, oy0))

    W0, H0 = real_pil.size
    return out.crop((pad, pad, pad + W0, pad + H0))


def load_font(path=DEFAULT_FONT_PATH, size=60):
    try:
        return ImageFont.truetype(path, size=size)
    except Exception as e:
        print(f"[warn] failed to load font {path}: {e}")
        return ImageFont.load_default()


def edit_text_region(pipe, font, image, box, text,
                     size=DEFAULT_SIZE, steps=20, guidance_scale=30.0, seed=0,
                     crop_pad_ratio=0.0):
    """Page + box + new text -> dict with the cropped result and a page preview."""
    image = image.convert("RGB")
    W, H = image.size
    x0, y0, x1, y1 = [float(v) for v in box]
    x0, x1 = max(0.0, min(x0, W - 1)), max(1.0, min(x1, W))
    y0, y1 = max(0.0, min(y0, H - 1)), max(1.0, min(y1, H))
    if x1 - x0 < 2 or y1 - y0 < 2:
        raise ValueError(f"box too small: {(x0, y0, x1, y1)}")
    polygon = box_to_polygon((x0, y0, x1, y1))

    gen_scene, poly_local = run_editing_inference(
        pipe, image, text, polygon, [], font,
        num_steps=int(steps), guidance_scale=float(guidance_scale), seed=int(seed),
        size=int(size), deterministic=True, allow_pad=True,
    )
    if gen_scene is None:
        raise ValueError("no valid square crop window; please draw a smaller box")

    crop, crop_box = crop_with_padding(gen_scene, poly_local, pad_ratio=float(crop_pad_ratio))
    page = paste_back(image, gen_scene, polygon, int(size))
    return {
        "crop": crop,
        "page": page,
        "crop_box": crop_box,
        "polygon_local": poly_local,
        "canvas": concat_canvas_size(int(size)),
        "prompt": generate_prompt_concat([text]),
    }