Update app.py

app.py

@@ -1,223 +1,233 @@
 # app.py
-import os, time, random
+import os, io, math, random, time
 from typing import Tuple
-import gradio as gr
 import numpy as np
-from PIL import Image, ImageFilter, ImageDraw
-
+from PIL import Image, ImageFilter, ImageOps
 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))
+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
+
 
-def make_square_canvas(img: Image.Image, pad_px: int, bg=(200,200,200)) -> Tuple[Image.Image, Tuple[int,int]]:
+# -----------------------------
+# 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
-    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))
+    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 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)
+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:
-        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"
-)
+        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="""
-#mini {font-size: 0.9em; opacity: 0.9}
-.caption {font-size: 0.9em; color: #aaa}
+#note {opacity:.8}
 """) as demo:
-    gr.Markdown("## Qwen Outpaint (SD2 Inpaint)\nКвадратная дорисовка краёв. Центр сохраняется 1:1. PNG в полном размере.")
-
+    gr.Markdown("## Qwen Image — Seamless Background Expand")
     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.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():
-                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]
+                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__":
-    # КРИТИЧЕСКОЕ ИСПРАВЛЕНИЕ: разрешаем /mnt/data
-    demo.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        inbrowser=False,
-        allowed_paths=["/mnt/data"]
-    )
+    demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)