249 lines
9.5 KiB
Python
249 lines
9.5 KiB
Python
import base64
|
|
import io
|
|
import json
|
|
import requests
|
|
from requests.adapters import HTTPAdapter
|
|
from urllib3.util.retry import Retry
|
|
from PIL import Image
|
|
import numpy as np
|
|
import torch
|
|
|
|
|
|
# Base ComfyUI defined nodes
|
|
class BaseImageNode:
|
|
BASE_INPUT = {
|
|
"required": {
|
|
"api_key": ("STRING", {"default": "", "multiline": False}),
|
|
"model": ([], {}),
|
|
"prompt": ("STRING", {"default": "", "multiline": True, "placeholder": "提示词"}),
|
|
"image_size": (["1536x1024", "1024x1024", "1024x1536", "自定义"], {"default": "1024x1024"}),
|
|
"custom_width": ("INT", {"default": 1024, "min": 512, "max": 4096, "step": 64}),
|
|
"custom_height": ("INT", {"default": 1024, "min": 512, "max": 4096, "step": 64}),
|
|
"retries": ("INT", {"default": 3, "min": 1, "max": 10}),
|
|
"timeout": ("INT", {"default": 120, "min": 5, "max": 180}),
|
|
},
|
|
"optional": {
|
|
"reference_image": ("IMAGE",),
|
|
"extra": ("STRING",)
|
|
}
|
|
}
|
|
EXTEND_INPUT = {} # replace it with your own input to extend
|
|
SUPPORTED_MODELS = [] # replace it with your own supported moldes
|
|
DELETED_INPUT = [] # replace it with your own deleted input
|
|
CATEGORY = "OneThingAI/image generation"
|
|
FUNCTION = "pre_input"
|
|
UA = "OneThingAI ComfyUI/1.1"
|
|
URL_GENERATIONS = "https://api-model.onethingai.com/v1/images/generations"
|
|
URL_EDIT = "https://api-model.onethingai.com/v1/images/edits"
|
|
RETURN_TYPES = ("IMAGE",)
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, **kwargs):
|
|
return float("NaN")
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
cls.BASE_INPUT["required"]["model"] = (cls.SUPPORTED_MODELS, {})
|
|
required = {**cls.BASE_INPUT.get("required", {}), **cls.EXTEND_INPUT.get("required", {})}
|
|
optional = {**cls.BASE_INPUT.get("optional", {}), **cls.EXTEND_INPUT.get("optional", {})}
|
|
return {
|
|
"required": {k: v for k, v in required.items() if k not in cls.DELETED_INPUT},
|
|
"optional": {k: v for k, v in optional.items() if k not in cls.DELETED_INPUT}
|
|
}
|
|
|
|
@staticmethod
|
|
def new_http_client(retries) -> requests.Session:
|
|
retry_strategy = Retry(
|
|
total=retries,
|
|
backoff_factor=1,
|
|
status_forcelist=[429, 500, 502, 503, 504],
|
|
)
|
|
session = requests.Session()
|
|
adapter = HTTPAdapter(max_retries=retry_strategy)
|
|
session.mount("https://", adapter)
|
|
return session
|
|
|
|
@staticmethod
|
|
def read_reference_image(image):
|
|
# Convert the first image if it's a batch
|
|
if len(image.shape) == 4:
|
|
image_to_encode = [img for img in image]
|
|
else:
|
|
image_to_encode = [image]
|
|
|
|
buffered = []
|
|
for img in image_to_encode:
|
|
# Convert from torch tensor to PIL Image
|
|
img = (img.cpu().numpy() * 255).astype(np.uint8)
|
|
if len(img.shape) == 3:
|
|
img = Image.fromarray(img)
|
|
|
|
# write buffer
|
|
buf = io.BytesIO()
|
|
img.save(buf, format="PNG")
|
|
buf.seek(0)
|
|
buffered.append(buf)
|
|
print("buffered:", buffered)
|
|
return buffered
|
|
|
|
@staticmethod
|
|
def read_image_result(response: requests.Response):
|
|
response.raise_for_status()
|
|
|
|
# Parse response
|
|
result = response.json()
|
|
|
|
if "data" not in result or result["data"] is None:
|
|
raise ValueError(f"Unexpected API response: {result}")
|
|
|
|
# Process images
|
|
if len(result["data"]) == 0:
|
|
raise ValueError("no image result")
|
|
if "b64_json" not in result["data"][0] or not result["data"][0]["b64_json"]:
|
|
raise ValueError("no b64 json in result")
|
|
image_bytes = base64.b64decode(result["data"][0]["b64_json"])
|
|
image = Image.open(io.BytesIO(image_bytes))
|
|
|
|
# Convert to RGB if necessary
|
|
if image.mode != "RGB":
|
|
image = image.convert("RGB")
|
|
final_image = np.array(image)
|
|
final_image = torch.from_numpy(final_image).float() / 255.0
|
|
final_image = final_image.unsqueeze(0)
|
|
return (final_image,)
|
|
|
|
def pre_input(self, api_key, model, prompt, retries, timeout, reference_image=None, extra=None, **kwargs):
|
|
client = self.new_http_client(retries)
|
|
|
|
headers = {
|
|
"Authorization": f"Bearer {api_key}",
|
|
"User-Agent": self.UA,
|
|
}
|
|
|
|
payload = {
|
|
"model": model,
|
|
"prompt": prompt,
|
|
}
|
|
|
|
if "image_size" in kwargs:
|
|
if kwargs["image_size"] == "自定义":
|
|
payload["size"] = f"{kwargs['custom_width']}x{kwargs['custom_height']}"
|
|
else:
|
|
payload["size"] = kwargs["image_size"]
|
|
|
|
# filter
|
|
input_types = self.INPUT_TYPES()
|
|
for key, settings in input_types["required"].items():
|
|
if len(settings) > 1 and "enabled" in settings[1] and model not in settings[1]["enabled"]:
|
|
if key in kwargs:
|
|
del kwargs[key]
|
|
# print(f"filtered input: {key} at models: {model}")
|
|
for key, settings in input_types["optional"].items():
|
|
if len(settings) > 1 and "enabled" in settings[1] and model not in settings[1]["enabled"]:
|
|
if key == "reference_image":
|
|
reference_image = None
|
|
# print("reference_image off")
|
|
continue
|
|
if key in kwargs:
|
|
del kwargs[key]
|
|
# print(f"filtered input: {key} at models: {model}")
|
|
|
|
extra_params = {}
|
|
|
|
if extra:
|
|
try:
|
|
extra_params = json.loads(extra)
|
|
except json.decoder.JSONDecodeError:
|
|
raise ValueError("extra not a json string")
|
|
|
|
payload, timeout, extra_params, reference_image, kwargs = self.post_input(payload, timeout, extra_params,
|
|
reference_image, **kwargs)
|
|
|
|
return self.generate(client, headers, payload, timeout, extra_params, reference_image, **kwargs)
|
|
|
|
def post_input(self, payload: dict, timeout: int, extra_params: dict, reference_image=None, **kwargs) -> tuple[
|
|
dict, int, dict, list, dict]:
|
|
return payload, timeout, extra_params, reference_image, kwargs
|
|
|
|
def generate(self, client: requests.Session, headers: dict[str, str], payload: dict, timeout: int,
|
|
extra_params: dict, reference_image=None, **kwargs):
|
|
# Implement it
|
|
raise NotImplementedError()
|
|
|
|
|
|
class OpenAICompatibleNode(BaseImageNode):
|
|
|
|
def generate(self, client, headers, payload, timeout, extra_params, reference_image=None, **kwargs):
|
|
payload.update(kwargs)
|
|
payload.update(extra_params)
|
|
payload["n"] = 1
|
|
try:
|
|
if reference_image is None:
|
|
# generate mode, use json api
|
|
headers["Content-Type"] = "application/json"
|
|
response = client.post(
|
|
self.URL_GENERATIONS,
|
|
headers=headers,
|
|
json=payload,
|
|
timeout=timeout,
|
|
verify=False
|
|
)
|
|
return self.read_image_result(response)
|
|
|
|
else:
|
|
# edit mode, use form api
|
|
# headers["Content-Type"] = "multipart/form-data"
|
|
img = self.read_reference_image(reference_image)
|
|
if len(img) == 1:
|
|
files = {"image": img[0]}
|
|
else:
|
|
files = [("image[]", i) for i in img]
|
|
response = client.post(
|
|
self.URL_EDIT,
|
|
headers=headers,
|
|
data=payload,
|
|
files=files,
|
|
timeout=timeout,
|
|
verify=False
|
|
)
|
|
return self.read_image_result(response)
|
|
except requests.exceptions.RequestException as e:
|
|
raise RuntimeError(f"API request failed: {str(e)}")
|
|
except json.JSONDecodeError:
|
|
raise RuntimeError("Failed to parse API response")
|
|
except Exception as e:
|
|
raise RuntimeError(f"Error generating images: {str(e)}")
|
|
finally:
|
|
client.close()
|
|
|
|
|
|
class VolcengineNode(BaseImageNode):
|
|
|
|
def generate(self, client, headers, payload, timeout, extra_params, reference_image=None, **kwargs):
|
|
payload.update(kwargs)
|
|
payload.update(extra_params)
|
|
headers["Content-Type"] = "application/json"
|
|
if reference_image is not None:
|
|
img = self.read_reference_image(reference_image)
|
|
if len(img) == 1:
|
|
payload["image"] = "data:image/png;base64," + base64.b64encode(img[0].getvalue()).decode()
|
|
else:
|
|
payload["image"] = ["data:image/png;base64,"+base64.b64encode(b64_img.getvalue()).decode() for b64_img in img]
|
|
try:
|
|
response = client.post(
|
|
self.URL_GENERATIONS,
|
|
headers=headers,
|
|
json=payload,
|
|
timeout=timeout,
|
|
verify=False
|
|
)
|
|
return self.read_image_result(response)
|
|
except requests.exceptions.RequestException as e:
|
|
raise RuntimeError(f"API request failed: {str(e)}")
|
|
except json.JSONDecodeError:
|
|
raise RuntimeError("Failed to parse API response")
|
|
except Exception as e:
|
|
raise RuntimeError(f"Error generating images: {str(e)}")
|
|
finally:
|
|
client.close()
|