diff --git a/invokeai/backend/ip_adapter/ip_adapter.py b/invokeai/backend/ip_adapter/ip_adapter.py index 7f320a2c35..94f3202ba0 100644 --- a/invokeai/backend/ip_adapter/ip_adapter.py +++ b/invokeai/backend/ip_adapter/ip_adapter.py @@ -9,6 +9,8 @@ from diffusers.models import UNet2DConditionModel from PIL import Image from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection +from invokeai.backend.model_management.models.base import calc_model_size_by_data + from .attention_processor import AttnProcessor2_0, IPAttnProcessor2_0 from .resampler import Resampler @@ -87,6 +89,20 @@ class IPAdapter: if self._attn_processors is not None: torch.nn.ModuleList(self._attn_processors.values()).to(device=self.device, dtype=self.dtype) + def calc_size(self): + if self._state_dict is not None: + image_proj_size = sum( + [tensor.nelement() * tensor.element_size() for tensor in self._state_dict["image_proj"].values()] + ) + ip_adapter_size = sum( + [tensor.nelement() * tensor.element_size() for tensor in self._state_dict["ip_adapter"].values()] + ) + return image_proj_size + ip_adapter_size + else: + return calc_model_size_by_data(self._image_proj_model) + calc_model_size_by_data( + torch.nn.ModuleList(self._attn_processors.values()) + ) + def _init_image_proj_model(self, state_dict): return ImageProjModel.from_state_dict(state_dict, self._num_tokens).to(self.device, dtype=self.dtype) diff --git a/invokeai/backend/model_management/models/ip_adapter.py b/invokeai/backend/model_management/models/ip_adapter.py index 8e1e97c9e0..63694af0c8 100644 --- a/invokeai/backend/model_management/models/ip_adapter.py +++ b/invokeai/backend/model_management/models/ip_adapter.py @@ -13,6 +13,7 @@ from invokeai.backend.model_management.models.base import ( ModelConfigBase, ModelType, SubModelType, + calc_model_size_by_fs, classproperty, ) @@ -30,7 +31,7 @@ class IPAdapterModel(ModelBase): assert model_type == ModelType.IPAdapter super().__init__(model_path, base_model, model_type) - self.model_size = os.path.getsize(self.model_path) + self.model_size = calc_model_size_by_fs(self.model_path) @classmethod def detect_format(cls, path: str) -> str: @@ -63,10 +64,13 @@ class IPAdapterModel(ModelBase): if child_type is not None: raise ValueError("There are no child models in an IP-Adapter model.") - return build_ip_adapter( + model = build_ip_adapter( ip_adapter_ckpt_path=os.path.join(self.model_path, "ip_adapter.bin"), device="cpu", dtype=torch_dtype ) + self.model_size = model.calc_size() + return model + @classmethod def convert_if_required( cls,