From b2a4994e188e331cb8f75aacf22213565cbfe14b Mon Sep 17 00:00:00 2001 From: Jason Hoku Date: Mon, 18 May 2026 12:36:08 -0700 Subject: [PATCH] =?UTF-8?q?fix:=20florence2=20=E2=80=94=20slash=20max=5Fne?= =?UTF-8?q?w=5Ftokens=201024->256=20+=20aggressive=20cleanup?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- florence2_hires.py | 22 ++++++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/florence2_hires.py b/florence2_hires.py index bccaedf..edc3c92 100644 --- a/florence2_hires.py +++ b/florence2_hires.py @@ -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