t # This is a combination of 2 commits.

Add model cache after loaded

comfy/cache.py

@@ -0,0 +1,98 @@
+import gc
+from collections import OrderedDict
+import os
+import logging
+
+from comfy.cli_args import args
+
+
+class ModelCache:
+    def __init__(self):
+        self._cache = OrderedDict()
+        # ignore gpu and highvram state, may cause OOM error
+        self._cache_state = False if (args.highvram or args.gpu_only) else True
+
+    def _get_item(self, key):
+        if self._cache.get(key) is None:
+            self._cache[key] = {}
+        return self._cache[key]
+
+    def cache_model(self, key, model):
+        if not self._cache_state:
+            return
+        item = self._get_item(key)
+        item['model'] = model
+
+    def cache_clip_vision(self, key, clip_vision):
+        if not self._cache_state:
+            return
+        item = self._get_item(key)
+        item['clip_vision'] = clip_vision
+
+    def cache_vae(self, key, vae):
+        if not self._cache_state:
+            return
+        item = self._get_item(key)
+        item['vae'] = vae
+
+    def cache_sd(self, key, sd):
+        if not self._cache_state:
+            return
+        assert isinstance(sd, dict)
+        keys = list(sd.keys())
+        values = list(sd.values())
+        item = self._get_item(key)
+        item['sd'] = (keys, values)
+
+    def cache_clip(self, key, clip_key, clip):
+        item = self._get_item(key)
+        item[clip_key] = clip
+
+    def refresh_cache(self, key):
+        if key in self._cache:
+            self._cache.move_to_end(key)
+
+    def get_item(self, key, prop):
+        item = self._cache.get(key)
+        if item is None:
+            return None
+        if prop == "sd":
+            if item.get('sd') is None:
+                return None
+            k, values = item.get("sd")
+            return dict(zip(k, values))
+        return item.get(prop)
+
+    def __len__(self):
+        return len(self._cache)
+
+    @property
+    def cache(self):
+        return self._cache
+
+    @staticmethod
+    def unpatch_offload_model(model):
+        model.model_patches_to(model.offload_device)
+
+    def free_one_model_cache(self):
+        if len(self) == 0:
+            return
+        cache_k, item = self._cache.popitem(last=False)
+        item.pop("sd", None)
+
+        for k in list(item.keys()):
+            model = item.pop(k, None)
+            if model is not None:
+                if hasattr(model, "patcher"):
+                    self.unpatch_offload_model(model.patcher)
+                else:
+                    self.unpatch_offload_model(model)
+
+        item.clear()
+        model_dir, model_name = os.path.split(cache_k)
+        dir_name = os.path.basename(model_dir)
+        gc.collect()
+        logging.info(f"Drop model cache: {model_name} ({dir_name})")
+
+
+model_cache = ModelCache()

comfy/model_management.py

@@ -5,6 +5,7 @@ from comfy.cli_args import args
 import torch
 import sys
 import platform
+from comfy.cache import model_cache
 
 class VRAMState(Enum):
     DISABLED = 0    #No vram present: no need to move models to vram
@@ -466,8 +467,24 @@ def load_models_gpu(models, memory_required=0, force_patch_weights=False):
 def load_model_gpu(model):
     return load_models_gpu([model])
 
+
+def get_cpu_memory():
+    return round(get_free_memory(torch.device("cpu")) / 1024 ** 3, 2)
+
+
+def check_and_free_cpu_memory():
+    # set free cpu memory size threshold flag to make free action
+    current_free_memory_cpu = get_cpu_memory()
+    while (current_free_memory_cpu <= 8) and len(model_cache) > 0:
+        model_cache.free_one_model_cache()   
+        soft_empty_cache(force=True)
+        current_free_memory_cpu = get_cpu_memory()
+
+
 def cleanup_models(keep_clone_weights_loaded=False):
+    check_and_free_cpu_memory()
     to_delete = []
+
     for i in range(len(current_loaded_models)):
         if sys.getrefcount(current_loaded_models[i].model) <= 2:
             if not keep_clone_weights_loaded:

comfy/sd.py

@@ -19,6 +19,7 @@ from . import model_detection
 from . import sd1_clip
 from . import sd2_clip
 from . import sdxl_clip
