From ba8f06c285475db528de18c49287d13d454cf586 Mon Sep 17 00:00:00 2001
From: Lincoln Stein <lstein@gmail.com>
Date: Sat, 20 Jul 2024 21:08:21 -0400
Subject: [PATCH] add memory usage calculations for controlnet, scheduler,
 tokenizer and upscaler

---
 .../load/model_cache/model_cache_default.py   |  2 +-
 .../backend/model_manager/load/model_util.py  | 24 +++++++++++++++----
 .../app/services/model_load/test_load_api.py  |  2 +-
 3 files changed, 22 insertions(+), 6 deletions(-)

diff --git a/invokeai/backend/model_manager/load/model_cache/model_cache_default.py b/invokeai/backend/model_manager/load/model_cache/model_cache_default.py
index d1a9740fbd..fedcff45df 100644
--- a/invokeai/backend/model_manager/load/model_cache/model_cache_default.py
+++ b/invokeai/backend/model_manager/load/model_cache/model_cache_default.py
@@ -293,7 +293,7 @@ class ModelCache(ModelCacheBase[AnyModel]):
                         new_dict[k] = v.to(target_device, copy=True)
                     cache_entry.model.load_state_dict(new_dict, assign=True)
             try:
-                cache_entry.model.to(target_device, non_blocking=TorchDevice.get_non_blocking(target_device))
+                cache_entry.model.to(target_device)
             except TypeError as e:
                 if "got an unexpected keyword argument 'non_blocking'" in str(e):
                     cache_entry.model.to(target_device)
diff --git a/invokeai/backend/model_manager/load/model_util.py b/invokeai/backend/model_manager/load/model_util.py
index f070a42965..503a02cfd7 100644
--- a/invokeai/backend/model_manager/load/model_util.py
+++ b/invokeai/backend/model_manager/load/model_util.py
@@ -3,8 +3,9 @@
 
 import json
 import logging
+import sys
 from pathlib import Path
-from typing import Optional
+from typing import Any, Optional
 
 import torch
 from diffusers.pipelines.pipeline_utils import DiffusionPipeline
@@ -30,12 +31,14 @@ def calc_model_size_by_data(logger: logging.Logger, model: AnyModel) -> int:
     elif isinstance(model, IAIOnnxRuntimeModel):
         return _calc_onnx_model_by_data(model)
     elif isinstance(model, SchedulerMixin):
-        return 0
+        assert hasattr(model, "config")  # size is dominated by config
+        return sys.getsizeof(model.config)
     elif isinstance(model, CLIPTokenizer):
-        # TODO(ryand): Accurately calculate the tokenizer's size. It's small enough that it shouldn't matter for now.
-        return 0
+        return sys.getsizeof(model.get_vocab())  # size is dominated by the vocab dict
     elif isinstance(model, (TextualInversionModelRaw, IPAdapter, LoRAModelRaw, SpandrelImageToImageModel)):
         return model.calc_size()
+    elif isinstance(model, dict):
+        return _calc_size_from_dict(model, logger)
     else:
         # TODO(ryand): Promote this from a log to an exception once we are confident that we are handling all of the
         # supported model types.
@@ -70,6 +73,19 @@ def _calc_onnx_model_by_data(model: IAIOnnxRuntimeModel) -> int:
     return mem
 
 
+def _calc_size_from_dict(model: dict[str, Any] | torch.Tensor | torch.nn.Module, logger: logging.Logger) -> int:
+    total = sys.getsizeof(model)  # get python overhead for object
+    if isinstance(model, dict):
+        total += sum(_calc_size_from_dict(model[x], logger) for x in model.keys())
+    elif isinstance(model, torch.Tensor):
+        total += model.element_size() * model.nelement()
+    elif isinstance(model, torch.nn.Module):
+        total += calc_module_size(model)
+    else:
+        logger.warning(f"Failed to calculate model size for unexpected model type: {type(model)}.")
+    return total
+
+
 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():
diff --git a/tests/app/services/model_load/test_load_api.py b/tests/app/services/model_load/test_load_api.py
index 17628cfc53..21018c8dbf 100644
--- a/tests/app/services/model_load/test_load_api.py
+++ b/tests/app/services/model_load/test_load_api.py
@@ -88,6 +88,7 @@ def test_download_diffusers_subfolder(mock_context: InvocationContext) -> None:
         model_path / "diffusion_pytorch_model.safetensors"
     ).exists()
 
+
 def test_download_diffusers_preserve_subfolders(mock_context: InvocationContext) -> None:
     model_path = mock_context.models.download_and_cache_model(
         "stabilityai/sdxl-turbo::/vae",
@@ -98,4 +99,3 @@ def test_download_diffusers_preserve_subfolders(mock_context: InvocationContext)
     assert (model_path / "diffusion_pytorch_model.fp16.safetensors").exists() or (
         model_path / "diffusion_pytorch_model.safetensors"
     ).exists()
-