From 0d353d7dddf772d2257e27f2855fbb993ab00886 Mon Sep 17 00:00:00 2001 From: "Mr.Chip" Date: Mon, 9 Dec 2024 13:21:37 +0800 Subject: [PATCH] FEAT: add basic function to save image to S3 --- __init__.py | 9 +++ common/__init__.py | 62 +++++++++++++++++++++ nodes/__init__.py | 1 + nodes/images.py | 83 ++++++++++++++++++++++++++++ server/__init__.py | 1 + server/setting.py | 100 ++++++++++++++++++++++++++++++++++ web/plugin.js | 133 +++++++++++++++++++++++++++++++++++++++++++++ 7 files changed, 389 insertions(+) create mode 100644 __init__.py create mode 100644 common/__init__.py create mode 100644 nodes/__init__.py create mode 100644 nodes/images.py create mode 100644 server/__init__.py create mode 100644 server/setting.py create mode 100644 web/plugin.js diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..3f4b326 --- /dev/null +++ b/__init__.py @@ -0,0 +1,9 @@ +from .nodes.images import SaveImageS3 +from .server import * + +NODE_CLASS_MAPPINGS = { + "SaveImageS3": SaveImageS3 +} + + +WEB_DIRECTORY = "./web" \ No newline at end of file diff --git a/common/__init__.py b/common/__init__.py new file mode 100644 index 0000000..e45b944 --- /dev/null +++ b/common/__init__.py @@ -0,0 +1,62 @@ +from minio import Minio +from minio.error import S3Error + +class Settings: + def __init__(self): + self.url = None + self.bucket = None + self.access_key = None + self.secret_key = None + +class S3Client(object): + def __init__(self, url, access_key, secret_key, bucket): + self._url = url + self._access_key = access_key + self._secret_key = secret_key + self._bucket = bucket + + self.client = Minio( + endpoint=self._url, + access_key=self._access_key, + secret_key=self._secret_key, + secure=False) + + def upload_image(self, object_name, data, metadata=None): + try: + self.client.put_object( + bucket_name=self._bucket, + object_name=object_name, + data=data, + length=data.getbuffer().nbytes, + content_type='image/png', + metadata=metadata + ) + except S3Error as err: + print(err) + + def upload_meta(self, object_name, data): + try: + self.client.put_object( + bucket_name=self._bucket, + object_name=object_name, + data=data, + length=data.getbuffer().nbytes, + content_type='application/json', + ) + except S3Error as err: + print(err) + +g_settings = Settings() +g_s3_client = None + + +def clear_s3_client(): + global g_s3_client + g_s3_client = None + +def get_s3_client(): + global g_s3_client + global g_settings + if g_s3_client is None: + g_s3_client = S3Client(g_settings.url, g_settings.access_key, g_settings.secret_key, g_settings.bucket) + return g_s3_client \ No newline at end of file diff --git a/nodes/__init__.py b/nodes/__init__.py new file mode 100644 index 0000000..0d24fc2 --- /dev/null +++ b/nodes/__init__.py @@ -0,0 +1 @@ +# intentionally left blank \ No newline at end of file diff --git a/nodes/images.py b/nodes/images.py new file mode 100644 index 0000000..ed7448b --- /dev/null +++ b/nodes/images.py @@ -0,0 +1,83 @@ +from nodes import SaveImage +import importlib +import numpy as np +from comfy.cli_args import args +import os +from io import BytesIO + +from PIL import Image, ImageOps +from PIL.ExifTags import TAGS, GPSTAGS, IFD +from PIL.PngImagePlugin import PngImageFile, PngInfo +from PIL.JpegImagePlugin import JpegImageFile + +import json + +import folder_paths + +from ..common import get_s3_client, S3Client +class SaveImageS3(SaveImage): + def __init__(self): + super().__init__() + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "images": ("IMAGE", {"tooltip": "The images to save."}), + "filename_prefix": ("STRING", {"default": "ComfyUI", "tooltip": "The prefix for the file to save. This may include formatting information such as %date:yyyy-MM-dd% or %Empty Latent Image.width% to include values from nodes."}) + }, + "hidden": { + "prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO" + }, + } + + RETURN_TYPES = () + FUNCTION = "save_images" + + OUTPUT_NODE = True + + CATEGORY = "image" + DESCRIPTION = "Saves the input images to your S3 storage." + + def save_images(self, images, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None): + s3_client = get_s3_client() + filename_prefix += self.prefix_append + full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0]) + results = list() + for (batch_number, image) in enumerate(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])) + + filename_with_batch_num = filename.replace("%batch_num%", str(batch_number)) + file = f"{filename_with_batch_num}_{counter:05}_.png" + # save meta info in json file + meta_file = f"{filename_with_batch_num}_{counter:05}_.json" + oss_meta = {} + oss_meta["prompt"] = prompt + if extra_pnginfo is not None: + for x in extra_pnginfo: + oss_meta[x] = extra_pnginfo[x] + data = BytesIO() + img.save(data, format="PNG", pnginfo=metadata, compress_level=self.compress_level) + data.seek(0) + oss_meta_json = json.dumps(oss_meta, indent=4).encode('utf-8') + s3_client.upload_image(file, data) + s3_client.upload_meta(meta_file, BytesIO(oss_meta_json)) + img.save(os.path.join(full_output_folder, file), pnginfo=metadata, compress_level=self.compress_level) + results.append({ + "filename": file, + "subfolder": subfolder, + "type": self.type + }) + counter += 1 + + return { "ui": { "images": results } } + \ No newline at end of file diff --git a/server/__init__.py b/server/__init__.py new file mode 100644 index 0000000..d4f1004 --- /dev/null +++ b/server/__init__.py @@ -0,0 +1 @@ +from .setting import * \ No newline at end of file diff --git a/server/setting.py b/server/setting.py new file mode 100644 index 0000000..d1db992 --- /dev/null +++ b/server/setting.py @@ -0,0 +1,100 @@ +from server import PromptServer +from aiohttp import web +from ..common import g_settings, clear_s3_client + +@PromptServer.instance.routes.patch("/saveimages3/setting") +async def new_setting(request): + try: + settings = await request.json() + if 'url' in settings is not None and 'bucket' in settings is not None \ + and 'accessKey' in settings is not None and'secretKey' in settings is not None: + url = settings['url'] + bucket = settings['bucket'] + accessKey = settings['access'] + secretKey = settings['secret'] + + if type(url) is not str: + raise Exception('url must be string') + if type(bucket) is not str: + raise Exception('bucket must be string') + if type(accessKey) is not str: + raise Exception('accessKey must be string') + if type(secretKey) is not str: + raise Exception('secretKey must be string') + + g_settings.url = url + g_settings.access_key = accessKey + g_settings.secret_key = secretKey + g_settings.bucket = bucket + clear_s3_client() + + return web.Response(status=200) + except Exception as e: + return web.Response(status=400, text=str(e)) + +@PromptServer.instance.routes.patch("/saveimages3/setting/url") +async def update_url(request): + try: + settings = await request.json() + if 'url' in settings is not None: + url = settings['url'] + if type(url) is not str: + raise Exception('url must be string') + + g_settings.url = url + clear_s3_client() + + + return web.Response(status=200) + except Exception as e: + return web.Response(status=400, text=str(e)) + + + +@PromptServer.instance.routes.patch("/saveimages3/setting/bucket") +async def update_bucket(request): + try: + settings = await request.json() + if 'bucket' in settings is not None: + bucket = settings['bucket'] + if type(bucket) is not str: + raise Exception('bucket must be string') + g_settings.bucket = bucket + clear_s3_client() + + + return web.Response(status=200) + except Exception as e: + return web.Response(status=400, text=str(e)) + + +@PromptServer.instance.routes.patch("/saveimages3/setting/access") +async def update_access(request): + try: + settings = await request.json() + if 'access' in settings is not None: + accessKey = settings['access'] + if type(accessKey) is not str: + raise Exception('bucket must be string') + g_settings.access_key = accessKey + clear_s3_client() + + return web.Response(status=200) + except Exception as e: + return web.Response(status=400, text=str(e)) + + +@PromptServer.instance.routes.patch("/saveimages3/setting/secret") +async def update_secret(request): + try: + settings = await request.json() + if 'secret' in settings is not None: + secretKey = settings['secret'] + if type(secretKey) is not str: + raise Exception('bucket must be string') + g_settings.secret_key = secretKey + clear_s3_client() + + return web.Response(status=200) + except Exception as e: + return web.Response(status=400, text=str(e)) \ No newline at end of file diff --git a/web/plugin.js b/web/plugin.js new file mode 100644 index 0000000..4d0a84b --- /dev/null +++ b/web/plugin.js @@ -0,0 +1,133 @@ +import { ComfyApp, app } from "../../scripts/app.js" +import { api } from "../../scripts/api.js"; + + +function debounce(func, wait) { + let timeout; + + return function executedFunction(...args) { + const later = () => { + clearTimeout(timeout); + func(...args); + }; + + clearTimeout(timeout); + timeout = setTimeout(later, wait); + }; +} + +const debounce_update_url_settings = debounce(update_settings, 1000); +const debounce_update_bucket_settings = debounce(update_settings, 1000); +const debounce_update_access_settings = debounce(update_settings, 1000); +const debounce_update_secret_settings = debounce(update_settings, 1000); + +async function update_settings(setting, value) { + console.log(`Setting ${setting}: ${value}`); + const data = {} + data[setting] = value; + const resp = await api.fetchApi(`/saveimages3/setting/${setting}`, { + method: 'PATCH', + body: JSON.stringify(data), + cache: 'no-store', + }); + if (resp.status === 200) { + return await resp.text(); + } + throw new Error(resp.statusText); +} + +class SaveImageS3Plugin +{ + extentionName = "Chipset.SaveImageS3"; + + createSettingsURL = () => { + this.settingsURL = { + id: "Chipset.SaveImageS3.URL", + type: "string", + name: "S3 URL", + category: ['SaveImageS3', 'Configuration', 'url'], + tooltip: 'S3 URL', + type: "text", + defaultValue: "", + onChange: (value) => { + debounce_update_url_settings("url", value); + } + } + } + + createSettingsAccessKey = () => { + this.settingsAccessKey = { + id: plugin.extentionName + ".AccessKey", + type: "string", + name: "Access Key", + category: ['SaveImageS3', 'Configuration', 'access'], + tooltip: 'Access Key', + type: "text", + defaultValue: "", + onChange: (value) => { + debounce_update_access_settings('access', value); + } + } + } + + createSettingsSecretKey = () => { + this.settingsSecretKey = { + id: plugin.extentionName + ".SecretKey", + type: "string", + name: "Secret Key", + category: ['SaveImageS3', 'Configuration', 'secret'], + tooltip: 'Secret Key', + type: "text", + defaultValue: "", + onChange: (value) => { + debounce_update_secret_settings('secret', value); + } + } + } + + createSettingsBucketName = () => { + this.settingsBucketName = { + id: plugin.extentionName + ".BucketName", + type: "string", + name: "Bucket Name", + category: ['SaveImageS3', 'Configuration', 'bucket'], + tooltip: 'Bucket Name', + type: "text", + defaultValue: "", + onChange: (value) => { + debounce_update_bucket_settings('bucket', value); + } + } + } + + addSettingsUI = () => { + app.ui.settings.addSetting(this.settingsSecretKey); + app.ui.settings.addSetting(this.settingsAccessKey); + app.ui.settings.addSetting(this.settingsBucketName); + app.ui.settings.addSetting(this.settingsURL); + + } + +} + +const plugin = new SaveImageS3Plugin(); + +app.registerExtension({ + name: plugin.extentionName, + async setup() { + console.log("Comfy.SaveImageS3 setup!!!"); + plugin.createSettingsURL(); + plugin.createSettingsAccessKey(); + plugin.createSettingsSecretKey(); + plugin.createSettingsBucketName(); + + plugin.addSettingsUI(); + }, + async beforeRegisterNodeDef(nodeType, nodeData, app) { + // console.log(nodeType.comfyClass); + if (nodeType.comfyClass=="SaveImageS3") { + console.log("SaveImageS3 node is registed") + } + } +}) +