From 5df47d60d57b358bbaab840559252e37a9e87e96 Mon Sep 17 00:00:00 2001 From: David Strand Date: Wed, 11 Dec 2024 23:56:07 -0600 Subject: [PATCH] Initial commit: upload_images and classifiers nodes --- LICENSE | 21 +++++++++++ __init__.py | 12 +++++++ classifiers.py | 67 ++++++++++++++++++++++++++++++++++ requirements.txt | 3 ++ upload_images.py | 94 ++++++++++++++++++++++++++++++++++++++++++++++++ 5 files changed, 197 insertions(+) create mode 100644 LICENSE create mode 100644 __init__.py create mode 100644 classifiers.py create mode 100644 requirements.txt create mode 100644 upload_images.py diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..c4e25c2 --- /dev/null +++ b/LICENSE @@ -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. diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..1fad1b4 --- /dev/null +++ b/__init__.py @@ -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" +} diff --git a/classifiers.py b/classifiers.py new file mode 100644 index 0000000..ae0638e --- /dev/null +++ b/classifiers.py @@ -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) + + diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..d6a5666 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,3 @@ +numpy +Pillow +Requests diff --git a/upload_images.py b/upload_images.py new file mode 100644 index 0000000..5817484 --- /dev/null +++ b/upload_images.py @@ -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,) }