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:
Jason Hoku
2026-05-18 12:36:08 -07:00
parent d4b88b0420
commit b2a4994e18
+18 -4
View File
@@ -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