From d18cdb18597f525ef8d613a0cb447080fbab8fce Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Tue, 5 May 2026 14:00:56 +0300
Subject: [PATCH] Fix offload on interrupt

---
 nodes_sampler.py | 12 ++++--------
 1 file changed, 4 insertions(+), 8 deletions(-)

diff --git a/nodes_sampler.py b/nodes_sampler.py
index dfe17ad..4dc208f 100644
--- a/nodes_sampler.py
+++ b/nodes_sampler.py
@@ -2593,10 +2593,10 @@ class WanVideoSampler:
 
             except Exception as e:
                 log.error(f"Error during sampling: {e}")
-                if force_offload:
-                    if not model["auto_cpu_offload"]:
-                        offload_transformer(transformer)
-                raise e
+                raise
+            finally:
+                if force_offload and not model["auto_cpu_offload"]:
+                    offload_transformer(transformer)
 
         if phantom_latents is not None:
             latent = latent[:,:-phantom_latents.shape[1]]
@@ -2620,10 +2620,6 @@ class WanVideoSampler:
                     "magcache_state": transformer.magcache_state,
                 }
 
-        if force_offload:
-            if not model["auto_cpu_offload"]:
-                offload_transformer(transformer)
-
         try:
             print_memory(device)
             torch.cuda.reset_peak_memory_stats(device)