fix: florence2 — slash max_new_tokens 1024->256 + aggressive cleanup
Same OOM persisted on 8 GiB card after duplicate-load fix. Florence2's beam-search KV cache scales linearly with max_new_tokens × num_beams × hidden_dim × num_layers. For Florence-2-large with num_beams=3 and max_new_tokens=1024, the KV cache alone is multiple GiB — and it's pre- allocated upfront regardless of actual generation length. For referring_expression_segmentation outputting a face polygon (~10-50 tokens), 1024 is 20x over-provisioned. Cut to 256: still 5x headroom for complex polygons, ~4x less VRAM for the KV cache. Also added gc.collect() + torch.cuda.empty_cache() at end of run_florence2_step alongside the existing soft_empty_cache, since Python- side dangling tensor references from the KSampler stage were preventing GPU memory from being reclaimed before the next item iterates.
This commit is contained in:
+18
-4
@@ -491,6 +491,11 @@ def run_florence2_step(
|
||||
|
||||
text_input = manifest_extras["florence2_text_input"]
|
||||
print(f"[Florence2HiResFix] Detecting '{text_input}' in image...")
|
||||
# max_new_tokens=256 (was 1024): Florence2's KV cache scales linearly with
|
||||
# sequence length × num_beams × hidden_dim × num_layers. For
|
||||
# referring_expression_segmentation outputting a face polygon (~10-50 tokens),
|
||||
# 1024 is wildly over-provisioned and wastes GiB of VRAM on dashboard batches
|
||||
# where the session checkpoint is already loaded.
|
||||
det_result = _call_node(
|
||||
f2r_cls,
|
||||
image=source_image,
|
||||
@@ -499,7 +504,7 @@ def run_florence2_step(
|
||||
task="referring_expression_segmentation",
|
||||
fill_mask=True,
|
||||
keep_model_loaded=True,
|
||||
max_new_tokens=1024,
|
||||
max_new_tokens=256,
|
||||
num_beams=3,
|
||||
do_sample=False,
|
||||
output_mask_select=step_config.get("output_mask_select", ""),
|
||||
@@ -705,12 +710,21 @@ def run_florence2_step(
|
||||
manifest_extras["duration"] = round(time.time() - t_start, 2)
|
||||
|
||||
# VRAM hygiene: release intermediate tensors before next item. The KSampler
|
||||
# latents, VAE-decoded crops, and Florence2 beam-search buffers can fragment
|
||||
# GPU memory across iterations. soft_empty_cache lets ComfyUI reclaim it.
|
||||
# Skip in tests (mm import would pull comfy stack at module load).
|
||||
# latents, VAE-decoded crops, and Florence2 beam-search KV cache buffers can
|
||||
# fragment GPU memory across iterations. gc.collect() drops Python-side
|
||||
# references; soft_empty_cache + torch.cuda.empty_cache let ComfyUI / PyTorch
|
||||
# reclaim the freed memory. Skip in tests (mm import would pull comfy stack).
|
||||
try:
|
||||
import gc as _gc
|
||||
_gc.collect()
|
||||
import comfy.model_management as _mm
|
||||
_mm.soft_empty_cache()
|
||||
try:
|
||||
import torch as _torch
|
||||
if _torch.cuda.is_available():
|
||||
_torch.cuda.empty_cache()
|
||||
except Exception:
|
||||
pass
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
Reference in New Issue
Block a user