nodes/florestefano1975-ComfyUI-StabilityAI-Suite
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Compare commits

merge into: nodes/florestefano1975-ComfyUI-StabilityAI-Suite:main
Branches Tags
nodes/florestefano1975-ComfyUI-StabilityAI-Suite:main nodes/florestefano1975-ComfyUI-StabilityAI-Suite:patch-1
..
pull from: nodes/florestefano1975-ComfyUI-StabilityAI-Suite:patch-1
Branches Tags
nodes/florestefano1975-ComfyUI-StabilityAI-Suite:main nodes/florestefano1975-ComfyUI-StabilityAI-Suite:patch-1
1 Commits
main .. patch-1
Author SHA1 Message Date
florestefano1975 bfc175b5da 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.
 2024-04-21 09:05:00 +02:00
12 changed files with 343 additions and 467 deletions
Expand all files Collapse all files
.github/workflows/publish.yml
-22
View File
@@ -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
View File
@@ -1,3 +1 @@
__pycache__/
sai_platform_key.txt
log.txt
__pycache__/
LICENSE.txt
+21 -21
View File
@@ -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
+35 -44
View File
@@ -1,44 +1,35 @@
# 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/)
# 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.
__init__.py
+11 -24
View File
@@ -1,24 +1,11 @@
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,
}
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,
}
api_cat_with_workflow.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 13 MiB

images/creative_upscale_recover_file.png
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 484 KiB

images/image_core_style.png
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 44 KiB

pyproject.toml
-14
View File
@@ -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 = ""
sai_api.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 2.4 MiB

sai_platform_key.txt
View File
stability_api.py
+275 -339
View File
@@ -1,339 +1,275 @@
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)
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)
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")
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}),
},
}