app.py

# app.py
import os
import time
from typing import Optional, Tuple

import gradio as gr
import numpy as np
from PIL import Image, ImageFilter, ImageOps, ImageFile

# лечим обрезанные файлы (webp/jpeg)
ImageFile.LOAD_TRUNCATED_IMAGES = True

import torch
from diffusers import (
    StableDiffusionXLInpaintPipeline,
    StableDiffusionInpaintPipeline,
)

# Опциональный латентный апскейлер (если доступен)
try:
    from diffusers import StableDiffusionLatentUpscalePipeline
    HAS_LATENT_UP = True
except Exception:
    HAS_LATENT_UP = False

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32

MODEL_PREFS = [
    ("sdxl", "diffusers/stable-diffusion-xl-1.0-inpainting-0.1"),
    ("sd2",  "stabilityai/stable-diffusion-2-inpainting"),
    ("sd15", "runwayml/stable-diffusion-inpainting"),
]

def load_inpaint() -> tuple[object, str, str]:
    last = None
    for family, repo in MODEL_PREFS:
        try:
            if family == "sdxl":
                pipe = StableDiffusionXLInpaintPipeline.from_pretrained(
                    repo, torch_dtype=DTYPE, use_safetensors=True
                )
            else:
                pipe = StableDiffusionInpaintPipeline.from_pretrained(
                    repo, torch_dtype=DTYPE, use_safetensors=True
                )
            pipe = pipe.to(DEVICE)
            if hasattr(pipe, "enable_attention_slicing"): pipe.enable_attention_slicing()
            if hasattr(pipe, "enable_vae_slicing"): pipe.enable_vae_slicing()
            if hasattr(pipe, "enable_model_cpu_offload"): pipe.enable_model_cpu_offload()
            return pipe, repo, family
        except Exception as e:
            print(f"[load_inpaint] fail {repo}: {e}")
            last = e
    raise RuntimeError(f"Не удалось загрузить ни одну модель: {last}")

PIPE, MODEL_ID, MODEL_FAMILY = load_inpaint()

LATENT_UP: Optional[StableDiffusionLatentUpscalePipeline] = None
if HAS_LATENT_UP:
    try:
        LATENT_UP = StableDiffusionLatentUpscalePipeline.from_pretrained(
            "stabilityai/sd-x2-latent-upscaler", torch_dtype=DTYPE, use_safetensors=True
        ).to(DEVICE)
        if hasattr(LATENT_UP, "enable_attention_slicing"): LATENT_UP.enable_attention_slicing()
        if hasattr(LATENT_UP, "enable_vae_slicing"): LATENT_UP.enable_vae_slicing()
        if hasattr(LATENT_UP, "enable_model_cpu_offload"): LATENT_UP.enable_model_cpu_offload()
        print("Latent upscaler loaded.")
    except Exception as e:
        LATENT_UP = None
        print("Latent upscaler not available:", e)

# ---------- утилиты ----------

def compute_target_size(
    w: int, h: int, L: int, R: int, T: int, B: int, force_square: bool
) -> Tuple[int, int, Tuple[int, int, int, int]]:
    W = w + L + R
    H = h + T + B
    if force_square:
        side = max(W, H)
        add_w = side - W
        add_h = side - H
        L += add_w // 2
        R += add_w - add_w // 2
        T += add_h // 2
        B += add_h - add_h // 2
        W = side
        H = side
    return W, H, (L, R, T, B)

def smart_canvas(img: Image.Image, L: int, R: int, T: int, B: int, blur_px: int = 24) -> tuple[Image.Image, tuple[int,int,int,int]]:
    w, h = img.size
    W, H = w + L + R, h + T + B
    arr = np.array(img)
    arr_pad = np.pad(arr, ((T, B), (L, R), (0, 0)), mode="edge")
    edge_img = Image.fromarray(arr_pad)
    blur_base = img.resize((W, H), Image.BICUBIC).filter(ImageFilter.GaussianBlur(blur_px))
    bg = Image.blend(blur_base, edge_img, alpha=0.5)
    bg.paste(img, (L, T))
    bbox = (L, T, L + w, T + h)  # (left, top, right, bottom) в координатах канваса
    return bg, bbox

