diff --git a/python_server/model/Inpaint/flux_fill/loader.py b/python_server/model/Inpaint/flux_fill/loader.py index cef841e..a03d2ec 100644 --- a/python_server/model/Inpaint/flux_fill/loader.py +++ b/python_server/model/Inpaint/flux_fill/loader.py @@ -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):