diff --git a/nodes/loader.py b/nodes/loader.py index ccb2d79..59232e1 100644 --- a/nodes/loader.py +++ b/nodes/loader.py @@ -3,6 +3,7 @@ from __future__ import annotations from comfy.comfy_types.node_typing import ComfyNodeABC, InputTypeDict, IO +from ..utils.helpers import pull_metadata, update_model_timestamp # Import specific utilities instead of wildcard import from ..utils import get_lora_stack_keywords @@ -49,6 +50,9 @@ class Sage_UNETLoaderFromInfo(ComfyNodeABC): unet_node = add_unet_node_from_info(graph, unet_info) if unet_node is None: raise ValueError("UNET info is missing or invalid.") + else: + update_model_timestamp(unet_info["path"]) + # Return the UNET node as a model component. print(f"UNET loaded: {unet_info['path']}") unet_out = unet_node.out(0) if unet_node else None @@ -86,6 +90,9 @@ class Sage_CLIPLoaderFromInfo(ComfyNodeABC): clip_node = add_clip_node_from_info(graph, clip_info) if clip_node is None: raise ValueError("CLIP info is missing or invalid.") + else: + update_model_timestamp(clip_info["path"]) + # Return the CLIP node as a model component. print(f"CLIP loaded: {clip_info['path']}") clip_out = clip_node.out(0) if clip_node else None @@ -123,6 +130,9 @@ class Sage_ChromaCLIPLoaderFromInfo(ComfyNodeABC): clip_node = add_clip_node_from_info(graph, clip_info) if clip_node is None: raise ValueError("CLIP info is missing or invalid.") + else: + update_model_timestamp(clip_info["path"]) + clip_node = graph.node("T5TokenizerOptions", clip=clip_node.out(0), min_padding=1, min_length=0) # Return the CLIP node as a model component. print(f"CLIP loaded: {clip_info['path']}") @@ -159,6 +169,9 @@ class Sage_VAELoaderFromInfo(ComfyNodeABC): vae_node = add_vae_node_from_info(graph, vae_info) if vae_node is None: raise ValueError("VAE info is missing or invalid.") + else: + update_model_timestamp(vae_info["path"]) + # Return the VAE node as a model component. print(f"VAE loaded: {vae_info['path']}") vae_out = vae_node.out(0) if vae_node else None @@ -286,18 +299,34 @@ class Sage_LoadModelFromInfo(ComfyNodeABC): if ckpt_info is not None: ckpt_node = add_ckpt_node_from_info(graph, ckpt_info) + if ckpt_node is None: + raise ValueError("Checkpoint info is missing or invalid.") + else: + update_model_timestamp(ckpt_info["path"]) # If we have a UNET, load it with the UNETLoader node. if unet_info is not None: unet_node = add_unet_node_from_info(graph, unet_info) + if unet_node is None: + raise ValueError("UNET info is missing or invalid.") + else: + update_model_timestamp(unet_info["path"]) # If we have a CLIP, load it with the appropriate CLIPLoader node. if clip_info is not None: clip_node = add_clip_node_from_info(graph, clip_info) + if clip_node is None: + raise ValueError("CLIP info is missing or invalid.") + else: + update_model_timestamp(clip_info["path"]) # If we have a VAE, load it with the VAELoader node. if vae_info is not None: vae_node = add_vae_node_from_info(graph, vae_info) + if vae_node is None: + raise ValueError("VAE info is missing or invalid.") + else: + update_model_timestamp(vae_info["path"]) # We need to determine which outputs to return from the nodes. # If there's a checkpoint, set all the outputs to use its outputs initially. @@ -406,6 +435,8 @@ class Sage_UNETLoRALoader(ComfyNodeABC): unet_out = unet_node.out(0) if unet_node else None if unet_node is None: raise ValueError("UNET info is missing or invalid.") + else: + update_model_timestamp(unet_info["path"]) if model_shifts is not None and unet_out is not None: unet_out = create_model_shift_nodes_v2(graph, unet_out, model_shifts) diff --git a/nodes/util.py b/nodes/util.py index 72b9fef..24b14b6 100644 --- a/nodes/util.py +++ b/nodes/util.py @@ -752,7 +752,7 @@ class Sage_CheckLorasForUpdates(ComfyNodeABC): if lora is not None: print(f"Checking {lora[0]} for updates...") lora_path = folder_paths.get_full_path_or_raise("loras", lora[0]) - pull_metadata(lora_path, timestamp=False, force_all=force) + pull_metadata(lora_path, force_all=force) print(f"Update check complete for {lora[0]}") if "update_available" in cache.by_path(lora_path): diff --git a/server_routes.py b/server_routes.py index 8d5fe1f..9eed0cd 100644 --- a/server_routes.py +++ b/server_routes.py @@ -501,7 +501,7 @@ try: # Call the pull_metadata function try: - pull_metadata(file_path, timestamp=True, force_all=force) + pull_metadata(file_path, force_all=force) return web.json_response({ "success": True, diff --git a/utils/common.py b/utils/common.py index 712bfd6..a5e38a6 100644 --- a/utils/common.py +++ b/utils/common.py @@ -19,6 +19,7 @@ from .helpers import ( update_cache_without_civitai_json, add_file_to_cache, recheck_hash, + update_model_timestamp, pull_metadata, lora_to_string, lora_to_prompt, @@ -122,7 +123,7 @@ __all__ = [ 'str_to_bool', 'bool_to_str', 'name_from_path', 'get_path_without_base', 'get_file_extension', 'get_file_sha256', 'get_files_in_dir', 'last_used', 'days_since_last_used', 'get_file_modification_date', 'update_cache_from_civitai_json', 'update_cache_without_civitai_json', - 'add_file_to_cache', 'recheck_hash', 'pull_metadata', + 'add_file_to_cache', 'recheck_hash', 'pull_metadata', 'update_model_timestamp', 'lora_to_string', 'lora_to_prompt', 'get_lora_hash', 'model_scan', 'get_recently_used_models', 'clean_keywords', 'clean_text', 'condition_text', 'get_save_file_path', 'unwrap_tuple', diff --git a/utils/helpers.py b/utils/helpers.py index 9b883f1..e98591d 100644 --- a/utils/helpers.py +++ b/utils/helpers.py @@ -222,6 +222,16 @@ def recheck_hash(file_path, hash): hash = new_hash return hash +def update_model_timestamp(file_path): + cache.load() + # If the file_path isn't a list, make it a list. + if not isinstance(file_path, (list, tuple)): + file_path = [file_path] + for path in file_path: + 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): pull_json = True metadata_days_recheck = 7 @@ -237,10 +247,7 @@ def pull_metadata(file_paths, timestamp = True, force_all = False, pbar = None): for file_path in file_paths: force = force_all - #print(f"Processing file: {file_path}") - #print(f"cache.hash: {cache.hash}") hash = cache.hash.get(str(file_path), None) - #print(f"Current hash: {hash}") if hash is None: print(f"Hash not found in cache for {file_path}. Adding to cache.") hash = add_file_to_cache(file_path) @@ -301,10 +308,6 @@ def pull_metadata(file_paths, timestamp = True, force_all = False, pbar = None): retries += 1 file_cache['civitai_failed_count'] = retries - if timestamp: - print("Updating timestamp.") - cache.update_last_used_by_path(file_path) - cache.hash[file_path] = hash cache.info[hash] = file_cache if pbar is not None: @@ -346,7 +349,7 @@ def model_scan(the_path, force = False): model_list = [str(x) for x in model_list] print(f"Scanning {len(model_list)} models for metadata.") pbar = comfy.utils.ProgressBar(len(model_list)) - pull_metadata(model_list, force_all=force, timestamp=False, pbar=pbar) + pull_metadata(model_list, force_all=force, pbar=pbar) def get_recently_used_models(model_type): model_list = list() diff --git a/utils/lora_stack.py b/utils/lora_stack.py index 303d9e1..641b83f 100644 --- a/utils/lora_stack.py +++ b/utils/lora_stack.py @@ -14,7 +14,7 @@ def get_lora_stack_keywords(lora_stack=None): # Collect unique lora names lora_names = {lora[0] for lora in lora_stack} lora_paths = [folder_paths.get_full_path_or_raise("loras", name) for name in lora_names] - pull_metadata(lora_paths, timestamp=True) + pull_metadata(lora_paths) # Gather all keywords into a set for uniqueness all_keywords = set() for name in lora_names: diff --git a/utils/model_cache.py b/utils/model_cache.py index 05a45ec..6b103a0 100644 --- a/utils/model_cache.py +++ b/utils/model_cache.py @@ -340,8 +340,12 @@ class SageCache: Update the 'lastUsed' field for a given file path. """ file_hash = self.hash.get(file_path) - if file_hash and file_hash in self.info: - self.info[file_hash]['lastUsed'] = datetime.datetime.now().isoformat() + if file_hash: + if file_hash in self.info: + self.info[file_hash]['lastUsed'] = datetime.datetime.now().isoformat() + else: + self.add_entry(file_path, file_hash) + self.info[file_hash]['lastUsed'] = datetime.datetime.now().isoformat() def update_last_used_by_hash(self, file_hash: str) -> None: """ diff --git a/utils/model_info.py b/utils/model_info.py index ce4f97c..fc853d9 100644 --- a/utils/model_info.py +++ b/utils/model_info.py @@ -23,7 +23,7 @@ def get_model_info_ckpt(ckpt_name: str) -> tuple: tuple: A tuple containing the model_info dictionary. """ model_info = {"type": "CKPT", "path": folder_paths.get_full_path_or_raise("checkpoints", ckpt_name)} - pull_metadata(model_info["path"], timestamp=True) + pull_metadata(model_info["path"]) model_info["hash"] = cache.hash[model_info["path"]] return (model_info,) @@ -43,7 +43,7 @@ def get_model_info_unet(unet_name: str, weight_dtype: str = "default") -> tuple: unet_name = unet_name[len(base):].lstrip("/\\") break model_info = {"type": "UNET", "path": folder_paths.get_full_path_or_raise("diffusion_models", unet_name)} - pull_metadata(model_info["path"], timestamp=True) + pull_metadata(model_info["path"]) model_info["hash"] = cache.hash[model_info["path"]] if weight_dtype and (weight_dtype in weight_dtype_options): model_info["weight_dtype"] = weight_dtype @@ -73,7 +73,7 @@ def get_model_info_clips(clip_names: list, clip_type: str = "") -> tuple: clip_paths = [] for key in clip_names: name = folder_paths.get_full_path_or_raise("text_encoders", key) - pull_metadata(name, timestamp=True) + pull_metadata(name) clip_paths.append(name) model_info = { @@ -97,7 +97,7 @@ def get_model_info_vae(vae_name: str) -> tuple: tuple: A tuple containing the model_info dictionary. """ model_info = {"type": "VAE", "path": folder_paths.get_full_path_or_raise("vae", vae_name)} - pull_metadata(model_info["path"], timestamp=True) + pull_metadata(model_info["path"]) model_info["hash"] = cache.hash[model_info["path"]] print(f"VAE model info: {model_info}") return (model_info,)