Zengtudor/Wan2.1
1
0
Fork 0
mirror of https://github.com/Wan-Video/Wan2.1.git synced 2025-06-07 07:44:53 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Check for cuda is available for macos

Browse Source
This commit is contained in:
Bakhtiyor Sulaymonov 2025-02-26 22:18:18 +05:00
parent 9ab8f963c8
commit 9d3d4d784f
1 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
wan/text2video.py
View File

@ -57,7 +57,7 @@ class WanT2V:
t5_cpu (`bool`, *optional*, defaults to False):
Whether to place T5 model on CPU. Only works without t5_fsdp.
"""
self.device = torch.device(f"cuda:{device_id}")
self.device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
self.config = config
self.rank = rank
self.t5_cpu = t5_cpu
@ -259,7 +259,8 @@ class WanT2V:
del sample_scheduler
if offload_model:
gc.collect()
torch.cuda.synchronize()
if torch.cuda.is_available():
torch.cuda.synchronize()
if dist.is_initialized():
dist.barrier()