mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2025-04-19 10:53:29 +00:00
64 lines
1.7 KiB
Python
64 lines
1.7 KiB
Python
import hashlib
|
|
import logging
|
|
import threading
|
|
import time
|
|
from comfy.cli_args import args
|
|
|
|
|
|
class ModelHasher:
|
|
|
|
def __init__(self):
|
|
self._thread = None
|
|
self._lock = threading.Lock()
|
|
self._model_entity = None
|
|
|
|
def start(self, model_entity):
|
|
if args.disable_model_hashing:
|
|
return
|
|
|
|
self._model_entity = model_entity
|
|
|
|
if self._thread is None:
|
|
# Lock to prevent multiple threads from starting
|
|
with self._lock:
|
|
if self._thread is None:
|
|
self._thread = threading.Thread(target=self._hash_models)
|
|
self._thread.daemon = True
|
|
self._thread.start()
|
|
|
|
def _get_models(self):
|
|
models = self._model_entity.get("WHERE hash IS NULL")
|
|
return models
|
|
|
|
def _hash_model(self, model_path):
|
|
h = hashlib.sha256()
|
|
b = bytearray(128 * 1024)
|
|
mv = memoryview(b)
|
|
with open(model_path, "rb", buffering=0) as f:
|
|
while n := f.readinto(mv):
|
|
h.update(mv[:n])
|
|
hash = h.hexdigest()
|
|
return hash
|
|
|
|
def _hash_models(self):
|
|
while True:
|
|
models = self._get_models()
|
|
|
|
if len(models) == 0:
|
|
break
|
|
|
|
for model in models:
|
|
time.sleep(0)
|
|
now = time.time()
|
|
logging.info(f"Hashing model {model['path']}")
|
|
hash = self._hash_model(model["path"])
|
|
logging.info(
|
|
f"Hashed model {model['path']} in {time.time() - now} seconds"
|
|
)
|
|
self._model_entity.update((model["id"],), {"hash": hash})
|
|
|
|
self._thread = None
|
|
|
|
|
|
model_hasher = ModelHasher()
|