theluyuan/DiffSynth-Studio
1
0
Fork 0
mirror of https://github.com/modelscope/DiffSynth-Studio.git synced 2026-03-18 22:08:13 +00:00
Code Issues Packages Projects Releases Wiki Activity

vram management support torch<2.6.0 (#613)

Browse Source 
support torch<2.6.0
This commit is contained in:
Zhongjie Duan
2025-06-16 13:08:29 +08:00
committed by GitHub
parent 8584e50309
commit c164519ef1
2 changed files with 5 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
diffsynth/pipelines/wan_video_new.py
View File

@@ -143,10 +143,8 @@ class BasePipeline(torch.nn.Module):
self.vram_management_enabled = True
def get_free_vram(self):
total_memory = torch.cuda.get_device_properties(self.device).total_memory
allocated_memory = torch.cuda.device_memory_used(self.device)
return (total_memory - allocated_memory) / (1024 ** 3)
def get_vram(self):
return torch.cuda.mem_get_info(self.device)[1] / (1024 ** 3)
def freeze_except(self, model_names):
@@ -247,7 +245,7 @@ class WanVideoPipeline(BasePipeline):
vram_limit = None
else:
if vram_limit is None:
vram_limit = self.get_free_vram()
vram_limit = self.get_vram()
vram_limit = vram_limit - vram_buffer
if self.text_encoder is not None:
dtype = next(iter(self.text_encoder.parameters())).dtype

3
diffsynth/vram_management/layers.py
View File

@@ -13,7 +13,8 @@ class AutoTorchModule(torch.nn.Module):
super().__init__()
def check_free_vram(self):
used_memory = torch.cuda.device_memory_used(self.computation_device) / (1024 ** 3)
gpu_mem_state = torch.cuda.mem_get_info(self.computation_device)
used_memory = (gpu_mem_state[1] - gpu_mem_state[0]) / (1024 ** 3)
return used_memory < self.vram_limit
def offload(self):