image_utils.py

"""Pure image helpers — no torch, no diffusers, no gradio state.

Owns: EXIF handling, dimension snapping, canvas fitting, editor-composite
extraction, HEIC decoding, PNG metadata embedding.
"""
from __future__ import annotations

import json
import tempfile
from typing import Any

import numpy as np
from PIL import Image, ImageOps, ImageFilter, ImageDraw, ImageColor
from PIL.PngImagePlugin import PngInfo
import gradio as gr

# ── HEIC / HEIF support ──────────────────────────────────────────────────────
try:
    from pillow_heif import register_heif_opener
    register_heif_opener()
except ImportError:
    print("pillow-heif not installed — HEIC/HEIF uploads will not work. "
          "Add `pillow-heif` to requirements.txt.")


# ── EXIF / dimension helpers ─────────────────────────────────────────────────

def fix_orientation(img: Image.Image | None) -> Image.Image | None:
    if img is None:
        return None
    return ImageOps.exif_transpose(img)


def _snap16(v: float) -> int:
    """Snap to a multiple of 16 — required by FLUX's VAE."""
    return max(16, (int(v) // 16) * 16)


def compute_base_dimensions(image: Image.Image | None) -> tuple[int, int]:
    if image is None:
        return 1024, 1024
    w, h = image.size
    scale = min(1024 / w, 1024 / h)
    return _snap16(w * scale), _snap16(h * scale)


update_dimensions_on_upload = compute_base_dimensions


def compute_canvas_dimensions(
    base_image: Image.Image | None,
    canvas_mode: str,
    custom_width: int,
    custom_height: int,
) -> tuple[int, int]:
    if canvas_mode == "Custom":
        return _snap16(custom_width), _snap16(custom_height)
    return compute_base_dimensions(base_image)


# ── Canvas fitting ──────────────────────────────────────────────────────────

def fit_to_canvas(
    img: Image.Image,
    width: int,
    height: int,
    mode: str = "Stretch",
    pad_color: str = "#000000",
) -> Image.Image:
    """Return `img` resized to exactly width×height using the given strategy.

    Modes:
      - "Stretch"        : resize ignoring aspect (current default, may distort)
      - "Pad (color)"    : scale to fit, pad with `pad_color`
      - "Pad (blur)"     : scale to fit, pad with a blurred cover of the image
      - "Crop (cover)"   : scale to cover, center-crop to canvas
    """
    img = img.convert("RGB")

    if mode == "Stretch":
        return img.resize((width, height), Image.LANCZOS)

    iw, ih = img.size

    if mode == "Pad (color)":
        scale = min(width / iw, height / ih)
        nw, nh = max(1, int(iw * scale)), max(1, int(ih * scale))
        resized = img.resize((nw, nh), Image.LANCZOS)
        canvas = Image.new("RGB", (width, height), pad_color)
        canvas.paste(resized, ((width - nw) // 2, (height - nh) // 2))
        return canvas

    if mode == "Pad (blur)":
        # Foreground: scale-to-fit
        scale = min(width / iw, height / ih)
        nw, nh = max(1, int(iw * scale)), max(1, int(ih * scale))
        fg = img.resize((nw, nh), Image.LANCZOS)
        # Background: scale-to-cover, center-crop, then blur heavily
        cscale = max(width / iw, height / ih)
        cw, ch = max(1, int(iw * cscale)), max(1, int(ih * cscale))
        bg = img.resize((cw, ch), Image.LANCZOS)
        bg = bg.crop(((cw - width) // 2, (ch - height) // 2,
                     (cw - width) // 2 + width, (ch - height) // 2 + height))
        bg = bg.filter(ImageFilter.GaussianBlur(radius=32))
        bg.paste(fg, ((width - nw) // 2, (height - nh) // 2))
        return bg

    if mode == "Crop (cover)":
        cscale = max(width / iw, height / ih)
        nw, nh = max(1, int(iw * cscale)), max(1, int(ih * cscale))
        resized = img.resize((nw, nh), Image.LANCZOS)
        left = (nw - width) // 2
        top = (nh - height) // 2
        return resized.crop((left, top, left + width, top + height))

    # Unknown mode → fall back to stretch rather than erroring during inference
    print(f"[fit_to_canvas] unknown mode {mode!r} — falling back to Stretch.")
    return img.resize((width, height), Image.LANCZOS)


# ── UI label updates ────────────────────────────────────────────────────────

def on_base_image_change(img) -> str:
    if img is None:
        return "*No base image uploaded yet*"
    try:
        pil_img = img if isinstance(img, Image.Image) else Image.open(img)
        ow, oh = pil_img.size
        bw, bh = compute_base_dimensions(pil_img)
        return (
            f"Input: **{ow} × {oh}** px → "
            f"Auto canvas (pre-upscale): **{bw} × {bh}** px"
        )
    except Exception as e:
        return f"*Could not read dimensions: {e}*"


def make_solid_base_image(
    color: str = "#FFFFFF",
    canvas_mode: str = "Auto (from base image)",
    custom_width: int = 1024,
    custom_height: int = 1024,
) -> Image.Image:
    """Unicolor base PNG for t2i-style runs that still require a base image.

    Uses Custom canvas W×H when canvas mode is Custom; otherwise 1024×1024
    (matches Auto longest-side target). Colour often becomes the background.
    """
    if canvas_mode == "Custom":
        w, h = _snap16(custom_width or 1024), _snap16(custom_height or 1024)
    else:
        w, h = 1024, 1024
    try:
        fill = ImageColor.getrgb(color or "#FFFFFF")
    except Exception:
        fill = (255, 255, 255)
        gr.Warning(f"Could not parse colour {color!r} — using white.")
    img = Image.new("RGB", (w, h), fill)
    gr.Info(f"Solid base {w}×{h} ({color or '#FFFFFF'})")
    return img


MAX_REFERENCE_IMAGES = 3
# base + 3 refs — used by reorder helpers
MAX_IMAGE_SLOTS = 1 + MAX_REFERENCE_IMAGES


def _is_empty_image(img) -> bool:
    return img is None


def _compact_refs(*refs) -> list:
    """Drop empty ref slots and pack remaining images upward."""
    packed = [r for r in refs if not _is_empty_image(r)]
    while len(packed) < MAX_REFERENCE_IMAGES:
        packed.append(None)
    return packed[:MAX_REFERENCE_IMAGES]


def collect_reference_images(ref1=None, ref2=None, ref3=None) -> list:
    """Ordered non-empty reference slots (gaps are compacted away)."""
    return [r for r in _compact_refs(ref1, ref2, ref3) if r is not None]


def reference_info_text(ref1=None, ref2=None, ref3=None) -> str:
    count = len(collect_reference_images(ref1, ref2, ref3))
    if count == 0:
        return "📷 No reference images"
    return f"📷 {count} reference image{'s' if count != 1 else ''} (max {MAX_REFERENCE_IMAGES})"


def _img_slot_update(new, old, *, visible=None):
    """Update a slot only when value identity or visibility actually changes.

    Avoids re-writing the same image object into the component that fired
    `.change`, which would loop forever.
    """
    kw = {}
    if new is not old:
        kw["value"] = new
    if visible is not None:
        kw["visible"] = visible
    return gr.update(**kw) if kw else gr.update()


def _ref_slot_updates(r1, r2, r3, *, old1=None, old2=None, old3=None):
    """Values + progressive visibility for the three ref boxes."""
    show2 = r1 is not None
    show3 = r1 is not None and r2 is not None
    return (
        _img_slot_update(r1, old1),
        _img_slot_update(r2, old2, visible=show2),
        _img_slot_update(r3, old3, visible=show3),
        reference_info_text(r1, r2, r3),
        gr.update(visible=show2),  # ref2 up
        gr.update(visible=show2),  # ref2 down
        gr.update(visible=show3),  # ref3 up
    )


def compact_reference_slots(ref1, ref2, ref3):
    """On any ref change: pack non-empty images into ref1..refN.

    Deleting Reference 1 with 2/3 filled shifts them up instead of wiping all.
    Returns (ref1, ref2, ref3, info, ref2_up_vis, ref2_down_vis, ref3_up_vis).
    """
    r1, r2, r3 = _compact_refs(ref1, ref2, ref3)
    return _ref_slot_updates(r1, r2, r3, old1=ref1, old2=ref2, old3=ref3)


# Back-compat names used by older call sites / imports
def on_reference_change(images) -> str:
    if not images:
        return "📷 No reference images"
    if isinstance(images, (list, tuple)):
        count = len([x for x in images if x is not None])
    else:
        count = 1
    if count == 0:
        return "📷 No reference images"
    return f"📷 {count} reference image{'s' if count != 1 else ''} (max {MAX_REFERENCE_IMAGES})"


def on_ref1_change(ref1, ref2, ref3):
    return compact_reference_slots(ref1, ref2, ref3)


def on_ref2_change(ref1, ref2, ref3):
    return compact_reference_slots(ref1, ref2, ref3)


def sync_reference_slots(ref1, ref2, ref3):
    return compact_reference_slots(ref1, ref2, ref3)


def _slot_tuple(base, ref1, ref2, ref3):
    return [base, ref1, ref2, ref3]


def _apply_slots(slots, *, old_slots=None):
    """Return base/ref updates + ref button visibility from a 4-slot list."""
    old = list(old_slots) if old_slots is not None else [None, None, None, None]
    base, r1, r2, r3 = slots[0], slots[1], slots[2], slots[3]
    # Compact only the ref portion so base can be empty while refs remain.
    r1, r2, r3 = _compact_refs(r1, r2, r3)
    show2 = r1 is not None
    show3 = r1 is not None and r2 is not None
    return (
        _img_slot_update(base, old[0]),
        _img_slot_update(r1, old[1]),
        _img_slot_update(r2, old[2], visible=show2),
        _img_slot_update(r3, old[3], visible=show3),
        reference_info_text(r1, r2, r3),
        gr.update(visible=show2),
        gr.update(visible=show2),
        gr.update(visible=show3),
    )


def move_image_slot(base, ref1, ref2, ref3, index: int, direction: int):
    """Swap base/ref slot `index` with neighbour (`direction` = -1 up / +1 down).

    Slot indices: 0=base, 1=ref1, 2=ref2, 3=ref3.
    Empty refs stay compacted after the move.
    """
    old = _slot_tuple(base, ref1, ref2, ref3)
    slots = list(old)
    j = int(index) + int(direction)
    if index < 0 or index >= len(slots) or j < 0 or j >= len(slots):
        return _apply_slots(slots, old_slots=old)
    slots[index], slots[j] = slots[j], slots[index]
    return _apply_slots(slots, old_slots=old)


def move_base_down(base, ref1, ref2, ref3):
    return move_image_slot(base, ref1, ref2, ref3, 0, +1)


def move_ref1_up(base, ref1, ref2, ref3):
    return move_image_slot(base, ref1, ref2, ref3, 1, -1)


def move_ref1_down(base, ref1, ref2, ref3):
    return move_image_slot(base, ref1, ref2, ref3, 1, +1)


def move_ref2_up(base, ref1, ref2, ref3):
    return move_image_slot(base, ref1, ref2, ref3, 2, -1)


def move_ref2_down(base, ref1, ref2, ref3):
    return move_image_slot(base, ref1, ref2, ref3, 2, +1)


def move_ref3_up(base, ref1, ref2, ref3):
    return move_image_slot(base, ref1, ref2, ref3, 3, -1)


def _place_in_ref_slots(img, ref1, ref2, ref3):
    """Put `img` into the first empty reference slot (or replace ref3 if full).

    Returns updates for (ref1, ref2, ref3, reference_info).
    """
    if img is None:
        raise gr.Error("Nothing to send.")
    if not isinstance(img, Image.Image):
        try:
            img = Image.open(img)
        except Exception as e:
            raise gr.Error(f"Could not open image: {e}")
    img = fix_orientation(img).convert("RGB")

    r1, r2, r3 = _compact_refs(ref1, ref2, ref3)
    # Full update includes reorder-button visibility (ref2↑↓, ref3↑).
    if r1 is None:
        gr.Info("Sent to Reference 1")
        return _ref_slot_updates(img, None, None, old1=ref1, old2=ref2, old3=ref3)
    if r2 is None:
        gr.Info("Sent to Reference 2")
        return _ref_slot_updates(r1, img, None, old1=ref1, old2=ref2, old3=ref3)
    if r3 is None:
        gr.Info("Sent to Reference 3")
        return _ref_slot_updates(r1, r2, img, old1=ref1, old2=ref2, old3=ref3)
    gr.Warning("All 3 reference slots full — replaced Reference 3.")
    return _ref_slot_updates(r1, r2, img, old1=ref1, old2=ref2, old3=ref3)


# ── Upload round-trip (fixes HEIC preview in main tab) ──────────────────────

def reencode_upload(img):
    if img is None:
        return None
    if not isinstance(img, Image.Image):
        try:
            img = Image.open(img)
        except Exception:
            return img
    return fix_orientation(img).convert("RGB")


# ── Inference input assembly ────────────────────────────────────────────────

def process_images(base_image, reference_images) -> list[Image.Image]:
    pil_images: list[Image.Image] = []
    if base_image is not None:
        try:
            img = base_image if isinstance(base_image, Image.Image) else Image.open(base_image)
            pil_images.append(fix_orientation(img).convert("RGB"))
        except Exception as e:
            print(f"Skipping invalid base image: {e}")
    for item in (reference_images or []):
        try:
            path_or_img = item[0] if isinstance(item, (tuple, list)) else item
            if path_or_img is None:
                continue
            if isinstance(path_or_img, Image.Image):
                img = path_or_img
            elif isinstance(path_or_img, str):
                img = Image.open(path_or_img)
            else:
                img = Image.open(path_or_img.name)
            pil_images.append(fix_orientation(img).convert("RGB"))
        except Exception as e:
            print(f"Skipping invalid reference image: {e}")
    return pil_images


# ── ImageEditor helpers ─────────────────────────────────────────────────────

def _editor_composite(editor_value) -> Image.Image:
    if not editor_value or editor_value.get("composite") is None:
        raise gr.Error("Upload and crop an image in the editor first.")
    composite = editor_value["composite"]
    if isinstance(composite, np.ndarray):
        composite = Image.fromarray(composite)
    return composite.convert("RGB")


def send_editor_to_base(editor_value) -> Image.Image:
    composite = fix_orientation(_editor_composite(editor_value))
    gr.Info("Sent to Base Image")
    return composite


def send_editor_to_reference(editor_value, ref1, ref2, ref3):
    composite = fix_orientation(_editor_composite(editor_value))
    return _place_in_ref_slots(composite, ref1, ref2, ref3)


def load_heic_to_editor(path):
    if not path:
        return gr.update()
    try:
        img = fix_orientation(Image.open(path)).convert("RGB")
    except Exception as e:
        raise gr.Error(f"Could not decode HEIC/HEIF: {e}")
    gr.Info("HEIC loaded into editor.")
    return img


# ── Send output → base / reference (gallery-aware) ──────────────────────────

def _gallery_item_path(item):
    """Normalize a Gallery entry / FileData-ish value to a filesystem path str."""
    if item is None:
        return None
    if isinstance(item, (list, tuple)):
        item = item[0] if item else None
    if item is None:
        return None
    if isinstance(item, dict):
        item = item.get("path") or item.get("name") or item.get("url")
    if hasattr(item, "path") and not isinstance(item, (str, bytes)):
        try:
            item = item.path
        except Exception:
            pass
    if isinstance(item, str) and item:
        return item
    return None


def _resolve_gallery_path(selected_path, gallery_value):
    """Pick the path the Send-to-* / Download controls should use.

    Bug we're guarding against: `selected_path` comes from a gr.State that
    persists across generation runs, so after a new batch has replaced the
    gallery it can still hold a stale path from a previous run — or even a
    path from the *first* image of the current batch, because some Gradio
    builds auto-fire .select(index=0) right after the gallery repopulates.

    Rule: only honour the selection if it's actually still one of the paths
    currently in the gallery. Otherwise use the *last* (most recent) item.
    """
    if not gallery_value:
        return None

    current_paths = []
    for item in gallery_value:
        p = _gallery_item_path(item)
        if p:
            current_paths.append(p)

    if not current_paths:
        return None

    selected = _gallery_item_path(selected_path)
    if selected and selected in current_paths:
        return selected
    return current_paths[-1]


def resolve_output_download_path(selected_path, gallery_value):
    """Path for the large Download button under the output gallery.

    Same selection rules as Send→*: honour a still-valid gallery selection,
    otherwise the latest generated image. Returns None when empty so the
    DownloadButton stays inactive.
    """
    return _resolve_gallery_path(selected_path, gallery_value)


def latest_gallery_download_path(gallery_value):
    """Always the newest gallery item — used right after generate yields."""
    return _resolve_gallery_path(None, gallery_value)


def send_output_to_base(selected_path, gallery_value):
    path = _resolve_gallery_path(selected_path, gallery_value)
    if not path:
        raise gr.Error("Nothing to send — generate an image first.")
    img = Image.open(path).convert("RGB")
    gr.Info("Output sent to Base Image.")
    return img


def send_output_to_reference(selected_path, gallery_value, ref1, ref2, ref3):
    path = _resolve_gallery_path(selected_path, gallery_value)
    if not path:
        raise gr.Error("Nothing to send — generate an image first.")
    img = Image.open(path).convert("RGB")
    return _place_in_ref_slots(img, ref1, ref2, ref3)


# ── PNG metadata embedding ──────────────────────────────────────────────────

def _format_parameters_string(meta: dict[str, Any]) -> str:
    prompt = meta.get("prompt", "") or ""
    fields = [
        ("Seed", meta.get("seed")),
        ("Steps", meta.get("steps")),
        ("CFG scale", meta.get("guidance_scale")),
        ("Size", f"{meta.get('width')}x{meta.get('height')}"),
        ("Model", meta.get("model")),
        ("Upscaler", meta.get("upscale_factor")),
        ("Canvas mode", meta.get("canvas_mode")),
        ("Fit mode", meta.get("canvas_fit_mode")),
    ]
    loras = meta.get("loras") or []
    if loras:
        fields.append(("LoRAs", ", ".join(f"{n}:{w:.2f}" for n, w in loras)))
    kv = ", ".join(f"{k}: {v}" for k, v in fields if v not in (None, "", "None"))
    return f"{prompt}\n{kv}".strip()


def build_pnginfo(meta: dict[str, Any]) -> PngInfo:
    """Public so bulk processing can reuse it for in-place saves."""
    info = PngInfo()
    info.add_text("parameters", _format_parameters_string(meta))
    for k in ("prompt", "seed", "steps", "guidance_scale", "width", "height",
              "model", "upscale_factor", "canvas_mode", "canvas_fit_mode",
              "lora_prompt", "custom_prompt"):
        info.add_text(k, str(meta.get(k, "")))
    info.add_text("loras", json.dumps(meta.get("loras") or []))
    return info


def save_with_metadata(image: Image.Image, meta: dict[str, Any],
                       path: str | None = None) -> str:
    """Save `image` as PNG with embedded generation metadata.

    If `path` is given, write there (used by bulk-process to keep predictable
    filenames inside its work directory). Otherwise allocate a temp PNG.
    """
    if path is None:
        tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False, prefix="flux2_klein_")
        tmp.close()
        path = tmp.name
    image.save(path, format="PNG", pnginfo=build_pnginfo(meta))
    return path

# ── Send arbitrary PIL → base / reference (used by the Depth/Pose tab) ──────

def push_pil_to_base(img):
    if img is None:
        raise gr.Error("Nothing to send — generate it first.")
    gr.Info("Sent to Base Image.")
    return img


def push_pil_to_reference(img, ref1, ref2, ref3):
    if img is None:
        raise gr.Error("Nothing to send — generate it first.")
    return _place_in_ref_slots(img, ref1, ref2, ref3)

# ── Canvas extension (used by the Crop / Fix tab) ───────────────────────────
# Percentages are relative to the *original* dimensions. Padding is filled
# with a flat colour. Returns a new PIL image; does NOT save to disk — the
# result is loaded straight back into the ImageEditor so the existing
# Send → Base / Send → Reference buttons handle export.

from PIL import ImageColor  # add near the other PIL imports at top of file


def _parse_fill(color_str: str, mode: str):
    rgba = ImageColor.getcolor(color_str or "#000000", "RGBA")
    if mode == "RGB":
        return rgba[:3]
    if mode == "L":
        r, g, b, _ = rgba
        return int(0.299 * r + 0.587 * g + 0.114 * b)
    return rgba  # RGBA


def compute_extend_padding(w: int, h: int, up, down, left, right) -> tuple[int, int, int, int]:
    """Return (pad_left, pad_top, pad_right, pad_bottom) in pixels for the
    given percentage inputs. Percentages are relative to original W/H so
    Down=100 on 960×960 → 960×1920 with the source at the top."""
    pl = int(round(w * (float(left  or 0) / 100.0)))
    pr = int(round(w * (float(right or 0) / 100.0)))
    pt = int(round(h * (float(up    or 0) / 100.0)))
    pb = int(round(h * (float(down  or 0) / 100.0)))
    return pl, pt, pr, pb


def extend_canvas(img: Image.Image, up, down, left, right, fill: str) -> Image.Image:
    """Extend `img`'s canvas by the given per-side percentages, filling the
    new area with `fill`. Original image mode is preserved so RGBA stays
    RGBA (no transparency loss)."""
    if img is None:
        raise gr.Error("Nothing to extend — upload an image into the editor first.")
    if img.mode not in ("RGB", "RGBA", "L"):
        img = img.convert("RGBA")
    for name, v in (("Up", up), ("Down", down), ("Left", left), ("Right", right)):
        if v is None or float(v) < 0:
            raise gr.Error(f"'{name} %' must be ≥ 0.")
    w, h = img.size
    pl, pt, pr, pb = compute_extend_padding(w, h, up, down, left, right)
    if pl == pt == pr == pb == 0:
        gr.Info("All percentages are 0 — image unchanged.")
        return img
    new_size = (w + pl + pr, h + pt + pb)
    canvas = Image.new(img.mode, new_size, _parse_fill(fill, img.mode))
    canvas.paste(img, (pl, pt))
    return canvas


def render_extend_schematic(
    img: Image.Image | None, up, down, left, right, fill: str,
    max_dim: int = 320,
) -> tuple[Image.Image | None, str]:
    """Live, to-scale preview of what extend_canvas() will produce, without
    doing the full render. Returns (schematic PIL, info markdown)."""
    if img is None:
        return None, "*Upload something into the editor first.*"
    w, h = img.size
    pl, pt, pr, pb = compute_extend_padding(w, h, up, down, left, right)
    nw, nh = w + pl + pr, h + pt + pb
    scale = min(max_dim / nw, max_dim / nh, 1.0)
    sw, sh = max(1, int(nw * scale)), max(1, int(nh * scale))
    spl, spt = int(pl * scale), int(pt * scale)
    sow, soh = max(1, int(w * scale)), max(1, int(h * scale))

    schem = Image.new("RGB", (sw, sh), _parse_fill(fill, "RGB"))
    d = ImageDraw.Draw(schem)  # ImageDraw already imported at top of module
    d.rectangle([spl, spt, spl + sow - 1, spt + soh - 1],
                fill=(200, 200, 200), outline=(255, 0, 0), width=2)
    info = (f"**Original:** {w} × {h}  \n"
            f"**Padding (L, T, R, B):** {pl}, {pt}, {pr}, {pb}  \n"
            f"**Final canvas:** {nw} × {nh}")
    return schem, info


def extend_editor_canvas(editor_value, up, down, left, right, fill) -> Image.Image:
    """Take the current editor composite, extend it, and return the result so
    it can be loaded straight back into the same ImageEditor. Any crop marks
    the user had placed are baked in before extending (that's what
    `_editor_composite` already does)."""
    composite = fix_orientation(_editor_composite(editor_value))
    out = extend_canvas(composite, up, down, left, right, fill)
    gr.Info(f"Canvas extended to {out.size[0]} × {out.size[1]}.")
    return out