From bade95b10d47bfec9636207662ef00a3b79d84db Mon Sep 17 00:00:00 2001 From: TemryL Date: Fri, 2 Feb 2024 12:25:40 +0100 Subject: [PATCH] save image to s3 node --- src/client_s3.py | 83 ++++++++++++++++++++++++++++++++------ src/nodes/save_image.py | 0 src/nodes/save_image_s3.py | 82 +++++++++++++++++++++++++++++++++++++ src/nodes_mappings.py | 8 +++- 4 files changed, 159 insertions(+), 14 deletions(-) delete mode 100644 src/nodes/save_image.py create mode 100644 src/nodes/save_image_s3.py diff --git a/src/client_s3.py b/src/client_s3.py index b398471..213d166 100644 --- a/src/client_s3.py +++ b/src/client_s3.py @@ -13,11 +13,17 @@ class S3: self.secret_key = secret_key self.bucket_name = bucket_name self.s3_client = self.get_client() + self.input_dir = os.getenv("S3_INPUT_DIR") + self.output_dir = os.getenv("S3_OUTPUT_DIR") + if not self.does_folder_exist(self.input_dir): + self.create_folder(self.output_dir) + if not self.does_folder_exist(self.output_dir): + self.create_folder(self.output_dir) def get_client(self): if not all([self.region, self.access_key, self.secret_key, self.bucket_name]): raise ValueError("Missing required S3 environment variables") - + try: s3 = boto3.resource( service_name='s3', @@ -30,13 +36,32 @@ class S3: raise RuntimeError(f"Failed to create S3 client: {e}") def get_files(self, prefix): + if self.does_folder_exist(prefix): + try: + bucket = self.s3_client.Bucket(self.bucket_name) + files = [obj.key for obj in bucket.objects.filter(Prefix=prefix)] + files = [f.replace(prefix, "") for f in files] + return files + except Exception as e: + raise RuntimeError(f"Failed to get files from S3: {e}") + else: + return [] + + def does_folder_exist(self, folder_name): try: bucket = self.s3_client.Bucket(self.bucket_name) - files = [obj.key for obj in bucket.objects.filter(Prefix=prefix)] - return files + response = bucket.objects.filter(Prefix=folder_name) + return any(obj.key.startswith(folder_name) for obj in response) except Exception as e: - raise RuntimeError(f"Failed to get files from S3: {e}") - + raise RuntimeError(f"Failed to check if folder exists in S3: {e}") + + def create_folder(self, folder_name): + try: + bucket = self.s3_client.Bucket(self.bucket_name) + bucket.put_object(Key=f"{folder_name}/") + except Exception as e: + raise RuntimeError(f"Failed to create folder in S3: {e}") + def download_file(self, s3_path, local_path): local_dir = os.path.dirname(local_path) if not os.path.exists(local_dir): @@ -59,6 +84,46 @@ class S3: raise RuntimeError("Credentials not available or not valid.") except Exception as e: raise RuntimeError(f"Failed to upload file to S3: {e}") + + def get_save_image_path(self, filename_prefix, image_width=0, image_height=0): + def map_filename(filename): + prefix_len = len(os.path.basename(filename_prefix)) + prefix = filename[:prefix_len + 1] + try: + digits = int(filename[prefix_len + 1:].split('_')[0]) + except: + digits = 0 + return (digits, prefix) + + def compute_vars(input, image_width, image_height): + input = input.replace("%width%", str(image_width)) + input = input.replace("%height%", str(image_height)) + return input + + filename_prefix = compute_vars(filename_prefix, image_width, image_height) + subfolder = os.path.dirname(os.path.normpath(filename_prefix)) + filename = os.path.basename(os.path.normpath(filename_prefix)) + + full_output_folder_s3 = os.path.join(self.output_dir, subfolder) + + # Check if the output folder exists, create it if it doesn't + if not self.does_folder_exist(full_output_folder_s3): + self.create_folder(full_output_folder_s3) + + try: + # Continue with the counter calculation + files = self.get_files(full_output_folder_s3) + print(files) + counter = max( + filter( + lambda a: a[1][:-1] == filename and a[1][-1] == "_", + map(map_filename, files) + ) + )[0] + 1 + except (ValueError, KeyError): + counter = 1 + + return full_output_folder_s3, filename, counter, subfolder, filename_prefix def get_s3_instance(): @@ -68,10 +133,4 @@ def get_s3_instance(): secret_key=os.getenv("S3_SECRET_KEY"), bucket_name=os.getenv("S3_BUCKET_NAME") ) - return s3_instance - - -# Example usage: -# s3_instance = get_s3_instance() -# s3_instance.upload_file(local_path="input/example_s3.jpeg", s3_path="input/example_s3.jpeg") -# s3_instance.download_file(s3_path="input/example_s3.jpeg", local_path="input/example_s3.jpeg") \ No newline at end of file + return s3_instance \ No newline at end of file diff --git a/src/nodes/save_image.py b/src/nodes/save_image.py deleted file mode 100644 index e69de29..0000000 diff --git a/src/nodes/save_image_s3.py b/src/nodes/save_image_s3.py new file mode 100644 index 0000000..6332e75 --- /dev/null +++ b/src/nodes/save_image_s3.py @@ -0,0 +1,82 @@ +import os +import json +import tempfile +import numpy as np +from PIL import Image +from PIL.PngImagePlugin import PngInfo +from comfy.cli_args import args + +from ..client_s3 import get_s3_instance +S3_INSTANCE = get_s3_instance() + + +class SaveImageS3: + def __init__(self): + base_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.realpath(__file__)))) + self.temp_dir = os.path.join(base_dir, "temp/") + self.s3_output_dir = os.getenv("S3_OUTPUT_DIR") + self.type = "output" + self.prefix_append = "" + self.compress_level = 4 + + @classmethod + def INPUT_TYPES(s): + return {"required": { + "images": ("IMAGE", ), + "filename_prefix": ("STRING", {"default": "ComfyUI"})}, + "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO" + }, + } + + RETURN_TYPES = () + FUNCTION = "save_images" + + OUTPUT_NODE = True + + CATEGORY = "image" + + def save_images(self, images, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None): + filename_prefix += self.prefix_append + full_output_folder, filename, counter, subfolder, filename_prefix = S3_INSTANCE.get_save_image_path(filename_prefix, images[0].shape[1], images[0].shape[0]) + results = list() + + for image in images: + i = 255. * image.cpu().numpy() + img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) + metadata = None + if not args.disable_metadata: + metadata = PngInfo() + if prompt is not None: + metadata.add_text("prompt", json.dumps(prompt)) + if extra_pnginfo is not None: + for x in extra_pnginfo: + metadata.add_text(x, json.dumps(extra_pnginfo[x])) + + file = f"{filename}_{counter:05}_.png" + temp_file = None + try: + # Create a temporary file + with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as temp_file: + temp_file_path = temp_file.name + + # Save the image to the temporary file + img.save(temp_file_path, pnginfo=metadata, compress_level=self.compress_level) + + # Upload the temporary file to S3 + s3_path = os.path.join(full_output_folder, file) + S3_INSTANCE.upload_file(temp_file_path, s3_path) + + # Add the result to the results list + results.append({ + "filename": file, + "subfolder": subfolder, + "type": self.type + }) + counter += 1 + + finally: + # Delete the temporary file + if temp_file_path and os.path.exists(temp_file_path): + os.remove(temp_file_path) + + return { "ui": { "images": results } } diff --git a/src/nodes_mappings.py b/src/nodes_mappings.py index 6fc6f26..8a535a7 100644 --- a/src/nodes_mappings.py +++ b/src/nodes_mappings.py @@ -1,9 +1,13 @@ from .nodes.load_image_s3 import LoadImageS3 +from .nodes.save_image_s3 import SaveImageS3 + NODE_CLASS_MAPPINGS = { - "LoadImageS3": LoadImageS3 + "LoadImageS3": LoadImageS3, + "SaveImageS3": SaveImageS3 } NODE_DISPLAY_NAME_MAPPINGS = { - "LoadImageS3": "Load Image from S3" + "LoadImageS3": "Load Image from S3", + "SaveImageS3": "Save Image to S3" } \ No newline at end of file