144 lines
5.5 KiB
Python
144 lines
5.5 KiB
Python
import hashlib
|
|
import difflib
|
|
import folder_paths
|
|
import comfy.samplers
|
|
import os
|
|
import numpy as np
|
|
from pathlib import Path
|
|
|
|
def get_model_hash(filename):
|
|
hash_sha256 = hashlib.sha256()
|
|
blksize = 1024 * 1024
|
|
|
|
is_link = os.path.islink(str(filename))
|
|
if is_link == True:
|
|
filename = Path(str(filename)).resolve()
|
|
|
|
if os.path.isfile(filename):
|
|
with open(filename, "rb") as f:
|
|
for chunk in iter(lambda: f.read(blksize), b""):
|
|
hash_sha256.update(chunk)
|
|
|
|
return hash_sha256.hexdigest()[0:10]
|
|
else:
|
|
return 'unknown'
|
|
|
|
def check_model_from_exif(model_hash_exif, model_name_exif, model_name, model_hash_check):
|
|
ckpt_path = folder_paths.get_full_path("checkpoints", model_name_exif)
|
|
if ckpt_path is not None:
|
|
if os.path.exists(ckpt_path):
|
|
return model_name_exif
|
|
|
|
checkpointpaths = folder_paths.get_folder_paths("checkpoints")[0]
|
|
allcheckpoints = folder_paths.get_filename_list("checkpoints")
|
|
|
|
source_model_name = model_name_exif
|
|
if '\\\\' not in model_name_exif:
|
|
source_model_name = model_name_exif.split('_', 1)[-1]
|
|
|
|
cutoff_list = list(np.around(np.arange(0.1, 1.05, 0.05).tolist(), 2))[::-1] # [1, 0.9, 0.8, 0.7, 0.6, 0.5, 0.4, 0.3, 0.2, 0.1]
|
|
is_found = []
|
|
for trycut in cutoff_list:
|
|
is_found = difflib.get_close_matches(source_model_name, allcheckpoints, cutoff=round(trycut, 2))
|
|
if len(is_found) == 1:
|
|
break
|
|
|
|
if len(is_found) != 1:
|
|
for trycut in cutoff_list:
|
|
is_found = difflib.get_close_matches(model_name_exif, allcheckpoints, cutoff=round(trycut, 2))
|
|
if len(is_found) == 1:
|
|
break
|
|
|
|
if len(is_found) == 1:
|
|
valid_model = is_found[0]
|
|
model_full_path = checkpointpaths + os.sep + valid_model
|
|
|
|
if model_hash_check == True:
|
|
is_link = os.path.islink(str(model_full_path))
|
|
match_model_hash = None
|
|
if is_link == False:
|
|
match_model_hash = get_model_hash(model_full_path)
|
|
if match_model_hash == model_hash_exif:
|
|
model_name = valid_model
|
|
else:
|
|
print('Model name:' + model_name_exif + ' not available by hashcheck, using system source: ' + model_name)
|
|
else:
|
|
model_name = valid_model
|
|
else:
|
|
if len(is_found) == 0:
|
|
print('Model name:' + model_name_exif + ' not available by diffcheck, using system source: ' + model_name)
|
|
if len(is_found) > 0:
|
|
print('Model name:' + model_name_exif + ' more than one results by diffcheck, using first found: ' + str(is_found[0]))
|
|
model_name = is_found[0]
|
|
|
|
return model_name
|
|
|
|
|
|
def change_exif_samplers(sampler_name_exif, comfy_schedulers):
|
|
lastchars = sampler_name_exif[-2:]
|
|
if lastchars == ' a':
|
|
sampler_name_exif = sampler_name_exif.rsplit(' a', 1)[0] + ' ancestral'
|
|
|
|
sampler_name_exif = sampler_name_exif.replace(' a ', ' ancestral ').replace(' ', '_').replace('++', 'pp').replace('dpm2', 'dpm_2').replace('unipc', 'uni_pc')
|
|
|
|
for comfy_scheduler in comfy_schedulers:
|
|
sampler_name_exif = sampler_name_exif.removesuffix(comfy_scheduler).removesuffix('_')
|
|
|
|
return sampler_name_exif
|
|
|
|
|
|
def check_sampler_from_exif(sampler_name_exif, sampler_name, scheduler_name):
|
|
comfy_samplers = comfy.samplers.KSampler.SAMPLERS
|
|
comfy_schedulers = comfy.samplers.KSampler.SCHEDULERS
|
|
|
|
sampler_name_exif_for_cutoff = change_exif_samplers(sampler_name_exif, comfy_schedulers)
|
|
is_found_sampler = []
|
|
is_found_scheduler = []
|
|
|
|
cutoff_list_samplers = list(np.around(np.arange(0.6, 1.05, 0.1).tolist(), 2))[::-1]
|
|
for trycut in cutoff_list_samplers:
|
|
is_found_sampler = difflib.get_close_matches(sampler_name_exif_for_cutoff, comfy_samplers, cutoff=trycut)
|
|
|
|
if len(is_found_sampler) >= 1:
|
|
sampler_name = is_found_sampler[0]
|
|
|
|
if " " in sampler_name_exif:
|
|
if any((match := substring) in sampler_name_exif for substring in comfy_schedulers):
|
|
scheduler_name = match
|
|
else:
|
|
cutoff_list_schedulers = list(np.around(np.arange(0.4, 0.8, 0.1).tolist(), 2))[::-1]
|
|
for trycut in cutoff_list_schedulers:
|
|
is_found_scheduler = difflib.get_close_matches(sampler_name_exif, comfy_schedulers, cutoff=trycut)
|
|
|
|
if len(is_found_scheduler) >= 1:
|
|
scheduler_name = is_found_scheduler[0]
|
|
|
|
if sampler_name not in comfy_samplers:
|
|
sampler_name = comfy_samplers[0]
|
|
|
|
if scheduler_name not in comfy_schedulers:
|
|
scheduler_name = comfy_schedulers[0]
|
|
|
|
return {'sampler': sampler_name, 'scheduler': scheduler_name}
|
|
|
|
def comfy_samplers2a11(comfy_sampler, comfy_scheduler):
|
|
sampler_name_a11 = comfy_sampler.capitalize().replace('_', ' ').replace('pp', '++').replace('Dpm 2', 'DPM2').replace('Uni pc', 'UniPC').replace(' ancestral', ' a') + ' ' + comfy_scheduler.capitalize()
|
|
sampler_name_a11 = sampler_name_a11.replace(' Normal', '')
|
|
sampler_name_a11 = sampler_name_a11.replace('Ddim', 'DDIM').replace('Dpm', 'DPM').replace('2m', '2M').replace('sde', 'SDE').replace('2s', '2S').replace('Lms', 'LMS')
|
|
|
|
return sampler_name_a11
|
|
|
|
def check_vae_exif(vae_name_exif, vae_name):
|
|
comfy_vaes = folder_paths.get_filename_list("vae")
|
|
cutoff_list = list(np.around(np.arange(0.7, 1.05, 0.1).tolist(), 2))[::-1]
|
|
is_found = []
|
|
|
|
for trycut in cutoff_list:
|
|
is_found = difflib.get_close_matches(vae_name_exif, comfy_vaes, cutoff=trycut)
|
|
if len(is_found) == 1:
|
|
break
|
|
|
|
if len(is_found) == 1:
|
|
vae_name = is_found[0]
|
|
|
|
return vae_name |