shenglin.net/ComfyUI-MochiWrapper-2
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix non-accelerate model loading

Browse Source
This commit is contained in:
kijai 2024-10-23 20:40:42 +03:00
parent fae591800c
commit a32064eefb
1 changed files with 2 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
mochi_preview/t2v_synth_mochi.py
View File

@ -2,7 +2,6 @@ import json
import random import random
from typing import Dict, List from typing import Dict, List
from safetensors.torch import load_file
import numpy as np import numpy as np
import torch import torch
import torch.nn as nn import torch.nn as nn
@ -163,7 +162,7 @@ class T2VSynthMochiModel:
else: else:
print("Loading state_dict without accelerate...") print("Loading state_dict without accelerate...")
model.load_state_dict(dit_sd) model.load_state_dict(dit_sd)
for name, param in self.dit.named_parameters(): for name, param in model.named_parameters():
if not any(keyword in name for keyword in params_to_keep): if not any(keyword in name for keyword in params_to_keep):
param.data = param.data.to(weight_dtype) param.data = param.data.to(weight_dtype)
else: else:
@ -180,7 +179,7 @@ class T2VSynthMochiModel:
self.vae_mean = torch.Tensor(vae_stats["mean"]).to(self.device) self.vae_mean = torch.Tensor(vae_stats["mean"]).to(self.device)
self.vae_std = torch.Tensor(vae_stats["std"]).to(self.device) self.vae_std = torch.Tensor(vae_stats["std"]).to(self.device)
#t.print_stats() t.print_stats()
def get_conditioning(self, prompts, *, zero_last_n_prompts: int): def get_conditioning(self, prompts, *, zero_last_n_prompts: int):
B = len(prompts) B = len(prompts)