113 lines
3.2 KiB
Python
113 lines
3.2 KiB
Python
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
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
|
|
|
|
root = Path(model_id.strip()).expanduser()
|
|
# 本机 diffusers 目录（含 model_index.json）时不访问 Hub，避免自动下载。
|
|
local_files_only = root.is_dir() and (root / "model_index.json").is_file()
|
|
|
|
pipe = FluxFillPipeline.from_pretrained(
|
|
str(root),
|
|
torch_dtype=torch_dtype,
|
|
token=None,
|
|
local_files_only=local_files_only,
|
|
)
|
|
|
|
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
|