# 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)