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
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 on_reference_change(images) -> str:
    if not images:
        return "📷 No reference images"
    count = len(images)
    return f"📷 {count} reference image{'s' if count != 1 else ''} uploaded"


# ── 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 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, current_gallery) -> list:
    composite = fix_orientation(_editor_composite(editor_value))
    current = list(current_gallery or [])
    current.append(composite)
    gr.Info("Added to Reference Images")
    return current


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 _resolve_gallery_path(selected_path, gallery_value):
    """Pick the path the Send-to-* buttons should use.

    Prefers the user's currently-selected gallery item; falls back to the most
    recent (last) item so single-image runs and "didn't click anything" cases
    both behave intuitively.
    """
    if selected_path:
        return selected_path
    if not gallery_value:
        return None
    item = gallery_value[-1]
    return item[0] if isinstance(item, (list, tuple)) else item


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, current_gallery):
    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")
    current = list(current_gallery or [])
    current.append(img)
    gr.Info("Output added to Reference Images.")
    return current


# ── 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, current_gallery):
    if img is None:
        raise gr.Error("Nothing to send — generate it first.")
    current = list(current_gallery or [])
    current.append(img)
    gr.Info("Added to Reference Images.")
    return current