passes.py

"""Препроцес G-буфер-пасів під ControlNet-Union + IP-Adapter.

Тут живе ГОЛОВНА пастка Фази 0 — конвенція нормалей і depth (див. docs/01).
three.js `MeshNormalMaterial` / рендер Blender кодують нормаль у своєму RGB;
ControlNet-normal чекає своє. Ці функції — те місце, де ми це вирівнюємо
дебаг-ручками (`normal_flip`, `normal_space`, `depth_invert`), не чіпаючи пайплайн.
"""
from __future__ import annotations

import io

import numpy as np
from PIL import Image


def _load_rgb(data: bytes) -> np.ndarray:
    return np.asarray(Image.open(io.BytesIO(data)).convert("RGB"))


def _load_gray(data: bytes) -> np.ndarray:
    return np.asarray(Image.open(io.BytesIO(data)).convert("L"))


def prep_normal(
    data: bytes,
    *,
    flip: dict | None = None,
    space: str = "view",
    size: tuple[int, int] | None = None,
) -> Image.Image:
    """view-space normal PNG (кодування (n+1)/2) → нормаль у конвенції ControlNet.

    `flip` = {"x":bool,"y":bool,"z":bool}: інвертує вісь ПІСЛЯ декоду в [-1,1].
    Типова пастка three.js↔ControlNet — перевернути Y (і/або Z). `space` поки
    лише мітка (world→view вимагає viewmat; у Фазі 0 працюємо у view-space).
    """
    flip = flip or {}
    rgb = _load_rgb(data).astype(np.float32) / 255.0
    n = rgb * 2.0 - 1.0  # [0,1] → [-1,1]

    if flip.get("x"):
        n[..., 0] *= -1.0
    if flip.get("y"):
        n[..., 1] *= -1.0
    if flip.get("z"):
        n[..., 2] *= -1.0

    # ре-нормалізуємо (фліпи довжину не міняють, але resize/JPEG-артефакти можуть)
    norm = np.linalg.norm(n, axis=-1, keepdims=True)
    n = n / np.clip(norm, 1e-6, None)

    enc = ((n + 1.0) * 0.5 * 255.0).clip(0, 255).astype(np.uint8)
    img = Image.fromarray(enc)
    if size is not None:
        img = img.resize(size, Image.BILINEAR)
    return img


def prep_depth(
    data: bytes,
    *,
    invert: bool = False,
    size: tuple[int, int] | None = None,
) -> Image.Image:
    """depth PNG → grayscale-мапа у конвенції ControlNet-depth (MiDaS: near=bright).

    Нормалізуємо [0,1] по фактичному діапазону кадру (робастно, 2..98 перцентиль),
    `invert` якщо джерело віддає near=dark.
    """
    d = _load_gray(data).astype(np.float32)
    lo, hi = np.percentile(d, [2, 98])
    d = np.clip((d - lo) / max(hi - lo, 1e-6), 0.0, 1.0)
    if invert:
        d = 1.0 - d

    img = Image.fromarray((d * 255).astype(np.uint8)).convert("RGB")
    if size is not None:
        img = img.resize(size, Image.BILINEAR)
    return img


def prep_beauty(data: bytes, *, size: tuple[int, int] | None = None) -> Image.Image:
    """lit beauty-пас → sRGB-картинка для CLIP-енкодера IP-Adapter-а."""
    img = Image.fromarray(_load_rgb(data))
    if size is not None:
        img = img.resize(size, Image.BILINEAR)
    return img




def build_regions(seg_bytes: bytes, parts: list, size: tuple[int, int]) -> list | None:
    """seg-маска (кожен регіон = свій palette-колір) + parts[{color,prompt}] →
    список {mask: float32 (H8,W8), prompt}. Колір пікселя → НАЙБЛИЖЧИЙ у палітрі
    (чорний фон = index 0), тож resize/JPEG-краї не ламають межі. Маска даунсемплиться
    до latent-роздільності (H/8×W/8) — саме її чекає RegionalCrossAttnProcessor.
    Частини без промпту або невидимі в кадрі пропускаємо."""
    W, H = size
    img = Image.open(io.BytesIO(seg_bytes)).convert("RGB").resize((W, H), Image.NEAREST)
    arr = np.asarray(img).astype(np.int16)                       # H,W,3

    def _hex(h):
        h = h.lstrip("#")
        return [int(h[0:2], 16), int(h[2:4], 16), int(h[4:6], 16)]

    pal = np.array([[0, 0, 0]] + [_hex(p["color"]) for p in parts], np.int16)   # (P+1,3)
    idx = ((arr[:, :, None, :] - pal[None, None, :, :]) ** 2).sum(-1).argmin(-1)  # H,W
    h8, w8 = H // 8, W // 8
    out = []
    for i, p in enumerate(parts):
        m = (idx == i + 1).astype(np.uint8) * 255
        m8 = np.asarray(Image.fromarray(m).resize((w8, h8), Image.NEAREST))
        m8 = (m8 > 127).astype(np.float32)
        if m8.sum() < 1:                                          # регіон не в кадрі
            continue
        out.append({"mask": m8, "prompt": p["prompt"]})
    return out or None