352 lines
12 KiB
Python
352 lines
12 KiB
Python
import base64
|
|
import io
|
|
import threading
|
|
import time
|
|
from datetime import datetime
|
|
from typing import List, Optional, Tuple
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
import torch
|
|
import comfy.model_management as model_management
|
|
|
|
|
|
def _get_first_image(image):
|
|
while isinstance(image, list):
|
|
if not image:
|
|
return None
|
|
image = image[0]
|
|
return image
|
|
|
|
|
|
def _tensor_to_pil(image: torch.Tensor) -> Image.Image:
|
|
if not isinstance(image, torch.Tensor):
|
|
raise TypeError("image must be torch.Tensor")
|
|
arr = image.detach().cpu().numpy()
|
|
if arr.ndim == 4:
|
|
arr = arr[0]
|
|
if arr.ndim == 3 and arr.shape[0] in (1, 3):
|
|
arr = np.transpose(arr, (1, 2, 0))
|
|
arr = np.clip(arr * 255.0, 0, 255).astype(np.uint8)
|
|
if arr.ndim == 2:
|
|
return Image.fromarray(arr, mode="L").convert("RGB")
|
|
return Image.fromarray(arr).convert("RGB")
|
|
|
|
|
|
def _pil_to_tensor(image: Image.Image) -> torch.Tensor:
|
|
if not isinstance(image, Image.Image):
|
|
raise TypeError("image must be PIL.Image")
|
|
rgb = image.convert("RGB")
|
|
arr = np.array(rgb).astype(np.float32) / 255.0
|
|
return torch.from_numpy(arr)[None,]
|
|
|
|
|
|
class Gemini_API_Image:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"API_Key": ("STRING", {"default": "<your_key>"}),
|
|
"prompt": ("STRING", {"default": "", "multiline": True, "rows": 4}),
|
|
"model_name": ("STRING", {"default": "gemini-3-pro-image-preview"}),
|
|
"aspect_ratio": (["1:1", "2:3", "3:2", "3:4", "4:3", "4:5", "5:4", "9:16", "16:9", "21:9"], {"default": "5:4"}),
|
|
"resolution": (["1K", "2K", "4K"], {"default": "2K"}),
|
|
"thinking_mode": ("BOOLEAN", {"default": False}),
|
|
"noise_seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "control_after_generate": True}),
|
|
},
|
|
"optional": {
|
|
"aspect_ratio_text": ("STRING", {"default": "", "forceInput": True}),
|
|
"image_1": ("IMAGE",),
|
|
"image_2": ("IMAGE",),
|
|
"image_3": ("IMAGE",),
|
|
"image_4": ("IMAGE",),
|
|
"image_5": ("IMAGE",),
|
|
"image_6": ("IMAGE",),
|
|
"image_7": ("IMAGE",),
|
|
"image_8": ("IMAGE",),
|
|
"image_9": ("IMAGE",),
|
|
"image_10": ("IMAGE",),
|
|
"image_11": ("IMAGE",),
|
|
"image_12": ("IMAGE",),
|
|
"image_13": ("IMAGE",),
|
|
"image_14": ("IMAGE",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "STRING")
|
|
RETURN_NAMES = ("image", "text")
|
|
FUNCTION = "generate"
|
|
CATEGORY = "DaNodes/API"
|
|
|
|
def _log(self, message: str) -> None:
|
|
ts = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
|
print(f"🌞[Gemini_API_Image {ts}] {message}", flush=True)
|
|
|
|
def _check_interrupted(self) -> None:
|
|
model_management.throw_exception_if_processing_interrupted()
|
|
|
|
def _pil_to_png_bytes(self, image: Image.Image) -> bytes:
|
|
buf = io.BytesIO()
|
|
image.save(buf, format="PNG")
|
|
return buf.getvalue()
|
|
|
|
def _build_contents(self, prompt: str, images: List[Image.Image], types_mod) -> list:
|
|
contents = []
|
|
if prompt:
|
|
contents.append(types_mod.Part.from_text(text=prompt))
|
|
for img in images:
|
|
data = self._pil_to_png_bytes(img)
|
|
contents.append(types_mod.Part.from_bytes(data=data, mime_type="image/png"))
|
|
return contents
|
|
|
|
def _build_generate_config(self, types_mod, image_cfg, thinking_mode: bool):
|
|
config_kwargs = {
|
|
"response_modalities": ["TEXT", "IMAGE"],
|
|
"image_config": image_cfg,
|
|
}
|
|
if thinking_mode and hasattr(types_mod, "ThinkingConfig"):
|
|
try:
|
|
config_kwargs["thinking_config"] = types_mod.ThinkingConfig(include_thoughts=False)
|
|
except Exception:
|
|
pass
|
|
return types_mod.GenerateContentConfig(**config_kwargs)
|
|
|
|
def _import_genai(self):
|
|
try:
|
|
from google import genai as genai_mod
|
|
from google.genai import types as types_mod
|
|
return genai_mod, types_mod, None
|
|
except Exception:
|
|
pass
|
|
try:
|
|
import importlib
|
|
|
|
genai_mod = importlib.import_module("google.genai")
|
|
types_mod = importlib.import_module("google.genai.types")
|
|
return genai_mod, types_mod, None
|
|
except Exception as e:
|
|
return None, None, e
|
|
|
|
def _collect_reference_images(self, *images) -> List[Image.Image]:
|
|
refs: List[Image.Image] = []
|
|
for entry in images:
|
|
if entry is None:
|
|
continue
|
|
entry = _get_first_image(entry)
|
|
if entry is None:
|
|
continue
|
|
if isinstance(entry, Image.Image):
|
|
refs.append(entry.convert("RGB"))
|
|
continue
|
|
if isinstance(entry, torch.Tensor):
|
|
refs.append(_tensor_to_pil(entry))
|
|
continue
|
|
raise TypeError("image inputs must be IMAGE tensors")
|
|
return refs[:14]
|
|
|
|
def _extract_response_parts(self, response) -> Tuple[Optional[Image.Image], str]:
|
|
parts = getattr(response, "parts", None)
|
|
if parts is None and hasattr(response, "candidates"):
|
|
try:
|
|
parts = response.candidates[0].content.parts
|
|
except Exception:
|
|
parts = []
|
|
if parts is None:
|
|
parts = []
|
|
|
|
text_parts: List[str] = []
|
|
output_image: Optional[Image.Image] = None
|
|
|
|
response_text = getattr(response, "text", None)
|
|
if isinstance(response_text, str) and response_text.strip():
|
|
text_parts.append(response_text.strip())
|
|
|
|
for part in parts:
|
|
text = getattr(part, "text", None)
|
|
if text:
|
|
text_parts.append(text)
|
|
continue
|
|
inline = getattr(part, "inline_data", None)
|
|
data = getattr(inline, "data", None) if inline is not None else None
|
|
if data:
|
|
try:
|
|
raw = base64.b64decode(data) if isinstance(data, str) else data
|
|
output_image = Image.open(io.BytesIO(raw)).convert("RGB")
|
|
except Exception:
|
|
output_image = None
|
|
continue
|
|
try:
|
|
image = part.as_image()
|
|
except Exception:
|
|
image = None
|
|
if image is None:
|
|
continue
|
|
if isinstance(image, Image.Image):
|
|
output_image = image.convert("RGB")
|
|
elif isinstance(image, (bytes, bytearray)):
|
|
try:
|
|
output_image = Image.open(io.BytesIO(image)).convert("RGB")
|
|
except Exception:
|
|
output_image = None
|
|
else:
|
|
try:
|
|
output_image = Image.open(image).convert("RGB")
|
|
except Exception:
|
|
output_image = None
|
|
return output_image, "\n".join(t for t in text_parts if t).strip()
|
|
|
|
def _empty_image(self) -> torch.Tensor:
|
|
return torch.zeros((1, 1, 1, 3), dtype=torch.float32)
|
|
|
|
def generate(
|
|
self,
|
|
API_Key: str,
|
|
prompt: str,
|
|
model_name: str,
|
|
aspect_ratio: str,
|
|
resolution: str,
|
|
thinking_mode: bool = False,
|
|
noise_seed: int = 0,
|
|
aspect_ratio_text: str = "",
|
|
image_1=None,
|
|
image_2=None,
|
|
image_3=None,
|
|
image_4=None,
|
|
image_5=None,
|
|
image_6=None,
|
|
image_7=None,
|
|
image_8=None,
|
|
image_9=None,
|
|
image_10=None,
|
|
image_11=None,
|
|
image_12=None,
|
|
image_13=None,
|
|
image_14=None,
|
|
) -> Tuple[torch.Tensor, str]:
|
|
try:
|
|
self._check_interrupted()
|
|
except model_management.InterruptProcessingException:
|
|
self._log("Interrupted before request started.")
|
|
return (self._empty_image(), "Task cancelled.")
|
|
|
|
if not API_Key or API_Key.strip() in ("<your_key>", "your_key", "api_key"):
|
|
return (self._empty_image(), "API key is missing.")
|
|
|
|
genai, types, import_err = self._import_genai()
|
|
if import_err is not None:
|
|
msg = (
|
|
"google-genai is not available or the 'google' package is shadowing it. "
|
|
f"Original error: {import_err}. "
|
|
"Try: pip uninstall -y google; pip install -U google-genai"
|
|
)
|
|
return (self._empty_image(), msg)
|
|
|
|
ref_images = self._collect_reference_images(
|
|
image_1,
|
|
image_2,
|
|
image_3,
|
|
image_4,
|
|
image_5,
|
|
image_6,
|
|
image_7,
|
|
image_8,
|
|
image_9,
|
|
image_10,
|
|
image_11,
|
|
image_12,
|
|
image_13,
|
|
image_14,
|
|
)
|
|
|
|
if not prompt and not ref_images:
|
|
return (self._empty_image(), "Prompt and reference images are both empty.")
|
|
|
|
contents = self._build_contents(prompt, ref_images, types)
|
|
|
|
try:
|
|
client = genai.Client(api_key=API_Key)
|
|
aspect_ratio_value = aspect_ratio_text.strip() or aspect_ratio
|
|
image_cfg = types.ImageConfig(
|
|
aspect_ratio=aspect_ratio_value,
|
|
image_size=resolution,
|
|
)
|
|
if noise_seed:
|
|
try:
|
|
setattr(image_cfg, "seed", int(noise_seed))
|
|
except Exception:
|
|
pass
|
|
self._log(
|
|
f"Request start: model={model_name}, refs={len(ref_images)}, aspect_ratio={aspect_ratio_value}, resolution={resolution}, thinking={bool(thinking_mode)}, seed={int(noise_seed) if noise_seed else 0}"
|
|
)
|
|
start_time = time.monotonic()
|
|
|
|
response_box = {"response": None, "error": None}
|
|
|
|
def _do_request():
|
|
try:
|
|
response_box["response"] = client.models.generate_content(
|
|
model=model_name,
|
|
contents=contents,
|
|
config=self._build_generate_config(types, image_cfg, bool(thinking_mode)),
|
|
)
|
|
except Exception as e: # defensive: capture worker errors
|
|
if not thinking_mode:
|
|
response_box["error"] = e
|
|
return
|
|
try:
|
|
response_box["response"] = client.models.generate_content(
|
|
model=model_name,
|
|
contents=contents,
|
|
config=self._build_generate_config(types, image_cfg, False),
|
|
)
|
|
except Exception as fallback_error:
|
|
response_box["error"] = fallback_error
|
|
|
|
worker = threading.Thread(
|
|
target=_do_request,
|
|
name="GeminiGenerateContent",
|
|
daemon=True,
|
|
)
|
|
worker.start()
|
|
|
|
while worker.is_alive():
|
|
try:
|
|
self._check_interrupted()
|
|
except model_management.InterruptProcessingException:
|
|
elapsed = time.monotonic() - start_time
|
|
self._log(f"Request interrupted after {elapsed:.2f}s.")
|
|
return (self._empty_image(), "Task cancelled.")
|
|
time.sleep(0.25)
|
|
|
|
if response_box["error"] is not None:
|
|
raise response_box["error"]
|
|
|
|
response = response_box["response"]
|
|
elapsed = time.monotonic() - start_time
|
|
self._log(f"Request finished in {elapsed:.2f}s.")
|
|
except Exception as e:
|
|
self._log(f"Request failed: {e}")
|
|
return (self._empty_image(), f"Gemini request failed: {e}")
|
|
|
|
try:
|
|
self._check_interrupted()
|
|
except model_management.InterruptProcessingException:
|
|
self._log("Interrupted after request finished.")
|
|
return (self._empty_image(), "Task cancelled.")
|
|
|
|
output_image, text = self._extract_response_parts(response)
|
|
if output_image is None:
|
|
self._log("No image returned from Gemini.")
|
|
return (self._empty_image(), text or "No image returned from Gemini.")
|
|
|
|
self._log("Image returned successfully.")
|
|
return (_pil_to_tensor(output_image), text)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"Gemini_API_Image": Gemini_API_Image,
|
|
}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"Gemini_API_Image": "Gemini_API_Image☀",
|
|
}
|