This commit is contained in:
2026-05-15 00:55:49 +08:00
parent f64e624a39
commit 12ea7273ef
+50 -13
View File
@@ -1,5 +1,6 @@
from __future__ import annotations
import os
from pathlib import Path
from typing import Any, Callable
@@ -61,16 +62,32 @@ def make_flux_fill_predictor(
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
pipe.to("cuda")
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)
@@ -99,6 +116,13 @@ def make_flux_fill_predictor(
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)
@@ -116,22 +140,35 @@ def make_flux_fill_predictor(
if device == "cuda":
try:
import gc
gc.collect()
torch.cuda.empty_cache()
except Exception:
pass
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]
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):