update
This commit is contained in:
@@ -1,5 +1,6 @@
|
|||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import os
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import Any, Callable
|
from typing import Any, Callable
|
||||||
|
|
||||||
@@ -61,16 +62,32 @@ def make_flux_fill_predictor(
|
|||||||
local_files_only=local_files_only,
|
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":
|
if device == "cuda":
|
||||||
|
offload = os.environ.get("FLUX_FILL_OFFLOAD", "sequential").strip().lower()
|
||||||
try:
|
try:
|
||||||
pipe.enable_vae_tiling()
|
pipe.enable_vae_tiling()
|
||||||
except Exception:
|
except Exception:
|
||||||
pass
|
pass
|
||||||
|
try:
|
||||||
|
pipe.enable_vae_slicing()
|
||||||
|
except Exception:
|
||||||
|
pass
|
||||||
try:
|
try:
|
||||||
pipe.enable_attention_slicing()
|
pipe.enable_attention_slicing()
|
||||||
except Exception:
|
except Exception:
|
||||||
pass
|
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:
|
else:
|
||||||
pipe.to(device)
|
pipe.to(device)
|
||||||
|
|
||||||
@@ -99,6 +116,13 @@ def make_flux_fill_predictor(
|
|||||||
run_w = _align_spatial_dim_for_flux_fill(run_w)
|
run_w = _align_spatial_dim_for_flux_fill(run_w)
|
||||||
run_h = _align_spatial_dim_for_flux_fill(run_h)
|
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):
|
if (run_w, run_h) != (orig_w, orig_h):
|
||||||
image_run = image.resize((run_w, run_h), resample=Image.BICUBIC)
|
image_run = image.resize((run_w, run_h), resample=Image.BICUBIC)
|
||||||
mask_run = mask.resize((run_w, run_h), resample=Image.NEAREST)
|
mask_run = mask.resize((run_w, run_h), resample=Image.NEAREST)
|
||||||
@@ -116,22 +140,35 @@ def make_flux_fill_predictor(
|
|||||||
|
|
||||||
if device == "cuda":
|
if device == "cuda":
|
||||||
try:
|
try:
|
||||||
|
import gc
|
||||||
|
|
||||||
|
gc.collect()
|
||||||
torch.cuda.empty_cache()
|
torch.cuda.empty_cache()
|
||||||
except Exception:
|
except Exception:
|
||||||
pass
|
pass
|
||||||
|
|
||||||
out = pipe(
|
try:
|
||||||
prompt=p_clip,
|
out = pipe(
|
||||||
prompt_2=p_t5,
|
prompt=p_clip,
|
||||||
image=image_run,
|
prompt_2=p_t5,
|
||||||
mask_image=mask_run,
|
image=image_run,
|
||||||
height=run_h,
|
mask_image=mask_run,
|
||||||
width=run_w,
|
height=run_h,
|
||||||
strength=st,
|
width=run_w,
|
||||||
guidance_scale=float(guidance_scale),
|
strength=st,
|
||||||
num_inference_steps=int(num_inference_steps),
|
guidance_scale=float(guidance_scale),
|
||||||
max_sequence_length=int(max_sequence_length),
|
num_inference_steps=int(num_inference_steps),
|
||||||
).images[0]
|
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")
|
out = out.convert("RGB")
|
||||||
if out.size != (orig_w, orig_h):
|
if out.size != (orig_w, orig_h):
|
||||||
|
|||||||
Reference in New Issue
Block a user