179 lines
6.1 KiB
Python
179 lines
6.1 KiB
Python
from __future__ import annotations
|
||
|
||
import os
|
||
from pathlib import Path
|
||
from typing import Any, Callable
|
||
|
||
import numpy as np
|
||
from PIL import Image
|
||
|
||
|
||
def _align_spatial_dim_for_flux_fill(n: int, multiple: int = 16) -> int:
|
||
"""FLUX Fill 要求 H/W 为 (vae_scale_factor*2) 的倍数，常见为 16。"""
|
||
n = int(max(n, multiple))
|
||
return n - (n % multiple)
|
||
|
||
|
||
def _flux_clip_prompt_and_t5_full(pipe: Any, text_full: str) -> tuple[str, str]:
|
||
"""
|
||
CLIP 仅约 77 token：用与管线一致的 tokenizer 截断得到 prompt（池化向量）。
|
||
完整文本走 prompt_2 → T5（max_sequence_length 可达 512），长中文不再整段喂给 CLIP。
|
||
"""
|
||
text_full = (text_full or "").strip()
|
||
if not text_full:
|
||
d = "Chinese ink painting background, natural texture, seamless fill, no figures"
|
||
return d, d
|
||
tok = getattr(pipe, "tokenizer", None)
|
||
if tok is None:
|
||
return text_full[:120], text_full
|
||
ml = int(getattr(tok, "model_max_length", 77) or 77)
|
||
enc = tok([text_full], padding="max_length", max_length=ml, truncation=True, return_tensors="pt")
|
||
clip_s = tok.batch_decode(enc.input_ids, skip_special_tokens=True)[0].strip()
|
||
if not clip_s:
|
||
clip_s = text_full[:80]
|
||
return clip_s, text_full
|
||
|
||
|
||
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,
|
||
)
|
||
|
||
# 显存：默认顺序 CPU offload（各子模块用时再上 GPU，峰值远低于 pipe.to(cuda)）。
|
||
# 整卡加载易占满 32GB 再在推理时 OOM。环境变量 FLUX_FILL_OFFLOAD=none 可改回整卡（需有余量）。
|
||
# 勿用 enable_model_cpu_offload（与 transformers 5.x CLIP 的 _hf_hook 问题相关）。
|
||
used_sequential = False
|
||
if device == "cuda":
|
||
offload = os.environ.get("FLUX_FILL_OFFLOAD", "sequential").strip().lower()
|
||
try:
|
||
pipe.enable_vae_tiling()
|
||
except Exception:
|
||
pass
|
||
try:
|
||
pipe.enable_vae_slicing()
|
||
except Exception:
|
||
pass
|
||
try:
|
||
pipe.enable_attention_slicing()
|
||
except Exception:
|
||
pass
|
||
if offload not in ("none", "gpu", "0", "false", "no"):
|
||
try:
|
||
pipe.enable_sequential_cpu_offload(gpu_id=0)
|
||
used_sequential = True
|
||
except Exception:
|
||
used_sequential = False
|
||
if not used_sequential:
|
||
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 = _align_spatial_dim_for_flux_fill(run_w)
|
||
run_h = _align_spatial_dim_for_flux_fill(run_h)
|
||
|
||
# 推理分辨率硬顶（像素域），防止极大 crop + 高 max_side 时注意力/VAE 爆显存；输出仍会缩回 orig。
|
||
run_cap = int(os.environ.get("FLUX_FILL_RUN_MAX_SIDE", "4096"))
|
||
if run_cap >= 256 and max(run_w, run_h) > run_cap:
|
||
s = run_cap / float(max(run_w, run_h))
|
||
run_w = _align_spatial_dim_for_flux_fill(int(round(run_w * s)))
|
||
run_h = _align_spatial_dim_for_flux_fill(int(round(run_h * s)))
|
||
|
||
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_full = (prompt or "").strip()
|
||
if not p_full:
|
||
p_full = "Chinese ink painting background, natural texture, seamless fill, no figures"
|
||
p_clip, p_t5 = _flux_clip_prompt_and_t5_full(pipe, p_full)
|
||
|
||
st = float(strength)
|
||
st = min(1.0, max(0.05, st))
|
||
|
||
if device == "cuda":
|
||
try:
|
||
import gc
|
||
|
||
gc.collect()
|
||
torch.cuda.empty_cache()
|
||
except Exception:
|
||
pass
|
||
|
||
try:
|
||
out = pipe(
|
||
prompt=p_clip,
|
||
prompt_2=p_t5,
|
||
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]
|
||
finally:
|
||
if device == "cuda":
|
||
try:
|
||
import gc
|
||
|
||
gc.collect()
|
||
torch.cuda.empty_cache()
|
||
except Exception:
|
||
pass
|
||
|
||
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
|