339 lines
13 KiB
Python
339 lines
13 KiB
Python
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) |