From 6978686272ffc2258b0cf3083b394414520163b9 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Fri, 29 Aug 2025 15:16:02 +0300
Subject: [PATCH] Update model.py

---
 wanvideo/modules/model.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/wanvideo/modules/model.py b/wanvideo/modules/model.py
index 3bde4c8..756f1c7 100644
--- a/wanvideo/modules/model.py
+++ b/wanvideo/modules/model.py
@@ -2516,7 +2516,7 @@ class WanModel(torch.nn.Module):
                     self.controlnet.to(self.offload_device)
 
             # Asynchronous block offloading with CUDA streams and events
-            cuda_stream = mm.get_offload_stream(device)
+            cuda_stream = torch.cuda.Stream(device=device, priority=0)
             events = [torch.cuda.Event() for _ in self.blocks]
             swap_start_idx = len(self.blocks) - self.blocks_to_swap if self.blocks_to_swap > 0 else len(self.blocks)