diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..723ef36 --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +.idea \ No newline at end of file diff --git a/__init__.py b/__init__.py index 0836072..440c9a6 100644 --- a/__init__.py +++ b/__init__.py @@ -1,3 +1,3 @@ -from .onethingai import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS +from .register import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS -__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] \ No newline at end of file +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] diff --git a/base.py b/base.py new file mode 100644 index 0000000..9b02102 --- /dev/null +++ b/base.py @@ -0,0 +1,202 @@ +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 + CATEGORY = "OneThingAI/image generation" + FUNCTION = "pre_input" + UA = "OneThingAI ComfyUI/1.0" + 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 INPUT_TYPES(cls): + cls.BASE_INPUT["required"]["model"] = (cls.SUPPORTED_MODELS, {}) + return { + "required": {**cls.BASE_INPUT.get("required", {}), **cls.EXTEND_INPUT.get("required", {})}, + "optional": {**cls.BASE_INPUT.get("optional", {}), **cls.EXTEND_INPUT.get("optional", {})} + } + + @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 = image[0] + else: + image_to_encode = image + + # Convert from torch tensor to PIL Image + image_to_encode = (image_to_encode.cpu().numpy() * 255).astype(np.uint8) + if len(image_to_encode.shape) == 3: + image_to_encode = Image.fromarray(image_to_encode) + + # write buffer + buffered = io.BytesIO() + image_to_encode.save(buffered, format="PNG") + return buffered + + @staticmethod + def read_image_result(response: requests.Response): + response.raise_for_status() + + # Parse response + result = response.json() + + if "data" not in result: + 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]: + 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, image_size, custom_width, custom_height, retries, timeout, + reference_image=None, extra=None, **kwargs): + client = self.new_http_client(retries) + if image_size == "自定义": + image_size = f"{custom_width}x{custom_height}" + + headers = { + "Authorization": f"Bearer {api_key}", + "User-Agent": self.UA, + } + + payload = { + "model": model, + "prompt": prompt, + "size": image_size + } + + extra_params = {} + + if extra: + try: + extra_params = json.loads(extra) + except json.decoder.JSONDecodeError: + raise ValueError("extra not a json string") + + return self.generate(client, headers, 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) + 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 + ) + 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) + files = {"image[]": 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) + payload["image"] = "data:image/png;base64," + base64.b64encode(img.getvalue()).decode() + 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() diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..3af3b81 --- /dev/null +++ b/nodes.py @@ -0,0 +1,39 @@ +from .base import OpenAICompatibleNode, VolcengineNode + + +class GeminiImage(OpenAICompatibleNode): + SUPPORTED_MODELS = ["gemini-2.5-flash-image"] + +class WanxiangImage(OpenAICompatibleNode): + SUPPORTED_MODELS = [ + "wanx2.1-t2i-plus", + "wanx2.1-t2i-turbo", + "wanx2.0-t2i-turbo", + "flux-dev", + ] + +class HunyuanImage(OpenAICompatibleNode): + SUPPORTED_MODELS = [ + "hunyuan-image" + ] + +class OpenAIImage(OpenAICompatibleNode): + SUPPORTED_MODELS = [ + "gpt-image-1" + ] + +class SeedreamImage(VolcengineNode): + SUPPORTED_MODELS = [ + "doubao-seedream-4-0-250828", + "doubao-seedream-3-0-t2i-250415", + "doubao-seededit-3-0-i2i-250628" + ] + EXTEND_INPUT = { + "required": { + "seed": ("INT", {"default": "-1", "min": -1, "max": 2147483647}), + "watermark": ("BOOLEAN", {"default": True}), + }, + "optional": { + "guidance_scale": ("FLOAT", {"min": 1, "max": 10}) + } + } \ No newline at end of file diff --git a/onethingai.py b/onethingai.py index 2f779d2..53ad539 100644 --- a/onethingai.py +++ b/onethingai.py @@ -53,7 +53,8 @@ class OneThingAILoader: # Request headers headers = { "Authorization": f"Bearer {api_key}", - "Content-Type": "application/json" + "Content-Type": "application/json", + "User-Agent": "OneThingAI ComfyUI/1.0" } # Determine image size based on selection @@ -143,13 +144,3 @@ class OneThingAILoader: raise RuntimeError(f"Error generating images: {str(e)}") finally: session.close() - -# Register the node class -NODE_CLASS_MAPPINGS = { - "OneThingAILoader": OneThingAILoader -} - -# Add descriptions -NODE_DISPLAY_NAME_MAPPINGS = { - "OneThingAILoader": "OneThingAI Image Generator" -} diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..622c92d --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,13 @@ +[project] +name = "comfyui_onething_image" +description = "A custom node for ComfyUI that integrates with OneThing AI's image generation API." +version = "1.0.0" + +[project.urls] +Repository = "https://github.com/OneThingAI/ComfyUI_Onething_Image" +# Used by Comfy Registry https://comfyregistry.org + +[tool.comfy] +PublisherId = "OneThingAI" +DisplayName = "ComfyUI_OneThingAI_ImageNode" +Icon = "" \ No newline at end of file diff --git a/register.py b/register.py new file mode 100644 index 0000000..51c6436 --- /dev/null +++ b/register.py @@ -0,0 +1,22 @@ +from .nodes import * +from .onethingai import * + +# Register the node class +NODE_CLASS_MAPPINGS = { + "OneThingAILoader": OneThingAILoader, + "OnethingAIImageGemini":GeminiImage, + "OnethingAIImageWanxiang":WanxiangImage, + "OnethingAIImageHunyuan":HunyuanImage, + "OnethingAIImageOpenAI":OpenAIImage, + "OnethingAIImageSeedream": SeedreamImage +} + +# Add descriptions +NODE_DISPLAY_NAME_MAPPINGS = { + "OneThingAILoader": "OneThingAI Image Generator", + "OnethingAIImageGemini": "OneThingAI Gemini Image", + "OnethingAIImageWanxiang": "OneThingAI Wanxiang Image", + "OnethingAIImageHunyuan": "OneThingAI Hunyuan Image", + "OnethingAIImageOpenAI": "OneThingAI OpenAI Image", + "OnethingAIImageSeedream": "OneThingAI Seedream Image" +}