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)
+The MIT License (MIT)
-
+
-Copyright (c) 2024 Stability AI
+Copyright (c) 2024 Stability AI
-
+
-Permission is hereby granted, free of charge, to any person obtaining a copy
+Permission is hereby granted, free of charge, to any person obtaining a copy
-of this software and associated documentation files (the "Software"), to deal
+of this software and associated documentation files (the "Software"), to deal
-in the Software without restriction, including without limitation the rights
+in the Software without restriction, including without limitation the rights
-to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-copies of the Software, and to permit persons to whom the Software is
+copies of the Software, and to permit persons to whom the Software is
-furnished to do so, subject to the following conditions:
+furnished to do so, subject to the following conditions:
-
+
-The above copyright notice and this permission notice shall be included in all
+The above copyright notice and this permission notice shall be included in all
-copies or substantial portions of the Software.
+copies or substantial portions of the Software.
-
+
-THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+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
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-SOFTWARE.
+SOFTWARE.

README.md

@@ -1,35 +1,44 @@
-# Stability API nodes for ComfyUI
+# ComfyUI StabilityAI Suite
-
+
-![alt text](api_cat_with_workflow.png)
+This fork of the official StabilityAI repository contains a number of enhancements and implementations.
-
+
-### Usage:
+**_If this project is useful to you and you like it, please consider a small donation to the author._**
-
+
-Add API key to environment variable "`SAI_API_KEY`"
+➡️ https://ko-fi.com/stefanoflore75
-
+
-Alternatively you can write your API key to file "`sai_platform_key.txt`"
+Buy my workflows:
-
+
-You can also use and/or override the above by entering your API key in the '`api_key_override`' field of each node.
+➡️ https://stefanoflore.it/download/
-
+
-# License
+## Image Core + Style Preset
-
+
-The MIT License (MIT)
+In this fork the node **Image Core** has been updated: it contains a new true/false switch to activate styles, and the corresponding selector.
-
+
-Copyright (c) 2024 Stability AI
+**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.**
-
+
-Permission is hereby granted, free of charge, to any person obtaining a copy
+![Image Core + Style Preset](/images/image_core_style.png)
-of this software and associated documentation files (the "Software"), to deal
+
-in the Software without restriction, including without limitation the rights
+## Creative Upscale Recover File 
-to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+
-copies of the Software, and to permit persons to whom the Software is
+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.
-furnished to do so, subject to the following conditions:
+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.
-
+
-The above copyright notice and this permission notice shall be included in all
+Follow these steps:
-copies or substantial portions of the Software.
+
-
+- Open the file and copy the desired `id`.
-THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+- Paste the `id` into the __Creative Upscale Recover File__ node to recover the lost image.
-IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+
-FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+![Creative Upscale Recover File](/images/creative_upscale_recover_file.png)
-AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+
-LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+## Usage:
-OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+
-SOFTWARE.
+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
+from .stability_api import StabilityCreativeUpscale, StabilityCreativeUpscaleRecover, StabilityRemoveBackground, StabilityInpainting, StabilityCore, StabilitySearchAndReplace, StabilityOutpainting, StabilitySD3
-
+
-NODE_CLASS_MAPPINGS = {
+NODE_CLASS_MAPPINGS = {
-    "Stability Creative Upscale": StabilityCreativeUpscale,
+    "StabilityAI Suite - Creative Upscale": StabilityCreativeUpscale,
-    "Stability Remove Background": StabilityRemoveBackground,
+    "StabilityAI Suite - Creative Upscale Recover File": StabilityCreativeUpscaleRecover,
-    "Stability Inpainting": StabilityInpainting,
+    "StabilityAI Suite - Remove Background": StabilityRemoveBackground,
-    "Stability Image Core": StabilityCore,
+    "StabilityAI Suite - Inpainting": StabilityInpainting,
-    "Stability Search and Replace": StabilitySearchAndReplace,
+    "StabilityAI Suite - Image Core + Style Preset": StabilityCore,
-    "Stability Outpainting": StabilityOutpainting,
+    "StabilityAI Suite - Search and Replace": StabilitySearchAndReplace,
-    "Stability SD3": StabilitySD3,
+    "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
+import requests
-from requests.models import PreparedRequest
+from requests.models import PreparedRequest
-from PIL import Image
+from PIL import Image
-import numpy as np
+import numpy as np
-import torch
+import torch
-from torchvision.transforms import ToPILImage
+from torchvision.transforms import ToPILImage
-from io import BytesIO
+from io import BytesIO
-import os
+import os
-import time
+import time
-
+from datetime import datetime
-API_KEY = os.environ.get("SAI_API_KEY")
+import base64
-
+import io
-# Check for API key in file as a backup, not recommended
+
-try:
+API_KEY = os.environ.get("SAI_API_KEY")
-    if not API_KEY:
+
-        dir_path = os.path.dirname(os.path.realpath(__file__))
+# Check for API key in file as a backup, not recommended
-        with open(os.path.join(dir_path, "sai_platform_key.txt"), "r") as f:
+try:
-            API_KEY = f.read().strip()
+    if not API_KEY:
-            print(f"API Key found in sai_platform_key.txt: {API_KEY}")
+        dir_path = os.path.dirname(os.path.realpath(__file__))
-        # Validate the key is not empty
+        with open(os.path.join(dir_path, "sai_platform_key.txt"), "r") as f:
-        if API_KEY.strip() == "":
+            API_KEY = f.read().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.")
+            print(f"API Key found in sai_platform_key.txt: {API_KEY}")
-        
+        # Validate the key is not empty
-except Exception as e:
+        if API_KEY.strip() == "":
-    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")
+            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.")
-
+        
-ROOT_API = "https://api.stability.ai/v2beta/"
+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")
-
+
-class StabilityBase:
+ROOT_API = "https://api.stability.ai/v2beta/"
-    API_ENDPOINT = ""
+
-    POLL_ENDPOINT = ""
+
-    ACCEPT = ""
+class StabilityBase:
-
+    API_ENDPOINT = ""
-    @classmethod
+    POLL_ENDPOINT = ""
-    def INPUT_TYPES(cls):
+    ACCEPT = ""
-        return cls.INPUT_SPEC
+
-
+    @classmethod
-    RETURN_TYPES = ("IMAGE",)
+    def INPUT_TYPES(cls):
-    FUNCTION = "call"
+        return cls.INPUT_SPEC
-    CATEGORY = "Stability"
+
-
+    RETURN_TYPES = ("IMAGE",)
-    def call(self, *args, **kwargs):
+    FUNCTION = "call"
-        
+    CATEGORY = "AI WizArt/Stability AI Suite"
-        buffered = BytesIO()
+
-        files = {'none': None}
+    def call(self, *args, **kwargs):
-        data = None
+        
-        
+        buffered = BytesIO()
-        image = kwargs.get('image', None)
+        files = {'none': None}
-        if image is not None:
+        data = None
-            kwargs["mode"] = "image-to-image"
+        
-            kwargs.pop("aspect_ratio", None)
+        image = kwargs.get('image', None)
-            image = ToPILImage()(image.squeeze(0).permute(2,0,1))
+        if image is not None:
-            image.save(buffered, format="PNG")
+            kwargs["mode"] = "image-to-image"
-            files = self._get_files(buffered, **kwargs)
+            kwargs.pop("aspect_ratio", None)
-        else:
+            image = ToPILImage()(image.squeeze(0).permute(2,0,1))
-            kwargs.pop("strength", None)
+            image.save(buffered, format="PNG")
-
+            files = self._get_files(buffered, **kwargs)
-        headers = {
+        else:
-            "Authorization": API_KEY,
+            kwargs.pop("strength", None)
-        }
+        
-
+        style = kwargs.get('style', False)
-        if kwargs.get("api_key_override"):
+        if style is False:
-            headers = {
+            kwargs.pop('style_preset', None)
-                "Authorization": kwargs.get("api_key_override"),
+
-            }
+        headers = {
-
+            "Authorization": API_KEY,
-        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'")
+
-
+        if kwargs.get("api_key_override"):
-        headers["Accept"] = self.ACCEPT
+            headers = {
-
+                "Authorization": kwargs.get("api_key_override"),
-        data = self._get_data(**kwargs)
+            }
-
+
-        req = PreparedRequest()
+        if headers.get("Authorization") is None:
-        req.prepare_method('POST')
+            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'")
-        req.prepare_url(f"{ROOT_API}{self.API_ENDPOINT}", None)
+
-        req.prepare_headers(headers)
+        headers["Accept"] = self.ACCEPT
-        req.prepare_body(data=data, files=files)
+
-        response = requests.Session().send(req)
+        data = self._get_data(**kwargs)
-
+
-        if response.status_code == 200:
+        req = PreparedRequest()
-            if self.POLL_ENDPOINT != "":
+        req.prepare_method('POST')
-                id = response.json().get("id")
+        req.prepare_url(f"{ROOT_API}{self.API_ENDPOINT}", None)
-                timeout = 240
+        req.prepare_headers(headers)
-                start_time = time.time()
+        req.prepare_body(data=data, files=files)
-                while True:
+        response = requests.Session().send(req)
-                    response = requests.get(f"{ROOT_API}{self.POLL_ENDPOINT}{id}", headers=headers)
+
-                    if response.status_code == 200:
+        if response.status_code == 200:
-                        if self.ACCEPT == "image/*":
+            if self.POLL_ENDPOINT != "":
-                            return self._return_image(response)
+                id = response.json().get("id")
-                        if self.ACCEPT == "video/*":
+                logFile(f"Image/video ID for recovery: {id}") # saving id for recovery in case of malfunction
-                            return self._return_video(response)
+                timeout = 240
-                        break
+                start_time = time.time()
-                    elif response.status_code == 202:
+                while True:
-                        time.sleep(10)
+                    response = requests.get(f"{ROOT_API}{self.POLL_ENDPOINT}{id}", headers=headers)
-                    elif time.time() - start_time > timeout:
+                    if response.status_code == 200:
-                        raise Exception("Stability API Timeout: Request took too long to complete")
+                        if self.ACCEPT == "image/*":
-                    else:
+                            return self._return_image(response)
-                        error_info = response.json()
+                        if self.ACCEPT == "video/*":
-                        raise Exception(f"Stability API Error: {error_info}")
+                            return self._return_video(response)
-            else:
+                        break
-                result_image = Image.open(BytesIO(response.content))
+                    elif response.status_code == 202:
-                result_image = result_image.convert("RGBA")
+                        time.sleep(10)
-                result_image = np.array(result_image).astype(np.float32) / 255.0
+                    elif time.time() - start_time > timeout:
-                result_image = torch.from_numpy(result_image)[None,]
+                        raise Exception("Stability API Timeout: Request took too long to complete")
-                return (result_image,)
+                    else:
-        else:
+                        error_info = response.json()
-            error_info = response.json()
+                        raise Exception(f"Stability API Error: {error_info}")
-            if error_info.get("name") == "unauthorized":
+            else:
-                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'")
+                result_image = Image.open(BytesIO(response.content))
-            if error_info.get("name") == "payment_required":
+                result_image = result_image.convert("RGBA")
-                raise Exception("Stability API Error: Not enough credits.\n\nPlease ensure your SAI API account has enough credits to complete this action.")
+                result_image = np.array(result_image).astype(np.float32) / 255.0
-            if error_info.get("name") == "bad_request":
+                result_image = torch.from_numpy(result_image)[None,]
-                errors = '\n'.join(error_info.get('errors'))
+                return (result_image,)
-                raise Exception(f"Stability API Error: Bad request.\n\n{errors}")
+        else:
-            else:
+            error_info = response.json()
-                raise Exception(f"Stability API Error: {error_info}")
+            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'")
-    def _return_image(self, response):
+            if error_info.get("name") == "payment_required":
-        result_image = Image.open(BytesIO(response.content))
+                raise Exception("Stability API Error: Not enough credits.\n\nPlease ensure your SAI API account has enough credits to complete this action.")
-        result_image = result_image.convert("RGBA")
+            if error_info.get("name") == "bad_request":
-        result_image = np.array(result_image).astype(np.float32) / 255.0
+                errors = '\n'.join(error_info.get('errors'))
-        result_image = torch.from_numpy(result_image)[None,]
+                raise Exception(f"Stability API Error: Bad request.\n\n{errors}")
-        return (result_image,)
+            else:
-
+                raise Exception(f"Stability API Error: {error_info}")
-    def _return_video(self, response):
+    
-        result_video = response.content
+    def _return_image(self, response):
-        return (result_video,)
+        result_image = Image.open(BytesIO(response.content))
-
+        result_image = result_image.convert("RGBA")
-    def _get_files(self, buffered, **kwargs):
+        result_image = np.array(result_image).astype(np.float32) / 255.0
-        return {
+        result_image = torch.from_numpy(result_image)[None,]
-            "image": buffered.getvalue()
+        return (result_image,)
-        }
+
-
+    def _return_video(self, response):
-    def _get_data(self, **kwargs):
+        result_video = response.content
-        return {k: v for k, v in kwargs.items() if k != "image"}
+        return (result_video,)
-
+
-
+    def _get_files(self, buffered, **kwargs):
-class StabilityCore(StabilityBase):
+        return {
-    API_ENDPOINT = "stable-image/generate/core"
+            "image": buffered.getvalue()
-    ACCEPT = "image/*"
+        }
-    INPUT_SPEC = {
+
-        "required": {
+    def _get_data(self, **kwargs):
-            "prompt": ("STRING", {"multiline": True}),
+        return {k: v for k, v in kwargs.items() if k != "image"}
-        },
+
-        "optional": {
+
-            "negative_prompt": ("STRING", {"multiline": True}),
+class StabilityCore(StabilityBase):
-            "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}),
+    API_ENDPOINT = "stable-image/generate/core"
-            "output_format": (["png", "webp", "jpeg"],),
+    ACCEPT = "image/*"
-            "aspect_ratio": (["16:9", "1:1", "21:9", "2:3", "3:2", "4:5", "5:4", "9:16", "9:21"],),
+    INPUT_SPEC = {
-            "api_key_override": ("STRING", {"multiline": False}),
+        "required": {
-        }
+            "prompt": ("STRING", {"multiline": True}),
-    }
+        },
-
+        "optional": {
-
+            "negative_prompt": ("STRING", {"multiline": True}),
-class StabilityCreativeUpscale(StabilityBase):
+            "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}),
-    API_ENDPOINT = "stable-image/upscale/creative"
+            "output_format": (["png", "webp", "jpeg"],),
-    POLL_ENDPOINT  = "stable-image/upscale/creative/result/"
+            "aspect_ratio": (["16:9", "1:1", "21:9", "2:3", "3:2", "4:5", "5:4", "9:16", "9:21"],),
-    ACCEPT = "image/*"
+            "style": ("BOOLEAN", {"default": False}),
-    INPUT_SPEC = {
+            "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"],),
-        "required": {
+            "api_key_override": ("STRING", {"multiline": False}),
-            "image": ("IMAGE",),
+        }
-            "prompt": ("STRING", {"multiline": True}),
+    }
-        },
+
-        "optional": {
+
-            "negative_prompt": ("STRING", {"multiline": True}),
+class StabilityCreativeUpscale(StabilityBase):
-            "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}),
+    API_ENDPOINT = "stable-image/upscale/creative"
-            "creativity": ("FLOAT", {"default": 0.3, "min": 0.01, "max": 0.35, "step": 0.01}),
+    POLL_ENDPOINT  = "stable-image/upscale/creative/result/"
-            "output_format": (["png", "webp", "jpeg"],),
+    ACCEPT = "image/*"
-            "api_key_override": ("STRING", {"multiline": False}),
+    INPUT_SPEC = {
-        }
+        "required": {
-    }
+            "image": ("IMAGE",),
-
+            "prompt": ("STRING", {"multiline": True}),
-
+        },
-class StabilityRemoveBackground(StabilityBase):
+        "optional": {
-    API_ENDPOINT = "stable-image/edit/remove-background"
+            "negative_prompt": ("STRING", {"multiline": True}),
-    ACCEPT = "image/*"
+            "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}),
-    INPUT_SPEC = {
+            "creativity": ("FLOAT", {"default": 0.3, "min": 0.01, "max": 0.35, "step": 0.01}),
-        "required": {
+            "output_format": (["png", "webp", "jpeg"],),
-            "image": ("IMAGE",),
+            "api_key_override": ("STRING", {"multiline": False}),
-        },
+        }
-    }
+    }
-
+
-class StabilityInpainting(StabilityBase):
+
-    API_ENDPOINT = "stable-image/edit/inpaint"
+class StabilityRemoveBackground(StabilityBase):
-    ACCEPT = "image/*"
+    API_ENDPOINT = "stable-image/edit/remove-background"
-    INPUT_SPEC = {
+    ACCEPT = "image/*"
-        "required": {
+    INPUT_SPEC = {
-            "image": ("IMAGE",),
+        "required": {
-            "mask": ("MASK",),
+            "image": ("IMAGE",),
-            "prompt": ("STRING", {"multiline": True, "default": ""}),\
+        },
-        },
+    }
-        "optional": {
+
-            "negative_prompt": ("STRING", {"multiline": True, "default": ""}),\
+class StabilityInpainting(StabilityBase):
-            "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}),
+    API_ENDPOINT = "stable-image/edit/inpaint"
-            "output_format": (["png", "webp", "jpeg"],),
+    ACCEPT = "image/*"
-            "api_key_override": ("STRING", {"multiline": False}),
+    INPUT_SPEC = {
-        }
+        "required": {
-    }
+            "image": ("IMAGE",),
-    def _get_files(self, buffered, **kwargs):
+            "mask": ("MASK",),
-        mask = kwargs.get("mask")
+            "prompt": ("STRING", {"multiline": True, "default": ""}),\
-        to_pil = ToPILImage()
+        },
-        mask = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3)
+        "optional": {
-        mask = to_pil(mask.squeeze(0).permute(2,0,1))
+            "negative_prompt": ("STRING", {"multiline": True, "default": ""}),\
-        buffered_mask = BytesIO()
+            "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}),
-        mask.save(buffered_mask, format="PNG")
+            "output_format": (["png", "webp", "jpeg"],),
-        return {
+            "api_key_override": ("STRING", {"multiline": False}),
-            "image": buffered.getvalue(),
+        }
-            "mask": buffered_mask.getvalue(),
+    }
-        }
+    def _get_files(self, buffered, **kwargs):
-
+        mask = kwargs.get("mask")
-
+        to_pil = ToPILImage()
-class StabilitySearchAndReplace(StabilityBase):
+        mask = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3)
-    API_ENDPOINT = "stable-image/edit/search-and-replace"
+        mask = to_pil(mask.squeeze(0).permute(2,0,1))
-    ACCEPT = "image/*"
+        buffered_mask = BytesIO()
-    INPUT_SPEC = {
+        mask.save(buffered_mask, format="PNG")
-        "required": {
+        return {
-            "image": ("IMAGE",),
+            "image": buffered.getvalue(),
-            "search_prompt": ("STRING", {"multiline": True}, "Search Prompt"),
+            "mask": buffered_mask.getvalue(),
-            "prompt": ("STRING", {"multiline": True}),
+        }
-        },
+
-        "optional": {
+
-            "negative_prompt": ("STRING", {"multiline": True}),
+class StabilitySearchAndReplace(StabilityBase):
-            "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}),
+    API_ENDPOINT = "stable-image/edit/search-and-replace"
-            "api_key_override": ("STRING", {"multiline": False}),
+    ACCEPT = "image/*"
-            "output_format": (["png", "webp", "jpeg"],),
+    INPUT_SPEC = {
-        },
+        "required": {
-    }
+            "image": ("IMAGE",),
-
+            "search_prompt": ("STRING", {"multiline": True}, "Search Prompt"),
-
+            "prompt": ("STRING", {"multiline": True}),
-class StabilitySD3(StabilityBase):
+        },
-    API_ENDPOINT = "stable-image/generate/sd3"
+        "optional": {
-    ACCEPT = "image/*"
+            "negative_prompt": ("STRING", {"multiline": True}),
-    INPUT_SPEC = {
+            "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}),
-        "required": {
+            "api_key_override": ("STRING", {"multiline": False}),
-            "model": (["sd3", "sd3-turbo"],),
+            "output_format": (["png", "webp", "jpeg"],),
-            "prompt": ("STRING", {"multiline": True}),
+        },
-        },
+    }
-        "optional": {
+
-            "image": ("IMAGE",),
+
-            "negative_prompt": ("STRING", {"multiline": True}),
+class StabilitySD3(StabilityBase):
-            "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}),
+    API_ENDPOINT = "stable-image/generate/sd3"
-            "strength": ("FLOAT", {"default": 0.5, "min": 0.01, "max": 1.0, "step": 0.01}),
+    ACCEPT = "image/*"
-            "aspect_ratio": (["16:9", "1:1", "21:9", "2:3", "3:2", "4:5", "5:4", "9:16", "9:21"],),
+    INPUT_SPEC = {
-            "output_format": (["png", "jpeg"],),
+        "required": {
-            "api_key_override": ("STRING", {"multiline": False}),
+            "model": (["sd3", "sd3-turbo"],),
-        },
+            "prompt": ("STRING", {"multiline": True}),
-    }
+        },
-
+        "optional": {
-
+            "image": ("IMAGE",),
-class StabilityOutpainting(StabilityBase):
+            "negative_prompt": ("STRING", {"multiline": True}),
-    API_ENDPOINT = "stable-image/edit/outpaint"
+            "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}),
-    ACCEPT = "image/*"
+            "strength": ("FLOAT", {"default": 0.5, "min": 0.01, "max": 1.0, "step": 0.01}),
-    INPUT_SPEC = {
+            "aspect_ratio": (["16:9", "1:1", "21:9", "2:3", "3:2", "4:5", "5:4", "9:16", "9:21"],),
-        "required": {
+            "output_format": (["png", "jpeg"],),
-            "image": ("IMAGE",),
+            "api_key_override": ("STRING", {"multiline": False}),
-            "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}),\
+
-        },
+class StabilityOutpainting(StabilityBase):
-        "optional": {
+    API_ENDPOINT = "stable-image/edit/outpaint"
-            "prompt": ("STRING", {"multiline": True}),
+    ACCEPT = "image/*"
-            "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}),
+    INPUT_SPEC = {
-            "output_format": (["png", "webp", "jpeg"],),
+        "required": {
-            "api_key_override": ("STRING", {"multiline": False}),
+            "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)