122 lines
4.1 KiB
Python
122 lines
4.1 KiB
Python
"""
|
|
Gemini 3 Pro Image (Nano Banana Pro) request builder.
|
|
Supports text-to-image and image editing with advanced features like
|
|
aspect ratio, resolution control, and Google Search grounding.
|
|
"""
|
|
import io
|
|
import base64
|
|
import numpy as np
|
|
from PIL import Image
|
|
from typing import Optional, List, Dict, Any
|
|
from pydantic import Field
|
|
from torch import Tensor
|
|
|
|
from ..utils import BaseRequest
|
|
|
|
|
|
def _tensor_to_base64(image: Tensor, mime_type: str = "image/png") -> Dict[str, str]:
|
|
"""Convert a ComfyUI image tensor (HWC or 4D with batch) to base64 without data URI."""
|
|
if image is None:
|
|
return None
|
|
|
|
# If batch, take first frame
|
|
if hasattr(image, 'shape') and len(image.shape) == 4:
|
|
image = image[0]
|
|
|
|
# Tensor (H, W, C) in 0..1 -> uint8
|
|
np_img = np.clip(255.0 * image.cpu().numpy(), 0, 255).astype(np.uint8)
|
|
pil_img = Image.fromarray(np_img)
|
|
|
|
fmt = 'PNG' if mime_type.lower().endswith('png') else 'JPEG'
|
|
with io.BytesIO() as bio:
|
|
pil_img.save(bio, format=fmt)
|
|
data = bio.getvalue()
|
|
|
|
return {
|
|
"mimeType": mime_type,
|
|
"data": base64.b64encode(data).decode("utf-8"),
|
|
}
|
|
|
|
|
|
class GeminiProImageRequest(BaseRequest):
|
|
"""
|
|
Request builder for Gemini 3 Pro Image (Nano Banana Pro) generateContent endpoint.
|
|
Supports text-to-image, image editing, and advanced features like aspect ratio,
|
|
resolution control (1K/2K/4K), and Google Search grounding.
|
|
"""
|
|
|
|
API_PATH = "/v1beta/models/gemini-3-pro-image:generateContent"
|
|
|
|
prompt: str = Field(..., description="Text prompt")
|
|
images: Optional[List[Tensor]] = Field(default=None, description="Optional input images (up to 14)")
|
|
mime_type: str = Field(default="image/png", description="MIME type for inline image data")
|
|
aspect_ratio: Optional[str] = Field(default=None, description="Aspect ratio: 1:1, 2:3, 3:2, 3:4, 4:3, 4:5, 5:4, 9:16, 16:9, 21:9")
|
|
image_size: Optional[str] = Field(default=None, description="Image resolution: 1K, 2K, 4K")
|
|
use_google_search: bool = Field(default=False, description="Enable Google Search grounding")
|
|
|
|
def __init__(
|
|
self,
|
|
prompt: str,
|
|
images: Optional[List[Tensor]] = None,
|
|
mime_type: str = "image/png",
|
|
aspect_ratio: Optional[str] = None,
|
|
image_size: Optional[str] = None,
|
|
use_google_search: bool = False,
|
|
**kwargs
|
|
):
|
|
super().__init__(**kwargs)
|
|
self.prompt = prompt
|
|
self.images = images
|
|
self.mime_type = mime_type
|
|
self.aspect_ratio = aspect_ratio
|
|
self.image_size = image_size
|
|
self.use_google_search = use_google_search
|
|
|
|
def build_payload(self) -> Dict[str, Any]:
|
|
parts: List[Dict[str, Any]] = []
|
|
|
|
# Add text prompt
|
|
if self.prompt:
|
|
parts.append({"text": self.prompt})
|
|
|
|
# Add images (up to 14)
|
|
if self.images:
|
|
for img in self.images[:14]: # Limit to 14 images
|
|
if isinstance(img, Tensor):
|
|
parts.append({
|
|
"inlineData": _tensor_to_base64(img, self.mime_type)
|
|
})
|
|
|
|
payload = {
|
|
"contents": [
|
|
{
|
|
"role": "user",
|
|
"parts": parts,
|
|
}
|
|
],
|
|
"generationConfig": {
|
|
"responseModalities": ["TEXT", "IMAGE"],
|
|
},
|
|
}
|
|
|
|
# Add imageConfig if aspect_ratio or image_size specified
|
|
if self.aspect_ratio or self.image_size:
|
|
image_config = {}
|
|
if self.aspect_ratio:
|
|
image_config["aspectRatio"] = self.aspect_ratio
|
|
if self.image_size:
|
|
image_config["imageSize"] = self.image_size
|
|
payload["generationConfig"]["imageConfig"] = image_config
|
|
|
|
# Add Google Search tool if enabled
|
|
if self.use_google_search:
|
|
payload["tools"] = [{"google_search": {}}]
|
|
|
|
return self._remove_empty_fields(payload)
|
|
|
|
def field_required(self):
|
|
return ["prompt"]
|
|
|
|
def field_order(self):
|
|
return ["prompt", "images", "mime_type", "aspect_ratio", "image_size", "use_google_search"]
|