+from .cache import model_cache
 
 import comfy.model_patcher
 import comfy.lora
@@ -464,20 +465,35 @@ def load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, o
 
     if model_config.clip_vision_prefix is not None:
         if output_clipvision:
+            if output_clipvision:
+                clipvision = model_cache.get_item(ckpt_path, 'clipvision')
+                if clipvision is None:
                     clipvision = clip_vision.load_clipvision_from_sd(sd, model_config.clip_vision_prefix, True)
 
     if output_model:
+        model_patcher = model_cache.get_item(ckpt_path, 'model')
         inital_load_device = model_management.unet_inital_load_device(parameters, unet_dtype)
+        if model_patcher is None:
             offload_device = model_management.unet_offload_device()
             model = model_config.get_model(sd, "model.diffusion_model.", device=inital_load_device)
             model.load_model_weights(sd, "model.diffusion_model.")
+            model_patcher = comfy.model_patcher.ModelPatcher(model, load_device=load_device, offload_device=offload_device, current_device=inital_load_device)
+
+        if inital_load_device != torch.device("cpu"):
+            logging.info("loaded straight to GPU")
+            model_management.load_model_gpu(model_patcher)
 
     if output_vae:
+        vae = model_cache.get_item(ckpt_path, 'vae')
+        if vae is None:
             vae_sd = comfy.utils.state_dict_prefix_replace(sd, {k: "" for k in model_config.vae_key_prefix}, filter_keys=True)
             vae_sd = model_config.process_vae_state_dict(vae_sd)
             vae = VAE(sd=vae_sd)
 
+    clip_key = f"clip_{embedding_directory}"
     if output_clip:
+        clip = model_cache.get_item(ckpt_path, clip_key)
+        if clip is None:
             clip_target = model_config.clip_target()
             if clip_target is not None:
                 clip_sd = model_config.process_clip_state_dict(sd)
@@ -499,13 +515,20 @@ def load_checkpoint_guess_config(ckpt_path, output_vae=True, output_clip=True, o
     left_over = sd.keys()
     if len(left_over) > 0:
         logging.debug("left over keys: {}".format(left_over))
+    if model:
+        logging.debug(f"cache model of : {ckpt_path}")
+        model_cache.cache_model(ckpt_path, model_patcher)
+    if clip:
+        logging.debug(f"cache clip of : {ckpt_path}")
+        model_cache.cache_clip(ckpt_path, clip_key, clip)
+    if vae:
+        logging.debug(f"cache vae of : {ckpt_path}")
+        model_cache.cache_vae(ckpt_path, vae)
+    if clipvision:
+        logging.debug(f"cache clipvision of : {ckpt_path}")
+        model_cache.cache_clip_vision(ckpt_path, clipvision)
 
-    if output_model:
+    model_cache.refresh_cache(clipvision)
-        model_patcher = comfy.model_patcher.ModelPatcher(model, load_device=load_device, offload_device=model_management.unet_offload_device(), current_device=inital_load_device)
-        if inital_load_device != torch.device("cpu"):
-            logging.info("loaded straight to GPU")
-            model_management.load_model_gpu(model_patcher)
-
     return (model_patcher, clip, vae, clipvision)
 
 

comfy/utils.py

@@ -6,10 +6,17 @@ import safetensors.torch
 import numpy as np
 from PIL import Image
 import logging
+from comfy.cache import model_cache
 
 def load_torch_file(ckpt, safe_load=False, device=None):
+    # read the sd from cache
+    cache_sd = model_cache.get_item(ckpt, 'sd')
+    if cache_sd:
+        return cache_sd
+    
     if device is None:
         device = torch.device("cpu")
+   
     if ckpt.lower().endswith(".safetensors"):
         sd = safetensors.torch.load_file(ckpt, device=device.type)
     else:
@@ -27,6 +34,8 @@ def load_torch_file(ckpt, safe_load=False, device=None):
             sd = pl_sd["state_dict"]
         else:
             sd = pl_sd
+    # save the references of Tensor to cache
+    model_cache.cache_sd(ckpt, sd)
     return sd
 
 def save_torch_file(sd, ckpt, metadata=None):