base.py

import base64
import io
import json
import random

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"]

        if "seed" in kwargs:
            if kwargs["seed"] == -1:
                kwargs["seed"] = random.randint(1, 2147483647)

        # 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()