From c5a48b15bddab089f0054af0ed72eea868263ef6 Mon Sep 17 00:00:00 2001
From: Thomas Ward <teward@thomas-ward.net>
Date: Tue, 16 Jul 2024 18:27:09 -0400
Subject: [PATCH] Make default hash lib configurable without code changes via
 CLI argument (#3947)

* cli_args: Add --duplicate-check-hash-function.

* server.py: compare_image_hash configurable hash function

Uses an argument added in cli_args to specify the type of hashing to default to for duplicate hash checking.  Uses an `eval()` to identify the specific hashlib class to utilize, but ultimately safely operates because we have specific options and only those options/choices in the arg parser.  So we don't have any unsafe input there.

* Add hasher() to node_helpers

* hashlib selection moved to node_helpers

* default-hashing-function instead of dupe checking hasher

This makes a default-hashing-function option instead of previous selected option.

* Use args.default_hashing_function

* Use safer handling for node_helpers.hasher()

Uses a safer handling method than `eval` to evaluate default hashing function.

* Stray parentheses are evil.

* Indentation fix.

Somehow when I hit save I didn't notice I missed a space to make indentation work proper.  Oops!
---
 comfy/cli_args.py |  1 +
 node_helpers.py   | 13 +++++++++++++
 server.py         |  7 +++++--
 3 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/comfy/cli_args.py b/comfy/cli_args.py
index 6251e3d2..2397de3d 100644
--- a/comfy/cli_args.py
+++ b/comfy/cli_args.py
@@ -112,6 +112,7 @@ vram_group.add_argument("--lowvram", action="store_true", help="Split the unet i
 vram_group.add_argument("--novram", action="store_true", help="When lowvram isn't enough.")
 vram_group.add_argument("--cpu", action="store_true", help="To use the CPU for everything (slow).")
 
+parser.add_argument("--default-hashing-function", type=str, choices=['md5', 'sha1', 'sha256', 'sha512'], default='sha256', help="Allows you to choose the hash function to use for duplicate filename / contents comparison. Default is sha256.")
 
 parser.add_argument("--disable-smart-memory", action="store_true", help="Force ComfyUI to agressively offload to regular ram instead of keeping models in vram when it can.")
 parser.add_argument("--deterministic", action="store_true", help="Make pytorch use slower deterministic algorithms when it can. Note that this might not make images deterministic in all cases.")
diff --git a/node_helpers.py b/node_helpers.py
index 43b9e829..fee62879 100644
--- a/node_helpers.py
+++ b/node_helpers.py
@@ -1,3 +1,7 @@
+import hashlib
+
+from comfy.cli_args import args
+
 from PIL import ImageFile, UnidentifiedImageError
 
 def conditioning_set_values(conditioning, values={}):
@@ -22,3 +26,12 @@ def pillow(fn, arg):
         if prev_value is not None:
             ImageFile.LOAD_TRUNCATED_IMAGES = prev_value
         return x
+
+def hasher():
+    hashfuncs = {
+        "md5": hashlib.md5,
+        "sha1": hashlib.sha1,
+        "sha256": hashlib.sha256,
+        "sha512": hashlib.sha512
+    }
+    return hashfuncs[args.default_hashing_function]
diff --git a/server.py b/server.py
index 03c4c421..23ca2fd3 100644
--- a/server.py
+++ b/server.py
@@ -25,6 +25,7 @@ import mimetypes
 from comfy.cli_args import args
 import comfy.utils
 import comfy.model_management
+import node_helpers
 from app.frontend_management import FrontendManager
 from app.user_manager import UserManager
 
@@ -161,10 +162,12 @@ class PromptServer():
             return type_dir, dir_type
 
         def compare_image_hash(filepath, image):
+            hasher = node_helpers.hasher()
+            
             # function to compare hashes of two images to see if it already exists, fix to #3465
             if os.path.exists(filepath):
-                a = hashlib.sha256()
-                b = hashlib.sha256()
+                a = hasher()
+                b = hasher()
                 with open(filepath, "rb") as f:
                     a.update(f.read())
                     b.update(image.file.read())