From 44f62944ee108b84fe228a422560fe71e294e84b Mon Sep 17 00:00:00 2001 From: Ryan Dick Date: Tue, 2 Jul 2024 11:55:05 -0400 Subject: [PATCH] Fix circular import caused by the organization the model size utils. --- invokeai/backend/ip_adapter/ip_adapter.py | 6 +- .../model_manager/load/model_size_utils.py | 79 ++++++++++++++++++ .../backend/model_manager/load/model_util.py | 82 +------------------ 3 files changed, 84 insertions(+), 83 deletions(-) create mode 100644 invokeai/backend/model_manager/load/model_size_utils.py diff --git a/invokeai/backend/ip_adapter/ip_adapter.py b/invokeai/backend/ip_adapter/ip_adapter.py index c33cb3f4ab..12ea0ee142 100644 --- a/invokeai/backend/ip_adapter/ip_adapter.py +++ b/invokeai/backend/ip_adapter/ip_adapter.py @@ -11,6 +11,7 @@ from PIL import Image from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection from invokeai.backend.ip_adapter.ip_attention_weights import IPAttentionWeights +from invokeai.backend.model_manager.load.model_size_utils import calc_module_size from ..raw_model import RawModel from .resampler import Resampler @@ -137,10 +138,7 @@ class IPAdapter(RawModel): 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 - - 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] diff --git a/invokeai/backend/model_manager/load/model_size_utils.py b/invokeai/backend/model_manager/load/model_size_utils.py new file mode 100644 index 0000000000..587c4e4c4c --- /dev/null +++ b/invokeai/backend/model_manager/load/model_size_utils.py @@ -0,0 +1,79 @@ +import json +from pathlib import Path +from typing import Optional + +import torch + + +def calc_module_size(model: torch.nn.Module) -> int: + """Estimate the size of a torch.nn.Module in bytes.""" + mem_params = sum([param.nelement() * param.element_size() for param in model.parameters()]) + mem_bufs = sum([buf.nelement() * buf.element_size() for buf in model.buffers()]) + mem: int = mem_params + mem_bufs # in bytes + return mem + + +def calc_model_size_by_fs(model_path: Path, subfolder: Optional[str] = None, variant: Optional[str] = None) -> int: + """Estimate the size of a model on disk in bytes.""" + if model_path.is_file(): + return model_path.stat().st_size + + if subfolder is not None: + model_path = model_path / subfolder + + # this can happen when, for example, the safety checker is not downloaded. + if not model_path.exists(): + return 0 + + all_files = [f for f in model_path.iterdir() if (model_path / f).is_file()] + + fp16_files = {f for f in all_files if ".fp16." in f.name or ".fp16-" in f.name} + bit8_files = {f for f in all_files if ".8bit." in f.name or ".8bit-" in f.name} + other_files = set(all_files) - fp16_files - bit8_files + + if not variant: # ModelRepoVariant.DEFAULT evaluates to empty string for compatability with HF + files = other_files + elif variant == "fp16": + files = fp16_files + elif variant == "8bit": + files = bit8_files + else: + raise NotImplementedError(f"Unknown variant: {variant}") + + # try read from index if exists + index_postfix = ".index.json" + if variant is not None: + index_postfix = f".index.{variant}.json" + + for file in files: + if not file.name.endswith(index_postfix): + continue + try: + with open(model_path / file, "r") as f: + index_data = json.loads(f.read()) + return int(index_data["metadata"]["total_size"]) + except Exception: + pass + + # calculate files size if there is no index file + formats = [ + (".safetensors",), # safetensors + (".bin",), # torch + (".onnx", ".pb"), # onnx + (".msgpack",), # flax + (".ckpt",), # tf + (".h5",), # tf2 + ] + + for file_format in formats: + model_files = [f for f in files if f.suffix in file_format] + if len(model_files) == 0: + continue + + model_size = 0 + for model_file in model_files: + file_stats = (model_path / model_file).stat() + model_size += file_stats.st_size + return model_size + + return 0 # scheduler/feature_extractor/tokenizer - models without loading to gpu diff --git a/invokeai/backend/model_manager/load/model_util.py b/invokeai/backend/model_manager/load/model_util.py index c55eee48fa..f35e78dda2 100644 --- a/invokeai/backend/model_manager/load/model_util.py +++ b/invokeai/backend/model_manager/load/model_util.py @@ -1,14 +1,11 @@ # Copyright (c) 2024 The InvokeAI Development Team """Various utility functions needed by the loader and caching system.""" -import json -from pathlib import Path -from typing import Optional - import torch from diffusers import DiffusionPipeline from invokeai.backend.model_manager.config import AnyModel +from invokeai.backend.model_manager.load.model_size_utils import calc_module_size from invokeai.backend.onnx.onnx_runtime import IAIOnnxRuntimeModel @@ -17,7 +14,7 @@ def calc_model_size_by_data(model: AnyModel) -> int: if isinstance(model, DiffusionPipeline): return _calc_pipeline_by_data(model) elif isinstance(model, torch.nn.Module): - return _calc_model_by_data(model) + return calc_module_size(model) elif isinstance(model, IAIOnnxRuntimeModel): return _calc_onnx_model_by_data(model) else: @@ -30,84 +27,11 @@ def _calc_pipeline_by_data(pipeline: DiffusionPipeline) -> int: for submodel_key in pipeline.components.keys(): submodel = getattr(pipeline, submodel_key) if submodel is not None and isinstance(submodel, torch.nn.Module): - res += _calc_model_by_data(submodel) + res += calc_module_size(submodel) return res -def _calc_model_by_data(model: torch.nn.Module) -> int: - mem_params = sum([param.nelement() * param.element_size() for param in model.parameters()]) - mem_bufs = sum([buf.nelement() * buf.element_size() for buf in model.buffers()]) - mem: int = mem_params + mem_bufs # in bytes - return mem - - def _calc_onnx_model_by_data(model: IAIOnnxRuntimeModel) -> int: tensor_size = model.tensors.size() * 2 # The session doubles this mem = tensor_size # in bytes return mem - - -def calc_model_size_by_fs(model_path: Path, subfolder: Optional[str] = None, variant: Optional[str] = None) -> int: - """Estimate the size of a model on disk in bytes.""" - if model_path.is_file(): - return model_path.stat().st_size - - if subfolder is not None: - model_path = model_path / subfolder - - # this can happen when, for example, the safety checker is not downloaded. - if not model_path.exists(): - return 0 - - all_files = [f for f in model_path.iterdir() if (model_path / f).is_file()] - - fp16_files = {f for f in all_files if ".fp16." in f.name or ".fp16-" in f.name} - bit8_files = {f for f in all_files if ".8bit." in f.name or ".8bit-" in f.name} - other_files = set(all_files) - fp16_files - bit8_files - - if not variant: # ModelRepoVariant.DEFAULT evaluates to empty string for compatability with HF - files = other_files - elif variant == "fp16": - files = fp16_files - elif variant == "8bit": - files = bit8_files - else: - raise NotImplementedError(f"Unknown variant: {variant}") - - # try read from index if exists - index_postfix = ".index.json" - if variant is not None: - index_postfix = f".index.{variant}.json" - - for file in files: - if not file.name.endswith(index_postfix): - continue - try: - with open(model_path / file, "r") as f: - index_data = json.loads(f.read()) - return int(index_data["metadata"]["total_size"]) - except Exception: - pass - - # calculate files size if there is no index file - formats = [ - (".safetensors",), # safetensors - (".bin",), # torch - (".onnx", ".pb"), # onnx - (".msgpack",), # flax - (".ckpt",), # tf - (".h5",), # tf2 - ] - - for file_format in formats: - model_files = [f for f in files if f.suffix in file_format] - if len(model_files) == 0: - continue - - model_size = 0 - for model_file in model_files: - file_stats = (model_path / model_file).stat() - model_size += file_stats.st_size - return model_size - - return 0 # scheduler/feature_extractor/tokenizer - models without loading to gpu