initial commit

python_server/model/Inpaint/__init__.py

@@ -0,0 +1 @@
+

python_server/model/Inpaint/inpaint_loader.py

@@ -0,0 +1,413 @@
+from __future__ import annotations
+
+"""
+统一的补全（Inpaint）模型加载入口。
+
+当前支持：
+- SDXL Inpaint（diffusers AutoPipelineForInpainting）
+- ControlNet（占位，暂未统一封装）
+"""
+
+from dataclasses import dataclass
+from enum import Enum
+from typing import Callable
+
+import numpy as np
+from PIL import Image
+
+from config_loader import (
+    load_app_config,
+    get_sdxl_base_model_from_app,
+    get_controlnet_base_model_from_app,
+    get_controlnet_model_from_app,
+)
+
+
+class InpaintBackend(str, Enum):
+    SDXL_INPAINT = "sdxl_inpaint"
+    CONTROLNET = "controlnet"
+
+
+@dataclass
+class UnifiedInpaintConfig:
+    backend: InpaintBackend = InpaintBackend.SDXL_INPAINT
+    device: str | None = None
+    # SDXL base model (HF id 或本地目录)，不填则用 config.py 的默认值
+    sdxl_base_model: str | None = None
+
+
+@dataclass
+class UnifiedDrawConfig:
+    """
+    统一绘图配置：
+    - 纯文生图：image=None
+    - 图生图（模仿输入图）：image=某张参考图
+    """
+    device: str | None = None
+    sdxl_base_model: str | None = None
+
+
+def _resolve_device_and_dtype(device: str | None):
+    import torch
+
+    app_cfg = load_app_config()
+    if device is None:
+        device = app_cfg.inpaint.device
+    device = "cuda" if device.startswith("cuda") and torch.cuda.is_available() else "cpu"
+    torch_dtype = torch.float16 if device == "cuda" else torch.float32
+    return device, torch_dtype
+
+
+def _enable_memory_opts(pipe, device: str) -> None:
+    if device == "cuda":
+        try:
+            pipe.enable_attention_slicing()
+        except Exception:
+            pass
+        try:
+            pipe.enable_vae_slicing()
+        except Exception:
+            pass
+        try:
+            pipe.enable_vae_tiling()
+        except Exception:
+            pass
+        try:
+            pipe.enable_model_cpu_offload()
+        except Exception:
+            pass
+
+
+def _align_size(orig_w: int, orig_h: int, max_side: int) -> tuple[int, int]:
+    run_w, run_h = orig_w, orig_h
+    if max_side > 0 and max(orig_w, orig_h) > max_side:
+        scale = max_side / float(max(orig_w, orig_h))
+        run_w = int(round(orig_w * scale))
+        run_h = int(round(orig_h * scale))
+    run_w = max(8, run_w - (run_w % 8))
+    run_h = max(8, run_h - (run_h % 8))
+    return run_w, run_h
+
+
+def _make_sdxl_inpaint_predictor(
+    cfg: UnifiedInpaintConfig,
+) -> Callable[[Image.Image, Image.Image, str, str], Image.Image]:
+    """
+    返回补全函数：
+    - 输入：image(PIL RGB), mask(PIL L/1), prompt, negative_prompt
+    - 输出：PIL RGB 结果图
+    """
+    import torch
+    from diffusers import AutoPipelineForText2Image, AutoPipelineForInpainting
+
+    app_cfg = load_app_config()
+    base_model = cfg.sdxl_base_model or get_sdxl_base_model_from_app(app_cfg)
+
+    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"
+
+    torch_dtype = torch.float16 if device == "cuda" else torch.float32
+
+    pipe_t2i = AutoPipelineForText2Image.from_pretrained(
+        base_model,
+        torch_dtype=torch_dtype,
+        variant="fp16" if device == "cuda" else None,
+        use_safetensors=True,
+    ).to(device)
+    pipe = AutoPipelineForInpainting.from_pipe(pipe_t2i).to(device)
+
+    # 省显存设置（尽量不改变输出语义）
+    # 注意：CPU offload 会明显变慢，但能显著降低显存占用。
+    if device == "cuda":
+        try:
+            pipe.enable_attention_slicing()
+        except Exception:
+            pass
+        try:
+            pipe.enable_vae_slicing()
+        except Exception:
+            pass
+        try:
+            pipe.enable_vae_tiling()
+        except Exception:
+            pass
+        try:
+            pipe.enable_model_cpu_offload()
+        except Exception:
+            pass
+
+    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,
+    ) -> Image.Image:
+        image = image.convert("RGB")
+        # diffusers 要求 mask 为单通道，白色区域为需要重绘
+        mask = mask.convert("L")
+
+        # SDXL / diffusers 通常要求宽高为 8 的倍数；同时为了避免 OOM，
+        # 推理时将图像按比例缩放到不超过 max_side（默认 1024）并对齐到 8 的倍数。
+        # 推理后再 resize 回原始尺寸，保证输出与原图分辨率一致。
+        orig_w, orig_h = image.size
+        run_w, run_h = orig_w, orig_h
+
+        if max(orig_w, orig_h) > max_side:
+            scale = max_side / float(max(orig_w, orig_h))
+            run_w = int(round(orig_w * scale))
+            run_h = int(round(orig_h * scale))
+
+        run_w = max(8, run_w - (run_w % 8))
+        run_h = max(8, run_h - (run_h % 8))
+        if run_w <= 0:
+            run_w = 8
+        if run_h <= 0:
+            run_h = 8
+
+        if (run_w, run_h) != (orig_w, orig_h):
+            image_run = image.resize((run_w, run_h), resample=Image.BICUBIC)
+            mask_run = mask.resize((run_w, run_h), resample=Image.NEAREST)
+        else:
+            image_run = image
+            mask_run = mask
+
+        if device == "cuda":
+            try:
+                torch.cuda.empty_cache()
+            except Exception:
+                pass
+
+        out = pipe(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            image=image_run,
+            mask_image=mask_run,
+            strength=strength,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_inference_steps,
+            width=run_w,
+            height=run_h,
+        ).images[0]
+        out = out.convert("RGB")
+        if out.size != (orig_w, orig_h):
+            out = out.resize((orig_w, orig_h), resample=Image.BICUBIC)
+        return out
+
+    return _predict
+
+
+def _make_controlnet_predictor(_: UnifiedInpaintConfig):
+    import torch
+    from diffusers import ControlNetModel, StableDiffusionControlNetInpaintPipeline
+
+    app_cfg = load_app_config()
+
+    device = _.device
+    if device is None:
+        device = app_cfg.inpaint.device
+    device = "cuda" if device.startswith("cuda") and torch.cuda.is_available() else "cpu"
+    torch_dtype = torch.float16 if device == "cuda" else torch.float32
+
+    base_model = get_controlnet_base_model_from_app(app_cfg)
+    controlnet_id = get_controlnet_model_from_app(app_cfg)
+
+    controlnet = ControlNetModel.from_pretrained(controlnet_id, torch_dtype=torch_dtype)
+    pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
+        base_model,
+        controlnet=controlnet,
+        torch_dtype=torch_dtype,
+        safety_checker=None,
+    )
+
+    if device == "cuda":
+        try:
+            pipe.enable_attention_slicing()
+        except Exception:
+            pass
+        try:
+            pipe.enable_vae_slicing()
+        except Exception:
+            pass
+        try:
+            pipe.enable_vae_tiling()
+        except Exception:
+            pass
+        try:
+            pipe.enable_model_cpu_offload()
+        except Exception:
+            pipe.to(device)
+    else:
+        pipe.to(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,
+        controlnet_conditioning_scale: float = 1.0,
+        max_side: int = 768,
+    ) -> Image.Image:
+        import cv2
+        import numpy as np
+
+        image = image.convert("RGB")
+        mask = mask.convert("L")
+
+        orig_w, orig_h = image.size
+        run_w, run_h = orig_w, orig_h
+        if max(orig_w, orig_h) > max_side:
+            scale = max_side / float(max(orig_w, orig_h))
+            run_w = int(round(orig_w * scale))
+            run_h = int(round(orig_h * scale))
+        run_w = max(8, run_w - (run_w % 8))
+        run_h = max(8, run_h - (run_h % 8))
+
+        if (run_w, run_h) != (orig_w, orig_h):
+            image_run = image.resize((run_w, run_h), resample=Image.BICUBIC)
+            mask_run = mask.resize((run_w, run_h), resample=Image.NEAREST)
+        else:
+            image_run = image
+            mask_run = mask
+
+        # control image：使用 canny 边缘作为约束（最通用）
+        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)
+
+        if device == "cuda":
+            try:
+                torch.cuda.empty_cache()
+            except Exception:
+                pass
+
+        out = pipe(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            image=image_run,
+            mask_image=mask_run,
+            control_image=control_image,
+            strength=strength,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_inference_steps,
+            controlnet_conditioning_scale=controlnet_conditioning_scale,
+            width=run_w,
+            height=run_h,
+        ).images[0]
+
+        out = out.convert("RGB")
+        if out.size != (orig_w, orig_h):
+            out = out.resize((orig_w, orig_h), resample=Image.BICUBIC)
+        return out
+
+    return _predict
+
+
+def build_inpaint_predictor(
+    cfg: UnifiedInpaintConfig | None = None,
+) -> tuple[Callable[..., Image.Image], InpaintBackend]:
+    """
+    统一构建补全预测函数。
+    """
+    cfg = cfg or UnifiedInpaintConfig()
+
+    if cfg.backend == InpaintBackend.SDXL_INPAINT:
+        return _make_sdxl_inpaint_predictor(cfg), InpaintBackend.SDXL_INPAINT
+
+    if cfg.backend == InpaintBackend.CONTROLNET:
+        return _make_controlnet_predictor(cfg), InpaintBackend.CONTROLNET
+
+    raise ValueError(f"不支持的补全后端: {cfg.backend}")
+
+
+def build_draw_predictor(
+    cfg: UnifiedDrawConfig | None = None,
+) -> Callable[..., Image.Image]:
+    """
+    构建统一绘图函数：
+    - 文生图：draw(prompt, image=None, ...)
+    - 图生图：draw(prompt, image=ref_image, strength=0.55, ...)
+    """
+    import torch
+    from diffusers import AutoPipelineForText2Image, AutoPipelineForImage2Image
+
+    cfg = cfg or UnifiedDrawConfig()
+    app_cfg = load_app_config()
+    base_model = cfg.sdxl_base_model or get_sdxl_base_model_from_app(app_cfg)
+    device, torch_dtype = _resolve_device_and_dtype(cfg.device)
+
+    pipe_t2i = AutoPipelineForText2Image.from_pretrained(
+        base_model,
+        torch_dtype=torch_dtype,
+        variant="fp16" if device == "cuda" else None,
+        use_safetensors=True,
+    ).to(device)
+    pipe_i2i = AutoPipelineForImage2Image.from_pipe(pipe_t2i).to(device)
+
+    _enable_memory_opts(pipe_t2i, device)
+    _enable_memory_opts(pipe_i2i, device)
+
+    def _draw(
+        prompt: str,
+        image: Image.Image | None = None,
+        negative_prompt: str = "",
+        strength: float = 0.55,
+        guidance_scale: float = 7.5,
+        num_inference_steps: int = 30,
+        width: int = 1024,
+        height: int = 1024,
+        max_side: int = 1024,
+    ) -> Image.Image:
+        prompt = prompt or ""
+        negative_prompt = negative_prompt or ""
+
+        if device == "cuda":
+            try:
+                torch.cuda.empty_cache()
+            except Exception:
+                pass
+
+        if image is None:
+            run_w, run_h = _align_size(width, height, max_side=max_side)
+            out = pipe_t2i(
+                prompt=prompt,
+                negative_prompt=negative_prompt,
+                guidance_scale=guidance_scale,
+                num_inference_steps=num_inference_steps,
+                width=run_w,
+                height=run_h,
+            ).images[0]
+            return out.convert("RGB")
+
+        image = image.convert("RGB")
+        orig_w, orig_h = image.size
+        run_w, run_h = _align_size(orig_w, orig_h, max_side=max_side)
+        if (run_w, run_h) != (orig_w, orig_h):
+            image_run = image.resize((run_w, run_h), resample=Image.BICUBIC)
+        else:
+            image_run = image
+
+        out = pipe_i2i(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            image=image_run,
+            strength=strength,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_inference_steps,
+        ).images[0].convert("RGB")
+
+        if out.size != (orig_w, orig_h):
+            out = out.resize((orig_w, orig_h), resample=Image.BICUBIC)
+        return out
+
+    return _draw
+