nodes/gateway-ComfyUI-Kie-API
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

wip: nanobanana refactor checkpoint

Browse Source
This commit is contained in:
Steve Tietze
2025-12-18 10:43:39 +07:00
parent af52345b90
commit 5ee8e6136e
2 changed files with 201 additions and 104 deletions
Download Patch File Download Diff File Expand all files Collapse all files
kie_api/nanobanana.py
+186 -4
View File
@@ -1,3 +1,15 @@
"""
Nano Banana Pro API implementation.
This module contains all logic specific to the KIE Nano Banana Pro model:
- task creation
- polling for completion
- result URL extraction
- image download and decoding
This file is intentionally model-specific and not generic.
"""
import json
import time
from io import BytesIO
@@ -6,7 +18,10 @@ from typing import Any
import torch
from PIL import Image
from .auth import _load_api_key
from .credits import _log_remaining_credits
from .http import TransientKieError, requests
from .upload import _image_tensor_to_png_bytes, _truncate_url, _upload_image
CREATE_TASK_URL = "https://api.kie.ai/api/v1/jobs/createTask"
@@ -16,6 +31,26 @@ ASPECT_RATIO_OPTIONS = ["1:1", "2:3", "3:2", "3:4", "4:3", "4:5", "5:4", "9:16",
RESOLUTION_OPTIONS = ["1K", "2K", "4K"]
OUTPUT_FORMAT_OPTIONS = ["png", "jpg"]
# -----------------------------
# Constants / enums
# -----------------------------
# -----------------------------
# Validation
# -----------------------------
# -----------------------------
# Request building (Nano Banana specific)
# -----------------------------
# -----------------------------
# Task lifecycle (submit + poll)
# -----------------------------
# -----------------------------
# Result handling (parse + download + decode)
# -----------------------------
def _log(enabled: bool, msg: str) -> None:
"""Log a message to stdout when verbose logging is enabled.
@@ -159,11 +194,11 @@ def _poll_task_until_complete(
timeout_s: int,
log: bool,
start_time: float,
) -> tuple[dict[str, Any], str, Any, str | None]:
) -> dict[str, Any]:
"""Poll recordInfo until the task completes, fails, or times out.
Returns:
A tuple of (data_dict, raw_response_text, message_field, elapsed_text).
The task record data dict returned by recordInfo.
Raises:
RuntimeError: If the task times out or returns a non-retryable failure.
TransientKieError: If the task fails with a retryable condition.
@@ -198,7 +233,9 @@ def _poll_task_until_complete(
last_state = state
if state == "success":
return data, raw_json, message_field, f"{elapsed:.1f}s"
if log:
_log(log, f"Task {task_id} completed (elapsed={elapsed:.1f}s)")
return data
if state == "fail":
fail_code = data.get("failCode")
fail_msg = data.get("failMsg") or data.get("msg")
@@ -221,8 +258,31 @@ def _poll_task_until_complete(
time.sleep(interval)
def _create_nanobanana_task(api_key: str, payload: dict[str, Any]) -> tuple[str, str]:
"""Backward-compatible alias for node usage."""
return _create_nano_banana_task(api_key, payload)
def _poll_nanobanana_until_complete(
api_key: str,
task_id: str,
poll_interval_s: float,
timeout_s: int,
log: bool,
start_time: float,
) -> dict[str, Any]:
"""Backward-compatible polling helper used by nodes.
Returns:
The task record data dict returned by recordInfo.
Raises:
RuntimeError: If the task times out or returns a non-retryable failure.
TransientKieError: If the task fails with a retryable condition.
"""
return _poll_task_until_complete(api_key, task_id, poll_interval_s, timeout_s, log, start_time)
def _extract_result_urls(record_data: dict[str, Any]) -> list[str]:
"""Parse the resultJson payload and return the result URL list."""
result_json = record_data.get("resultJson")
if not result_json:
raise RuntimeError("Task completed without resultJson.")
@@ -239,6 +299,11 @@ def _extract_result_urls(record_data: dict[str, Any]) -> list[str]:
return result_urls
def _extract_nanobanana_result_urls(record_data: dict[str, Any]) -> list[str]:
"""Backward-compatible alias for node usage."""
return _extract_result_urls(record_data)
def _download_image(url: str) -> bytes:
"""Download a result image and return its raw bytes.
@@ -258,6 +323,11 @@ def _download_image(url: str) -> bytes:
return response.content
def _download_nanobanana_image(url: str) -> bytes:
"""Backward-compatible alias for node usage."""
return _download_image(url)
def _image_bytes_to_tensor(image_bytes: bytes) -> torch.Tensor:
"""Convert image bytes into a normalized torch tensor.
@@ -279,3 +349,115 @@ def _image_bytes_to_tensor(image_bytes: bytes) -> torch.Tensor:
return tensor.unsqueeze(0)
except Exception as exc:
raise RuntimeError("Failed to decode result image.") from exc
def run_nanobanana_image_job(
prompt: str,
aspect_ratio: str = "auto",
resolution: str = "1K",
output_format: str = "png",
log: bool = True,
poll_interval_s: float = 1.0,
timeout_s: int = 300,
retry_on_fail: bool = True,
max_retries: int = 2,
retry_backoff_s: float = 3.0,
images: torch.Tensor | None = None,
) -> torch.Tensor:
"""Run a Nano Banana Pro image generation job end-to-end.
Orchestrates the same steps as the ComfyUI node:
validation → optional reference image upload → task creation → polling → result download → decode.
Returns:
A torch tensor of shape (1, H, W, 3) with float values in [0, 1].
Raises:
RuntimeError: For validation errors, non-retryable API failures, timeouts, or decoding failures.
TransientKieError: For retryable API/task failures (used to trigger retries when enabled).
"""
_validate_prompt(prompt)
if aspect_ratio not in ASPECT_RATIO_OPTIONS:
raise RuntimeError("Invalid aspect_ratio. Use the pinned enum options.")
if resolution not in RESOLUTION_OPTIONS:
raise RuntimeError("Invalid resolution. Use the pinned enum options.")
if output_format not in OUTPUT_FORMAT_OPTIONS:
raise RuntimeError("Invalid output_format. Use the pinned enum options.")
attempts = max_retries + 1 if retry_on_fail else 1
attempts = max(attempts, 1)
backoff = retry_backoff_s if retry_backoff_s >= 0 else 0.0
for attempt in range(1, attempts + 1):
start_time = time.time()
try:
api_key = _load_api_key()
image_urls: list[str] = []
if images is not None:
if not isinstance(images, torch.Tensor):
raise RuntimeError("images input must be a tensor batch.")
if images.dim() != 4 or images.shape[-1] != 3:
raise RuntimeError("images input must have shape [B, H, W, 3].")
total_images = images.shape[0]
if total_images > 8 and log:
_log(log, f"More than 8 images provided ({total_images}); only first 8 will be used.")
upload_count = min(total_images, 8)
if upload_count > 0:
_log(log, f"Uploading {upload_count} images...")
for idx in range(upload_count):
try:
png_bytes = _image_tensor_to_png_bytes(images[idx])
url = _upload_image(api_key, png_bytes)
image_urls.append(url)
_log(log, f"Image {idx + 1} upload success: {_truncate_url(url)}")
except Exception as exc:
_log(log, f"Image {idx + 1} upload failed: {exc}")
raise
payload = {
"model": MODEL_NAME,
"input": {
"prompt": prompt,
"aspect_ratio": aspect_ratio,
"resolution": resolution,
"output_format": output_format,
"image_input": image_urls,
},
}
_log(log, f"Sending {len(image_urls)} image URLs to createTask")
_log(log, "Creating Nano Banana Pro task...")
task_id, create_response_text = _create_nano_banana_task(api_key, payload)
_log(log, f"createTask response (elapsed={time.time() - start_time:.1f}s): {create_response_text}")
_log(log, f"Task created with ID {task_id}. Polling for completion...")
record_data = _poll_task_until_complete(
api_key,
task_id,
poll_interval_s,
timeout_s,
log,
start_time,
)
result_urls = _extract_result_urls(record_data)
_log(log, f"Result URLs: {result_urls}")
_log(log, f"Downloading result image from {result_urls[0]}...")
image_bytes = _download_image(result_urls[0])
image_tensor = _image_bytes_to_tensor(image_bytes)
_log(log, "Image downloaded and decoded.")
_log_remaining_credits(log, record_data, api_key, _log)
return image_tensor
except TransientKieError:
if not retry_on_fail or attempt >= attempts:
raise
_log(log, f"Retrying (attempt {attempt + 1}/{attempts}) after {backoff}s")
time.sleep(backoff)
continue
nodes.py
+15 -100
View File
@@ -1,27 +1,14 @@
import time
import torch
from .kie_api.auth import _load_api_key
from .kie_api.credits import _fetch_remaining_credits, _log_remaining_credits
from .kie_api.http import TransientKieError
from .kie_api.nanobanana import (
ASPECT_RATIO_OPTIONS,
MODEL_NAME,
OUTPUT_FORMAT_OPTIONS,
RESOLUTION_OPTIONS,
_create_nano_banana_task,
_download_image,
_extract_result_urls,
_image_bytes_to_tensor,
_log,
_poll_task_until_complete,
_validate_prompt,
)
from .kie_api.upload import (
_image_tensor_to_png_bytes,
_truncate_url,
_upload_image,
run_nanobanana_image_job,
)
@@ -81,92 +68,20 @@ class KIE_NanoBananaPro_Image:
retry_backoff_s: float = 3.0,
images: torch.Tensor | None = None,
):
_validate_prompt(prompt)
if aspect_ratio not in ASPECT_RATIO_OPTIONS:
raise RuntimeError("Invalid aspect_ratio. Use the pinned enum options.")
if resolution not in RESOLUTION_OPTIONS:
raise RuntimeError("Invalid resolution. Use the pinned enum options.")
if output_format not in OUTPUT_FORMAT_OPTIONS:
raise RuntimeError("Invalid output_format. Use the pinned enum options.")
attempts = max_retries + 1 if retry_on_fail else 1
attempts = max(attempts, 1)
backoff = retry_backoff_s if retry_backoff_s >= 0 else 0.0
for attempt in range(1, attempts + 1):
start_time = time.time()
try:
api_key = _load_api_key()
image_urls: list[str] = []
if images is not None:
if not isinstance(images, torch.Tensor):
raise RuntimeError("images input must be a tensor batch.")
if images.dim() != 4 or images.shape[-1] != 3:
raise RuntimeError("images input must have shape [B, H, W, 3].")
total_images = images.shape[0]
if total_images > 8 and log:
_log(log, f"More than 8 images provided ({total_images}); only first 8 will be used.")
upload_count = min(total_images, 8)
if upload_count > 0:
_log(log, f"Uploading {upload_count} images...")
for idx in range(upload_count):
try:
png_bytes = _image_tensor_to_png_bytes(images[idx])
url = _upload_image(api_key, png_bytes)
image_urls.append(url)
_log(log, f"Image {idx + 1} upload success: {_truncate_url(url)}")
except Exception as exc:
_log(log, f"Image {idx + 1} upload failed: {exc}")
raise
payload = {
"model": MODEL_NAME,
"input": {
"prompt": prompt,
"aspect_ratio": aspect_ratio,
"resolution": resolution,
"output_format": output_format,
"image_input": image_urls,
},
}
_log(log, f"Sending {len(image_urls)} image URLs to createTask")
_log(log, "Creating Nano Banana Pro task...")
task_id, create_response_text = _create_nano_banana_task(api_key, payload)
_log(log, f"createTask response (elapsed={time.time() - start_time:.1f}s): {create_response_text}")
_log(log, f"Task created with ID {task_id}. Polling for completion...")
record_data, _raw_json, _message_field, elapsed_text = _poll_task_until_complete(
api_key,
task_id,
poll_interval_s,
timeout_s,
log,
start_time,
)
result_urls = _extract_result_urls(record_data)
_log(log, f"Result URLs: {result_urls}")
_log(log, f"Downloading result image from {result_urls[0]} (elapsed={elapsed_text})...")
image_bytes = _download_image(result_urls[0])
image_tensor = _image_bytes_to_tensor(image_bytes)
_log(log, "Image downloaded and decoded.")
_log_remaining_credits(log, record_data, api_key, _log)
return (image_tensor,)
except TransientKieError as exc:
if not retry_on_fail or attempt >= attempts:
raise
_log(log, f"Retrying (attempt {attempt + 1}/{attempts}) after {backoff}s")
time.sleep(backoff)
continue
image_tensor = run_nanobanana_image_job(
prompt=prompt,
aspect_ratio=aspect_ratio,
resolution=resolution,
output_format=output_format,
log=log,
poll_interval_s=poll_interval_s,
timeout_s=timeout_s,
retry_on_fail=retry_on_fail,
max_retries=max_retries,
retry_backoff_s=retry_backoff_s,
images=images,
)
return (image_tensor,)
NODE_CLASS_MAPPINGS = {