diff --git a/kie_api/nanobanana.py b/kie_api/nanobanana.py index 8ac2a48..7415c19 100644 --- a/kie_api/nanobanana.py +++ b/kie_api/nanobanana.py @@ -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 diff --git a/nodes.py b/nodes.py index 997d74a..0cc9462 100644 --- a/nodes.py +++ b/nodes.py @@ -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 = {