| 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") |
|
|
| |
| 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() |
| |
| 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://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}), |
| } |
| } |
|
|
