mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2025-06-09 07:37:14 +00:00
123 lines
4.0 KiB
Python
123 lines
4.0 KiB
Python
import hashlib
|
|
import os
|
|
import logging
|
|
import time
|
|
from app.database.models import Model
|
|
from app.database.db import create_session
|
|
from folder_paths import get_relative_path
|
|
|
|
|
|
class ModelProcessor:
|
|
def _validate_path(self, model_path):
|
|
try:
|
|
if not os.path.exists(model_path):
|
|
logging.error(f"Model file not found: {model_path}")
|
|
return None
|
|
|
|
result = get_relative_path(model_path)
|
|
if not result:
|
|
logging.error(
|
|
f"Model file not in a recognized model directory: {model_path}"
|
|
)
|
|
return None
|
|
|
|
return result
|
|
except Exception as e:
|
|
logging.error(f"Error validating model path {model_path}: {str(e)}")
|
|
return None
|
|
|
|
def _hash_file(self, model_path):
|
|
try:
|
|
h = hashlib.sha256()
|
|
with open(model_path, "rb", buffering=0) as f:
|
|
b = bytearray(128 * 1024)
|
|
mv = memoryview(b)
|
|
while n := f.readinto(mv):
|
|
h.update(mv[:n])
|
|
return h.hexdigest()
|
|
except Exception as e:
|
|
logging.error(f"Error hashing file {model_path}: {str(e)}")
|
|
return None
|
|
|
|
def _get_existing_model(self, session, model_type, model_relative_path):
|
|
return (
|
|
session.query(Model)
|
|
.filter(Model.type == model_type)
|
|
.filter(Model.path == model_relative_path)
|
|
.first()
|
|
)
|
|
|
|
def _update_database(
|
|
self, session, model_type, model_relative_path, model_hash, model=None
|
|
):
|
|
try:
|
|
if not model:
|
|
model = self._get_existing_model(
|
|
session, model_type, model_relative_path
|
|
)
|
|
|
|
if not model:
|
|
model = Model(
|
|
path=model_relative_path,
|
|
type=model_type,
|
|
)
|
|
session.add(model)
|
|
|
|
model.hash = model_hash
|
|
session.commit()
|
|
return model
|
|
except Exception as e:
|
|
logging.error(
|
|
f"Error updating database for {model_relative_path}: {str(e)}"
|
|
)
|
|
|
|
def process_file(self, model_path):
|
|
try:
|
|
result = self._validate_path(model_path)
|
|
if not result:
|
|
return
|
|
model_type, model_relative_path = result
|
|
|
|
with create_session() as session:
|
|
existing_model = self._get_existing_model(
|
|
session, model_type, model_relative_path
|
|
)
|
|
if existing_model and existing_model.hash:
|
|
# File exists with hash, no need to process
|
|
return existing_model
|
|
|
|
start_time = time.time()
|
|
logging.info(f"Hashing model {model_relative_path}")
|
|
model_hash = self._hash_file(model_path)
|
|
if not model_hash:
|
|
return
|
|
logging.info(
|
|
f"Model hash: {model_hash} (duration: {time.time() - start_time} seconds)"
|
|
)
|
|
|
|
return self._update_database(session, model_type, model_relative_path, model_hash)
|
|
except Exception as e:
|
|
logging.error(f"Error processing model file {model_path}: {str(e)}")
|
|
|
|
def retrieve_hash(self, model_path, model_type=None):
|
|
try:
|
|
if model_type is not None:
|
|
result = self._validate_path(model_path)
|
|
if not result:
|
|
return None
|
|
model_type, model_relative_path = result
|
|
|
|
with create_session() as session:
|
|
model = self._get_existing_model(
|
|
session, model_type, model_relative_path
|
|
)
|
|
if model and model.hash:
|
|
return model.hash
|
|
return None
|
|
except Exception as e:
|
|
logging.error(f"Error retrieving hash for {model_path}: {str(e)}")
|
|
return None
|
|
|
|
|
|
model_processor = ModelProcessor()
|