commit 00da9d3e0873d97295be73e3f487e1eb34d5fea7 Author: TwoPerCent Date: Wed Apr 17 17:26:54 2024 -0400 Initial Version Initial Upload diff --git a/README.md b/README.md new file mode 100644 index 0000000..e2b5625 --- /dev/null +++ b/README.md @@ -0,0 +1,9 @@ +![alt text](api_cat_with_workflow.png) + +**Usage:** + +Add API key to environment variable "SAI_API_KEY" + +Alternatively you can write your API key to file "sai_platform_key.txt" + +You can also use and/or override the above by entering your API key in the 'api_override' field of each node diff --git a/api_cat_with_workflow.png b/api_cat_with_workflow.png new file mode 100644 index 0000000..fc1aa62 Binary files /dev/null and b/api_cat_with_workflow.png differ diff --git a/sai_api.png b/sai_api.png new file mode 100644 index 0000000..679d9d2 Binary files /dev/null and b/sai_api.png differ diff --git a/stability_api.py b/stability_api.py new file mode 100644 index 0000000..066c39a --- /dev/null +++ b/stability_api.py @@ -0,0 +1,266 @@ +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(f"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'") + 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}), + "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": "a winter wonderland"}), + "negative_prompt": ("STRING", {"multiline": True, "default": "a post-apocalyptic world"}), + "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}), + "output_format": (["png", "webp", "jpeg"],), + }, + "optional": { + "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}), + "negative_prompt": ("STRING", {"multiline": True}), + "api_key override": ("STRING", {"multiline": False}), + }, + } + + +class StabilitySD3(StabilityBase): + API_ENDPOINT = "stable-image/generate/sd3" + ACCEPT = "image/*" + INPUT_SPEC = { + "required": { + "model": (["sd3", "sd3-turbo"],), + "prompt": ("STRING", {"multiline": True}), + "negative_prompt": ("STRING", {"multiline": True}), + "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}), + "output_format": (["png", "jpeg"],), + }, + "optional": { + "image": ("IMAGE",), + "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"],), + "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}), + }, + } + +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, +}