Update calc_model_size_by_data(...) to handle all expected model types, and to log an error if an unexpected model type is received.

This commit is contained in:
Ryan Dick
2024-07-02 21:14:12 -04:00
parent 0fe92cd406
commit 414750a45d
4 changed files with 40 additions and 10 deletions

View File

@ -136,11 +136,11 @@ class IPAdapter(RawModel):
self._image_proj_model.to(device=self.device, dtype=self.dtype, non_blocking=non_blocking)
self.attn_weights.to(device=self.device, dtype=self.dtype, non_blocking=non_blocking)
def calc_size(self):
# workaround for circular import
from invokeai.backend.model_manager.load.model_util import calc_model_size_by_data
def calc_size(self) -> int:
# HACK(ryand): Fix this issue with circular imports.
from invokeai.backend.model_manager.load.model_util import calc_module_size
return calc_model_size_by_data(self._image_proj_model) + calc_model_size_by_data(self.attn_weights)
return calc_module_size(self._image_proj_model) + calc_module_size(self.attn_weights)
def _init_image_proj_model(
self, state_dict: dict[str, torch.Tensor]