fix: guarantee offload cleanup on exception with try/finally

Both nodes moved models to GPU before work then back to CPU after.
Any exception (OOM, cancellation, bad input) would skip the cleanup,
leaving models on GPU permanently until ComfyUI restarts.

Wrap the entire work block in try/finally so offload_to_cpu cleanup
always runs regardless of how the node exits. Also removes the unused
`mode` variable in SelvaSampler.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
2026-04-05 08:40:39 +02:00
parent 8bb2fb7015
commit 3dd6badfd9
2 changed files with 100 additions and 97 deletions
+2
View File
@@ -158,6 +158,7 @@ class SelvaFeatureExtractor:
print(f"[SelVA] Extracting features: duration={duration:.2f}s fps={fps:.3f} prompt='{prompt[:60]}'", flush=True) print(f"[SelVA] Extracting features: duration={duration:.2f}s fps={fps:.3f} prompt='{prompt[:60]}'", flush=True)
pbar = comfy.utils.ProgressBar(3) pbar = comfy.utils.ProgressBar(3)
try:
with torch.no_grad(): with torch.no_grad():
# --- CLIP frames: [1, N, C, 384, 384] float32 [0,1] --- # --- CLIP frames: [1, N, C, 384, 384] float32 [0,1] ---
clip_frames = _sample_frames(video, fps, _CLIP_FPS, duration) # [N, H, W, C] clip_frames = _sample_frames(video, fps, _CLIP_FPS, duration) # [N, H, W, C]
@@ -198,6 +199,7 @@ class SelvaFeatureExtractor:
print(f"[SelVA] clip_features: {tuple(clip_features.shape)}", flush=True) print(f"[SelVA] clip_features: {tuple(clip_features.shape)}", flush=True)
print(f"[SelVA] sync_features: {tuple(sync_features.shape)}", flush=True) print(f"[SelVA] sync_features: {tuple(sync_features.shape)}", flush=True)
finally:
if strategy == "offload_to_cpu": if strategy == "offload_to_cpu":
feature_utils.to(get_offload_device()) feature_utils.to(get_offload_device())
net_video_enc.to(get_offload_device()) net_video_enc.to(get_offload_device())
+2 -1
View File
@@ -54,7 +54,6 @@ class SelvaSampler:
strategy = model["strategy"] strategy = model["strategy"]
net_generator = model["generator"] net_generator = model["generator"]
feature_utils = model["feature_utils"] feature_utils = model["feature_utils"]
mode = model["mode"]
# Validate that features were extracted with the same model variant # Validate that features were extracted with the same model variant
feat_variant = features.get("variant") feat_variant = features.get("variant")
@@ -88,6 +87,7 @@ class SelvaSampler:
feature_utils.to(device) feature_utils.to(device)
soft_empty_cache() soft_empty_cache()
try:
clip_f = features["clip_features"].to(device, dtype) # [1, T_clip, 1024] clip_f = features["clip_features"].to(device, dtype) # [1, T_clip, 1024]
sync_f = features["sync_features"].to(device, dtype) # [1, T_sync, 768] sync_f = features["sync_features"].to(device, dtype) # [1, T_sync, 768]
@@ -154,6 +154,7 @@ class SelvaSampler:
"to 'offload_to_cpu', using a smaller variant, or reducing duration." "to 'offload_to_cpu', using a smaller variant, or reducing duration."
) )
finally:
if strategy == "offload_to_cpu": if strategy == "offload_to_cpu":
net_generator.to(get_offload_device()) net_generator.to(get_offload_device())
feature_utils.to(get_offload_device()) feature_utils.to(get_offload_device())