Merge pull request #2 from ComfyNodePRs/pyproject

Add pyproject.toml for Custom Node Registry

.github/workflows/publish.yml

@@ -0,0 +1,22 @@
+name: Publish to Comfy registry
+on:
+  workflow_dispatch:
+  push:
+    branches:
+      - main
+      - master
+    paths:
+      - "pyproject.toml"
+
+jobs:
+  publish-node:
+    name: Publish Custom Node to registry
+    runs-on: ubuntu-latest
+    steps:
+      - name: Check out code
+        uses: actions/checkout@v4
+      - name: Publish Custom Node
+        uses: Comfy-Org/publish-node-action@main
+        with:
+          ## Add your own personal access token to your Github Repository secrets and reference it here.
+          personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }}

.gitignore

@@ -1 +1,3 @@
-__pycache__/
+__pycache__/
+sai_platform_key.txt
+log.txt

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.

README.md

@@ -1,35 +1,44 @@
-# Stability API nodes for ComfyUI
-
-![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_key_override`' field of each node.
-
-# License
-
-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.
+# ComfyUI StabilityAI Suite
+
+This fork of the official StabilityAI repository contains a number of enhancements and implementations.
+
+**_If this project is useful to you and you like it, please consider a small donation to the author._**
+
+➡️ https://ko-fi.com/stefanoflore75
+
+Buy my workflows:
+
+➡️ https://stefanoflore.it/download/
+
+## Image Core + Style Preset
+
+In this fork the node **Image Core** has been updated: it contains a new true/false switch to activate styles, and the corresponding selector.
+
+**I submitted a change request in the original repository to implement the styles. The change was accepted, so find the settings for the styles in the official node as well.**
+
+![Image Core + Style Preset](/images/image_core_style.png)
+
+## Creative Upscale Recover File 
+
+The node for the creative uspcale performs two steps: generation and recovery of the generated image. If unfortunately your PC has a crash or suddenly lacks an Internet connection, the second step cannot complete and you risk losing the generated image and the credits used. Fortunately, the images are archived for 24 hours by StabilityAI.
+Thanks to a change in the general code, the  unique __id__ are automatically saved within the __log.txt__ file. This file is automatically created and saved in the same folder as the custom node.
+
+Follow these steps:
+
+- Open the file and copy the desired `id`.
+- Paste the `id` into the __Creative Upscale Recover File__ node to recover the lost image.
+
+![Creative Upscale Recover File](/images/creative_upscale_recover_file.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_key_override`' field of each node.
+
+## Other projects
+
+- [ComfyUI Portrait Master](https://github.com/florestefano1975/comfyui-portrait-master/)
+- [ComfyUI Prompt Composer](https://github.com/florestefano1975/comfyui-prompt-composer/)

__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,
+}

pyproject.toml

@@ -0,0 +1,14 @@
+[project]
+name = "comfyui-stabilityai-suite"
+description = "This fork of the official StabilityAI repository contains a number of enhancements and implementations."
+version = "1.0.0"
+license = "LICENSE"
+
+[project.urls]
+Repository = "https://github.com/florestefano1975/ComfyUI-StabilityAI-Suite"
+#  Used by Comfy Registry https://comfyregistry.org
+
+[tool.comfy]
+PublisherId = ""
+DisplayName = "ComfyUI-StabilityAI-Suite"
+Icon = ""

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)