From d9d20018be34fed0ebfeb48bda60ca0187296d57 Mon Sep 17 00:00:00 2001
From: wl-zhao <zhaowl20@mails.tsinghua.edu.cn>
Date: Thu, 14 Nov 2024 10:19:59 +0800
Subject: [PATCH] improve abs_path of file dependency

---
 dependency_checker.py |  5 ++++-
 file_upload.py        | 21 +++++++++++++++------
 utils/utils.py        | 17 ++++++++++++++++-
 3 files changed, 35 insertions(+), 8 deletions(-)

diff --git a/dependency_checker.py b/dependency_checker.py
index ecdfe6c..613325c 100644
--- a/dependency_checker.py
+++ b/dependency_checker.py
@@ -293,7 +293,10 @@ def resolve_dependencies(prompt, custom_dependencies): # resolve custom nodes an
 
     # step 3: handle local files
     process_local_file_path_async(file_mapping_dict, max_workers=20)
-    files_dict = {v[0]: {"filename": windows_to_linux_path(os.path.relpath(v[2], BASE_PATH)), "urls": [v[1]]} for v in file_mapping_dict.values()}
+    files_dict = {
+        v[0]: {
+            "filename": windows_to_linux_path(os.path.relpath(v[2], BASE_PATH)) if not v[3] else v[2], 
+            "urls": [v[1]]} for v in file_mapping_dict.values()}
     
     depencencies = {
         "comfyui_version": comfyui_version,
diff --git a/file_upload.py b/file_upload.py
index 3c4afbd..aa13ff8 100644
--- a/file_upload.py
+++ b/file_upload.py
@@ -5,7 +5,7 @@ import time
 from concurrent.futures import ThreadPoolExecutor, as_completed
 import folder_paths
 
-from .utils.utils import compute_sha256
+from .utils.utils import compute_sha256, get_alphanumeric_hash
 
 ext_to_type = {
     # image
@@ -27,7 +27,7 @@ ext_to_type = {
     '.m4a': 'audio/mp4',
 }
 
-def upload_file_to_myshell(local_file: str) -> str:
+def upload_file_to_myshell(local_file: str, target_path: str, is_abs) -> str:
     ''' Now we only support upload file one-by-one
     '''
     MYSHELL_KEY = os.environ.get('MYSHELL_KEY', "OPENSOURCE_FIXED")
@@ -51,8 +51,8 @@ def upload_file_to_myshell(local_file: str) -> str:
     response = requests.request("POST", server_url, headers=headers, files=files)
     if response.status_code == 200:
         end_time = time.time()
-        logging.info(f"{local_file} uploaded, time elapsed: {end_time - start_time}")
-        return [sha256sum, response.json()['url'], local_file]
+        logging.info(f"{local_file} uploaded, time elapsed: {end_time - start_time}, will be saved to {target_path}")
+        return [sha256sum, response.json()['url'], target_path, is_abs]
     else:
         raise Exception(
             f"[HTTP ERROR] {response.status_code} - {response.text} \n"
@@ -66,8 +66,11 @@ def collect_local_file(item, mapping_dict={}):
     abspath = os.path.abspath(item)
     input_abspath = os.path.join(input_dir, item)
     # required file type
+    is_abs = False
     if os.path.isfile(abspath):
         fpath = abspath
+        is_abs = True
+        
     elif os.path.isfile(input_abspath):
         fpath = input_abspath
     else:
@@ -75,7 +78,13 @@ def collect_local_file(item, mapping_dict={}):
     if fpath is not None:
         ext = os.path.splitext(fpath)[1]
         if ext.lower() in ext_to_type.keys():
-            mapping_dict[item] = fpath
+            if is_abs: # if use abs path, replace it
+                filename_hash = get_alphanumeric_hash(abspath)[:16]
+                count = len(mapping_dict)
+                target_path = f"/ShellAgentDeploy/ComfyUI/input/{filename_hash}_{count:06d}{ext}"
+                mapping_dict[item] = (fpath, target_path, is_abs)
+            else:
+                mapping_dict[item] = (fpath, fpath, is_abs)
             return
         else:
             return
@@ -86,7 +95,7 @@ def process_local_file_path_async(mapping_dict, max_workers=10):
     start_time = time.time()
     with ThreadPoolExecutor(max_workers=max_workers) as executor:
         # Submit tasks to the executor
-        futures = {executor.submit(upload_file_to_myshell, full_path): filename for filename, full_path in mapping_dict.items()}
+        futures = {executor.submit(upload_file_to_myshell, source_path, target_path, is_abs): filename for filename, (source_path, target_path, is_abs) in mapping_dict.items()}
         logging.info("submit done")
         # Collect the results as they complete
         for future in as_completed(futures):
diff --git a/utils/utils.py b/utils/utils.py
index 4854509..13e1216 100644
--- a/utils/utils.py
+++ b/utils/utils.py
@@ -1,6 +1,8 @@
 import hashlib
 import time
 from pathlib import PurePosixPath, Path, PureWindowsPath
+import base64
+import re
 
 def windows_to_linux_path(windows_path):
     return PureWindowsPath(windows_path).as_posix()
@@ -17,4 +19,17 @@ def compute_sha256(file_path, chunk_size=1024 ** 2):
             sha256.update(chunk)
     print("finish compute sha256 for", file_path, f"time: {time.time() - start}")
     # Return the hexadecimal digest of the hash
-    return sha256.hexdigest()
\ No newline at end of file
+    return sha256.hexdigest()
+
+
+def get_alphanumeric_hash(input_string: str) -> str:
+    # Generate a SHA-256 hash of the input string
+    sha256_hash = hashlib.sha256(input_string.encode()).digest()
+    
+    # Encode the hash in base64 to get a string with [A-Za-z0-9+/=]
+    base64_hash = base64.b64encode(sha256_hash).decode('ascii')
+    
+    # Remove any non-alphanumeric characters (+, /, =)
+    alphanumeric_hash = re.sub(r'[^a-zA-Z0-9]', '', base64_hash)
+    
+    return alphanumeric_hash
\ No newline at end of file