From 5770c7034b0338dd8d68ae8471f06675e36e2492 Mon Sep 17 00:00:00 2001 From: Dustin <6962246+djdarcy@users.noreply.github.com> Date: Mon, 4 May 2026 02:17:17 -0400 Subject: [PATCH] Guard against None offload_stream in prefetch_queue_pop cast_modules_with_vbar can return None when all modules in a prefetched block are already resident. Mirrors the existing None check in sync_stream. --- comfy/model_prefetch.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/comfy/model_prefetch.py b/comfy/model_prefetch.py index 0ad35deb5..72e11dec6 100644 --- a/comfy/model_prefetch.py +++ b/comfy/model_prefetch.py @@ -37,7 +37,8 @@ def prefetch_queue_pop(queue, device, module): consumed = queue.pop(0) if consumed is not None: offload_stream, prefetch_state = consumed - offload_stream.wait_stream(comfy.model_management.current_stream(device)) + if offload_stream is not None: + offload_stream.wait_stream(comfy.model_management.current_stream(device)) _, comfy_modules = prefetch_state if comfy_modules is not None: cleanup_prefetched_modules(comfy_modules)