95 lines
3.2 KiB
Python
95 lines
3.2 KiB
Python
import torch
|
|
|
|
from .kie_api.auth import _load_api_key
|
|
from .kie_api.credits import _fetch_remaining_credits, _log_remaining_credits
|
|
from .kie_api.nanobanana import (
|
|
ASPECT_RATIO_OPTIONS,
|
|
MODEL_NAME,
|
|
OUTPUT_FORMAT_OPTIONS,
|
|
RESOLUTION_OPTIONS,
|
|
_log,
|
|
run_nanobanana_image_job,
|
|
)
|
|
|
|
|
|
class KIE_GetRemainingCredits:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {"required": {"log": ("BOOLEAN", {"default": True})}}
|
|
|
|
RETURN_TYPES = ("STRING", "INT")
|
|
RETURN_NAMES = ("raw_json", "credits_remaining")
|
|
FUNCTION = "get_remaining_credits"
|
|
CATEGORY = "kie/api"
|
|
|
|
def get_remaining_credits(self, log: bool):
|
|
_log(log, "Requesting remaining credits...")
|
|
api_key = _load_api_key()
|
|
raw_json, credits_remaining = _fetch_remaining_credits(api_key)
|
|
_log(log, f"Credits remaining: {credits_remaining}")
|
|
return (raw_json, credits_remaining)
|
|
|
|
|
|
class KIE_NanoBananaPro_Image:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {"prompt": ("STRING", {"multiline": True})},
|
|
"optional": {
|
|
"images": ("IMAGE",),
|
|
"aspect_ratio": ("COMBO", {"options": ASPECT_RATIO_OPTIONS, "default": "auto"}),
|
|
"resolution": ("COMBO", {"options": RESOLUTION_OPTIONS, "default": "1K"}),
|
|
"output_format": ("COMBO", {"options": OUTPUT_FORMAT_OPTIONS, "default": "png"}),
|
|
"log": ("BOOLEAN", {"default": True}),
|
|
"poll_interval_s": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 60.0, "step": 0.1}),
|
|
"timeout_s": ("INT", {"default": 300, "min": 1, "max": 3600, "step": 1}),
|
|
"retry_on_fail": ("BOOLEAN", {"default": True}),
|
|
"max_retries": ("INT", {"default": 2, "min": 0, "max": 10, "step": 1}),
|
|
"retry_backoff_s": ("FLOAT", {"default": 3.0, "min": 0.0, "max": 300.0, "step": 0.5}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
FUNCTION = "generate"
|
|
CATEGORY = "kie/api"
|
|
|
|
def generate(
|
|
self,
|
|
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,
|
|
):
|
|
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 = {
|
|
"KIE_GetRemainingCredits": KIE_GetRemainingCredits,
|
|
"KIE_NanoBananaPro_Image": KIE_NanoBananaPro_Image,
|
|
}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"KIE_GetRemainingCredits": "KIE Get Remaining Credits",
|
|
"KIE_NanoBananaPro_Image": "KIE Nano Banana Pro (Image)",
|
|
}
|