# app.py import os import time from typing import Optional, Tuple import gradio as gr import numpy as np from PIL import Image, ImageFilter, ImageOps, ImageFile # лечим обрезанные файлы (webp/jpeg) ImageFile.LOAD_TRUNCATED_IMAGES = True import torch from diffusers import ( StableDiffusionXLInpaintPipeline, StableDiffusionInpaintPipeline, ) # Опциональный латентный апскейлер (если доступен) try: from diffusers import StableDiffusionLatentUpscalePipeline HAS_LATENT_UP = True except Exception: HAS_LATENT_UP = False DEVICE = "cuda" if torch.cuda.is_available() else "cpu" DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32 MODEL_PREFS = [ ("sdxl", "diffusers/stable-diffusion-xl-1.0-inpainting-0.1"), ("sd2", "stabilityai/stable-diffusion-2-inpainting"), ("sd15", "runwayml/stable-diffusion-inpainting"), ] def load_inpaint() -> tuple[object, str, str]: last = None for family, repo in MODEL_PREFS: try: if family == "sdxl": pipe = StableDiffusionXLInpaintPipeline.from_pretrained( repo, torch_dtype=DTYPE, use_safetensors=True ) else: pipe = StableDiffusionInpaintPipeline.from_pretrained( repo, torch_dtype=DTYPE, use_safetensors=True ) pipe = pipe.to(DEVICE) if hasattr(pipe, "enable_attention_slicing"): pipe.enable_attention_slicing() if hasattr(pipe, "enable_vae_slicing"): pipe.enable_vae_slicing() if hasattr(pipe, "enable_model_cpu_offload"): pipe.enable_model_cpu_offload() return pipe, repo, family except Exception as e: print(f"[load_inpaint] fail {repo}: {e}") last = e raise RuntimeError(f"Не удалось загрузить ни одну модель: {last}") PIPE, MODEL_ID, MODEL_FAMILY = load_inpaint() LATENT_UP: Optional[StableDiffusionLatentUpscalePipeline] = None if HAS_LATENT_UP: try: LATENT_UP = StableDiffusionLatentUpscalePipeline.from_pretrained( "stabilityai/sd-x2-latent-upscaler", torch_dtype=DTYPE, use_safetensors=True ).to(DEVICE) if hasattr(LATENT_UP, "enable_attention_slicing"): LATENT_UP.enable_attention_slicing() if hasattr(LATENT_UP, "enable_vae_slicing"): LATENT_UP.enable_vae_slicing() if hasattr(LATENT_UP, "enable_model_cpu_offload"): LATENT_UP.enable_model_cpu_offload() print("Latent upscaler loaded.") except Exception as e: LATENT_UP = None print("Latent upscaler not available:", e) # ---------- утилиты ---------- def compute_target_size( w: int, h: int, L: int, R: int, T: int, B: int, force_square: bool ) -> Tuple[int, int, Tuple[int, int, int, int]]: W = w + L + R H = h + T + B if force_square: side = max(W, H) add_w = side - W add_h = side - H L += add_w // 2 R += add_w - add_w // 2 T += add_h // 2 B += add_h - add_h // 2 W = side H = side return W, H, (L, R, T, B) def smart_canvas(img: Image.Image, L: int, R: int, T: int, B: int, blur_px: int = 24) -> tuple[Image.Image, tuple[int,int,int,int]]: w, h = img.size W, H = w + L + R, h + T + B arr = np.array(img) arr_pad = np.pad(arr, ((T, B), (L, R), (0, 0)), mode="edge") edge_img = Image.fromarray(arr_pad) blur_base = img.resize((W, H), Image.BICUBIC).filter(ImageFilter.GaussianBlur(blur_px)) bg = Image.blend(blur_base, edge_img, alpha=0.5) bg.paste(img, (L, T)) bbox = (L, T, L + w, T + h) # (left, top, right, bottom) в координатах канваса return bg, bbox def build_mask(size: Tuple[int,int], bbox: Tuple[int,int,int,int], feather_px: int) -> Image.Image: W, H = size L, T, R, B = bbox m = Image.new("L", (W, H), 255) # белое = редактировать m.paste(0, (L, T, R, B)) # чёрный центр = защищаем if feather_px > 0: m = m.filter(ImageFilter.GaussianBlur(radius=feather_px)) return m def resize_to_gen_side(im: Image.Image, gen_side: int) -> Tuple[Image.Image, float]: if gen_side <= 0: return im, 1.0 w, h = im.size m = max(w, h) if m <= gen_side: return im, 1.0 s = gen_side / m return im.resize((int(round(w*s)), int(round(h*s))), Image.LANCZOS), s def gaussian_alpha_mask(size: Tuple[int,int], rect: Tuple[int,int,int,int], feather: int) -> Image.Image: """Альфа такого же размера, как итог: внутри прямоугольника 255, края 0, с размытием.""" W, H = size L, T, R, B = rect a = Image.new("L", (W, H), 0) a.paste(255, (L, T, R, B)) if feather > 0: a = a.filter(ImageFilter.GaussianBlur(radius=feather)) return a def paste_center_preserving(original: Image.Image, outpaint_upscaled: Image.Image, L: int, T: int, feather: int) -> Image.Image: """ Складываем: base = outpaint_upscaled (RGB) overlay = пустой RGB того же размера, куда пастим original в (L,T) alpha = маска того же размера, 255 в bbox центра (с мягким краем) Результат = composite(overlay, base, alpha) — где 255 берём overlay (оригинал), 0 — base. """ base = outpaint_upscaled.convert("RGB") Wf, Hf = base.size w0, h0 = original.size overlay = Image.new("RGB", (Wf, Hf), (0, 0, 0)) overlay.paste(original, (L, T)) alpha = gaussian_alpha_mask((Wf, Hf), (L, T, L + w0, T + h0), feather=max(1, feather // 2)) final = Image.composite(overlay, base, alpha) return final # ---------- основной процесс ---------- def run_outpaint( image: Image.Image, left: int, right: int, top: int, bottom: int, force_square: bool, prompt: str, negative: str, steps: int, guidance: float, strength: float, feather_px: int, seed: int, gen_side: int, # размер для генерации (обычно 1024) ): if image is None: return None, "", "", None, {"error":"no input"}, None img = image.convert("RGB") w, h = img.size # 1) целевой размер Wt, Ht, (L, R, T, B) = compute_target_size(w, h, int(left), int(right), int(top), int(bottom), bool(force_square)) # 2) канвас и маска в целевом размере canvas_full, bbox_full = smart_canvas(img, L, R, T, B, blur_px=24) mask_full = build_mask(canvas_full.size, bbox_full, feather_px=int(feather_px)) # 3) уменьшаем канвас и маску до gen_side одной и той же шкалой canvas_gen, scale = resize_to_gen_side(canvas_full, int(gen_side)) new_w, new_h = canvas_gen.size mask_gen = mask_full.resize((new_w, new_h), Image.LANCZOS) # 4) инференс p = (prompt or "extend the image naturally").strip() n = (negative or "frames, borders, pillars, poster, collage tiles, mockup, UI elements, captions, text, watermark").strip() g = None if isinstance(seed, int) and seed >= 0: g = torch.Generator(device=DEVICE).manual_seed(int(seed)) out_small = PIPE( prompt=p, negative_prompt=n, image=canvas_gen, mask_image=mask_gen, num_inference_steps=int(steps), guidance_scale=float(guidance), strength=float(strength), inpaint_full_res=True, inpaint_full_res_padding=64, generator=g, ).images[0] # 5) апскейл к целевому размеру if LATENT_UP is not None and out_small.size[0] * 2 <= max(Wt, Ht) + 64: try: up_img = LATENT_UP(prompt=p, image=out_small).images[0] if up_img.size != (Wt, Ht): up_img = up_img.resize((Wt, Ht), Image.LANCZOS) except Exception as e: print("Latent upscaler failed, fallback to LANCZOS:", e) up_img = out_small.resize((Wt, Ht), Image.LANCZOS) else: up_img = out_small.resize((Wt, Ht), Image.LANCZOS) # 6) возвращаем оригинальный центр без потерь final = paste_center_preserving(original=img, outpaint_upscaled=up_img, L=L, T=T, feather=feather_px) # 7) сохраняем полный PNG ts = int(time.time()) out_path = f"/tmp/outpaint_{Wt}x{Ht}_{ts}.png" final.save(out_path, format="PNG") meta = { "loaded_model": MODEL_ID, "family": MODEL_FAMILY, "device": DEVICE, "target_size": [Wt, Ht], "gen_side_used": int(gen_side), "pads_px": {"left": L, "right": R, "top": T, "bottom": B}, "note": "Скачивай через кнопку File — это полный PNG, не превью.", } return final, p, n, mask_full, meta, out_path # ---------- интерфейс ---------- with gr.Blocks(title="Seamless Outpainting — hires PNG, preserved center") as demo: gr.Markdown( f"## Seamless Outpainting (hires)\n" f"**Model:** `{MODEL_ID}` • **Device:** `{DEVICE}`\n\n" "— Итоговый PNG = точное целевое разрешение.\n" "— Генерация на контролируемой стороне (по умолчанию 1024), затем апскейл.\n" "— Центр исходника возвращается пиксель-в-пиксель с мягким швом.\n" "— Дефолтный промт: **extend the image naturally**." ) with gr.Row(): with gr.Column(scale=1): img_in = gr.Image(type="pil", label="Исходная картинка") force_square = gr.Checkbox(value=True, label="Force square (сделать результат квадратным)") with gr.Row(): left = gr.Number(value=256, label="Left (px)", precision=0) right = gr.Number(value=256, label="Right (px)", precision=0) with gr.Row(): top = gr.Number(value=256, label="Top (px)", precision=0) bottom= gr.Number(value=256, label="Bottom (px)", precision=0) prompt = gr.Textbox(label="Prompt", value="extend the image naturally") negative = gr.Textbox( label="Negative prompt", value="frames, borders, pillars, poster, collage tiles, mockup, UI elements, captions, text, watermark", ) with gr.Row(): steps = gr.Slider(10, 60, value=30, step=1, label="Steps") guidance = gr.Slider(0.5, 12, value=5.0, step=0.1, label="Guidance") strength = gr.Slider(0.3, 1.0, value=0.7, step=0.05, label="Strength") with gr.Row(): feather = gr.Slider(0, 160, value=96, step=2, label="Feather (px, шов)") seed = gr.Number(value=-1, label="Seed (-1=random)", precision=0) gen_side = gr.Number(value=1024, label="Generation side (px)", precision=0) run_btn = gr.Button("Expand", variant="primary") with gr.Column(scale=1): img_out = gr.Image(label="Результат (превью = финал)", format="png") used_p = gr.Textbox(label="Использованный prompt", interactive=False) used_n = gr.Textbox(label="Использованный negative", interactive=False) dbg_mask = gr.Image(label="Маска (отладка)") meta = gr.JSON(label="Инфо") download = gr.File(label="Скачать полный PNG") run_btn.click( run_outpaint, inputs=[ img_in, left, right, top, bottom, force_square, prompt, negative, steps, guidance, strength, feather, seed, gen_side ], outputs=[img_out, used_p, used_n, dbg_mask, meta, download], ) if __name__ == "__main__": demo.launch()