diff --git a/nodes.py b/nodes.py index 7aef7ed..23c3d17 100644 --- a/nodes.py +++ b/nodes.py @@ -22,7 +22,15 @@ def parse_checkpoint_name(ckpt_name: str) -> str: return os.path.basename(ckpt_name) def parse_checkpoint_name_without_extension(ckpt_name: str) -> str: - return os.path.splitext(parse_checkpoint_name(ckpt_name))[0] + filename = parse_checkpoint_name(ckpt_name) + name_without_ext, ext = os.path.splitext(filename) + supported_extensions = folder_paths.supported_pt_extensions | {".gguf"} + + # Only remove extension if it's a known model file extension + if ext.lower() in supported_extensions: + return name_without_ext + else: + return filename # Keep full name if extension isn't recognized def get_timestamp(time_format: str) -> str: now = datetime.now() @@ -186,7 +194,32 @@ class ImageSaverMetadata: if download_civitai_data and civitai_resources: a111_params += f", Civitai resources: {json.dumps(civitai_resources, separators=(',', ':'))}" - final_hashes = ",".join(f"{Path(name.split(':')[-1]).stem + ':' if name else ''}{hash}{':' + str(weight) if weight is not None and download_civitai_data else ''}" for name, (_, weight, hash) in ({ modelname: ( ckpt_path, None, modelhash ) } | loras | embeddings | manual_entries).items()) + # Combine all resources (model, loras, embeddings, manual entries) for final hash string + all_resources = { modelname: ( ckpt_path, None, modelhash ) } | loras | embeddings | manual_entries + + hash_parts = [] + for name, (_, weight, hash) in all_resources.items(): + # Format: "name:hash" or "name:hash:weight" depending on download_civitai_data + if name: + # Extract clean name (only remove actual model file extensions, preserve dots in model names) + filename = name.split(':')[-1] + name_without_ext, ext = os.path.splitext(filename) + supported_extensions = folder_paths.supported_pt_extensions | {".gguf"} + + # Only remove extension if it's a known model file extension + if ext.lower() in supported_extensions: + clean_name = name_without_ext + else: + clean_name = filename # Keep full name if extension isn't recognized + + name_part = f"{clean_name}:" + else: + name_part = "" + + weight_part = f":{weight}" if weight is not None and download_civitai_data else "" + hash_parts.append(f"{name_part}{hash}{weight_part}") + + final_hashes = ",".join(hash_parts) metadata = Metadata(modelname, positive, negative, width, height, seed_value, steps, cfg, sampler_name, scheduler_name, denoise, clip_skip, custom, additional_hashes, ckpt_path, a111_params, final_hashes) return metadata diff --git a/utils.py b/utils.py index 6f2e871..9e02e85 100644 --- a/utils.py +++ b/utils.py @@ -55,7 +55,7 @@ def get_sha256(file_path: str) -> str: Based on a embedding name, eg: EasyNegative, finds the path as known in comfy, including extension """ def full_embedding_path_for(embedding: str) -> Optional[str]: - matching_embedding = next((x for x in __list_embeddings() if x.startswith(embedding)), None) + matching_embedding = next((x for x in folder_paths.get_filename_list("embeddings") if x.startswith(embedding)), None) if matching_embedding == None: return None return folder_paths.get_full_path("embeddings", matching_embedding) @@ -70,44 +70,42 @@ def full_lora_path_for(lora: str) -> Optional[str]: lora += ".safetensors" # Find the matching lora path - matching_lora = next((x for x in __list_loras() if x.endswith(lora)), None) + matching_lora = next((x for x in folder_paths.get_filename_list("loras") if x.endswith(lora)), None) if matching_lora is None: print(f'ComfyUI-Image-Saver: could not find full path to lora "{lora}"') return None return folder_paths.get_full_path("loras", matching_lora) -def __list_loras() -> list[str]: - return folder_paths.get_filename_list("loras") - -def __list_embeddings() -> list[str]: - return folder_paths.get_filename_list("embeddings") - def full_checkpoint_path_for(model_name: str) -> str: if not model_name: return '' last_dot_position = model_name.rfind('.') extension = model_name[last_dot_position:] if last_dot_position != -1 else "" - if extension not in folder_paths.supported_pt_extensions: + if extension not in folder_paths.supported_pt_extensions | {".gguf"}: model_name += ".safetensors" - matching_checkpoint = next((x for x in __list_checkpoints() if x.endswith(model_name)), None) + matching_checkpoint = next((x for x in folder_paths.get_filename_list("checkpoints") if x.endswith(model_name)), None) if matching_checkpoint: return folder_paths.get_full_path("checkpoints", matching_checkpoint) - matching_model = next((x for x in __list_diffusion_models() if x.endswith(model_name)), None) + matching_model = next((x for x in folder_paths.get_filename_list("diffusion_models") if x.endswith(model_name)), None) if matching_model: return folder_paths.get_full_path("diffusion_models", matching_model) + # Handle GGUF files in diffusion_models directories since they're not in supported_pt_extensions + if model_name.endswith('.gguf'): + diffusion_model_paths = folder_paths.folder_names_and_paths.get("diffusion_models", [[], set()])[0] + for path in diffusion_model_paths: + if os.path.exists(path): + for root, dirs, files in os.walk(path): + for file in files: + if file.endswith(model_name): + return os.path.join(root, file) + print(f'Could not find full path to checkpoint "{model_name}"') return '' -def __list_checkpoints() -> list[str]: - return folder_paths.get_filename_list("checkpoints") - -def __list_diffusion_models() -> list[str]: - return folder_paths.get_filename_list("diffusion_models") - def http_get_json(url: str) -> dict[str, Any] | None: try: response = requests.get(url, timeout=300)