From c33ba8b05cbed7eb62dd6468924c5277a84b292b Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Wed, 2 Apr 2025 18:27:40 +0300
Subject: [PATCH] Update model.py

---
 wanvideo/modules/model.py | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/wanvideo/modules/model.py b/wanvideo/modules/model.py
index fd4d6e3..8a211fb 100644
--- a/wanvideo/modules/model.py
+++ b/wanvideo/modules/model.py
@@ -820,15 +820,15 @@ class WanModel(ModelMixin, ConfigMixin):
         if vace_blocks_to_swap is not None:
             self.vace_blocks_to_swap = vace_blocks_to_swap
 
-        for b, block in tqdm(enumerate(self.vace_blocks), total=len(self.vace_blocks), desc="Initializing vace block swap"):
-            block_memory = get_module_memory_mb(block)
-            
-            if b > self.vace_blocks_to_swap:
-                block.to(self.main_device)
-                total_main_memory += block_memory
-            else:
-                block.to(self.offload_device, non_blocking=self.use_non_blocking)
-                total_offload_memory += block_memory
+            for b, block in tqdm(enumerate(self.vace_blocks), total=len(self.vace_blocks), desc="Initializing vace block swap"):
+                block_memory = get_module_memory_mb(block)
+                
+                if b > self.vace_blocks_to_swap:
+                    block.to(self.main_device)
+                    total_main_memory += block_memory
+                else:
+                    block.to(self.offload_device, non_blocking=self.use_non_blocking)
+                    total_offload_memory += block_memory
 
         mm.soft_empty_cache()
         gc.collect()