update

python_server/model/Inpaint/flux_fill/__init__.py

@@ -0,0 +1,3 @@
+from .loader import make_flux_fill_predictor
+
+__all__ = ["make_flux_fill_predictor"]

python_server/model/Inpaint/flux_fill/loader.py

@@ -0,0 +1,109 @@
+from __future__ import annotations
+
+import os
+from typing import Any, Callable
+
+import numpy as np
+from PIL import Image
+
+
+def make_flux_fill_predictor(
+    model_id: str,
+    device: str | None = None,
+) -> Callable[[Image.Image, Image.Image, str, str], Image.Image]:
+    import torch
+    from diffusers import FluxFillPipeline
+
+    if device is None:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    if device == "cuda" and torch.cuda.is_bf16_supported():
+        torch_dtype = torch.bfloat16
+    elif device == "cuda":
+        torch_dtype = torch.float16
+    else:
+        torch_dtype = torch.float32
+
+    hf_token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
+
+    pipe = FluxFillPipeline.from_pretrained(
+        model_id,
+        torch_dtype=torch_dtype,
+        token=hf_token,
+    )
+
+    if device == "cuda":
+        try:
+            pipe.enable_model_cpu_offload()
+        except Exception:
+            try:
+                pipe.enable_vae_tiling()
+            except Exception:
+                pass
+            pipe.to("cuda")
+    else:
+        pipe.to(device)
+
+    def _predict(
+        image: Image.Image,
+        mask: Image.Image,
+        prompt: str,
+        negative_prompt: str = "",
+        strength: float = 1.0,
+        guidance_scale: float = 30.0,
+        num_inference_steps: int = 28,
+        max_side: int = 1024,
+        max_sequence_length: int = 512,
+        **_: Any,
+    ) -> Image.Image:
+        _ = negative_prompt
+        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
+
+        p = (prompt or "").strip()
+        if not p:
+            p = "Chinese ink painting background, natural texture, seamless fill, no figures"
+
+        st = float(strength)
+        st = min(1.0, max(0.05, st))
+
+        if device == "cuda":
+            try:
+                torch.cuda.empty_cache()
+            except Exception:
+                pass
+
+        out = pipe(
+            prompt=p,
+            image=image_run,
+            mask_image=mask_run,
+            height=run_h,
+            width=run_w,
+            strength=st,
+            guidance_scale=float(guidance_scale),
+            num_inference_steps=int(num_inference_steps),
+            max_sequence_length=int(max_sequence_length),
+        ).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