Implementation of styles for Image Core

- On/off switch for applying a style.
- Style selector.
- Modification of the "StabilityBase" class to check for the presence of the style among the dispatch parameters with the API.

.github/workflows/publish.yml

@@ -1,22 +0,0 @@
-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,3 +1 @@
-__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,44 +1,35 @@
-# ComfyUI StabilityAI Suite
+# Stability API nodes for ComfyUI
-
+
-This fork of the official StabilityAI repository contains a number of enhancements and implementations.
+![alt text](api_cat_with_workflow.png)
-
+
-**_If this project is useful to you and you like it, please consider a small donation to the author._**
+### Usage:
-
+
-➡️ https://ko-fi.com/stefanoflore75
+Add API key to environment variable "`SAI_API_KEY`"
-
+
-Buy my workflows:
+Alternatively you can write your API key to file "`sai_platform_key.txt`"
-
+
-➡️ https://stefanoflore.it/download/
+You can also use and/or override the above by entering your API key in the '`api_key_override`' field of each node.
-
+
-## Image Core + Style Preset
+# License
-
+
-In this fork the node **Image Core** has been updated: it contains a new true/false switch to activate styles, and the corresponding selector.
+The MIT License (MIT)
-
+
-**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.**
+Copyright (c) 2024 Stability AI
-
+
-![Image Core + Style Preset](/images/image_core_style.png)
+Permission is hereby granted, free of charge, to any person obtaining a copy
-
+of this software and associated documentation files (the "Software"), to deal
-## Creative Upscale Recover File 
+in the Software without restriction, including without limitation the rights
-
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-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.
+copies of the Software, and to permit persons to whom the Software is
-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.
+furnished to do so, subject to the following conditions:
-
+
-Follow these steps:
+The above copyright notice and this permission notice shall be included in all
-
+copies or substantial portions of the Software.
-- Open the file and copy the desired `id`.
+
-- Paste the `id` into the __Creative Upscale Recover File__ node to recover the lost image.
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-![Creative Upscale Recover File](/images/creative_upscale_recover_file.png)
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-## Usage:
+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
-Add API key to environment variable "`SAI_API_KEY`".
+SOFTWARE.
-
-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,24 +1,11 @@
-from .stability_api import StabilityCreativeUpscale, StabilityCreativeUpscaleRecover, StabilityRemoveBackground, StabilityInpainting, StabilityCore, StabilitySearchAndReplace, StabilityOutpainting, StabilitySD3
+from .stability_api import StabilityCreativeUpscale, StabilityRemoveBackground, StabilityInpainting, StabilityCore, StabilitySearchAndReplace, StabilityOutpainting, StabilitySD3
-
+
-NODE_CLASS_MAPPINGS = {
+NODE_CLASS_MAPPINGS = {
-    "StabilityAI Suite - Creative Upscale": StabilityCreativeUpscale,
+    "Stability Creative Upscale": StabilityCreativeUpscale,
-    "StabilityAI Suite - Creative Upscale Recover File": StabilityCreativeUpscaleRecover,
+    "Stability Remove Background": StabilityRemoveBackground,
-    "StabilityAI Suite - Remove Background": StabilityRemoveBackground,
+    "Stability Inpainting": StabilityInpainting,
-    "StabilityAI Suite - Inpainting": StabilityInpainting,
+    "Stability Image Core": StabilityCore,
-    "StabilityAI Suite - Image Core + Style Preset": StabilityCore,
+    "Stability Search and Replace": StabilitySearchAndReplace,
-    "StabilityAI Suite - Search and Replace": StabilitySearchAndReplace,
+    "Stability Outpainting": StabilityOutpainting,
-    "StabilityAI Suite - Outpainting": StabilityOutpainting,
+    "Stability SD3": StabilitySD3,
-    "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

@@ -1,14 +0,0 @@
-[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,339 +1,275 @@
-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
+
-import base64
+API_KEY = os.environ.get("SAI_API_KEY")
-import io
+
-
+# Check for API key in file as a backup, not recommended
-API_KEY = os.environ.get("SAI_API_KEY")
+try:
-
+    if not API_KEY:
-# Check for API key in file as a backup, not recommended
+        dir_path = os.path.dirname(os.path.realpath(__file__))
-try:
+        with open(os.path.join(dir_path, "sai_platform_key.txt"), "r") as f:
-    if not API_KEY:
+            API_KEY = f.read().strip()
-        dir_path = os.path.dirname(os.path.realpath(__file__))
+            print(f"API Key found in sai_platform_key.txt: {API_KEY}")
-        with open(os.path.join(dir_path, "sai_platform_key.txt"), "r") as f:
+        # Validate the key is not empty
-            API_KEY = f.read().strip()
+        if API_KEY.strip() == "":
-            print(f"API Key found in sai_platform_key.txt: {API_KEY}")
+            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.")
-        # Validate the key is not empty
+        
-        if API_KEY.strip() == "":
+except Exception as e:
-            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"\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")
-        
+
-except Exception as e:
+ROOT_API = "https://api.stability.ai/v2beta/"
-    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 = ""
-class StabilityBase:
+    ACCEPT = ""
-    API_ENDPOINT = ""
+
-    POLL_ENDPOINT = ""
+    @classmethod
-    ACCEPT = ""
+    def INPUT_TYPES(cls):
-
+        return cls.INPUT_SPEC
-    @classmethod
+
-    def INPUT_TYPES(cls):
+    RETURN_TYPES = ("IMAGE",)
-        return cls.INPUT_SPEC
+    FUNCTION = "call"
-
+    CATEGORY = "Stability"
-    RETURN_TYPES = ("IMAGE",)
+
-    FUNCTION = "call"
+    def call(self, *args, **kwargs):
-    CATEGORY = "AI WizArt/Stability AI Suite"
+        
-
+        buffered = BytesIO()
-    def call(self, *args, **kwargs):
+        files = {'none': None}
-        
+        data = None
-        buffered = BytesIO()
+        
-        files = {'none': None}
+        image = kwargs.get('image', None)
-        data = None
+        if image is not None:
-        
+            kwargs["mode"] = "image-to-image"
-        image = kwargs.get('image', None)
+            kwargs.pop("aspect_ratio", None)
-        if image is not None:
+            image = ToPILImage()(image.squeeze(0).permute(2,0,1))
-            kwargs["mode"] = "image-to-image"
+            image.save(buffered, format="PNG")
-            kwargs.pop("aspect_ratio", None)
+            files = self._get_files(buffered, **kwargs)
-            image = ToPILImage()(image.squeeze(0).permute(2,0,1))
+        else:
-            image.save(buffered, format="PNG")
+            kwargs.pop("strength", None)
-            files = self._get_files(buffered, **kwargs)
+
-        else:
+        style = kwargs.get('style', False)
-            kwargs.pop("strength", None)
+        if style is False:
-        
+            kwargs.pop('style_preset', None)
-        style = kwargs.get('style', False)
+        
-        if style is False:
+        headers = {
-            kwargs.pop('style_preset', None)
+            "Authorization": API_KEY,
-
+        }
-        headers = {
+
-            "Authorization": API_KEY,
+        if kwargs.get("api_key_override"):
-        }
+            headers = {
-
+                "Authorization": kwargs.get("api_key_override"),
-        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'")
-
+
-        if headers.get("Authorization") is None:
+        headers["Accept"] = self.ACCEPT
-            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'")
+
-
+        data = self._get_data(**kwargs)
-        headers["Accept"] = self.ACCEPT
+
-
+        req = PreparedRequest()
-        data = self._get_data(**kwargs)
+        req.prepare_method('POST')
-
+        req.prepare_url(f"{ROOT_API}{self.API_ENDPOINT}", None)
-        req = PreparedRequest()
+        req.prepare_headers(headers)
-        req.prepare_method('POST')
+        req.prepare_body(data=data, files=files)
-        req.prepare_url(f"{ROOT_API}{self.API_ENDPOINT}", None)
+        response = requests.Session().send(req)
-        req.prepare_headers(headers)
+
-        req.prepare_body(data=data, files=files)
+        if response.status_code == 200:
-        response = requests.Session().send(req)
+            if self.POLL_ENDPOINT != "":
-
+                id = response.json().get("id")
-        if response.status_code == 200:
+                timeout = 240
-            if self.POLL_ENDPOINT != "":
+                start_time = time.time()
-                id = response.json().get("id")
+                while True:
-                logFile(f"Image/video ID for recovery: {id}") # saving id for recovery in case of malfunction
+                    response = requests.get(f"{ROOT_API}{self.POLL_ENDPOINT}{id}", headers=headers)
-                timeout = 240
+                    if response.status_code == 200:
-                start_time = time.time()
+                        if self.ACCEPT == "image/*":
-                while True:
+                            return self._return_image(response)
-                    response = requests.get(f"{ROOT_API}{self.POLL_ENDPOINT}{id}", headers=headers)
+                        if self.ACCEPT == "video/*":
-                    if response.status_code == 200:
+                            return self._return_video(response)
-                        if self.ACCEPT == "image/*":
+                        break
-                            return self._return_image(response)
+                    elif response.status_code == 202:
-                        if self.ACCEPT == "video/*":
+                        time.sleep(10)
-                            return self._return_video(response)
+                    elif time.time() - start_time > timeout:
-                        break
+                        raise Exception("Stability API Timeout: Request took too long to complete")
-                    elif response.status_code == 202:
+                    else:
-                        time.sleep(10)
+                        error_info = response.json()
-                    elif time.time() - start_time > timeout:
+                        raise Exception(f"Stability API Error: {error_info}")
-                        raise Exception("Stability API Timeout: Request took too long to complete")
+            else:
-                    else:
+                result_image = Image.open(BytesIO(response.content))
-                        error_info = response.json()
+                result_image = result_image.convert("RGBA")
-                        raise Exception(f"Stability API Error: {error_info}")
+                result_image = np.array(result_image).astype(np.float32) / 255.0
-            else:
+                result_image = torch.from_numpy(result_image)[None,]
-                result_image = Image.open(BytesIO(response.content))
+                return (result_image,)
-                result_image = result_image.convert("RGBA")
+        else:
-                result_image = np.array(result_image).astype(np.float32) / 255.0
+            error_info = response.json()
-                result_image = torch.from_numpy(result_image)[None,]
+            if error_info.get("name") == "unauthorized":
-                return (result_image,)
+                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'")
-        else:
+            if error_info.get("name") == "payment_required":
-            error_info = response.json()
+                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") == "unauthorized":
+            if error_info.get("name") == "bad_request":
-                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'")
+                errors = '\n'.join(error_info.get('errors'))
-            if error_info.get("name") == "payment_required":
+                raise Exception(f"Stability API Error: Bad request.\n\n{errors}")
-                raise Exception("Stability API Error: Not enough credits.\n\nPlease ensure your SAI API account has enough credits to complete this action.")
+            else:
-            if error_info.get("name") == "bad_request":
+                raise Exception(f"Stability API Error: {error_info}")
-                errors = '\n'.join(error_info.get('errors'))
+    
-                raise Exception(f"Stability API Error: Bad request.\n\n{errors}")
+    def _return_image(self, response):
-            else:
+        result_image = Image.open(BytesIO(response.content))
-                raise Exception(f"Stability API Error: {error_info}")
+        result_image = result_image.convert("RGBA")
-    
+        result_image = np.array(result_image).astype(np.float32) / 255.0
-    def _return_image(self, response):
+        result_image = torch.from_numpy(result_image)[None,]
-        result_image = Image.open(BytesIO(response.content))
+        return (result_image,)
-        result_image = result_image.convert("RGBA")
+
-        result_image = np.array(result_image).astype(np.float32) / 255.0
+    def _return_video(self, response):
-        result_image = torch.from_numpy(result_image)[None,]
+        result_video = response.content
-        return (result_image,)
+        return (result_video,)
-
+
-    def _return_video(self, response):
+    def _get_files(self, buffered, **kwargs):
-        result_video = response.content
+        return {
-        return (result_video,)
+            "image": buffered.getvalue()
-
+        }
-    def _get_files(self, buffered, **kwargs):
+
-        return {
+    def _get_data(self, **kwargs):
-            "image": buffered.getvalue()
+        return {k: v for k, v in kwargs.items() if k != "image"}
-        }
+
-
+
-    def _get_data(self, **kwargs):
+class StabilityCore(StabilityBase):
-        return {k: v for k, v in kwargs.items() if k != "image"}
+    API_ENDPOINT = "stable-image/generate/core"
-
+    ACCEPT = "image/*"
-
+    INPUT_SPEC = {
-class StabilityCore(StabilityBase):
+        "required": {
-    API_ENDPOINT = "stable-image/generate/core"
+            "prompt": ("STRING", {"multiline": True}),
-    ACCEPT = "image/*"
+        },
-    INPUT_SPEC = {
+        "optional": {
-        "required": {
+            "negative_prompt": ("STRING", {"multiline": True}),
-            "prompt": ("STRING", {"multiline": True}),
+            "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}),
-        },
+            "output_format": (["png", "webp", "jpeg"],),
-        "optional": {
+            "aspect_ratio": (["16:9", "1:1", "21:9", "2:3", "3:2", "4:5", "5:4", "9:16", "9:21"],),
-            "negative_prompt": ("STRING", {"multiline": True}),
+            "style": ("BOOLEAN", {"default": False}),
-            "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}),
+            "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"],),
-            "output_format": (["png", "webp", "jpeg"],),
+            "api_key_override": ("STRING", {"multiline": False}),
-            "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/*"
-class StabilityCreativeUpscale(StabilityBase):
+    INPUT_SPEC = {
-    API_ENDPOINT = "stable-image/upscale/creative"
+        "required": {
-    POLL_ENDPOINT  = "stable-image/upscale/creative/result/"
+            "image": ("IMAGE",),
-    ACCEPT = "image/*"
+            "prompt": ("STRING", {"multiline": True}),
-    INPUT_SPEC = {
+        },
-        "required": {
+        "optional": {
-            "image": ("IMAGE",),
+            "negative_prompt": ("STRING", {"multiline": True}),
-            "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}),
-        "optional": {
+            "output_format": (["png", "webp", "jpeg"],),
-            "negative_prompt": ("STRING", {"multiline": True}),
+            "api_key_override": ("STRING", {"multiline": False}),
-            "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 = {
-class StabilityRemoveBackground(StabilityBase):
+        "required": {
-    API_ENDPOINT = "stable-image/edit/remove-background"
+            "image": ("IMAGE",),
-    ACCEPT = "image/*"
+        },
-    INPUT_SPEC = {
+    }
-        "required": {
+
-            "image": ("IMAGE",),
+class StabilityInpainting(StabilityBase):
-        },
+    API_ENDPOINT = "stable-image/edit/inpaint"
-    }
+    ACCEPT = "image/*"
-
+    INPUT_SPEC = {
-class StabilityInpainting(StabilityBase):
+        "required": {
-    API_ENDPOINT = "stable-image/edit/inpaint"
+            "image": ("IMAGE",),
-    ACCEPT = "image/*"
+            "mask": ("MASK",),
-    INPUT_SPEC = {
+            "prompt": ("STRING", {"multiline": True, "default": ""}),\
-        "required": {
+        },
-            "image": ("IMAGE",),
+        "optional": {
-            "mask": ("MASK",),
+            "negative_prompt": ("STRING", {"multiline": True, "default": ""}),\
-            "prompt": ("STRING", {"multiline": True, "default": ""}),\
+            "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}),
-        },
+            "output_format": (["png", "webp", "jpeg"],),
-        "optional": {
+            "api_key_override": ("STRING", {"multiline": False}),
-            "negative_prompt": ("STRING", {"multiline": True, "default": ""}),\
+        }
-            "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}),
+    }
-            "output_format": (["png", "webp", "jpeg"],),
+    def _get_files(self, buffered, **kwargs):
-            "api_key_override": ("STRING", {"multiline": False}),
+        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)
-    def _get_files(self, buffered, **kwargs):
+        mask = to_pil(mask.squeeze(0).permute(2,0,1))
-        mask = kwargs.get("mask")
+        buffered_mask = BytesIO()
-        to_pil = ToPILImage()
+        mask.save(buffered_mask, format="PNG")
-        mask = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3)
+        return {
-        mask = to_pil(mask.squeeze(0).permute(2,0,1))
+            "image": buffered.getvalue(),
-        buffered_mask = BytesIO()
+            "mask": buffered_mask.getvalue(),
-        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 = {
-class StabilitySearchAndReplace(StabilityBase):
+        "required": {
-    API_ENDPOINT = "stable-image/edit/search-and-replace"
+            "image": ("IMAGE",),
-    ACCEPT = "image/*"
+            "search_prompt": ("STRING", {"multiline": True}, "Search Prompt"),
-    INPUT_SPEC = {
+            "prompt": ("STRING", {"multiline": True}),
-        "required": {
+        },
-            "image": ("IMAGE",),
+        "optional": {
-            "search_prompt": ("STRING", {"multiline": True}, "Search Prompt"),
+            "negative_prompt": ("STRING", {"multiline": True}),
-            "prompt": ("STRING", {"multiline": True}),
+            "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}),
-        },
+            "api_key_override": ("STRING", {"multiline": False}),
-        "optional": {
+            "output_format": (["png", "webp", "jpeg"],),
-            "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 = {
-class StabilitySD3(StabilityBase):
+        "required": {
-    API_ENDPOINT = "stable-image/generate/sd3"
+            "model": (["sd3", "sd3-turbo"],),
-    ACCEPT = "image/*"
+            "prompt": ("STRING", {"multiline": True}),
-    INPUT_SPEC = {
+        },
-        "required": {
+        "optional": {
-            "model": (["sd3", "sd3-turbo"],),
+            "image": ("IMAGE",),
-            "prompt": ("STRING", {"multiline": True}),
+            "negative_prompt": ("STRING", {"multiline": True}),
-        },
+            "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}),
-        "optional": {
+            "strength": ("FLOAT", {"default": 0.5, "min": 0.01, "max": 1.0, "step": 0.01}),
-            "image": ("IMAGE",),
+            "aspect_ratio": (["16:9", "1:1", "21:9", "2:3", "3:2", "4:5", "5:4", "9:16", "9:21"],),
-            "negative_prompt": ("STRING", {"multiline": True}),
+            "output_format": (["png", "jpeg"],),
-            "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}),
+            "api_key_override": ("STRING", {"multiline": False}),
-            "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 = {
-class StabilityOutpainting(StabilityBase):
+        "required": {
-    API_ENDPOINT = "stable-image/edit/outpaint"
+            "image": ("IMAGE",),
-    ACCEPT = "image/*"
+            "left": ("INT", {"default": 0, "min": 0, "max": 512}),
-    INPUT_SPEC = {
+            "right": ("INT", {"default": 0, "min": 0, "max": 512}),
-        "required": {
+            "up": ("INT", {"default": 0, "min": 0, "max": 512}),
-            "image": ("IMAGE",),
+            "down": ("INT", {"default": 0, "min": 0, "max": 512}),\
-            "left": ("INT", {"default": 0, "min": 0, "max": 512}),
+        },
-            "right": ("INT", {"default": 0, "min": 0, "max": 512}),
+        "optional": {
-            "up": ("INT", {"default": 0, "min": 0, "max": 512}),
+            "prompt": ("STRING", {"multiline": True}),
-            "down": ("INT", {"default": 0, "min": 0, "max": 512}),
+            "seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}),
-        },
+            "output_format": (["png", "webp", "jpeg"],),
-        "optional": {
+            "api_key_override": ("STRING", {"multiline": False}),
-            "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)