app.py

# app.py
import os, time, random
from typing import Tuple
import gradio as gr
import numpy as np
from PIL import Image, ImageFilter, ImageDraw

import torch
from diffusers import StableDiffusionInpaintPipeline, DPMSolverMultistepScheduler

# ------------------------------
# Config & model
# ------------------------------
MODEL_ID = os.getenv("MODEL_ID", "stabilityai/stable-diffusion-2-inpainting")
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32

OUT_DIR = "/mnt/data"                     # оставляем /mnt/data
os.makedirs(OUT_DIR, exist_ok=True)

print(f"Loading model: {MODEL_ID}  (device={DEVICE}, dtype={DTYPE})")
pipe = StableDiffusionInpaintPipeline.from_pretrained(
    MODEL_ID,
    torch_dtype=DTYPE,
    safety_checker=None,                  # для локального теста
).to(DEVICE)
# Более «послушный» шагатель, лучше держит NEG и фон
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
pipe.enable_attention_slicing()
if DEVICE == "cuda":
    try:
        pipe.enable_model_cpu_offload()
    except Exception:
        pass

# ------------------------------
# Helpers
# ------------------------------
def round_to_eight(x: int) -> int:
    return int(max(64, (x // 8) * 8))

def make_square_canvas(img: Image.Image, pad_px: int, bg=(200,200,200)) -> Tuple[Image.Image, Tuple[int,int]]:
    w, h = img.size
    base = max(w, h)
    W = base + 2*pad_px
    H = base + 2*pad_px
    canvas = Image.new("RGB", (W, H), bg)
    x0 = (W - w)//2
    y0 = (H - h)//2
    canvas.paste(img, (x0, y0))
    return canvas, (x0, y0)

def make_ring_mask(canvas_size: Tuple[int,int], inner_rect: Tuple[int,int,int,int], feather_px:int) -> Image.Image:
    W, H = canvas_size
    x0,y0,x1,y1 = inner_rect
    mask = Image.new("L", (W,H), 255)
    draw = ImageDraw.Draw(mask)
    draw.rectangle([x0,y0,x1,y1], fill=0)
    if feather_px > 0:
        mask = mask.filter(ImageFilter.GaussianBlur(radius=feather_px))
    return mask

def paste_original_back(generated: Image.Image, original: Image.Image, offset: Tuple[int,int]) -> Image.Image:
    out = generated.copy()
    out.paste(original, offset)
    return out

def preview_frame(image: Image.Image, pad_px:int):
    w,h = image.size
    base = max(w,h)
    W = base + 2*pad_px
    H = base + 2*pad_px
    canvas, offset = make_square_canvas(image, pad_px)
    draw = ImageDraw.Draw(canvas)
    x0,y0 = offset
    x1,y1 = x0 + w, y0 + h
    draw.rectangle([x0-1,y0-1,x1+1,y1+1], outline=(255, 90, 90), width=3)
    prev = canvas.copy()
    prev.thumbnail((512, 512), Image.LANCZOS)
    top = y0
    left = x0
    right = W - (x0 + w)
    bottom = H - (y0 + h)
    info = f"Final canvas: {W}×{H}px  •  pad: {pad_px}px  •  add: top {top}px, bottom {bottom}px, left {left}px, right {right}px"
    return prev, info

# ------------------------------
# Generation
# ------------------------------
def outpaint_generate(
    input_image: Image.Image,
    pad_px: int,
    prompt: str,
    negative_prompt: str,
    steps: int,
    cfg: float,
    feather_px: int,
    seed: int,
):
    if input_image is None:
        raise gr.Error("Сначала загрузите изображение.")

    canvas, offset = make_square_canvas(input_image.convert("RGB"), pad_px)
    w, h = input_image.size
    x0, y0 = offset
    x1, y1 = x0 + w, y0 + h

    mask = make_ring_mask(canvas.size, (x0,y0,x1,y1), feather_px)

    W = round_to_eight(canvas.size[0])
    H = round_to_eight(canvas.size[1])
    MAX_SIDE = 1536 if DEVICE == "cuda" else 1024
    scale = min(1.0, MAX_SIDE / max(W,H))
    if scale < 1.0:
        newW = round_to_eight(int(W*scale))
        newH = round_to_eight(int(H*scale))
        canvas_small = canvas.resize((newW,newH), Image.LANCZOS)
        mask_small = mask.resize((newW,newH), Image.LANCZOS)
        sx = int(x0*scale); sy = int(y0*scale)
        sw = int(w*scale);   sh = int(h*scale)
        inner_rect_small = (sx,sy,sx+sw,sy+sh)
    else:
        canvas_small, mask_small = canvas, mask
        inner_rect_small = (x0,y0,x1,y1)

    g = torch.Generator(device=DEVICE)
    seed_val = random.randint(0, 2**32 - 1) if (seed is None or int(seed) < 0) else int(seed)
    g.manual_seed(seed_val)

    with torch.autocast(device_type=DEVICE if DEVICE!="mps" else "cpu"):
        out = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            image=canvas_small,
            mask_image=mask_small,
            guidance_scale=float(cfg),
            num_inference_steps=int(steps),
            generator=g,
        ).images[0]

    if out.size != canvas.size:
        out = out.resize(canvas.size, Image.LANCZOS)

    final = paste_original_back(out, input_image.convert("RGB"), offset)

    fname = f"outpaint_{canvas.size[0]}x{canvas.size[1]}_{int(time.time()*1000)}.png"
    out_path = os.path.join(OUT_DIR, fname)
    os.makedirs(os.path.dirname(out_path), exist_ok=True)
    final.save(out_path, "PNG")

    return final, out_path, f"Seed: {seed_val} • Size: {canvas.size[0]}×{canvas.size[1]}"

# ------------------------------
# Gradio UI
# ------------------------------
DEFAULT_PROMPT = "extend the image naturally, seamless realistic background, consistent lighting, matching style"
DEFAULT_NEG = (
    "text, letters, words, caption, typography, logo, watermark, "
    "numbers, digits, signboard, poster text, "
    "lowres, blurry, artifacts, deformed, distorted, oversaturated, "
    "frame, border, mosaic, collage, extra limbs"
)

with gr.Blocks(css="""
#mini {font-size: 0.9em; opacity: 0.9}
.caption {font-size: 0.9em; color: #aaa}
""") as demo:
    gr.Markdown("## Qwen Outpaint (SD2 Inpaint)\nКвадратная дорисовка краёв. Центр сохраняется 1:1. PNG в полном размере.")

    with gr.Row():
        with gr.Column(scale=6):
            in_img = gr.Image(type="pil", label="Input image", height=560)
            pad = gr.Slider(0, 2048, value=256, step=1, label="Padding (px) around square")
            feather = gr.Slider(0, 64, value=16, step=1, label="Feather border (px)")
            prmpt = gr.Textbox(value=DEFAULT_PROMPT, label="Prompt")
            nprmpt = gr.Textbox(value=DEFAULT_NEG, label="Negative prompt")
            with gr.Row():
                steps = gr.Slider(10, 60, value=30, step=1, label="Steps")
                cfg = gr.Slider(1.0, 9.0, value=6.5, step=0.5, label="CFG (guidance)")
                seed = gr.Number(value=-1, precision=0, label="Seed (-1 = random)")
            go_btn = gr.Button("Generate", variant="primary")

        with gr.Column(scale=6):
            prev = gr.Image(label="Preview (outpaint region)", height=560)
            info = gr.Markdown(elem_id="mini")
            with gr.Tab("Result"):
                out_img = gr.Image(label="Result", height=560)
                meta = gr.Markdown("")
                file_out = gr.File(label="Download PNG")

    def _update_preview(img, pad_px):
        if img is None:
            return None, ""
        p, t = preview_frame(img, int(pad_px))
        return p, t

    in_img.change(_update_preview, [in_img, pad], [prev, info])
    pad.release(_update_preview, [in_img, pad], [prev, info])

    def go(image, pad_px, feather_px, prompt, negative_prompt, steps, cfg, seed):
        if image is None:
            raise gr.Error("Загрузите изображение.")
        res, path, meta_text = outpaint_generate(
            image, int(pad_px),
            prompt, negative_prompt,
            int(steps), float(cfg), int(feather_px), int(seed)
        )
        return res, meta_text, path

    go_btn.click(
        go,
        [in_img, pad, feather, prmpt, nprmpt, steps, cfg, seed],
        [out_img, meta, file_out]
    )

if __name__ == "__main__":
    # КРИТИЧЕСКОЕ ИСПРАВЛЕНИЕ: разрешаем /mnt/data
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        inbrowser=False,
        allowed_paths=["/mnt/data"]
    )