mirror of
https://git.datalinker.icu/comfyanonymous/ComfyUI
synced 2026-11-12 06:00:04 +08:00
101 lines
2.7 KiB
Python
101 lines
2.7 KiB
Python
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):
|
|
if not self._cache_state:
|
|
return
|
|
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()
|