FEAT: add basic function to save image to S3
This commit is contained in:
Mr.Chip
@@ -0,0 +1,9 @@
|
||||
from .nodes.images import SaveImageS3
|
||||
from .server import *
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"SaveImageS3": SaveImageS3
|
||||
}
|
||||
|
||||
|
||||
WEB_DIRECTORY = "./web"
|
||||
@@ -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
|
||||
@@ -0,0 +1 @@
|
||||
# intentionally left blank
|
||||
@@ -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 } }
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
from .setting import *
|
||||
@@ -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))
|
||||
+133
@@ -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")
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
Reference in New Issue
Block a user