From f40c0e8c024fa89bc87ab43ca52b97c61922876f Mon Sep 17 00:00:00 2001
From: Mattabyte <matthew.inger@intelligentpathways.com.au>
Date: Sat, 8 Mar 2025 09:13:31 +1000
Subject: [PATCH] Add SaveLatentToS3 node

---
 __init__.py             |  8 ++++--
 nodes/save_latent_s3.py | 61 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 66 insertions(+), 3 deletions(-)
 create mode 100644 nodes/save_latent_s3.py

diff --git a/__init__.py b/__init__.py
index 5c9cede..7158bc8 100644
--- a/__init__.py
+++ b/__init__.py
@@ -1,16 +1,18 @@
 from .nodes.secure_api_call import SecureApiCall
 from .nodes.secure_api_call_aws import SecureApiCallAws
 from .nodes.save_video_files_s3 import SaveVideoFilesS3
-
+from .nodes.save_latent_s3 import SaveLatentToS3
 NODE_CLASS_MAPPINGS = {
     "SecureAPI-SecureAPI": SecureApiCall,
     "SecureAPI-SecureAPI-AWS": SecureApiCallAws,
-    "SaveVideoFilesS3": SaveVideoFilesS3
+    "SaveVideoFilesS3": SaveVideoFilesS3,
+    "SaveLatentToS3": SaveLatentToS3
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
     "SecureAPI-SecureAPI": "Secure API Call",
     "SecureAPI-SecureAPI-AWS": "Secure API Call (AWS)",
-    "SaveVideoFilesS3": "Save Video Files to S3"
+    "SaveVideoFilesS3": "Save Video Files to S3",
+    "SaveLatentToS3": "Save Latent to S3"
 }
 
diff --git a/nodes/save_latent_s3.py b/nodes/save_latent_s3.py
new file mode 100644
index 0000000..40f396b
--- /dev/null
+++ b/nodes/save_latent_s3.py
@@ -0,0 +1,61 @@
+import os
+
+from ..client_s3 import get_s3_instance
+S3_INSTANCE = get_s3_instance()
+
+class SaveLatentToS3:
+    def __init__(self):
+        self.output_dir = folder_paths.get_output_directory()
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": { "samples": ("LATENT", ),
+                              "filename_prefix": ("STRING", {"default": "latents/ComfyUI"})},
+                "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
+                }
+    RETURN_TYPES = ("STRING")
+    RETURN_NAMES = ("latent_file_path",)
+    FUNCTION = "save"
+
+    OUTPUT_NODE = True
+
+    CATEGORY = "SaveVideoFilesS3"
+
+    def save(self, samples, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None):
+        full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir)
+        # support save metadata for latent sharing
+        prompt_info = ""
+        if prompt is not None:
+            prompt_info = json.dumps(prompt)
+
+        metadata = None
+        if not args.disable_metadata:
+            metadata = {"prompt": prompt_info}
+            if extra_pnginfo is not None:
+                for x in extra_pnginfo:
+                    metadata[x] = json.dumps(extra_pnginfo[x])
+
+        file = f"{filename}_.latent"
+
+        results: list[FileLocator] = []
+        results.append({
+            "filename": file,
+            "subfolder": subfolder,
+            "type": "output"
+        })
+
+        file = os.path.join(full_output_folder, file)
+
+        output = {}
+        output["latent_tensor"] = samples["samples"]
+        output["latent_format_version_0"] = torch.tensor([])
+
+        comfy.utils.save_torch_file(output, file, metadata=metadata)
+
+        #Upload to S3
+        full_output_folder, filename, counter, _, filename_prefix = S3_INSTANCE.get_save_path(filename_prefix)
+        s3_path = os.path.join(full_output_folder, file)
+        file_path = S3_INSTANCE.upload_file(file, s3_path)
+
+
+        return {file_path}