diff --git a/.gitignore b/.gitignore index c18dd8d..e17f5ed 100644 --- a/.gitignore +++ b/.gitignore @@ -1 +1,3 @@ -__pycache__/ +__pycache__/ +sai_platform_key.txt +log.txt \ No newline at end of file diff --git a/LICENSE.txt b/LICENSE.txt index c58c5b1..5b82316 100644 --- a/LICENSE.txt +++ b/LICENSE.txt @@ -1,21 +1,21 @@ -The MIT License (MIT) - -Copyright (c) 2024 Stability AI - -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. +The MIT License (MIT) + +Copyright (c) 2024 Stability AI + +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/README.md b/README.md index 8b13789..d3f5a12 100644 --- a/README.md +++ b/README.md @@ -1 +1 @@ - + diff --git a/__init__.py b/__init__.py index 8aad30c..b1df004 100644 --- a/__init__.py +++ b/__init__.py @@ -1,11 +1,24 @@ -from .stability_api import StabilityCreativeUpscale, StabilityRemoveBackground, StabilityInpainting, StabilityCore, StabilitySearchAndReplace, StabilityOutpainting, StabilitySD3 - -NODE_CLASS_MAPPINGS = { - "Stability Creative Upscale": StabilityCreativeUpscale, - "Stability Remove Background": StabilityRemoveBackground, - "Stability Inpainting": StabilityInpainting, - "Stability Image Core": StabilityCore, - "Stability Search and Replace": StabilitySearchAndReplace, - "Stability Outpainting": StabilityOutpainting, - "Stability SD3": StabilitySD3, -} +from .stability_api import StabilityCreativeUpscale, StabilityCreativeUpscaleRecover, StabilityRemoveBackground, StabilityInpainting, StabilityCore, StabilitySearchAndReplace, StabilityOutpainting, StabilitySD3 + +NODE_CLASS_MAPPINGS = { + "StabilityAI Suite - Creative Upscale": StabilityCreativeUpscale, + "StabilityAI Suite - Creative Upscale Recover File": StabilityCreativeUpscaleRecover, + "StabilityAI Suite - Remove Background": StabilityRemoveBackground, + "StabilityAI Suite - Inpainting": StabilityInpainting, + "StabilityAI Suite - Image Core + Style Preset": StabilityCore, + "StabilityAI Suite - Search and Replace": StabilitySearchAndReplace, + "StabilityAI Suite - Outpainting": StabilityOutpainting, + "StabilityAI Suite - SD3": StabilitySD3, +} +from .stability_api import StabilityCreativeUpscale, StabilityCreativeUpscaleRecover, StabilityRemoveBackground, StabilityInpainting, StabilityCore, StabilitySearchAndReplace, StabilityOutpainting, StabilitySD3 + +NODE_CLASS_MAPPINGS = { + "StabilityAI Suite - Creative Upscale": StabilityCreativeUpscale, + "StabilityAI Suite - Creative Upscale Recover File": StabilityCreativeUpscaleRecover, + "StabilityAI Suite - Remove Background": StabilityRemoveBackground, + "StabilityAI Suite - Inpainting": StabilityInpainting, + "StabilityAI Suite - Image Core + Style Preset": StabilityCore, + "StabilityAI Suite - Search and Replace": StabilitySearchAndReplace, + "StabilityAI Suite - Outpainting": StabilityOutpainting, + "StabilityAI Suite - SD3": StabilitySD3, +} diff --git a/api_cat_with_workflow.png b/api_cat_with_workflow.png deleted file mode 100644 index fc1aa62..0000000 Binary files a/api_cat_with_workflow.png and /dev/null differ diff --git a/images/creative_upscale_recover_file.png b/images/creative_upscale_recover_file.png new file mode 100644 index 0000000..86b45ae Binary files /dev/null and b/images/creative_upscale_recover_file.png differ diff --git a/log.txt b/log.txt new file mode 100644 index 0000000..df88fdc --- /dev/null +++ b/log.txt @@ -0,0 +1 @@ +2024-04-18_17-18-12 - Image/video ID for recovery: 1c72e698fb5090591c2d44b5cface686ed4fbac97d6e971b97134a8e63325e2e diff --git a/sai_api.png b/sai_api.png deleted file mode 100644 index 679d9d2..0000000 Binary files a/sai_api.png and /dev/null differ diff --git a/sai_platform_key.txt b/sai_platform_key.txt new file mode 100644 index 0000000..e69de29 diff --git a/stability_api.py b/stability_api.py index 6ee1985..b268149 100644 --- a/stability_api.py +++ b/stability_api.py @@ -1,269 +1,339 @@ -import requests -from requests.models import PreparedRequest -from PIL import Image -import numpy as np -import torch -from torchvision.transforms import ToPILImage -from io import BytesIO -import os -import time - -API_KEY = os.environ.get("SAI_API_KEY") - -# Check for API key in file as a backup, not recommended -try: - if not API_KEY: - dir_path = os.path.dirname(os.path.realpath(__file__)) - with open(os.path.join(dir_path, "sai_platform_key.txt"), "r") as f: - API_KEY = f.read().strip() - print(f"API Key found in sai_platform_key.txt: {API_KEY}") - # Validate the key is not empty - if API_KEY.strip() == "": - raise Exception(f"API Key is required to use the Stability API. \nPlease set the SAI_API_KEY environment variable to your API key or place in {dir_path}/sai_platform_key.txt.") - -except Exception as e: - print(f"\n\n***API Key is required to use the Stability API. Please set the SAI_API_KEY environment variable to your API key or place in {dir_path}/sai_platform_key.txt.***\n\n") - -ROOT_API = "https://api.stability.ai/v2beta/" - - -class StabilityBase: - API_ENDPOINT = "" - POLL_ENDPOINT = "" - ACCEPT = "" - - @classmethod - def INPUT_TYPES(cls): - return cls.INPUT_SPEC - - RETURN_TYPES = ("IMAGE",) - FUNCTION = "call" - CATEGORY = "Stability" - - def call(self, *args, **kwargs): - - buffered = BytesIO() - files = {'none': None} - data = None - - image = kwargs.get('image', None) - if image is not None: - kwargs["mode"] = "image-to-image" - kwargs.pop("aspect_ratio", None) - image = ToPILImage()(image.squeeze(0).permute(2,0,1)) - image.save(buffered, format="PNG") - files = self._get_files(buffered, **kwargs) - else: - kwargs.pop("strength", None) - - headers = { - "Authorization": API_KEY, - } - - if kwargs.get("api_key_override"): - headers = { - "Authorization": kwargs.get("api_key_override"), - } - - if headers.get("Authorization") is None: - raise Exception(f"No Stability key set.\n\nUse your Stability AI API key by:\n1. Setting the SAI_API_KEY environment variable to your API key\n3. Placing inside sai_platform_key.txt\n4. Passing the API key as an argument to the function with the key 'api_key_override'") - - headers["Accept"] = self.ACCEPT - - data = self._get_data(**kwargs) - - req = PreparedRequest() - req.prepare_method('POST') - req.prepare_url(f"{ROOT_API}{self.API_ENDPOINT}", None) - req.prepare_headers(headers) - req.prepare_body(data=data, files=files) - response = requests.Session().send(req) - - if response.status_code == 200: - if self.POLL_ENDPOINT != "": - id = response.json().get("id") - timeout = 240 - start_time = time.time() - while True: - response = requests.get(f"{ROOT_API}{self.POLL_ENDPOINT}{id}", headers=headers) - if response.status_code == 200: - if self.ACCEPT == "image/*": - return self._return_image(response) - if self.ACCEPT == "video/*": - return self._return_video(response) - break - elif response.status_code == 202: - time.sleep(10) - elif time.time() - start_time > timeout: - raise Exception("Stability API Timeout: Request took too long to complete") - else: - error_info = response.json() - raise Exception(f"Stability API Error: {error_info}") - else: - result_image = Image.open(BytesIO(response.content)) - result_image = result_image.convert("RGBA") - result_image = np.array(result_image).astype(np.float32) / 255.0 - result_image = torch.from_numpy(result_image)[None,] - return (result_image,) - else: - error_info = response.json() - if error_info.get("name") == "unauthorized": - raise Exception("Stability API Error: Unauthorized.\n\nUse your Stability AI API key by:\n1. Setting the SAI_API_KEY environment variable to your API key\n3. Placing inside sai_platform_key.txt\n4. Passing the API key as an argument to the function with the key 'api_key_override'") - if error_info.get("name") == "payment_required": - raise Exception("Stability API Error: Not enough credits.\n\nPlease ensure your SAI API account has enough credits to complete this action.") - if error_info.get("name") == "bad_request": - errors = '\n'.join(error_info.get('errors')) - raise Exception(f"Stability API Error: Bad request.\n\n{errors}") - else: - raise Exception(f"Stability API Error: {error_info}") - - def _return_image(self, response): - result_image = Image.open(BytesIO(response.content)) - result_image = result_image.convert("RGBA") - result_image = np.array(result_image).astype(np.float32) / 255.0 - result_image = torch.from_numpy(result_image)[None,] - return (result_image,) - - def _return_video(self, response): - result_video = response.content - return (result_video,) - - def _get_files(self, buffered, **kwargs): - return { - "image": buffered.getvalue() - } - - def _get_data(self, **kwargs): - return {k: v for k, v in kwargs.items() if k != "image"} - - -class StabilityCore(StabilityBase): - API_ENDPOINT = "stable-image/generate/core" - ACCEPT = "image/*" - INPUT_SPEC = { - "required": { - "prompt": ("STRING", {"multiline": True}), - }, - "optional": { - "negative_prompt": ("STRING", {"multiline": True}), - "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}), - "output_format": (["png", "webp", "jpeg"],), - "aspect_ratio": (["16:9", "1:1", "21:9", "2:3", "3:2", "4:5", "5:4", "9:16", "9:21"],), - "api_key_override": ("STRING", {"multiline": False}), - } - } - - -class StabilityCreativeUpscale(StabilityBase): - API_ENDPOINT = "stable-image/upscale/creative" - POLL_ENDPOINT = "stable-image/upscale/creative/result/" - ACCEPT = "image/*" - INPUT_SPEC = { - "required": { - "image": ("IMAGE",), - "prompt": ("STRING", {"multiline": True}), - }, - "optional": { - "negative_prompt": ("STRING", {"multiline": True}), - "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}), - "creativity": ("FLOAT", {"default": 0.3, "min": 0.01, "max": 0.35, "step": 0.01}), - "output_format": (["png", "webp", "jpeg"],), - "api_key_override": ("STRING", {"multiline": False}), - } - } - - -class StabilityRemoveBackground(StabilityBase): - API_ENDPOINT = "stable-image/edit/remove-background" - ACCEPT = "image/*" - INPUT_SPEC = { - "required": { - "image": ("IMAGE",), - }, - } - -class StabilityInpainting(StabilityBase): - API_ENDPOINT = "stable-image/edit/inpaint" - ACCEPT = "image/*" - INPUT_SPEC = { - "required": { - "image": ("IMAGE",), - "mask": ("MASK",), - "prompt": ("STRING", {"multiline": True, "default": ""}),\ - }, - "optional": { - "negative_prompt": ("STRING", {"multiline": True, "default": ""}),\ - "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}), - "output_format": (["png", "webp", "jpeg"],), - "api_key_override": ("STRING", {"multiline": False}), - } - } - def _get_files(self, buffered, **kwargs): - mask = kwargs.get("mask") - to_pil = ToPILImage() - mask = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3) - mask = to_pil(mask.squeeze(0).permute(2,0,1)) - buffered_mask = BytesIO() - mask.save(buffered_mask, format="PNG") - return { - "image": buffered.getvalue(), - "mask": buffered_mask.getvalue(), - } - - -class StabilitySearchAndReplace(StabilityBase): - API_ENDPOINT = "stable-image/edit/search-and-replace" - ACCEPT = "image/*" - INPUT_SPEC = { - "required": { - "image": ("IMAGE",), - "search_prompt": ("STRING", {"multiline": True}, "Search Prompt"), - "prompt": ("STRING", {"multiline": True}), - }, - "optional": { - "negative_prompt": ("STRING", {"multiline": True}), - "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}), - "api_key_override": ("STRING", {"multiline": False}), - "output_format": (["png", "webp", "jpeg"],), - }, - } - - -class StabilitySD3(StabilityBase): - API_ENDPOINT = "stable-image/generate/sd3" - ACCEPT = "image/*" - INPUT_SPEC = { - "required": { - "model": (["sd3", "sd3-turbo"],), - "prompt": ("STRING", {"multiline": True}), - }, - "optional": { - "image": ("IMAGE",), - "negative_prompt": ("STRING", {"multiline": True}), - "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}), - "strength": ("FLOAT", {"default": 0.5, "min": 0.01, "max": 1.0, "step": 0.01}), - "aspect_ratio": (["16:9", "1:1", "21:9", "2:3", "3:2", "4:5", "5:4", "9:16", "9:21"],), - "output_format": (["png", "jpeg"],), - "api_key_override": ("STRING", {"multiline": False}), - }, - } - - -class StabilityOutpainting(StabilityBase): - API_ENDPOINT = "stable-image/edit/outpaint" - ACCEPT = "image/*" - INPUT_SPEC = { - "required": { - "image": ("IMAGE",), - "left": ("INT", {"default": 0, "min": 0, "max": 512}), - "right": ("INT", {"default": 0, "min": 0, "max": 512}), - "up": ("INT", {"default": 0, "min": 0, "max": 512}), - "down": ("INT", {"default": 0, "min": 0, "max": 512}),\ - }, - "optional": { - "prompt": ("STRING", {"multiline": True}), - "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}), - "output_format": (["png", "webp", "jpeg"],), - "api_key_override": ("STRING", {"multiline": False}), - }, - } +import requests +from requests.models import PreparedRequest +from PIL import Image +import numpy as np +import torch +from torchvision.transforms import ToPILImage +from io import BytesIO +import os +import time +from datetime import datetime +import base64 +import io + +API_KEY = os.environ.get("SAI_API_KEY") + +# Check for API key in file as a backup, not recommended +try: + if not API_KEY: + dir_path = os.path.dirname(os.path.realpath(__file__)) + with open(os.path.join(dir_path, "sai_platform_key.txt"), "r") as f: + API_KEY = f.read().strip() + print(f"API Key found in sai_platform_key.txt: {API_KEY}") + # Validate the key is not empty + if API_KEY.strip() == "": + raise Exception(f"API Key is required to use the Stability API. \nPlease set the SAI_API_KEY environment variable to your API key or place in {dir_path}/sai_platform_key.txt.") + +except Exception as e: + print(f"\n\n***API Key is required to use the Stability API. Please set the SAI_API_KEY environment variable to your API key or place in {dir_path}/sai_platform_key.txt.***\n\n") + +ROOT_API = "https://api.stability.ai/v2beta/" + + +class StabilityBase: + API_ENDPOINT = "" + POLL_ENDPOINT = "" + ACCEPT = "" + + @classmethod + def INPUT_TYPES(cls): + return cls.INPUT_SPEC + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "call" + CATEGORY = "AI WizArt/Stability AI Suite" + + def call(self, *args, **kwargs): + + buffered = BytesIO() + files = {'none': None} + data = None + + image = kwargs.get('image', None) + if image is not None: + kwargs["mode"] = "image-to-image" + kwargs.pop("aspect_ratio", None) + image = ToPILImage()(image.squeeze(0).permute(2,0,1)) + image.save(buffered, format="PNG") + files = self._get_files(buffered, **kwargs) + else: + kwargs.pop("strength", None) + + style = kwargs.get('style', False) + if style is False: + kwargs.pop('style_preset', None) + + headers = { + "Authorization": API_KEY, + } + + if kwargs.get("api_key_override"): + headers = { + "Authorization": kwargs.get("api_key_override"), + } + + if headers.get("Authorization") is None: + raise Exception(f"No Stability key set.\n\nUse your Stability AI API key by:\n1. Setting the SAI_API_KEY environment variable to your API key\n3. Placing inside sai_platform_key.txt\n4. Passing the API key as an argument to the function with the key 'api_key_override'") + + headers["Accept"] = self.ACCEPT + + data = self._get_data(**kwargs) + + req = PreparedRequest() + req.prepare_method('POST') + req.prepare_url(f"{ROOT_API}{self.API_ENDPOINT}", None) + req.prepare_headers(headers) + req.prepare_body(data=data, files=files) + response = requests.Session().send(req) + + if response.status_code == 200: + if self.POLL_ENDPOINT != "": + id = response.json().get("id") + logFile(f"Image/video ID for recovery: {id}") # saving id for recovery in case of malfunction + timeout = 240 + start_time = time.time() + while True: + response = requests.get(f"{ROOT_API}{self.POLL_ENDPOINT}{id}", headers=headers) + if response.status_code == 200: + if self.ACCEPT == "image/*": + return self._return_image(response) + if self.ACCEPT == "video/*": + return self._return_video(response) + break + elif response.status_code == 202: + time.sleep(10) + elif time.time() - start_time > timeout: + raise Exception("Stability API Timeout: Request took too long to complete") + else: + error_info = response.json() + raise Exception(f"Stability API Error: {error_info}") + else: + result_image = Image.open(BytesIO(response.content)) + result_image = result_image.convert("RGBA") + result_image = np.array(result_image).astype(np.float32) / 255.0 + result_image = torch.from_numpy(result_image)[None,] + return (result_image,) + else: + error_info = response.json() + if error_info.get("name") == "unauthorized": + raise Exception("Stability API Error: Unauthorized.\n\nUse your Stability AI API key by:\n1. Setting the SAI_API_KEY environment variable to your API key\n3. Placing inside sai_platform_key.txt\n4. Passing the API key as an argument to the function with the key 'api_key_override'") + if error_info.get("name") == "payment_required": + raise Exception("Stability API Error: Not enough credits.\n\nPlease ensure your SAI API account has enough credits to complete this action.") + if error_info.get("name") == "bad_request": + errors = '\n'.join(error_info.get('errors')) + raise Exception(f"Stability API Error: Bad request.\n\n{errors}") + else: + raise Exception(f"Stability API Error: {error_info}") + + def _return_image(self, response): + result_image = Image.open(BytesIO(response.content)) + result_image = result_image.convert("RGBA") + result_image = np.array(result_image).astype(np.float32) / 255.0 + result_image = torch.from_numpy(result_image)[None,] + return (result_image,) + + def _return_video(self, response): + result_video = response.content + return (result_video,) + + def _get_files(self, buffered, **kwargs): + return { + "image": buffered.getvalue() + } + + def _get_data(self, **kwargs): + return {k: v for k, v in kwargs.items() if k != "image"} + + +class StabilityCore(StabilityBase): + API_ENDPOINT = "stable-image/generate/core" + ACCEPT = "image/*" + INPUT_SPEC = { + "required": { + "prompt": ("STRING", {"multiline": True}), + }, + "optional": { + "negative_prompt": ("STRING", {"multiline": True}), + "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}), + "output_format": (["png", "webp", "jpeg"],), + "aspect_ratio": (["16:9", "1:1", "21:9", "2:3", "3:2", "4:5", "5:4", "9:16", "9:21"],), + "style": ("BOOLEAN", {"default": False}), + "style_preset": (["3d-model", "analog-film", "anime", "cinematic", "comic-book", "digital-art", "enhance", "fantasy-art", "isometric", "line-art", "low-poly", "modeling-compound", "neon-punk", "origami", "photographic", "pixel-art", "tile-texture"],), + "api_key_override": ("STRING", {"multiline": False}), + } + } + + +class StabilityCreativeUpscale(StabilityBase): + API_ENDPOINT = "stable-image/upscale/creative" + POLL_ENDPOINT = "stable-image/upscale/creative/result/" + ACCEPT = "image/*" + INPUT_SPEC = { + "required": { + "image": ("IMAGE",), + "prompt": ("STRING", {"multiline": True}), + }, + "optional": { + "negative_prompt": ("STRING", {"multiline": True}), + "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}), + "creativity": ("FLOAT", {"default": 0.3, "min": 0.01, "max": 0.35, "step": 0.01}), + "output_format": (["png", "webp", "jpeg"],), + "api_key_override": ("STRING", {"multiline": False}), + } + } + + +class StabilityRemoveBackground(StabilityBase): + API_ENDPOINT = "stable-image/edit/remove-background" + ACCEPT = "image/*" + INPUT_SPEC = { + "required": { + "image": ("IMAGE",), + }, + } + +class StabilityInpainting(StabilityBase): + API_ENDPOINT = "stable-image/edit/inpaint" + ACCEPT = "image/*" + INPUT_SPEC = { + "required": { + "image": ("IMAGE",), + "mask": ("MASK",), + "prompt": ("STRING", {"multiline": True, "default": ""}),\ + }, + "optional": { + "negative_prompt": ("STRING", {"multiline": True, "default": ""}),\ + "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}), + "output_format": (["png", "webp", "jpeg"],), + "api_key_override": ("STRING", {"multiline": False}), + } + } + def _get_files(self, buffered, **kwargs): + mask = kwargs.get("mask") + to_pil = ToPILImage() + mask = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3) + mask = to_pil(mask.squeeze(0).permute(2,0,1)) + buffered_mask = BytesIO() + mask.save(buffered_mask, format="PNG") + return { + "image": buffered.getvalue(), + "mask": buffered_mask.getvalue(), + } + + +class StabilitySearchAndReplace(StabilityBase): + API_ENDPOINT = "stable-image/edit/search-and-replace" + ACCEPT = "image/*" + INPUT_SPEC = { + "required": { + "image": ("IMAGE",), + "search_prompt": ("STRING", {"multiline": True}, "Search Prompt"), + "prompt": ("STRING", {"multiline": True}), + }, + "optional": { + "negative_prompt": ("STRING", {"multiline": True}), + "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}), + "api_key_override": ("STRING", {"multiline": False}), + "output_format": (["png", "webp", "jpeg"],), + }, + } + + +class StabilitySD3(StabilityBase): + API_ENDPOINT = "stable-image/generate/sd3" + ACCEPT = "image/*" + INPUT_SPEC = { + "required": { + "model": (["sd3", "sd3-turbo"],), + "prompt": ("STRING", {"multiline": True}), + }, + "optional": { + "image": ("IMAGE",), + "negative_prompt": ("STRING", {"multiline": True}), + "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}), + "strength": ("FLOAT", {"default": 0.5, "min": 0.01, "max": 1.0, "step": 0.01}), + "aspect_ratio": (["16:9", "1:1", "21:9", "2:3", "3:2", "4:5", "5:4", "9:16", "9:21"],), + "output_format": (["png", "jpeg"],), + "api_key_override": ("STRING", {"multiline": False}), + }, + } + + +class StabilityOutpainting(StabilityBase): + API_ENDPOINT = "stable-image/edit/outpaint" + ACCEPT = "image/*" + INPUT_SPEC = { + "required": { + "image": ("IMAGE",), + "left": ("INT", {"default": 0, "min": 0, "max": 512}), + "right": ("INT", {"default": 0, "min": 0, "max": 512}), + "up": ("INT", {"default": 0, "min": 0, "max": 512}), + "down": ("INT", {"default": 0, "min": 0, "max": 512}), + }, + "optional": { + "prompt": ("STRING", {"multiline": True}), + "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}), + "output_format": (["png", "webp", "jpeg"],), + "api_key_override": ("STRING", {"multiline": False}), + }, + } + +# ======================================================== +# FILE RECOVER +# ======================================================== + +class StabilityCreativeUpscaleRecover(StabilityBase): + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "image_id": ("STRING", { + "multiline": False + }) + } + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image_out",) + FUNCTION = "creativeUpscaleRecover" + + def creativeUpscaleRecover(self, image_id): + get_code = 202 + while get_code == 202: + response_get = requests.request( + "GET", + f"https://api.stability.ai/v2beta/stable-image/upscale/creative/result/{image_id}", + headers={ + "accept": "application/json", + "authorization": f"Bearer {API_KEY}" + }, + ) + get_code = response_get.status_code + time.sleep(10) + print("Waiting image...") + if response_get.status_code == 200: + json_data = response_get.json() + image_base64 = json_data['image'] + image_bytes = base64.b64decode(image_base64) + image_data = Image.open(io.BytesIO(image_bytes)) + output_t = pil2tensor(image_data) + return (output_t,) + else: + print(response_get.json()) + +# ======================================================== +# UTILITIES +# ======================================================== + +def logFile(text): + now = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + new_entry = f"{now} - {text}" + logfile = os.path.join(dir_path, 'log.txt') + with open(logfile, "a") as file: + file.write(new_entry + "\n") + +def pil2tensor(image): + return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0) \ No newline at end of file