Song367-ComfyUI-KieAPI/lib/client.py

from __future__ import annotations

import asyncio
import sys
import time
import uuid
from pathlib import Path

import aiohttp


PLUGIN_ROOT = Path(__file__).resolve().parents[1]
if str(PLUGIN_ROOT) not in sys.path:
    sys.path.insert(0, str(PLUGIN_ROOT))

from lib.models import (  # noqa: E402
    KieBase64UploadRequest,
    KieBase64UploadResponse,
    KieCreateTaskRequest,
    KieCreateTaskResponse,
    KieNanoBanana2Input,
    KieTaskRecordData,
    KieTaskRecordResponse,
)


class KieClient:
    def __init__(
        self,
        api_key: str,
        *,
        api_base_url: str = "https://api.kie.ai",
        upload_base_url: str = "https://kieai.redpandaai.co",
    ):
        self.api_key = api_key
        self.api_base_url = api_base_url.rstrip("/")
        self.upload_base_url = upload_base_url.rstrip("/")

    def build_headers(self) -> dict[str, str]:
        return {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
        }

    def _now(self) -> float:
        return time.monotonic()

    async def _sleep(self, seconds: float) -> None:
        await asyncio.sleep(seconds)

    def _poll_interval(self, elapsed_seconds: float) -> float:
        if elapsed_seconds < 30:
            return 2.0
        if elapsed_seconds < 120:
            return 5.0
        return 15.0

    def build_reference_file_name(self, index: int) -> str:
        return f"reference_{index}_{uuid.uuid4().hex}.png"

    async def _request_json(
        self,
        method: str,
        url: str,
        *,
        params: dict | None = None,
        json_body: dict | None = None,
    ) -> dict:
        timeout = aiohttp.ClientTimeout(total=120)
        async with aiohttp.ClientSession(timeout=timeout) as session:
            async with session.request(
                method,
                url,
                params=params,
                json=json_body,
                headers=self.build_headers(),
            ) as response:
                response.raise_for_status()
                return await response.json()

    async def _download_bytes(self, url: str) -> bytes:
        timeout = aiohttp.ClientTimeout(total=120)
        async with aiohttp.ClientSession(timeout=timeout) as session:
            async with session.get(url) as response:
                response.raise_for_status()
                return await response.read()

    async def upload_reference_image(self, base64_data: str, *, file_name: str = "reference.png") -> str:
        request = KieBase64UploadRequest(
            base64Data=base64_data,
            uploadPath="images/base64",
            fileName=file_name,
        )
        response_json = await self._request_json(
            "POST",
            f"{self.upload_base_url}/api/file-base64-upload",
            json_body=request.model_dump(),
        )
        response = KieBase64UploadResponse.model_validate(response_json)
        if not response.data.downloadUrl:
            raise RuntimeError("Kie upload succeeded but did not return a downloadUrl.")
        return response.data.downloadUrl

    async def create_nano_banana_2_task(
        self,
        *,
        prompt: str,
        aspect_ratio: str,
        resolution: str,
        output_format: str,
        google_search: bool,
        image_input: list[str] | None = None,
    ) -> str:
        request = KieCreateTaskRequest(
            input=KieNanoBanana2Input(
                prompt=prompt,
                aspect_ratio=aspect_ratio,
                resolution=resolution,
                output_format=output_format,
                google_search=google_search,
                image_input=image_input or [],
            )
        )
        response_json = await self._request_json(
            "POST",
            f"{self.api_base_url}/api/v1/jobs/createTask",
            json_body=request.model_dump(exclude_none=True),
        )
        response = KieCreateTaskResponse.model_validate(response_json)
        if not response.data.taskId:
            raise RuntimeError("Kie createTask response did not include a taskId.")
        return response.data.taskId

    async def get_task_details(self, task_id: str) -> KieTaskRecordData:
        response_json = await self._request_json(
            "GET",
            f"{self.api_base_url}/api/v1/jobs/recordInfo",
            params={"taskId": task_id},
        )
        response = KieTaskRecordResponse.model_validate(response_json)
        return response.data

    async def wait_for_task(self, task_id: str, *, timeout_seconds: int = 600) -> KieTaskRecordData:
        started_at = self._now()
        while True:
            task = await self.get_task_details(task_id)
            state = task.state.lower().strip()

            if state == "success":
                return task
            if state == "fail":
                reason = task.failure_reason() or "Task failed without a detailed error message."
                raise RuntimeError(f"Kie task {task_id} failed: {reason}")

            elapsed = self._now() - started_at
            if elapsed >= timeout_seconds:
                raise TimeoutError(f"Kie task {task_id} did not complete within {timeout_seconds} seconds.")

            await self._sleep(self._poll_interval(elapsed))

    async def download_result_images(self, task: KieTaskRecordData) -> list[bytes]:
        urls = task.result_urls()
        if not urls:
            raise RuntimeError(f"Kie task {task.taskId} completed without result URLs.")
        return [await self._download_bytes(url) for url in urls]

    async def generate_images(
        self,
        *,
        prompt: str,
        aspect_ratio: str,
        resolution: str,
        output_format: str,
        google_search: bool,
        poll_timeout_seconds: int,
        reference_image_data_uris: list[str] | None = None,
    ) -> tuple[list[bytes], str]:
        image_input: list[str] = []
        for index, reference_image_data_uri in enumerate(reference_image_data_uris or []):
            image_input.append(
                await self.upload_reference_image(
                    reference_image_data_uri,
                    file_name=self.build_reference_file_name(index),
                )
            )

        task_id = await self.create_nano_banana_2_task(
            prompt=prompt,
            aspect_ratio=aspect_ratio,
            resolution=resolution,
            output_format=output_format,
            google_search=google_search,
            image_input=image_input,
        )
        task = await self.wait_for_task(task_id, timeout_seconds=poll_timeout_seconds)
        return await self.download_result_images(task), task_id