236 lines
8.5 KiB
Python
236 lines
8.5 KiB
Python
"""Model metadata/cache maintenance helpers extracted from helpers facade."""
|
|
|
|
import datetime
|
|
|
|
from .helpers_civitai import (
|
|
get_civitai_model_version_json_by_hash,
|
|
get_civitai_model_version_json_by_id,
|
|
get_latest_model_version,
|
|
)
|
|
from .logger import get_logger
|
|
from .model_cache import cache
|
|
from .file_utils import days_since_last_used, get_file_modification_date, get_file_sha256
|
|
from .type_utils import str_to_bool
|
|
|
|
logger = get_logger('model.metadata')
|
|
|
|
|
|
def update_cache_from_civitai_json(file_path, json_data, timestamp=True):
|
|
"""Update cache entry from a successful CivitAI model-version payload."""
|
|
the_files = json_data.get("files", [])
|
|
hashes = {}
|
|
|
|
if len(the_files) > 0:
|
|
hashes = the_files[0].get("hashes", {})
|
|
update_available = True
|
|
|
|
latest_model = None
|
|
if json_data.get("modelId", None) is not None:
|
|
latest_model = get_latest_model_version(json_data["modelId"])
|
|
if latest_model == json_data["id"] or latest_model is None:
|
|
update_available = False
|
|
|
|
if latest_model is None:
|
|
latest_model = ""
|
|
|
|
file_cache = cache.by_path(file_path)
|
|
file_cache.update({
|
|
'civitai': "True",
|
|
'civitai_failed_count': 0,
|
|
'model': json_data.get("model", {}),
|
|
'name': json_data.get("name", ""),
|
|
'baseModel': json_data.get("baseModel", ""),
|
|
'id': json_data.get("id", ""),
|
|
'modelId': json_data.get("modelId", ""),
|
|
'update_available': update_available,
|
|
'update_version_id': latest_model,
|
|
'trainedWords': json_data.get("trainedWords", []),
|
|
'downloadUrl': json_data.get("downloadUrl", ""),
|
|
'hashes': hashes,
|
|
})
|
|
if timestamp:
|
|
logger.info("Updating timestamp.")
|
|
cache.update_last_used_by_path(file_path)
|
|
|
|
logger.info("Successfully pulled metadata.")
|
|
|
|
|
|
def update_cache_without_civitai_json(file_path, hash_value, timestamp=True):
|
|
"""Update cache entry when CivitAI metadata is unavailable."""
|
|
file_cache = cache.by_path(file_path)
|
|
logger.info("Unable to find metadata on CivitAI.")
|
|
file_cache['civitai'] = "False"
|
|
file_cache['civitai_failed_count'] = file_cache.get('civitai_failed_count', 0) + 1
|
|
file_cache['hash'] = hash_value
|
|
if timestamp:
|
|
cache.update_last_used_by_path(file_path)
|
|
|
|
|
|
def add_file_to_cache(file_path, hash_value=None):
|
|
"""Ensure a file path has hash/info cache entries and return its hash."""
|
|
file_path = str(file_path)
|
|
logger.info(f"Adding {file_path} to cache.")
|
|
if hash_value is None:
|
|
hash_value = get_file_sha256(file_path)
|
|
|
|
if file_path not in cache.hash:
|
|
cache.hash[file_path] = hash_value
|
|
|
|
if cache.info.get(hash_value, None) is None:
|
|
cache.info[hash_value] = {
|
|
'civitai': "False",
|
|
'update_available': False,
|
|
'update_version_id': "",
|
|
'hash': hash_value,
|
|
'lastUsed': datetime.datetime.now().isoformat(),
|
|
}
|
|
|
|
logger.info(f"Added {file_path} to cache with hash {hash_value}.")
|
|
return hash_value
|
|
|
|
|
|
def recheck_hash(file_path, hash_value):
|
|
"""Recompute a file hash and migrate cache references if it changed."""
|
|
new_hash = get_file_sha256(file_path)
|
|
if new_hash != hash_value:
|
|
logger.info("Hash mismatch detected. Using new hash.")
|
|
if file_path in cache.hash:
|
|
logger.info(f"Updating cache for {file_path} with new hash {new_hash}.")
|
|
if new_hash not in cache.info and hash_value in cache.info:
|
|
cache.info[new_hash] = cache.info[hash_value]
|
|
else:
|
|
logger.info(f"File {file_path} not in cache. Adding with new hash {new_hash}.")
|
|
add_file_to_cache(file_path, new_hash)
|
|
hash_value = new_hash
|
|
return hash_value
|
|
|
|
|
|
def pull_and_update_model_timestamp(file_paths, model_type):
|
|
"""Pull metadata for one-or-many model paths and update last-used timestamps."""
|
|
if not isinstance(file_paths, (list, tuple)):
|
|
file_paths = [file_paths]
|
|
|
|
for path in file_paths:
|
|
pull_metadata(path, model_type=model_type)
|
|
|
|
update_model_timestamp(file_paths)
|
|
|
|
|
|
def update_model_timestamp(file_paths):
|
|
"""Update last-used timestamps for one-or-many model paths in cache."""
|
|
cache.load()
|
|
if not isinstance(file_paths, (list, tuple)):
|
|
file_paths = [file_paths]
|
|
|
|
for path in file_paths:
|
|
if path in cache.hash:
|
|
cache.update_last_used_by_path(path)
|
|
cache.save()
|
|
|
|
|
|
def pull_metadata(file_paths, timestamp=True, force_all=False, pbar=None, model_type=None):
|
|
"""Pull model metadata from CivitAI and update cache entries."""
|
|
pull_json = True
|
|
metadata_days_recheck = 7
|
|
|
|
cache.load()
|
|
cache.backup_counter += 1
|
|
if cache.backup_counter >= cache.num_of_backups_to_keep:
|
|
cache.prune_all_backups()
|
|
cache.backup_counter = 0
|
|
|
|
if isinstance(file_paths, str):
|
|
file_paths = [file_paths]
|
|
|
|
if not file_paths:
|
|
logger.warning("No file paths provided.")
|
|
return
|
|
|
|
num_not_pulled = 0
|
|
|
|
for file_path in file_paths:
|
|
force = force_all
|
|
hash_value = cache.hash.get(str(file_path), None)
|
|
if hash_value is None:
|
|
logger.debug(f"Hash not found in cache for {file_path}. Adding to cache.")
|
|
hash_value = add_file_to_cache(file_path)
|
|
|
|
file_cache = cache.by_path(file_path)
|
|
|
|
last_used_date = datetime.datetime.fromisoformat(file_cache['lastUsed']) if 'lastUsed' in file_cache else None
|
|
|
|
modified = get_file_modification_date(file_path)
|
|
if last_used_date is not None and modified is not None and modified > last_used_date:
|
|
logger.info("File was modified after last used. Pulling metadata.")
|
|
force = True
|
|
|
|
civitai_val = False
|
|
try:
|
|
civitai_val = str_to_bool(file_cache.get('civitai', False))
|
|
except (TypeError, ValueError):
|
|
civitai_val = False
|
|
|
|
if not force and civitai_val is True:
|
|
if days_since_last_used(file_path) <= metadata_days_recheck:
|
|
num_not_pulled += 1
|
|
pull_json = False
|
|
|
|
if file_cache.get('blacklist'):
|
|
if not force:
|
|
logger.info(f"File {file_path} is blacklisted (previously not found). Skipping metadata pull.")
|
|
pull_json = False
|
|
|
|
if force:
|
|
logger.debug(f"Force flag is set. Recalculating hash for {file_path}.")
|
|
hash_value = recheck_hash(file_path, hash_value)
|
|
|
|
if pull_json or force:
|
|
logger.debug(f"Currently pulling metadata for {file_path}.")
|
|
json_data = get_civitai_model_version_json_by_hash(hash_value)
|
|
|
|
if 'error' in json_data:
|
|
retried = False
|
|
dead_model = False
|
|
if 'civitai_error' in json_data:
|
|
civitai_error = json_data['civitai_error']
|
|
if 'Model not found' in civitai_error or 'No model with id' in civitai_error:
|
|
dead_model = True
|
|
|
|
if dead_model is False:
|
|
if 'modelId' in file_cache:
|
|
logger.debug(f"Using cached model id {file_cache.get('id', None)}")
|
|
json_data = get_civitai_model_version_json_by_id(file_cache['id'])
|
|
retried = True
|
|
else:
|
|
logger.debug("No cached model id.")
|
|
|
|
if 'error' in json_data:
|
|
if retried:
|
|
logger.error(f"Error: {json_data['error']}")
|
|
if dead_model:
|
|
file_cache['blacklist'] = True
|
|
logger.info("Unable to find metadata on CivitAI.")
|
|
file_cache['civitai'] = "False"
|
|
file_cache['civitai_failed_count'] = file_cache.get('civitai_failed_count', 0) + 1
|
|
update_cache_without_civitai_json(file_path, hash_value, timestamp=timestamp)
|
|
|
|
if 'error' not in json_data:
|
|
update_cache_from_civitai_json(file_path, json_data, timestamp=timestamp)
|
|
else:
|
|
retries = file_cache.get('civitai_failed_count', 0) + 1
|
|
file_cache['civitai_failed_count'] = retries
|
|
|
|
cache.hash[file_path] = hash_value
|
|
cache.info[hash_value] = file_cache
|
|
if model_type is not None:
|
|
file_cache['model_type'] = model_type
|
|
|
|
if pbar is not None:
|
|
pbar.update(1)
|
|
|
|
if num_not_pulled > 0:
|
|
logger.info(
|
|
f"Metadata pull complete. Skipped {num_not_pulled} files checked within the last {metadata_days_recheck} days."
|
|
)
|
|
cache.save()
|