app.py

# app.py
import os, io, math, random, time
from typing import Tuple
import numpy as np
from PIL import Image, ImageFilter, ImageOps
import torch
from diffusers import StableDiffusionXLInpaintPipeline
import gradio as gr

# -----------------------------
# Settings
# -----------------------------
DEFAULT_PROMPT = "extend the image naturally"
DEFAULT_NEG = (
    "text, letters, words, numbers, caption, watermark, logo, signature, frame, border, "
    "collage tiles, UI elements, artifacts, deformed, duplicate, blurry"
)
SAFE_MAX = int(os.getenv("SAFE_MAX", 1792))   # макс. сторона генерации
HF_TOKEN = os.getenv("HF_TOKEN", None)

# -----------------------------
# Load pipeline (no xformers)
# -----------------------------
device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = StableDiffusionXLInpaintPipeline.from_pretrained(
    "diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
    use_safetensors=True,
    variant="fp16" if device == "cuda" else None,
    token=HF_TOKEN,
)
pipe = pipe.to(device)

# включаем SDPA/attention slicing вместо xformers
try:
    if device == "cuda":
        torch.backends.cuda.matmul.allow_tf32 = True
    else:
        pipe.enable_attention_slicing("max")
except Exception:
    pass


# -----------------------------
# Helpers
# -----------------------------
def _to_pil(img) -> Image.Image:
    return img if isinstance(img, Image.Image) else Image.fromarray(img)

def _round64(x: int) -> int:
    return int(math.ceil(x / 64) * 64)

def fit_to_safe_max(w: int, h: int, safe_max: int) -> Tuple[int, int, float]:
    """масштабирует (w,h) до safe_max, сохраняя пропорции; возвращает новые w,h и scale"""
    s = 1.0
    max_side = max(w, h)
    if max_side > safe_max:
        s = safe_max / max_side
        w = int(round(w * s))
        h = int(round(h * s))
    # SDXL требует кратность 64
    w = max(512, _round64(w))
    h = max(512, _round64(h))
    return w, h, s

def make_square_canvas(img: Image.Image, L: int, R: int, T: int, B: int) -> Tuple[int, int, Tuple[int,int,int,int]]:
    """Считает квадратную сторону и box для вставки оригинала"""
    w, h = img.size
    tgt_w = w + L + R
    tgt_h = h + T + B
    side = max(tgt_w, tgt_h)
    # центровка оригинала
    x0 = (side - w) // 2
    y0 = (side - h) // 2
    return side, side, (x0, y0, x0 + w, y0 + h)

def build_bg_from_edges(img: Image.Image, side: int) -> Image.Image:
    """Создаёт «умный» фон из повторённых краёв, слегка блюрит."""
    arr = np.asarray(img.convert("RGB"))
    pad_top = (side - img.height) // 2
    pad_bottom = side - img.height - pad_top
    pad_left = (side - img.width) // 2
    pad_right = side - img.width - pad_left
    padded = np.pad(arr, ((pad_top, pad_bottom), (pad_left, pad_right), (0, 0)), mode="edge")
    bg = Image.fromarray(padded, mode="RGB").filter(ImageFilter.GaussianBlur(radius=8))
    return bg

def build_mask(side: int, orig_box: Tuple[int,int,int,int], feather: int) -> Image.Image:
    """Белое = генерировать, чёрное = сохраняем оригинал"""
    mask = Image.new("L", (side, side), 255)
    x0, y0, x1, y1 = orig_box
    # чёрный прямоугольник по оригиналу
    ImageDraw = ImageDraw_ensure()
    d = ImageDraw.Draw(mask)
    d.rectangle([x0, y0, x1, y1], fill=0)
    if feather > 0:
        mask = mask.filter(ImageFilter.GaussianBlur(radius=feather))
    return mask

def ImageDraw_ensure():
    from PIL import ImageDraw
    return ImageDraw

def feather_alpha_mask(size: Tuple[int,int], feather: int) -> Image.Image:
    """Белый центр, мягкий край (для вставки оригинала поверх результата)"""
    w, h = size
    m = Image.new("L", (w, h), 255)
    if feather <= 0:
        return m
    # делаем тёмную рамку и размазываем
    d = ImageDraw_ensure().Draw(m)
    d.rectangle([0,0,w-1,h-1], outline=0, width=feather*2)
    m = m.filter(ImageFilter.GaussianBlur(radius=feather))
    # нормализация – центр белее
    return m

def paste_center_preserving(base: Image.Image, original: Image.Image, box: Tuple[int,int,int,int], feather: int) -> Image.Image:
    """Вставляет оригинал по box с мягким пером по краям."""
    x0, y0, x1, y1 = box
    orig_rgba = original.convert("RGBA")
    alpha = feather_alpha_mask(original.size, feather)
    orig_rgba.putalpha(alpha)
    base = base.convert("RGBA")
    base.alpha_composite(orig_rgba, dest=(x0, y0))
    return base.convert("RGB")

