save image to s3 node
This commit is contained in:
TemryL
+71
-12
@@ -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")
|
||||
return s3_instance
|
||||
@@ -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 } }
|
||||
@@ -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"
|
||||
}
|
||||
Reference in New Issue
Block a user