|
import json |
|
import os |
|
from concurrent.futures import ThreadPoolExecutor |
|
from multiprocessing import cpu_count |
|
|
|
import args_manager |
|
from modules.util import sha256, HASH_SHA256_LENGTH, get_file_from_folder_list |
|
|
|
hash_cache_filename = 'hash_cache.txt' |
|
hash_cache = {} |
|
|
|
|
|
def sha256_from_cache(filepath): |
|
global hash_cache |
|
if filepath not in hash_cache: |
|
print(f"[Cache] Calculating sha256 for {filepath}") |
|
hash_value = sha256(filepath) |
|
print(f"[Cache] sha256 for {filepath}: {hash_value}") |
|
hash_cache[filepath] = hash_value |
|
save_cache_to_file(filepath, hash_value) |
|
|
|
return hash_cache[filepath] |
|
|
|
|
|
def load_cache_from_file(): |
|
global hash_cache |
|
|
|
try: |
|
if os.path.exists(hash_cache_filename): |
|
with open(hash_cache_filename, 'rt', encoding='utf-8') as fp: |
|
for line in fp: |
|
entry = json.loads(line) |
|
for filepath, hash_value in entry.items(): |
|
if not os.path.exists(filepath) or not isinstance(hash_value, str) and len(hash_value) != HASH_SHA256_LENGTH: |
|
print(f'[Cache] Skipping invalid cache entry: {filepath}') |
|
continue |
|
hash_cache[filepath] = hash_value |
|
except Exception as e: |
|
print(f'[Cache] Loading failed: {e}') |
|
|
|
|
|
def save_cache_to_file(filename=None, hash_value=None): |
|
global hash_cache |
|
|
|
if filename is not None and hash_value is not None: |
|
items = [(filename, hash_value)] |
|
mode = 'at' |
|
else: |
|
items = sorted(hash_cache.items()) |
|
mode = 'wt' |
|
|
|
try: |
|
with open(hash_cache_filename, mode, encoding='utf-8') as fp: |
|
for filepath, hash_value in items: |
|
json.dump({filepath: hash_value}, fp) |
|
fp.write('\n') |
|
except Exception as e: |
|
print(f'[Cache] Saving failed: {e}') |
|
|
|
|
|
def init_cache(model_filenames, paths_checkpoints, lora_filenames, paths_loras): |
|
load_cache_from_file() |
|
|
|
if args_manager.args.rebuild_hash_cache: |
|
max_workers = args_manager.args.rebuild_hash_cache if args_manager.args.rebuild_hash_cache > 0 else cpu_count() |
|
rebuild_cache(lora_filenames, model_filenames, paths_checkpoints, paths_loras, max_workers) |
|
|
|
|
|
save_cache_to_file() |
|
|
|
|
|
def rebuild_cache(lora_filenames, model_filenames, paths_checkpoints, paths_loras, max_workers=cpu_count()): |
|
def thread(filename, paths): |
|
filepath = get_file_from_folder_list(filename, paths) |
|
sha256_from_cache(filepath) |
|
|
|
print('[Cache] Rebuilding hash cache') |
|
with ThreadPoolExecutor(max_workers=max_workers) as executor: |
|
for model_filename in model_filenames: |
|
executor.submit(thread, model_filename, paths_checkpoints) |
|
for lora_filename in lora_filenames: |
|
executor.submit(thread, lora_filename, paths_loras) |
|
print('[Cache] Done') |
|
|
