nodes/smthemex-ComfyUI_MS_Diffusion
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Update model.py

Browse Source 
fix some
This commit is contained in:
smthemex
2024-07-25 15:10:13 +08:00
committed by GitHub
parent 7fbcade8b1
commit 08aa7c6fe0
1 changed files with 4 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
msdiffusion/models/model.py
+4 -6
View File
@@ -43,13 +43,11 @@ class MSAdapter(torch.nn.Module):
orig_ip_proj_sum = torch.sum(torch.stack([torch.sum(p) for p in self.image_proj_model.parameters()]))
orig_adapter_sum = torch.sum(torch.stack([torch.sum(p) for p in self.adapter_modules.parameters()]))
state_dict = torch.load(ckpt_path, map_location="cpu")
#state_dict = torch.load(ckpt_path, map_location="cpu")
state_dict = torch.load(ckpt_path, map_location="cuda")
# Load state dict for image_proj_model and adapter_modules when using resampler
# self.image_proj_model.load_state_dict(state_dict["image_proj"], strict=False)
# self.adapter_modules.load_state_dict(state_dict["ms_adapter"], strict=False)
self.image_proj_model.load_state_dict(state_dict["image_proj"], strict=True)
self.adapter_modules.load_state_dict(state_dict["ms_adapter"], strict=True)
self.image_proj_model.load_state_dict(state_dict["image_proj"], strict=False)
self.adapter_modules.load_state_dict(state_dict["ms_adapter"], strict=False)
self.load_state_dict({"dummy_image_tokens": state_dict["dummy_image_tokens"]}, strict=False)
# Calculate new checksums