111 lines
4.0 KiB
Python
111 lines
4.0 KiB
Python
import base64
|
|
import torch
|
|
from typing import Optional, List, Dict, Any
|
|
from comfy.comfy_types.node_typing import IO
|
|
|
|
from .modelverse_api.client import ModelverseClient
|
|
from .modelverse_api.requests.gemini_flash_image import GeminiFlashImageRequest
|
|
from .modelverse_api.utils import decode_image, images2tensor
|
|
|
|
|
|
MODELS = ["gemini-3.1-flash-image", "gemini-2.5-flash-image"]
|
|
|
|
|
|
def _extract_images_from_gemini_response(resp: Dict[str, Any]) -> List[torch.Tensor]:
|
|
"""Parse Gemini API response and return list of tensors for any inline images."""
|
|
images = []
|
|
try:
|
|
candidates = resp.get("candidates", []) or []
|
|
for cand in candidates:
|
|
content = cand.get("content", {})
|
|
parts = content.get("parts", [])
|
|
for part in parts:
|
|
inline = part.get("inlineData")
|
|
if inline and isinstance(inline, dict):
|
|
data_b64 = inline.get("data")
|
|
if data_b64:
|
|
try:
|
|
img_bytes = base64.b64decode(data_b64)
|
|
pil_img = decode_image(img_bytes)
|
|
images.append(pil_img)
|
|
except Exception:
|
|
continue
|
|
except Exception:
|
|
pass
|
|
if not images:
|
|
return []
|
|
return [images2tensor(images)]
|
|
|
|
|
|
class GeminiFlashImageNode:
|
|
"""
|
|
Gemini Flash Image (text-to-image and image-edit) via Modelverse API.
|
|
|
|
Endpoint: /v1beta/models/{model}:generateContent
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"client": ("MODELVERSE_API_CLIENT",),
|
|
"model": (MODELS, {"default": "gemini-3.1-flash-image", "tooltip": "Gemini Flash Image model"}),
|
|
"prompt": (IO.STRING, {"multiline": True, "default": "Create a picture of a nano banana dish in a fancy restaurant with a Gemini theme"}),
|
|
"mime_type": (["image/png", "image/jpeg"], {"default": "image/png"}),
|
|
"num_requests": (IO.INT, {"default": 1, "min": 1, "max": 10, "step": 1, "display": "number"}),
|
|
},
|
|
"optional": {
|
|
"image": (IO.IMAGE, {"default": None, "tooltip": "Optional input image for edit"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = (IO.IMAGE,)
|
|
RETURN_NAMES = ("image",)
|
|
CATEGORY = "UCLOUD_MODELVERSE/Gemini"
|
|
FUNCTION = "execute"
|
|
|
|
async def execute(self,
|
|
client,
|
|
model: str,
|
|
prompt: str,
|
|
mime_type: str = "image/png",
|
|
num_requests: int = 1,
|
|
image=None):
|
|
if not prompt:
|
|
raise ValueError("Prompt is required")
|
|
|
|
mv_client = ModelverseClient(client["api_key"])
|
|
|
|
outputs: List[torch.Tensor] = []
|
|
for i in range(num_requests):
|
|
req = GeminiFlashImageRequest(prompt=prompt, model=model, image=image, mime_type=mime_type)
|
|
payload = req.build_payload()
|
|
resp = mv_client.post(req.API_PATH, payload)
|
|
|
|
if isinstance(resp, dict) and resp.get("error"):
|
|
err = resp.get("error")
|
|
raise Exception(f"GeminiFlashImage error: {err.get('message', 'Unknown error')}")
|
|
|
|
tensors = _extract_images_from_gemini_response(resp)
|
|
if not tensors:
|
|
print("WARN:", "No image data found in Gemini response; check console for details.")
|
|
continue
|
|
# tensors is a list with one batch tensor; append that tensor
|
|
outputs.append(tensors[0])
|
|
|
|
if not outputs:
|
|
# Return an empty black image tensor (1x3x1x1) to avoid breaking graph
|
|
return (torch.zeros((1, 3, 1, 1)),)
|
|
|
|
# Concatenate along batch dimension
|
|
return (torch.cat(outputs, dim=0),)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"NanoBanana": GeminiFlashImageNode,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"NanoBanana": "Modelverse Gemini Flash Image",
|
|
}
|