add memory usage calculations for controlnet, scheduler, tokenizer and upscaler

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)

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():

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()
-