# 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 # ------------------------------ # 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) 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, watermark, signature, logo, lowres, blurry, artifacts, deformed, distorted, oversaturated, extra limbs, frame, border" 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, 12.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"] )