Skip to content

Commit 5d31df0

Browse files
authored
Fix IP-Adapter calculation of memory footprint (#4692)
## What type of PR is this? (check all applicable) - [ ] Refactor - [ ] Feature - [x] Bug Fix - [ ] Optimization - [ ] Documentation Update - [ ] Community Node Submission ## Have you discussed this change with the InvokeAI team? - [ ] Yes - [x] No, because: ## Have you updated all relevant documentation? - [x] Yes - [ ] No ## Description The IP-Adapter memory footprint was not being calculated correctly. I think we could put checks in place to catch this type of error in the future, but for now I'm just fixing the bug. ## QA Instructions, Screenshots, Recordings I tested manually in a debugger. There are 3 pathways for calculating the model size. All were tested: - From file - From state_dict - From model weights ## Added/updated tests? - [ ] Yes - [x] No : This would require the ability to run tests that depend on models. I'm working on this in another branch, but not ready quite yet.
2 parents 062df07 + bd63454 commit 5d31df0

File tree

2 files changed

+22
-2
lines changed

2 files changed

+22
-2
lines changed

invokeai/backend/ip_adapter/ip_adapter.py

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,8 @@
99
from PIL import Image
1010
from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection
1111

12+
from invokeai.backend.model_management.models.base import calc_model_size_by_data
13+
1214
from .attention_processor import AttnProcessor2_0, IPAttnProcessor2_0
1315
from .resampler import Resampler
1416

@@ -87,6 +89,20 @@ def to(self, device: torch.device, dtype: Optional[torch.dtype] = None):
8789
if self._attn_processors is not None:
8890
torch.nn.ModuleList(self._attn_processors.values()).to(device=self.device, dtype=self.dtype)
8991

92+
def calc_size(self):
93+
if self._state_dict is not None:
94+
image_proj_size = sum(
95+
[tensor.nelement() * tensor.element_size() for tensor in self._state_dict["image_proj"].values()]
96+
)
97+
ip_adapter_size = sum(
98+
[tensor.nelement() * tensor.element_size() for tensor in self._state_dict["ip_adapter"].values()]
99+
)
100+
return image_proj_size + ip_adapter_size
101+
else:
102+
return calc_model_size_by_data(self._image_proj_model) + calc_model_size_by_data(
103+
torch.nn.ModuleList(self._attn_processors.values())
104+
)
105+
90106
def _init_image_proj_model(self, state_dict):
91107
return ImageProjModel.from_state_dict(state_dict, self._num_tokens).to(self.device, dtype=self.dtype)
92108

invokeai/backend/model_management/models/ip_adapter.py

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
ModelConfigBase,
1414
ModelType,
1515
SubModelType,
16+
calc_model_size_by_fs,
1617
classproperty,
1718
)
1819

@@ -30,7 +31,7 @@ def __init__(self, model_path: str, base_model: BaseModelType, model_type: Model
3031
assert model_type == ModelType.IPAdapter
3132
super().__init__(model_path, base_model, model_type)
3233

33-
self.model_size = os.path.getsize(self.model_path)
34+
self.model_size = calc_model_size_by_fs(self.model_path)
3435

3536
@classmethod
3637
def detect_format(cls, path: str) -> str:
@@ -63,10 +64,13 @@ def get_model(
6364
if child_type is not None:
6465
raise ValueError("There are no child models in an IP-Adapter model.")
6566

66-
return build_ip_adapter(
67+
model = build_ip_adapter(
6768
ip_adapter_ckpt_path=os.path.join(self.model_path, "ip_adapter.bin"), device="cpu", dtype=torch_dtype
6869
)
6970

71+
self.model_size = model.calc_size()
72+
return model
73+
7074
@classmethod
7175
def convert_if_required(
7276
cls,

0 commit comments

Comments
 (0)