def build_mask(size: Tuple[int,int], bbox: Tuple[int,int,int,int], feather_px: int) -> Image.Image:
    W, H = size
    L, T, R, B = bbox
    m = Image.new("L", (W, H), 255)  # белое = редактировать
    m.paste(0, (L, T, R, B))         # чёрный центр = защищаем
    if feather_px > 0:
        m = m.filter(ImageFilter.GaussianBlur(radius=feather_px))
    return m

def resize_to_gen_side(im: Image.Image, gen_side: int) -> Tuple[Image.Image, float]:
    if gen_side <= 0:
        return im, 1.0
    w, h = im.size
    m = max(w, h)
    if m <= gen_side:
        return im, 1.0
    s = gen_side / m
    return im.resize((int(round(w*s)), int(round(h*s))), Image.LANCZOS), s

def gaussian_alpha_mask(size: Tuple[int,int], rect: Tuple[int,int,int,int], feather: int) -> Image.Image:
    """Альфа такого же размера, как итог: внутри прямоугольника 255, края 0, с размытием."""
    W, H = size
    L, T, R, B = rect
    a = Image.new("L", (W, H), 0)
    a.paste(255, (L, T, R, B))
    if feather > 0:
        a = a.filter(ImageFilter.GaussianBlur(radius=feather))
    return a

