Update app.py

app.py

@@ -1,241 +1,251 @@
 # app.py
 import os
 import io
-import math
+import time
 import random
-from dataclasses import dataclass
-from typing import Tuple
-
-import gradio as gr
-import numpy as np
-from PIL import Image, ImageOps, ImageFilter
+from typing import Tuple, Optional
 
 import torch
-from diffusers import StableDiffusionInpaintPipeline
+from PIL import Image, ImageDraw, ImageFile, ImageFilter
+ImageFile.LOAD_TRUNCATED_IMAGES = True
 
-# =========================
+import gradio as gr
+from diffusers import StableDiffusionXLInpaintPipeline
+
+# -------------------------
 # Config
-# =========================
-MODEL_ID = os.getenv("MODEL_ID", "runwayml/stable-diffusion-inpainting")  # SD 1.5 Inpainting
-HF_TOKEN = os.getenv("HF_TOKEN", None)
+# -------------------------
+MODEL_ID = os.getenv("MODEL_ID", "stabilityai/stable-diffusion-xl-1.0-inpainting-0.1")
+HF_TOKEN = os.getenv("HF_TOKEN", None)  # добавь в Settings, если репо gated
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32
 
-DEFAULT_PROMPT = "extend the image naturally"
+DEFAULT_PROMPT = "extend the image naturally, seamless realistic continuation, preserve original center"
 DEFAULT_NEG = (
-    "frames, borders, phone mockup, collage tiles, UI elements, captions, watermark, "
-    "logo, text, letters, words, signatures, artifacts, duplicate objects, extra limbs"
+    "text, watermark, signature, logo, lowres, blurry, artifacts, deformed, distorted, "
+    "oversaturated, extra limbs, frames, borders, bad perspective"
 )
 
