mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
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.
This commit is contained in:
commit
5d31df0cb7
@ -9,6 +9,8 @@ from diffusers.models import UNet2DConditionModel
|
|||||||
from PIL import Image
|
from PIL import Image
|
||||||
from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection
|
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 .attention_processor import AttnProcessor2_0, IPAttnProcessor2_0
|
||||||
from .resampler import Resampler
|
from .resampler import Resampler
|
||||||
|
|
||||||
@ -87,6 +89,20 @@ class IPAdapter:
|
|||||||
if self._attn_processors is not None:
|
if self._attn_processors is not None:
|
||||||
torch.nn.ModuleList(self._attn_processors.values()).to(device=self.device, dtype=self.dtype)
|
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):
|
def _init_image_proj_model(self, state_dict):
|
||||||
return ImageProjModel.from_state_dict(state_dict, self._num_tokens).to(self.device, dtype=self.dtype)
|
return ImageProjModel.from_state_dict(state_dict, self._num_tokens).to(self.device, dtype=self.dtype)
|
||||||
|
|
||||||
|
@ -13,6 +13,7 @@ from invokeai.backend.model_management.models.base import (
|
|||||||
ModelConfigBase,
|
ModelConfigBase,
|
||||||
ModelType,
|
ModelType,
|
||||||
SubModelType,
|
SubModelType,
|
||||||
|
calc_model_size_by_fs,
|
||||||
classproperty,
|
classproperty,
|
||||||
)
|
)
|
||||||
|
|
||||||
@ -30,7 +31,7 @@ class IPAdapterModel(ModelBase):
|
|||||||
assert model_type == ModelType.IPAdapter
|
assert model_type == ModelType.IPAdapter
|
||||||
super().__init__(model_path, base_model, model_type)
|
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
|
@classmethod
|
||||||
def detect_format(cls, path: str) -> str:
|
def detect_format(cls, path: str) -> str:
|
||||||
@ -63,10 +64,13 @@ class IPAdapterModel(ModelBase):
|
|||||||
if child_type is not None:
|
if child_type is not None:
|
||||||
raise ValueError("There are no child models in an IP-Adapter model.")
|
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
|
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
|
@classmethod
|
||||||
def convert_if_required(
|
def convert_if_required(
|
||||||
cls,
|
cls,
|
||||||
|
Loading…
Reference in New Issue
Block a user