mirror of
https://git.datalinker.icu/comfyanonymous/ComfyUI
synced 2026-11-12 05:06:37 +08:00
t # This is a combination of 2 commits.
Add model cache after loaded
This commit is contained in:
david
parent
71ec5b144e
commit
e821321c6c
98
comfy/cache.py
Normal file
98
comfy/cache.py
Normal file
@ -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()
|
||||
@ -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,14 +467,30 @@ 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:
|
||||
to_delete = [i] + to_delete
|
||||
#TODO: find a less fragile way to do this.
|
||||
elif sys.getrefcount(current_loaded_models[i].real_model) <= 3: #references from .real_model + the .model
|
||||
# TODO: find a less fragile way to do this.
|
||||
elif sys.getrefcount(current_loaded_models[i].real_model) <= 3: # references from .real_model + the .model
|
||||
to_delete = [i] + to_delete
|
||||
|
||||
for i in to_delete:
|
||||
|
||||
91
comfy/sd.py
91
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,48 +465,70 @@ 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:
|
||||
clipvision = clip_vision.load_clipvision_from_sd(sd, model_config.clip_vision_prefix, True)
|
||||
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)
|
||||
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.")
|
||||
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 output_vae:
|
||||
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)
|
||||
|
||||
if output_clip:
|
||||
clip_target = model_config.clip_target()
|
||||
if clip_target is not None:
|
||||
clip_sd = model_config.process_clip_state_dict(sd)
|
||||
if len(clip_sd) > 0:
|
||||
clip = CLIP(clip_target, embedding_directory=embedding_directory)
|
||||
m, u = clip.load_sd(clip_sd, full_model=True)
|
||||
if len(m) > 0:
|
||||
m_filter = list(filter(lambda a: ".logit_scale" not in a and ".transformer.text_projection.weight" not in a, m))
|
||||
if len(m_filter) > 0:
|
||||
logging.warning("clip missing: {}".format(m))
|
||||
else:
|
||||
logging.debug("clip missing: {}".format(m))
|
||||
|
||||
if len(u) > 0:
|
||||
logging.debug("clip unexpected {}:".format(u))
|
||||
else:
|
||||
logging.warning("no CLIP/text encoder weights in checkpoint, the text encoder model will not be loaded.")
|
||||
|
||||
left_over = sd.keys()
|
||||
if len(left_over) > 0:
|
||||
logging.debug("left over keys: {}".format(left_over))
|
||||
|
||||
if output_model:
|
||||
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)
|
||||
|
||||
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)
|
||||
if len(clip_sd) > 0:
|
||||
clip = CLIP(clip_target, embedding_directory=embedding_directory)
|
||||
m, u = clip.load_sd(clip_sd, full_model=True)
|
||||
if len(m) > 0:
|
||||
m_filter = list(filter(lambda a: ".logit_scale" not in a and ".transformer.text_projection.weight" not in a, m))
|
||||
if len(m_filter) > 0:
|
||||
logging.warning("clip missing: {}".format(m))
|
||||
else:
|
||||
logging.debug("clip missing: {}".format(m))
|
||||
|
||||
if len(u) > 0:
|
||||
logging.debug("clip unexpected {}:".format(u))
|
||||
else:
|
||||
logging.warning("no CLIP/text encoder weights in checkpoint, the text encoder model will not be loaded.")
|
||||
|
||||
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)
|
||||
|
||||
model_cache.refresh_cache(clipvision)
|
||||
return (model_patcher, clip, vae, clipvision)
|
||||
|
||||
|
||||
|
||||
@ -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):
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user