Initial Version
Initial Upload
This commit is contained in:
TwoPerCent
@@ -0,0 +1,9 @@
|
||||
|
||||
|
||||
**Usage:**
|
||||
|
||||
Add API key to environment variable "SAI_API_KEY"
|
||||
|
||||
Alternatively you can write your API key to file "sai_platform_key.txt"
|
||||
|
||||
You can also use and/or override the above by entering your API key in the 'api_override' field of each node
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 13 MiB |
BIN
Binary file not shown.
|
After Width: | Height: | Size: 2.4 MiB |
@@ -0,0 +1,266 @@
|
||||
import requests
|
||||
from requests.models import PreparedRequest
|
||||
from PIL import Image
|
||||
import numpy as np
|
||||
import torch
|
||||
from torchvision.transforms import ToPILImage
|
||||
from io import BytesIO
|
||||
import os
|
||||
import time
|
||||
|
||||
API_KEY = os.environ.get("SAI_API_KEY")
|
||||
|
||||
# Check for API key in file as a backup, not recommended
|
||||
try:
|
||||
if not API_KEY:
|
||||
dir_path = os.path.dirname(os.path.realpath(__file__))
|
||||
with open(os.path.join(dir_path, "sai_platform_key.txt"), "r") as f:
|
||||
API_KEY = f.read().strip()
|
||||
print(f"API Key found in sai_platform_key.txt: {API_KEY}")
|
||||
# Validate the key is not empty
|
||||
if API_KEY.strip() == "":
|
||||
raise Exception(f"API Key is required to use the Stability API. \nPlease set the SAI_API_KEY environment variable to your API key or place in {dir_path}/sai_platform_key.txt.")
|
||||
|
||||
except Exception as e:
|
||||
print(f"\n\n***API Key is required to use the Stability API. Please set the SAI_API_KEY environment variable to your API key or place in {dir_path}/sai_platform_key.txt.***\n\n")
|
||||
|
||||
ROOT_API = "https://api.stability.ai/v2beta/"
|
||||
|
||||
class StabilityBase:
|
||||
API_ENDPOINT = ""
|
||||
POLL_ENDPOINT = ""
|
||||
ACCEPT = ""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return cls.INPUT_SPEC
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "call"
|
||||
CATEGORY = "Stability"
|
||||
|
||||
def call(self, *args, **kwargs):
|
||||
|
||||
buffered = BytesIO()
|
||||
files = {'none': None}
|
||||
data = None
|
||||
|
||||
image = kwargs.get('image', None)
|
||||
if image is not None:
|
||||
kwargs["mode"] = "image-to-image"
|
||||
kwargs.pop("aspect_ratio", None)
|
||||
image = ToPILImage()(image.squeeze(0).permute(2,0,1))
|
||||
image.save(buffered, format="PNG")
|
||||
files = self._get_files(buffered, **kwargs)
|
||||
else:
|
||||
kwargs.pop("strength", None)
|
||||
|
||||
|
||||
headers = {
|
||||
"Authorization": API_KEY,
|
||||
}
|
||||
|
||||
|
||||
if kwargs.get("api_key override"):
|
||||
headers = {
|
||||
"Authorization": kwargs.get("api_key override"),
|
||||
}
|
||||
|
||||
if headers.get("Authorization") is None:
|
||||
raise Exception(f"No Stability key set.\n\nUse your Stability AI API key by:\n1. Setting the SAI_API_KEY environment variable to your API key\n3. Placing inside sai_platform_key.txt\n4. Passing the API key as an argument to the function with the key 'api_key override'")
|
||||
|
||||
headers["Accept"] = self.ACCEPT
|
||||
|
||||
data = self._get_data(**kwargs)
|
||||
|
||||
req = PreparedRequest()
|
||||
req.prepare_method('POST')
|
||||
req.prepare_url(f"{ROOT_API}{self.API_ENDPOINT}", None)
|
||||
req.prepare_headers(headers)
|
||||
req.prepare_body(data=data, files=files)
|
||||
response = requests.Session().send(req)
|
||||
|
||||
if response.status_code == 200:
|
||||
if self.POLL_ENDPOINT != "":
|
||||
id = response.json().get("id")
|
||||
timeout = 240
|
||||
start_time = time.time()
|
||||
while True:
|
||||
response = requests.get(f"{ROOT_API}{self.POLL_ENDPOINT}{id}", headers=headers)
|
||||
if response.status_code == 200:
|
||||
if self.ACCEPT == "image/*":
|
||||
return self._return_image(response)
|
||||
if self.ACCEPT == "video/*":
|
||||
return self._return_video(response)
|
||||
break
|
||||
elif response.status_code == 202:
|
||||
time.sleep(10)
|
||||
elif time.time() - start_time > timeout:
|
||||
raise Exception("Stability API Timeout: Request took too long to complete")
|
||||
else:
|
||||
error_info = response.json()
|
||||
raise Exception(f"Stability API Error: {error_info}")
|
||||
else:
|
||||
result_image = Image.open(BytesIO(response.content))
|
||||
result_image = result_image.convert("RGBA")
|
||||
result_image = np.array(result_image).astype(np.float32) / 255.0
|
||||
result_image = torch.from_numpy(result_image)[None,]
|
||||
return (result_image,)
|
||||
else:
|
||||
error_info = response.json()
|
||||
if error_info.get("name") == "unauthorized":
|
||||
raise Exception(f"Stability API Error: Unauthorized.\n\nUse your Stability AI API key by:\n1. Setting the SAI_API_KEY environment variable to your API key\n3. Placing inside sai_platform_key.txt\n4. Passing the API key as an argument to the function with the key 'api_key override'")
|
||||
else:
|
||||
raise Exception(f"Stability API Error: {error_info}")
|
||||
|
||||
def _return_image(self, response):
|
||||
result_image = Image.open(BytesIO(response.content))
|
||||
result_image = result_image.convert("RGBA")
|
||||
result_image = np.array(result_image).astype(np.float32) / 255.0
|
||||
result_image = torch.from_numpy(result_image)[None,]
|
||||
return (result_image,)
|
||||
|
||||
def _return_video(self, response):
|
||||
result_video = response.content
|
||||
return (result_video,)
|
||||
|
||||
def _get_files(self, buffered, **kwargs):
|
||||
return {
|
||||
"image": buffered.getvalue()
|
||||
}
|
||||
|
||||
def _get_data(self, **kwargs):
|
||||
return {k: v for k, v in kwargs.items() if k != "image"}
|
||||
|
||||
class StabilityCore(StabilityBase):
|
||||
API_ENDPOINT = "stable-image/generate/core"
|
||||
ACCEPT = "image/*"
|
||||
INPUT_SPEC = {
|
||||
"required": {
|
||||
"prompt": ("STRING", {"multiline": True}),
|
||||
"negative_prompt": ("STRING", {"multiline": True}),
|
||||
"seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}),
|
||||
"output_format": (["png", "webp", "jpeg"],),
|
||||
"aspect_ratio": (["16:9", "1:1", "21:9", "2:3", "3:2", "4:5", "5:4", "9:16", "9:21"],),
|
||||
"api_key override": ("STRING", {"multiline": False}),
|
||||
},
|
||||
}
|
||||
|
||||
class StabilityCreativeUpscale(StabilityBase):
|
||||
API_ENDPOINT = "stable-image/upscale/creative"
|
||||
POLL_ENDPOINT = "stable-image/upscale/creative/result/"
|
||||
ACCEPT = "image/*"
|
||||
INPUT_SPEC = {
|
||||
"required": {
|
||||
"image": ("IMAGE",),
|
||||
"prompt": ("STRING", {"multiline": True}),
|
||||
},
|
||||
"optional": {
|
||||
"negative_prompt": ("STRING", {"multiline": True}),
|
||||
"seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}),
|
||||
"creativity": ("FLOAT", {"default": 0.3, "min": 0.01, "max": 0.35, "step": 0.01}),
|
||||
"output_format": (["png", "webp", "jpeg"],),
|
||||
"api_key override": ("STRING", {"multiline": False}),
|
||||
}
|
||||
}
|
||||
|
||||
class StabilityRemoveBackground(StabilityBase):
|
||||
API_ENDPOINT = "stable-image/edit/remove-background"
|
||||
ACCEPT = "image/*"
|
||||
INPUT_SPEC = {
|
||||
"required": {
|
||||
"image": ("IMAGE",),
|
||||
},
|
||||
}
|
||||
|
||||
class StabilityInpainting(StabilityBase):
|
||||
API_ENDPOINT = "stable-image/edit/inpaint"
|
||||
ACCEPT = "image/*"
|
||||
INPUT_SPEC = {
|
||||
"required": {
|
||||
"image": ("IMAGE",),
|
||||
"mask": ("MASK",),
|
||||
"prompt": ("STRING", {"multiline": True, "default": "a winter wonderland"}),
|
||||
"negative_prompt": ("STRING", {"multiline": True, "default": "a post-apocalyptic world"}),
|
||||
"seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}),
|
||||
"output_format": (["png", "webp", "jpeg"],),
|
||||
"api_key override": ("STRING", {"multiline": False}),
|
||||
},
|
||||
}
|
||||
def _get_files(self, buffered, **kwargs):
|
||||
mask = kwargs.get("mask")
|
||||
to_pil = ToPILImage()
|
||||
mask = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])).movedim(1, -1).expand(-1, -1, -1, 3)
|
||||
mask = to_pil(mask.squeeze(0).permute(2,0,1))
|
||||
buffered_mask = BytesIO()
|
||||
mask.save(buffered_mask, format="PNG")
|
||||
return {
|
||||
"image": buffered.getvalue(),
|
||||
"mask": buffered_mask.getvalue(),
|
||||
}
|
||||
|
||||
class StabilitySearchAndReplace(StabilityBase):
|
||||
API_ENDPOINT = "stable-image/edit/search-and-replace"
|
||||
ACCEPT = "image/*"
|
||||
INPUT_SPEC = {
|
||||
"required": {
|
||||
"image": ("IMAGE",),
|
||||
"search_prompt": ("STRING", {"multiline": True}, "Search Prompt"),
|
||||
"prompt": ("STRING", {"multiline": True}),
|
||||
"output_format": (["png", "webp", "jpeg"],),
|
||||
},
|
||||
"optional": {
|
||||
"seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}),
|
||||
"negative_prompt": ("STRING", {"multiline": True}),
|
||||
"api_key override": ("STRING", {"multiline": False}),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
class StabilitySD3(StabilityBase):
|
||||
API_ENDPOINT = "stable-image/generate/sd3"
|
||||
ACCEPT = "image/*"
|
||||
INPUT_SPEC = {
|
||||
"required": {
|
||||
"model": (["sd3", "sd3-turbo"],),
|
||||
"prompt": ("STRING", {"multiline": True}),
|
||||
"negative_prompt": ("STRING", {"multiline": True}),
|
||||
"seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}),
|
||||
"output_format": (["png", "jpeg"],),
|
||||
},
|
||||
"optional": {
|
||||
"image": ("IMAGE",),
|
||||
"strength": ("FLOAT", {"default": 0.5, "min": 0.01, "max": 1.0, "step": 0.01}),
|
||||
"aspect_ratio": (["16:9", "1:1", "21:9", "2:3", "3:2", "4:5", "5:4", "9:16", "9:21"],),
|
||||
"api_key override": ("STRING", {"multiline": False}),
|
||||
},
|
||||
}
|
||||
|
||||
class StabilityOutpainting(StabilityBase):
|
||||
API_ENDPOINT = "stable-image/edit/outpaint"
|
||||
ACCEPT = "image/*"
|
||||
INPUT_SPEC = {
|
||||
"required": {
|
||||
"image": ("IMAGE",),
|
||||
"left": ("INT", {"default": 0, "min": 0, "max": 512}),
|
||||
"right": ("INT", {"default": 0, "min": 0, "max": 512}),
|
||||
"up": ("INT", {"default": 0, "min": 0, "max": 512}),
|
||||
"down": ("INT", {"default": 0, "min": 0, "max": 512}),\
|
||||
},
|
||||
"optional": {
|
||||
"prompt": ("STRING", {"multiline": True}),
|
||||
"seed": ("INT", {"default": 0, "min": 0, "max": 4294967294}),
|
||||
"output_format": (["png", "webp", "jpeg"],),
|
||||
"api_key override": ("STRING", {"multiline": False}),
|
||||
},
|
||||
}
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"Stability Creative Upscale": StabilityCreativeUpscale,
|
||||
"Stability Remove Background": StabilityRemoveBackground,
|
||||
"Stability Inpainting": StabilityInpainting,
|
||||
"Stability Image Core": StabilityCore,
|
||||
"Stability Search and Replace": StabilitySearchAndReplace,
|
||||
"Stability Outpainting": StabilityOutpainting,
|
||||
"Stability SD3": StabilitySD3,
|
||||
}
|
||||
Reference in New Issue
Block a user