From 08aa7c6fe04ed262ccbc5caf36ff25ef8a00c7b9 Mon Sep 17 00:00:00 2001 From: smthemex <138738845+smthemex@users.noreply.github.com> Date: Thu, 25 Jul 2024 15:10:13 +0800 Subject: [PATCH] Update model.py fix some --- msdiffusion/models/model.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/msdiffusion/models/model.py b/msdiffusion/models/model.py index 04a9078..5afdfd6 100644 --- a/msdiffusion/models/model.py +++ b/msdiffusion/models/model.py @@ -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