comi-logic/Wan2GP
7
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

RAM optimizations and faster launch

Browse Source
This commit is contained in:
DeepBeepMeep
2025-03-25 00:14:14 +01:00
parent 7576ccc6ba
commit 1949b61a20
10 changed files with 103 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
wan/modules/vae.py
View File

@@ -744,11 +744,12 @@ def _video_vae(pretrained_path=None, z_dim=None, device='cpu', **kwargs):
with torch.device('meta'):
model = WanVAE_(**cfg)
from mmgp import offload
# load checkpoint
logging.info(f'loading {pretrained_path}')
model.load_state_dict(
torch.load(pretrained_path, map_location=device), assign=True)
# model.load_state_dict(
# torch.load(pretrained_path, map_location=device), assign=True)
offload.load_model_data(model, pretrained_path.replace(".pth", "_bf16.safetensors"), writable_tensors= False)
return model
@@ -778,7 +779,7 @@ class WanVAE:
self.model = _video_vae(
pretrained_path=vae_pth,
z_dim=z_dim,
).eval().requires_grad_(False).to(device)
).eval() #.requires_grad_(False).to(device)
def encode(self, videos, tile_size = 256, any_end_frame = False):
"""