Update app.py

app.py

@@ -1,5 +1,5 @@
 # app.py
-import os, io, time, math, random
+import os, time, random
 from typing import Tuple
 import gradio as gr
 import numpy as np
@@ -15,19 +15,19 @@ 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"
+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,            # вы уже тестируете локально
+    safety_checker=None,                  # для локального теста
 ).to(DEVICE)
 pipe.enable_attention_slicing()
 if DEVICE == "cuda":
     try:
-        pipe.enable_model_cpu_offload()  # экономия VRAM, если доступно
+        pipe.enable_model_cpu_offload()
     except Exception:
         pass
 
@@ -38,57 +38,43 @@ 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]]:
-    """Возвращает квадратный холст + координаты вставки исходника (x0,y0)."""
     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:
-    """Белый (генерировать) за пределами inner_rect, чёрный внутри.
-       Дополнительно — перья по границе."""
     W, H = canvas_size
     x0,y0,x1,y1 = inner_rect
-    mask = Image.new("L", (W,H), 255)  # по умолчанию всё белое — дорисовывать
+    mask = Image.new("L", (W,H), 255)
     draw = ImageDraw.Draw(mask)
-    draw.rectangle([x0,y0,x1,y1], fill=0)  # центр защищаем от изменений
+    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()
-    x0,y0 = offset
-    out.paste(original, (x0,y0))
+    out.paste(original, offset)
     return out
 
-def preview_frame(image: Image.Image, pad_px:int) -> Tuple[Image.Image, str]:
-    """Готовит превью квадрата с рамкой и текстом про пиксели по сторонам."""
+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
-    # уменьшаем для UI
-    show_W, show_H = 512, 512
     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((show_W, show_H), Image.LANCZOS)
-
-    # текст
+    prev.thumbnail((512, 512), Image.LANCZOS)
     top = y0
     left = x0
     right = W - (x0 + w)
@@ -112,16 +98,13 @@ def outpaint_generate(
     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)
 
-    # безопасные размеры (многократность 8, ограничение по памяти)
     W = round_to_eight(canvas.size[0])
     H = round_to_eight(canvas.size[1])
     MAX_SIDE = 1536 if DEVICE == "cuda" else 1024
@@ -131,7 +114,6 @@ def outpaint_generate(
         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)
@@ -140,13 +122,9 @@ def outpaint_generate(
         inner_rect_small = (x0,y0,x1,y1)
 
     g = torch.Generator(device=DEVICE)
-    if seed is None or int(seed) < 0:
-        seed_val = random.randint(0, 2**32 - 1)
-    else:
-        seed_val = int(seed)
+    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,
@@ -158,14 +136,11 @@ def outpaint_generate(
             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)
@@ -206,7 +181,6 @@ with gr.Blocks(css="""
                 meta = gr.Markdown("")
                 file_out = gr.File(label="Download PNG")
 
-    # live preview
     def _update_preview(img, pad_px):
         if img is None:
             return None, ""
@@ -233,4 +207,10 @@ with gr.Blocks(css="""
     )
 
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860, inbrowser=False)  # без ssr_mode
+    # КРИТИЧЕСКОЕ ИСПРАВЛЕНИЕ: разрешаем /mnt/data
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        inbrowser=False,
+        allowed_paths=["/mnt/data"]
+    )