-# Генерим при необходимости в даунскейле (для скорости/VRAM)
-SAFE_MAX = int(os.getenv("SAFE_MAX", "1024"))  # макс. сторона для генерации
-FEATHER_PX = 24  # мягкое склеивание по границе оригинала в финальном композе
-
-# =========================
-# Helpers
-# =========================
-def to_pil(img) -> Image.Image:
-    if isinstance(img, Image.Image):
-        return img.convert("RGB")
-    return Image.fromarray(img).convert("RGB")
-
-def make_square_with_border(im: Image.Image, border_px: int) -> Tuple[Image.Image, Image.Image]:
-    """Делаем квадратный холст и маску (белым — где дорисовывать)."""
-    w, h = im.size
-    side = max(w, h)
-    side = side + border_px * 2
-    canvas = Image.new("RGB", (side, side), (128, 128, 128))
-
-    # Центрируем оригинал
-    ox = (side - w)//2
-    oy = (side - h)//2
-    canvas.paste(im, (ox, oy))
-
-    # Маска: чёрный (0) — оригинал, белый (255) — где рисуем
-    mask = Image.new("L", (side, side), color=255)
-    mask.paste(0, (ox, oy, ox + w, oy + h))
-    return canvas, mask
-
-def fit_to_safe_max(w: int, h: int, safe_max: int) -> Tuple[int, int, float]:
-    """Укладываемся в SAFE_MAX, возвращаем целевые размеры и масштаб."""
-    scale = 1.0
-    if max(w, h) > safe_max:
-        scale = safe_max / float(max(w, h))
-        w = int(round(w * scale / 8.0)) * 8
-        h = int(round(h * scale / 8.0)) * 8
-    else:
-        # SD любит кратность 8
-        w = int(round(w / 8.0)) * 8
-        h = int(round(h / 8.0)) * 8
-    return w, h, scale
-
-def upscale_back(img: Image.Image, scale: float, target_wh: Tuple[int, int]) -> Image.Image:
-    """Возвращаемся к исходному (большему) размеру после генерации в сниж. разрешении."""
-    if scale == 1.0:
-        return img
-    tw, th = target_wh
-    return img.resize((tw, th), Image.LANCZOS)
-
-def paste_center_preserving(original: Image.Image, outpaint: Image.Image, feather: int = FEATHER_PX) -> Image.Image:
-    """Кладём исходник поверх результата (чтобы он остался чётким), мягко смешивая границы."""
-    w, h = outpaint.size
-    ow, oh = original.size
-    ox = (w - ow)//2
-    oy = (h - oh)//2
-
-    # делаем маску с мягкими краями
-    alpha = Image.new("L", (ow, oh), color=255)
-    if feather > 0:
-        # равномерное размывание по краю
-        alpha = ImageOps.expand(alpha, border=feather, fill=0)
-        alpha = alpha.filter(ImageFilter.GaussianBlur(radius=feather/2))
-        alpha = alpha.crop((feather, feather, feather + ow, feather + oh))
-
-    base = outpaint.convert("RGBA")
-    top = original.convert("RGBA")
-    top.putalpha(alpha)
-    base.alpha_composite(top, (ox, oy))
-    return base.convert("RGB")
-
-@dataclass
-class GenParams:
-    prompt: str
-    negative_prompt: str
-    steps: int
-    guidance: float
-    seed: int
-
-# =========================
-# Pipeline
-# =========================
-device = "cuda" if torch.cuda.is_available() else "cpu"
-dtype = torch.float16 if (device == "cuda") else torch.float32
-
-pipe = StableDiffusionInpaintPipeline.from_pretrained(
+# -------------------------
+# Pipeline (SDXL Inpaint)
+# -------------------------
+pipe = StableDiffusionXLInpaintPipeline.from_pretrained(
     MODEL_ID,
-    torch_dtype=dtype,
-    use_auth_token=HF_TOKEN
+    torch_dtype=DTYPE,
+    use_safetensors=True,
+    token=HF_TOKEN,        # если None — ок; если задана — пройдёт с токеном
 )
-pipe = pipe.to(device)
-pipe.safety_checker = None  # без блюра
-pipe.enable_attention_slicing("auto")
+pipe = pipe.to(DEVICE)
 
-# =========================
-# Core generate
-# =========================
-def outpaint_generate(
+# Без xformers
+# pipe.enable_xformers_memory_efficient_attention()  # НЕ ВКЛЮЧАЕМ
+
+
+# -------------------------
+# Helpers
+# -------------------------
+def center_paste_coords(
+    base_w: int, base_h: int, obj_w: int, obj_h: int
+) -> Tuple[int, int]:
+    """Вычисляет координаты для центрирования."""
+    return ( (base_w - obj_w) // 2, (base_h - obj_h) // 2 )
+
+def make_square_with_border(
     img: Image.Image,
     border_px: int,
-    params: GenParams
-):
-    img = to_pil(img)
-    # 1) расширяем до квадрата и строим маску
-    canvas, mask = make_square_with_border(img, border_px)
-
-    # 2) Поджимаем под SAFE_MAX для генерации
-    W0, H0 = canvas.size
-    gw, gh, scale = fit_to_safe_max(W0, H0, SAFE_MAX)
-    if (gw, gh) != (W0, H0):
-        canvas_small = canvas.resize((gw, gh), Image.LANCZOS)
-        mask_small = mask.resize((gw, gh), Image.NEAREST)
-    else:
-        canvas_small, mask_small = canvas, mask
-
-    # 3) Сид
-    generator = None
-    if params.seed is not None and params.seed != -1:
-        generator = torch.Generator(device=device).manual_seed(int(params.seed))
-
-    # 4) Генерация
-    out = pipe(
-        prompt=params.prompt or DEFAULT_PROMPT,
-        negative_prompt=params.negative_prompt or DEFAULT_NEG,
-        image=canvas_small,
-        mask_image=mask_small,
-        num_inference_steps=params.steps,
-        guidance_scale=params.guidance,
-        generator=generator,
+    bg=(128, 128, 128)
+) -> Tuple[Image.Image, Image.Image, Tuple[int, int, int, int]]:
+    """
+    Создаёт квадратный холст (max(w,h)+2*border), кладёт исходник по центру.
+    Возвращает:
+      padded_image (RGB), mask (L, 0 центр/255 поля), bbox исходника (x,y,w,h).
+    """
+    img = img.convert("RGB")
+    w, h = img.size
+    side = max(w, h) + 2 * border_px
+    canvas = Image.new("RGB", (side, side), bg)
+
+    x, y = center_paste_coords(side, side, w, h)
+    canvas.paste(img, (x, y))
+
+    # Маска: чёрный (0) = защитить (центр), белый (255) = дорисовать (поля)
+    mask = Image.new("L", (side, side), 255)
+    draw = ImageDraw.Draw(mask)
+    draw.rectangle([x, y, x + w - 1, y + h - 1], fill=0)
+
+    return canvas, mask, (x, y, w, h)
+
+def feather_mask(mask: Image.Image, feather_px: int) -> Image.Image:
+    """Смягчаем границу маски GaussianBlur-ом для более естественного склеивания."""
+    if feather_px <= 0:
+        return mask
+    return mask.filter(ImageFilter.GaussianBlur(radius=feather_px))
+
+def overlay_original_center(
+    out: Image.Image, orig: Image.Image, bbox: Tuple[int, int, int, int], feather_px: int
+) -> Image.Image:
+    """
+    Поверх результата накладываем оригинал (чтобы центр гарантированно остался без изменений).
+    Делаем мягкую маску по краям вставки, чтобы не было резкого стыка.
+    """
+    out = out.convert("RGB")
+    img = orig.convert("RGB")
+    x, y, w, h = bbox
+
+    # Маска для вставки: белый = виден оригинал, чёрный = виден out; мягкий край
+    m = Image.new("L", (out.width, out.height), 0)
+    draw = ImageDraw.Draw(m)
+    draw.rectangle([x, y, x + w - 1, y + h - 1], fill=255)
+    if feather_px > 0:
+        m = m.filter(ImageFilter.GaussianBlur(radius=feather_px))
+
+    out.paste(img, (x, y), m.crop((x, y, x + w, y + h)))
+    return out
+
+def format_dims(w: int, h: int) -> str:
+    return f"{w}×{h}px"
+
+# -------------------------
+# Core outpaint
+# -------------------------
+@torch.inference_mode()
+def outpaint_generate(
+    image: Image.Image,
+    border_px: int,
+    prompt: str,
+    neg_prompt: str,
+    steps: int,
+    guidance: float,
+    seed: Optional[int],
+    feather_px: int,
+) -> Tuple[Image.Image, str]:
+    """
+    Делает аутпейнтинг с SDXL Inpaint:
+     - готовит квадратный канвас и маску (белые поля = дорисовка),
+     - запускает пайплайн,
+     - мягко вставляет исходник по центру поверх результата,
+     - сохраняет PNG полноразмерный.
+    """
+    if not prompt or not prompt.strip():
+        prompt = DEFAULT_PROMPT
+    if not neg_prompt or not neg_prompt.strip():
+        neg_prompt = DEFAULT_NEG
+
+    # 1) Квадрат + маска
+    padded, mask, bbox = make_square_with_border(image, border_px, bg=(128, 128, 128))
+    mask = feather_mask(mask, max(0, feather_px))
+
+    # 2) Генератор по seed
+    if seed in (None, -1, ""):
+        seed = random.randint(0, 2**31 - 1)
+    g = torch.Generator(device=DEVICE)
+    g.manual_seed(int(seed))
+
+    # 3) SDXL Inpaint
+    # В SDXL-inpaint ничего дополнительно передавать не нужно (никаких added_cond_kwargs).
+    result = pipe(
+        prompt=prompt,
+        negative_prompt=neg_prompt,
+        image=padded,            # PIL RGB
+        mask_image=mask,         # PIL L (255 = дорисовать; 0 = сохраняем)
+        num_inference_steps=int(steps),
+        guidance_scale=float(guidance),
+        generator=g,
     ).images[0]
 
-    # 5) Возвращаемся к исходному квадрату
-    out_full = upscale_back(out, scale, (W0, H0))
-
-    # 6) Вставляем оригинал по центру с мягкой границей
-    final = paste_center_preserving(original=img, outpaint=out_full, feather=FEATHER_PX)
-
-    # 7) Сохраняем PNG без даунскейла
-    side = final.size[0]
-    save_name = f"outpaint_{side}x{side}_{random.randint(0, 2**31-1)}.png"
-    save_path = os.path.join("/home/user/app", save_name)
-    final.save(save_path, "PNG")
-
-    return final, save_path
-
-# =========================
-# Preview (без модели)
-# =========================
-def preview_frame(img, border_px):
-    if img is None:
-        return None
-    im = to_pil(img)
-    w, h = im.size
-    side = max(w, h) + border_px * 2
-    canvas = Image.new("RGB", (side, side), (40, 40, 40))
-    ox = (side - w)//2
-    oy = (side - h)//2
-    canvas.paste(im, (ox, oy))
-    # рисуем тонкую рамку границы оригинала
-    frame = Image.new("RGBA", (w, h), (255, 255, 255, 0))
-    outline = ImageOps.expand(Image.new("L", (w-2, h-2), 0), border=1, fill=255)
-    frame.putalpha(outline)
-    canvas = canvas.convert("RGBA")
-    canvas.alpha_composite(frame, (ox, oy))
-    return canvas.convert("RGB")
-
-# =========================
+    # 4) Поверх — снова центр, чтоб наверняка был 1:1
+    final = overlay_original_center(result, image, bbox, feather_px=max(0, feather_px // 2))
+
+    # 5) Сохранение PNG в полном размере
+    outdir = "outputs"
+    os.makedirs(outdir, exist_ok=True)
+    side = final.width  # квадрат
+    fname = f"outpaint_{side}x{side}_{seed}.png"
+    path = os.path.join(outdir, fname)
+    final.save(path, "PNG")
+
+    return final, path
+
+# -------------------------
+# Preview: рисуем чёрную рамку в том месте, где будет дорисовка
+# -------------------------
+def make_preview(image: Image.Image, border_px: int) -> Tuple[Image.Image, str, str]:
+    if image is None:
+        return None, "", ""
+
+    w, h = image.size
+    side = max(w, h) + 2 * int(border_px)
+    padded = Image.new("RGB", (side, side), (240, 240, 240))
+    x, y = center_paste_coords(side, side, w, h)
+    padded.paste(image.convert("RGB"), (x, y))
+
+    # Рамка места дорисовки
+    draw = ImageDraw.Draw(padded)
+    # Внутренняя окантовка исходника (покажем контур)
+    draw.rectangle([x, y, x + w - 1, y + h - 1], outline=(0, 0, 0), width=2)
+
+    return padded, f"Original: {format_dims(w, h)}", f"Result: {format_dims(side, side)}"
+
+# -------------------------
 # Gradio UI
-# =========================
-with gr.Blocks(fill_height=True) as demo:
-    gr.Markdown("# Qwen Image — Seamless Background Expand\nЗагрузи картинку, укажи, на сколько расширить (вокруг по периметру), и жми **Expand**.\nПо умолчанию промпт: *extend the image naturally*. Оригинал всегда остаётся по центру и сохраняет качество.")
+# -------------------------
+with gr.Blocks(css="""
+#dims {font-size: 14px; opacity: 0.8}
+""") as demo:
+    gr.Markdown("## Qwen Outpaint (SDXL Inpaint)\nЗагрузите изображение, выберите рамку (px) — модель дорисует края естественно. Центр сохраняется без изменений.")
 
     with gr.Row():
-        with gr.Column(scale=1):
-            inp_img = gr.Image(label="Картинка", type="pil")
-            border = gr.Slider(0, 2000, value=256, step=8, label="Border (px) – добавить по краю")
-            with gr.Accordion("Промпты и параметры", open=False):
-                prompt = gr.Textbox(label="Prompt", value=DEFAULT_PROMPT)
-                neg = gr.Textbox(label="Negative prompt", value=DEFAULT_NEG)
-                steps = gr.Slider(10, 60, value=30, step=1, label="Steps")
-                guidance = gr.Slider(1.0, 12.0, value=5.5, step=0.1, label="Guidance scale")
-                seed = gr.Number(value=-1, precision=0, label="Seed (-1 = random)")
-            preview = gr.Image(label="Превью (рамка расширения)", interactive=False)
-
-            # автопревью
-            inp_img.change(preview_frame, [inp_img, border], preview, api_name=False)
-            border.release(preview_frame, [inp_img, border], preview, api_name=False)
-
-        with gr.Column(scale=1):
-            out_img = gr.Image(label="Результат", interactive=False)
-            file_out = gr.File(label="Скачать PNG", interactive=False)
-
-            def go(image, border_px, prompt_txt, neg_txt, steps_v, guidance_v, seed_v):
-                if image is None:
-                    raise gr.Error("Сначала загрузите изображение.")
-                params = GenParams(
-                    prompt=prompt_txt or DEFAULT_PROMPT,
-                    negative_prompt=neg_txt or DEFAULT_NEG,
-                    steps=int(steps_v),
-                    guidance=float(guidance_v),
-                    seed=int(seed_v),
-                )
-                res, path = outpaint_generate(image, int(border_px), params)
-                return res, path
-
-            btn = gr.Button("Expand", variant="primary")
-            btn.click(
-                go,
-                [inp_img, border, prompt, neg, steps, guidance, seed],
-                [out_img, file_out]
-            )
-
-# SSR включён в Spaces; локально можно так:
+        with gr.Column():
+            image_in = gr.Image(label="Input image", type="pil")
+            border = gr.Slider(0, 2048, value=256, step=8, label="Padding (px) around (square)")
+            feather = gr.Slider(0, 64, value=16, step=1, label="Feather border (px)")
+            prompt = gr.Textbox(label="Prompt", value=DEFAULT_PROMPT)
+            neg_prompt = gr.Textbox(label="Negative prompt", value=DEFAULT_NEG)
+            steps = gr.Slider(10, 60, value=30, step=1, label="Steps")
+            guidance = gr.Slider(1.0, 12.0, value=5.5, step=0.1, label="CFG (guidance)")
+            seed = gr.Number(value=-1, precision=0, label="Seed (-1 random)")
+            run_btn = gr.Button("Generate", variant="primary")
+
+            gr.Markdown('<div id="dims"></div>')
+            dims_orig = gr.Markdown()
+            dims_final = gr.Markdown()
+
+        with gr.Column():
+            preview = gr.Image(label="Preview (where outpaint will happen)", interactive=False)
+            result = gr.Image(label="Result", interactive=False)
+            file_out = gr.File(label="Download PNG")
+
+    # Preview update
+    def _update_preview(img, b):
+        if img is None:
+            return None, "", ""
+        p, d1, d2 = make_preview(img, int(b))
+        return p, f"<span id='dims'>{d1}</span>", f"<span id='dims'>{d2}</span>"
+
+    image_in.change(_update_preview, [image_in, border], [preview, dims_orig, dims_final])
+    border.release(_update_preview, [image_in, border], [preview, dims_orig, dims_final])
+
+    # Generate
+    def go(img, b, fthr, prm, nprm, stp, cfg, sd):
+        if img is None:
+            raise gr.Error("Загрузите изображение.")
+        final, path = outpaint_generate(
+            image=img,
+            border_px=int(b),
+            prompt=(prm or DEFAULT_PROMPT),
+            neg_prompt=(nprm or DEFAULT_NEG),
+            steps=int(stp),
+            guidance=float(cfg),
+            seed=int(sd) if sd is not None else -1,
+            feather_px=int(fthr),
+        )
+        return final, path
+
+    run_btn.click(
+        go,
+        [image_in, border, feather, prompt, neg_prompt, steps, guidance, seed],
+        [result, file_out]
+    )
+
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=True)