def paste_center_preserving(original: Image.Image, outpaint_upscaled: Image.Image, L: int, T: int, feather: int) -> Image.Image:
    """
    Складываем:
      base = outpaint_upscaled (RGB)
      overlay = пустой RGB того же размера, куда пастим original в (L,T)
      alpha = маска того же размера, 255 в bbox центра (с мягким краем)
    Результат = composite(overlay, base, alpha) — где 255 берём overlay (оригинал), 0 — base.
    """
    base = outpaint_upscaled.convert("RGB")
    Wf, Hf = base.size
    w0, h0 = original.size

    overlay = Image.new("RGB", (Wf, Hf), (0, 0, 0))
    overlay.paste(original, (L, T))

    alpha = gaussian_alpha_mask((Wf, Hf), (L, T, L + w0, T + h0), feather=max(1, feather // 2))
    final = Image.composite(overlay, base, alpha)
    return final

# ---------- основной процесс ----------

def run_outpaint(
    image: Image.Image,
    left: int, right: int, top: int, bottom: int,
    force_square: bool,
    prompt: str, negative: str,
    steps: int, guidance: float, strength: float, feather_px: int,
    seed: int,
    gen_side: int,    # размер для генерации (обычно 1024)
):
    if image is None:
        return None, "", "", None, {"error":"no input"}, None

    img = image.convert("RGB")
    w, h = img.size

    # 1) целевой размер
    Wt, Ht, (L, R, T, B) = compute_target_size(w, h, int(left), int(right), int(top), int(bottom), bool(force_square))

    # 2) канвас и маска в целевом размере
    canvas_full, bbox_full = smart_canvas(img, L, R, T, B, blur_px=24)
    mask_full = build_mask(canvas_full.size, bbox_full, feather_px=int(feather_px))

    # 3) уменьшаем канвас и маску до gen_side одной и той же шкалой
    canvas_gen, scale = resize_to_gen_side(canvas_full, int(gen_side))
    new_w, new_h = canvas_gen.size
    mask_gen = mask_full.resize((new_w, new_h), Image.LANCZOS)

    # 4) инференс
    p = (prompt or "extend the image naturally").strip()
    n = (negative or "frames, borders, pillars, poster, collage tiles, mockup, UI elements, captions, text, watermark").strip()

    g = None
    if isinstance(seed, int) and seed >= 0:
        g = torch.Generator(device=DEVICE).manual_seed(int(seed))

    out_small = PIPE(
        prompt=p,
        negative_prompt=n,
        image=canvas_gen,
        mask_image=mask_gen,
        num_inference_steps=int(steps),
        guidance_scale=float(guidance),
        strength=float(strength),
        inpaint_full_res=True,
        inpaint_full_res_padding=64,
        generator=g,
    ).images[0]

    # 5) апскейл к целевому размеру
    if LATENT_UP is not None and out_small.size[0] * 2 <= max(Wt, Ht) + 64:
        try:
            up_img = LATENT_UP(prompt=p, image=out_small).images[0]
            if up_img.size != (Wt, Ht):
                up_img = up_img.resize((Wt, Ht), Image.LANCZOS)
        except Exception as e:
            print("Latent upscaler failed, fallback to LANCZOS:", e)
            up_img = out_small.resize((Wt, Ht), Image.LANCZOS)
    else:
        up_img = out_small.resize((Wt, Ht), Image.LANCZOS)

    # 6) возвращаем оригинальный центр без потерь
    final = paste_center_preserving(original=img, outpaint_upscaled=up_img, L=L, T=T, feather=feather_px)

    # 7) сохраняем полный PNG
    ts = int(time.time())
    out_path = f"/tmp/outpaint_{Wt}x{Ht}_{ts}.png"
    final.save(out_path, format="PNG")

    meta = {
        "loaded_model": MODEL_ID,
        "family": MODEL_FAMILY,
        "device": DEVICE,
        "target_size": [Wt, Ht],
        "gen_side_used": int(gen_side),
        "pads_px": {"left": L, "right": R, "top": T, "bottom": B},
        "note": "Скачивай через кнопку File — это полный PNG, не превью.",
    }

    return final, p, n, mask_full, meta, out_path

# ---------- интерфейс ----------
with gr.Blocks(title="Seamless Outpainting — hires PNG, preserved center") as demo:
    gr.Markdown(
        f"## Seamless Outpainting (hires)\n"
        f"**Model:** `{MODEL_ID}` • **Device:** `{DEVICE}`\n\n"
        "— Итоговый PNG = точное целевое разрешение.\n"
        "— Генерация на контролируемой стороне (по умолчанию 1024), затем апскейл.\n"
        "— Центр исходника возвращается пиксель-в-пиксель с мягким швом.\n"
        "— Дефолтный промт: **extend the image naturally**."
    )

    with gr.Row():
        with gr.Column(scale=1):
            img_in = gr.Image(type="pil", label="Исходная картинка")
            force_square = gr.Checkbox(value=True, label="Force square (сделать результат квадратным)")

            with gr.Row():
                left  = gr.Number(value=256, label="Left (px)", precision=0)
                right = gr.Number(value=256, label="Right (px)", precision=0)
            with gr.Row():
                top   = gr.Number(value=256, label="Top (px)", precision=0)
                bottom= gr.Number(value=256, label="Bottom (px)", precision=0)

            prompt = gr.Textbox(label="Prompt", value="extend the image naturally")
            negative = gr.Textbox(
                label="Negative prompt",
                value="frames, borders, pillars, poster, collage tiles, mockup, UI elements, captions, text, watermark",
            )

            with gr.Row():
                steps    = gr.Slider(10, 60, value=30, step=1,   label="Steps")
                guidance = gr.Slider(0.5, 12, value=5.0, step=0.1, label="Guidance")
                strength = gr.Slider(0.3, 1.0, value=0.7, step=0.05, label="Strength")
            with gr.Row():
                feather  = gr.Slider(0, 160, value=96, step=2, label="Feather (px, шов)")
                seed     = gr.Number(value=-1, label="Seed (-1=random)", precision=0)
                gen_side = gr.Number(value=1024, label="Generation side (px)", precision=0)

            run_btn = gr.Button("Expand", variant="primary")

        with gr.Column(scale=1):
            img_out   = gr.Image(label="Результат (превью = финал)", format="png")
            used_p    = gr.Textbox(label="Использованный prompt", interactive=False)
            used_n    = gr.Textbox(label="Использованный negative", interactive=False)
            dbg_mask  = gr.Image(label="Маска (отладка)")
            meta      = gr.JSON(label="Инфо")
            download  = gr.File(label="Скачать полный PNG")

    run_btn.click(
        run_outpaint,
        inputs=[
            img_in, left, right, top, bottom, force_square,
            prompt, negative, steps, guidance, strength, feather,
            seed, gen_side
        ],
        outputs=[img_out, used_p, used_n, dbg_mask, meta, download],
    )

if __name__ == "__main__":
    demo.launch()