Initial commit: upload_images and classifiers nodes
This commit is contained in:
David Strand
@@ -0,0 +1,21 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2024 Kudelabs.com
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
+12
@@ -0,0 +1,12 @@
|
||||
from .upload_images import UploadImages
|
||||
from .classifiers import *
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"UploadImagesToInstantStudio": UploadImages,
|
||||
'HuggingFace Classify': HFClassify,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"UploadImagesToInstantStudio": "Upload Images to Instant Studio"
|
||||
}
|
||||
@@ -0,0 +1,67 @@
|
||||
import os
|
||||
import folder_paths
|
||||
import torch
|
||||
from torchvision import transforms
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
def _tensor_check_image(image):
|
||||
if image.ndim != 4:
|
||||
raise ValueError(f"Expected NHWC tensor, but found {image.ndim} dimensions")
|
||||
if image.shape[-1] not in (1, 3, 4):
|
||||
raise ValueError(f"Expected 1, 3 or 4 channels for image, but found {image.shape[-1]} channels")
|
||||
return
|
||||
|
||||
def tensor2pil(image):
|
||||
_tensor_check_image(image)
|
||||
return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(0), 0, 255).astype(np.uint8))
|
||||
|
||||
def numpy2pil(image):
|
||||
return Image.fromarray(np.clip(255. * image.squeeze(0), 0, 255).astype(np.uint8))
|
||||
|
||||
def to_pil(image):
|
||||
if isinstance(image, Image.Image):
|
||||
return image
|
||||
if isinstance(image, torch.Tensor):
|
||||
return tensor2pil(image)
|
||||
if isinstance(image, np.ndarray):
|
||||
return numpy2pil(image)
|
||||
raise ValueError(f"Cannot convert {type(image)} to PIL.Image")
|
||||
|
||||
class HFClassify:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"hf_classifier_model": ("TRANSFORMERS_CLASSIFIER",),
|
||||
"image_to_classify": ("IMAGE",),
|
||||
"target_label": (["male", "female"],), # Toggle for selecting 'male' or 'female'
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("FLOAT", "INT") # Return the raw score as a float and the rounded score as an int
|
||||
RETURN_NAMES = ("raw_score", "rounded_score")
|
||||
FUNCTION = "classify"
|
||||
|
||||
def classify(self, hf_classifier_model, image_to_classify, target_label):
|
||||
# Convert image to PIL format
|
||||
image = to_pil(image_to_classify)
|
||||
|
||||
# Get classification result
|
||||
result = hf_classifier_model(image)
|
||||
|
||||
# Extract the score for the target label
|
||||
target_result = next((entry for entry in result if entry["label"] == target_label), None)
|
||||
if target_result:
|
||||
score = target_result["score"]
|
||||
else:
|
||||
# If the target label isn't in the results, default score to 0
|
||||
score = 0.0
|
||||
|
||||
# Round the score to 0 or 1
|
||||
rounded_score = int(round(score))
|
||||
|
||||
# Return the raw score as a float and the rounded score as an integer
|
||||
return (score, rounded_score)
|
||||
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
numpy
|
||||
Pillow
|
||||
Requests
|
||||
@@ -0,0 +1,94 @@
|
||||
import os
|
||||
import json
|
||||
from comfy.cli_args import args
|
||||
from PIL import Image
|
||||
from PIL.PngImagePlugin import PngInfo
|
||||
import numpy as np
|
||||
import requests
|
||||
|
||||
class UploadImages:
|
||||
def __init__(self):
|
||||
self.type = "output"
|
||||
self.compress_level = 4
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"images": ("IMAGE", {"tooltip": "The images to upload."}),
|
||||
"note": ("STRING", {"tooltip": "A note to associate with the image."}),
|
||||
"email": ("STRING", {"tooltip": "The email to use for authentication."}),
|
||||
"url": ("STRING", {"default": "https://toolkit.instantstudio.ai/api/v1/creations", "tooltip": "The URL to post the image to."}),
|
||||
"api_key": ("STRING", {"default": "TODO", "tooltip": "The API Key to use for authentication."})
|
||||
},
|
||||
"hidden": {
|
||||
"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ()
|
||||
FUNCTION = "upload_images"
|
||||
OUTPUT_NODE = True
|
||||
CATEGORY = "Instant Studio"
|
||||
DESCRIPTION = "Uploads the input images to Instant Studio."
|
||||
|
||||
def upload_images(self, images, note=None, email=None, url=None, api_key=None, prompt=None, extra_pnginfo=None):
|
||||
uuids = None
|
||||
|
||||
succeeded_count = 0
|
||||
failed_count = 0
|
||||
temp_file_paths = list()
|
||||
|
||||
try:
|
||||
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:
|
||||
for key in prompt:
|
||||
if "inputs" in prompt[key] and "api_key" in prompt[key]["inputs"]:
|
||||
prompt[key]["inputs"]["api_key"] = "TODO"
|
||||
metadata.add_text("prompt", json.dumps(prompt))
|
||||
if extra_pnginfo is not None:
|
||||
if "workflow" in extra_pnginfo and "nodes" in extra_pnginfo["workflow"]:
|
||||
for node in extra_pnginfo['workflow']['nodes']:
|
||||
if node['type'] == 'UploadImagesToInstantStudio':
|
||||
node['widgets_values'] = []
|
||||
for x in extra_pnginfo:
|
||||
metadata.add_text(x, json.dumps(extra_pnginfo[x]))
|
||||
|
||||
temp_file = None
|
||||
|
||||
import tempfile
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as temp_file:
|
||||
temp_file_path = temp_file.name
|
||||
img.save(temp_file_path, pnginfo=metadata, compress_level=self.compress_level)
|
||||
|
||||
temp_file_paths.append(temp_file_path)
|
||||
|
||||
files = [("files[]", open(temp_file_path, "rb")) for temp_file_path in temp_file_paths]
|
||||
headers = {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"X-User-Email": email
|
||||
}
|
||||
params = {
|
||||
"note": note,
|
||||
}
|
||||
response = requests.post(url, headers=headers, params=params, files=files)
|
||||
if response.status_code != 200:
|
||||
failed_count += 1
|
||||
raise Exception(f"Failed to post image: {response.status_code} - {response.text}")
|
||||
else:
|
||||
succeeded_count += 1
|
||||
response_data = response.json()
|
||||
uuids = response_data.get("uuids")
|
||||
|
||||
finally:
|
||||
for temp_file_path in temp_file_paths:
|
||||
if temp_file_path and os.path.exists(temp_file_path):
|
||||
os.remove(temp_file_path)
|
||||
|
||||
print(f"Successfully uploaded {succeeded_count} images to Instant Studio: {uuids}")
|
||||
return { "ui": { "uuids": uuids }, "result": (uuids,) }
|
||||
Reference in New Issue
Block a user