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:
@@ -158,6 +158,7 @@ class SelvaFeatureExtractor:
|
||||
print(f"[SelVA] Extracting features: duration={duration:.2f}s fps={fps:.3f} prompt='{prompt[:60]}'", flush=True)
|
||||
pbar = comfy.utils.ProgressBar(3)
|
||||
|
||||
try:
|
||||
with torch.no_grad():
|
||||
# --- CLIP frames: [1, N, C, 384, 384] float32 [0,1] ---
|
||||
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] sync_features: {tuple(sync_features.shape)}", flush=True)
|
||||
|
||||
finally:
|
||||
if strategy == "offload_to_cpu":
|
||||
feature_utils.to(get_offload_device())
|
||||
net_video_enc.to(get_offload_device())
|
||||
|
||||
@@ -54,7 +54,6 @@ class SelvaSampler:
|
||||
strategy = model["strategy"]
|
||||
net_generator = model["generator"]
|
||||
feature_utils = model["feature_utils"]
|
||||
mode = model["mode"]
|
||||
|
||||
# Validate that features were extracted with the same model variant
|
||||
feat_variant = features.get("variant")
|
||||
@@ -88,6 +87,7 @@ class SelvaSampler:
|
||||
feature_utils.to(device)
|
||||
soft_empty_cache()
|
||||
|
||||
try:
|
||||
clip_f = features["clip_features"].to(device, dtype) # [1, T_clip, 1024]
|
||||
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."
|
||||
)
|
||||
|
||||
finally:
|
||||
if strategy == "offload_to_cpu":
|
||||
net_generator.to(get_offload_device())
|
||||
feature_utils.to(get_offload_device())
|
||||
|
||||
Reference in New Issue
Block a user