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)