update

python_server/model/Inpaint/inpaint_loader.py

@@ -5,7 +5,8 @@ from __future__ import annotations
 
 当前支持：
 - SDXL Inpaint（diffusers AutoPipelineForInpainting）
-- ControlNet（占位，暂未统一封装）
+- ControlNet（Stable Diffusion + ControlNet Inpaint）
+- LaMa（torchscript big-lama，本地封装）
 """
 
 from dataclasses import dataclass
@@ -26,6 +27,7 @@ from config_loader import (
 class InpaintBackend(str, Enum):
     SDXL_INPAINT = "sdxl_inpaint"
     CONTROLNET = "controlnet"
+    LAMA = "lama"
 
 
 @dataclass
@@ -278,11 +280,13 @@ def _make_controlnet_predictor(_: UnifiedInpaintConfig):
             image_run = image
             mask_run = mask
 
-        # control image：使用 canny 边缘作为约束（最通用）
+        # control image：使用 Canny 边缘作为条件（更通用、更稳定）。
+        # 备注：当前项目的 ControlNet 权重未必是 inpaint 专用模型，
+        # 若直接把洞内置零作为条件，容易把输出“拉黑”。先回退到稳定方案。
         rgb = np.array(image_run, dtype=np.uint8)
         edges = cv2.Canny(rgb, 100, 200)
         edges3 = np.stack([edges, edges, edges], axis=-1)
-        control_image = Image.fromarray(edges3)
+        control_image = Image.fromarray(edges3, mode="RGB")
 
         if device == "cuda":
             try:
@@ -312,6 +316,37 @@ def _make_controlnet_predictor(_: UnifiedInpaintConfig):
     return _predict
 
 
+def _make_lama_predictor(cfg: UnifiedInpaintConfig):
+    import torch
+
+    app_cfg = load_app_config()
+    device = cfg.device
+    if device is None:
+        device = app_cfg.inpaint.device
+    device = "cuda" if device.startswith("cuda") and torch.cuda.is_available() else "cpu"
+
+    from .lama_torchscript import LaMaTorchscript
+
+    lama = LaMaTorchscript(device=device)
+
+    def _predict(
+        image: Image.Image,
+        mask: Image.Image,
+        _prompt: str,
+        _negative_prompt: str = "",
+        strength: float = 0.8,
+        guidance_scale: float = 7.5,
+        num_inference_steps: int = 30,
+        max_side: int = 1024,
+        **_kwargs,
+    ) -> Image.Image:
+        # LaMa 不使用 prompt/negative_prompt，也不使用 strength/steps 等扩散参数
+        _ = (strength, guidance_scale, num_inference_steps, max_side)
+        return lama(image, mask)
+
+    return _predict
+
+
 def build_inpaint_predictor(
     cfg: UnifiedInpaintConfig | None = None,
 ) -> tuple[Callable[..., Image.Image], InpaintBackend]:
@@ -326,6 +361,9 @@ def build_inpaint_predictor(
     if cfg.backend == InpaintBackend.CONTROLNET:
         return _make_controlnet_predictor(cfg), InpaintBackend.CONTROLNET
 
+    if cfg.backend == InpaintBackend.LAMA:
+        return _make_lama_predictor(cfg), InpaintBackend.LAMA
+
     raise ValueError(f"不支持的补全后端: {cfg.backend}")
 
 

python_server/model/Inpaint/lama_torchscript.py

@@ -0,0 +1,105 @@
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Optional
+
+import numpy as np
+from PIL import Image
+
+
+DEFAULT_LAMA_URL = os.environ.get(
+    "LAMA_MODEL_URL",
+    "https://github.com/Sanster/models/releases/download/add_big_lama/big-lama.pt",
+)
+
+
+def _default_cache_path() -> Path:
+    # 放在 python_server/outputs/cache 下，避免污染用户家目录
+    here = Path(__file__).resolve()
+    cache_dir = here.parents[2] / "outputs" / "cache"
+    cache_dir.mkdir(parents=True, exist_ok=True)
+    return cache_dir / "big-lama.pt"
+
+
+def _download(url: str, dst: Path) -> None:
+    import requests
+
+    dst.parent.mkdir(parents=True, exist_ok=True)
+    resp = requests.get(url, stream=True, timeout=120)
+    resp.raise_for_status()
+    tmp = dst.with_suffix(dst.suffix + ".tmp")
+    with tmp.open("wb") as f:
+        for chunk in resp.iter_content(chunk_size=1024 * 1024):
+            if not chunk:
+                continue
+            f.write(chunk)
+    tmp.replace(dst)
+
+
+def _pad_to_mod8(img: np.ndarray) -> tuple[np.ndarray, tuple[int, int]]:
+    h, w = img.shape[:2]
+    pad_h = (8 - (h % 8)) % 8
+    pad_w = (8 - (w % 8)) % 8
+    if pad_h == 0 and pad_w == 0:
+        return img, (0, 0)
+    out = np.pad(img, ((0, pad_h), (0, pad_w), (0, 0)), mode="reflect")
+    return out, (pad_h, pad_w)
+
+
+def _unpad(img: np.ndarray, pad_hw: tuple[int, int]) -> np.ndarray:
+    pad_h, pad_w = pad_hw
+    if pad_h == 0 and pad_w == 0:
+        return img
+    h, w = img.shape[:2]
+    return img[: h - pad_h, : w - pad_w]
+
+
+class LaMaTorchscript:
+    """
+    轻量 LaMa（big-lama.pt torchscript）推理封装。
+    - 输入：RGB 图 + L mask（白=需要修补）
+    - 输出：RGB 图（同尺寸）
+    """
+
+    def __init__(self, device: str = "cpu", model_path: Optional[str] = None, model_url: Optional[str] = None):
+        import torch
+
+        self.device = device
+        self.model_url = model_url or DEFAULT_LAMA_URL
+        p = Path(model_path) if model_path else _default_cache_path()
+        if not p.exists():
+            _download(self.model_url, p)
+        self.model_path = p
+        self.model = torch.jit.load(str(self.model_path), map_location=device).eval()
+
+    @staticmethod
+    def _norm_img_u8(rgb: np.ndarray) -> np.ndarray:
+        # [H,W,3] uint8 -> float32 [0,1]
+        return (rgb.astype(np.float32) / 255.0).clip(0.0, 1.0)
+
+    def __call__(self, image: Image.Image, mask: Image.Image) -> Image.Image:
+        import torch
+
+        img = np.asarray(image.convert("RGB"), dtype=np.uint8)
+        m = np.asarray(mask.convert("L"), dtype=np.uint8)
+        m = (m >= 128).astype(np.float32)  # 1=hole
+
+        img_f = self._norm_img_u8(img)
+        # pad to multiple of 8
+        img_pad, pad_hw = _pad_to_mod8(img_f)
+        m_pad, _ = _pad_to_mod8(np.repeat(m[:, :, None], 3, axis=2))
+        m1 = m_pad[:, :, 0]  # [H,W]
+
+        # torch: [1,C,H,W]
+        it = torch.from_numpy(img_pad).permute(2, 0, 1).unsqueeze(0).to(self.device)
+        mt = torch.from_numpy(m1).unsqueeze(0).unsqueeze(0).to(self.device)
+
+        with torch.no_grad():
+            out = self.model(it, mt)
+
+        out_np = out[0].permute(1, 2, 0).detach().cpu().numpy()
+        out_np = np.clip(out_np * 255.0, 0.0, 255.0).astype(np.uint8)
+        out_np = _unpad(out_np, pad_hw)
+        return Image.fromarray(out_np, mode="RGB")
+