File size: 6,825 Bytes
7906fb3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 |
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("CAI_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, "cai_platform_key.txt"), "r") as f:
API_KEY = f.read().strip()
# Validate the key is not empty
if API_KEY.strip() == "":
raise Exception(f"API Key is required to use Clarity AI. \nPlease set the CAI_API_KEY environment variable to your API key or place in {dir_path}/cai_platform_key.txt.")
except Exception as e:
print(f"\n\n***API Key is required to use Clarity AI. Please set the CAI_API_KEY environment variable to your API key or place in {dir_path}/cai_platform_key.txt.***\n\n")
#ROOT_API = "https://api.clarityai.cc/v1/upscale"
ROOT_API = "https://v1-upscale-endpoint-oak26mtdga-ey.a.run.app"
class ClarityBase:
API_ENDPOINT = ""
POLL_ENDPOINT = ""
ACCEPT = ""
@classmethod
def INPUT_TYPES(cls):
return cls.INPUT_SPEC
RETURN_TYPES = ("IMAGE",)
FUNCTION = "call"
CATEGORY = "Clarity AI"
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)
kwargs['comfyui'] = True
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 Clarity AI key set.\n\nUse your Clarity AI API key by:\n1. Setting the CAI_API_KEY environment variable to your API key\n3. Placing inside cai_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 = 550
start_time = time.time()
while True:
response = requests.get(f"{ROOT_API}{self.POLL_ENDPOINT}{id}", headers=headers, timeout=timeout)
if response.status_code == 200:
print("took time: ", time.time() - start_time)
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("Clarity AI API Timeout: Request took too long to complete")
else:
error_info = response.json()
raise Exception(f"Clarity AI 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:
print("Fehler!! Status Code:", response.status_code)
error_info = response.text
print("error_info: " + error_info)
if response.status_code == 401:
raise Exception("Clarity AI API Error: Unauthorized.\n\nUse your Clarity AI API key by:\n1. Setting the CAI_API_KEY environment variable to your API key\n3. Placing inside cai_platform_key.txt\n4. Passing the API key as an argument to the function with the key 'api_key_override' \n\n \n\n")
if response.status_code == 402:
raise Exception("Clarity AI API Error: Not enough credits.\n\nPlease ensure your Clarity AI API account has enough credits to complete this action. \n\n \n\n")
if response.status_code == 400:
raise Exception(f"Clarity AI API Error: Bad request.\n\n{error_info} \n\n \n\n")
else:
raise Exception(f"Clarity AI 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 ClarityAIUpscaler(ClarityBase):
API_ENDPOINT = ""
POLL_ENDPOINT = ""
ACCEPT = "image/*"
INPUT_SPEC = {
"required": {
"image": ("IMAGE",),
},
"optional": {
"prompt": ("STRING", {"multiline": True}),
"creativity": ("FLOAT", {"default": 0, "min": -10, "max": 10, "step": 1}),
"resemblance": ("FLOAT", {"default": 0, "min": -10, "max": 10, "step": 1}),
"dynamic": ("FLOAT", {"default": 0, "min": -10, "max": 10, "step": 1}),
"fractality": ("FLOAT", {"default": 0, "min": -10, "max": 10, "step": 1}),
"style": (["default", "portrait", "anime"],),
"scale_factor": (["2", "4", "6", "8", "10", "12", "14", "16"],),
"api_key_override": ("STRING", {"multiline": False}),
}
}
|