def outpaint_once(
    img: Image.Image,
    L: int, R: int, T: int, B: int,
    prompt: str, neg: str,
    steps: int, cfg: float, seed: int,
    feather: int, safe_max: int
) -> Tuple[Image.Image, Image.Image, Tuple[int,int,int,int], int]:
    """Генерит квадратное полотно, возвращает (outpaint, original, orig_box, used_seed)"""
    img = _to_pil(img).convert("RGB")
    # квадрат и позиция оригинала
    W, H, orig_box = make_square_canvas(img, L, R, T, B)
    # размер генерации с ограничением
    gen_w, gen_h, _ = fit_to_safe_max(W, H, safe_max)
    side = max(gen_w, gen_h)
    # фон и маска
    bg = build_bg_from_edges(img, side)
    # пересчитываем orig_box под side (если side!=W/H, центруем)
    cx, cy = side//2, side//2
    x0 = cx - img.width//2
    y0 = cy - img.height//2
    orig_box_side = (x0, y0, x0 + img.width, y0 + img.height)
    # кладём оригинал на фон (как исходник для inpaint)
    cond = bg.copy()
    cond.paste(img, (x0, y0))
    mask = build_mask(side, orig_box_side, feather=max(1, feather//2))
    # генерация
    if seed is None or seed < 0:
        seed = random.randint(0, 2**31-1)
    g = torch.Generator(device=device)
    if device == "cuda":
        g = g.manual_seed(seed)
    else:
        random.seed(seed)
    out = pipe(
        prompt=prompt if prompt.strip() else DEFAULT_PROMPT,
        negative_prompt=(neg.strip() or DEFAULT_NEG),
        image=cond,
        mask_image=mask,
        num_inference_steps=int(steps),
        guidance_scale=float(cfg),
        generator=g,
        width=side,
        height=side,
    ).images[0]
    return out, img, orig_box_side, seed

# -----------------------------
# Gradio core
# -----------------------------
def run(
    img, prompt, neg, L, R, T, B,
    steps, cfg, seed, feather, safe_max
):
    if img is None:
        raise gr.Error("Загрузи изображение.")
    out, orig, box, used_seed = outpaint_once(
        img=img, L=L, R=R, T=T, B=B,
        prompt=prompt or DEFAULT_PROMPT,
        neg=neg or DEFAULT_NEG,
        steps=steps, cfg=cfg, seed=seed,
        feather=feather, safe_max=safe_max
    )
    # поверх возвращаем оригинал (чтобы его качество не менялось)
    final = paste_center_preserving(out, orig, box, feather=max(8, feather))
    # сохраняем полный PNG
    side = final.size[0]
    ts = int(time.time())
    fname = f"/tmp/outpaint_{side}x{side}_{used_seed}.png"
    final.save(fname, "PNG")
    return final, fname, used_seed

with gr.Blocks(css="""
#note {opacity:.8}
""") as demo:
    gr.Markdown("## Qwen Image — Seamless Background Expand")
    with gr.Row():
        with gr.Column():
            in_img = gr.Image(label="Картинка", type="pil")
            prompt = gr.Textbox(label="Prompt", value=DEFAULT_PROMPT, lines=2)
            neg = gr.Textbox(label="Negative prompt", value=DEFAULT_NEG, lines=2)
            with gr.Row():
                L = gr.Number(label="Left (px)", value=256, precision=0)
                R = gr.Number(label="Right (px)", value=256, precision=0)
            with gr.Row():
                T = gr.Number(label="Top (px)", value=256, precision=0)
                B = gr.Number(label="Bottom (px)", value=256, precision=0)
            with gr.Row():
                steps = gr.Slider(10, 60, value=28, step=1, label="Steps")
                cfg = gr.Slider(1.0, 12.0, value=5.5, step=0.1, label="Guidance scale")
            with gr.Row():
                seed = gr.Number(label="Seed (-1=random)", value=-1, precision=0)
                feather = gr.Slider(0, 64, value=24, step=1, label="Feather (px)")
            safe_max = gr.Slider(1024, 2048, value=1792, step=64, label="Max side for generation (px)")
            run_btn = gr.Button("Expand", variant="primary")
        with gr.Column():
            out_img = gr.Image(label="Результат", interactive=False)
            file_out = gr.File(label="Скачать PNG (полный размер)")
            used_seed = gr.Number(label="Used seed", interactive=False)

    run_btn.click(
        fn=run,
        inputs=[in_img, prompt, neg, L, R, T, B, steps, cfg, seed, feather, safe_max],
        outputs=[out_img, file_out, used_seed],
    )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)