84 lines
2.7 KiB
Python
84 lines
2.7 KiB
Python
from typing import Optional, List, Dict, Any
|
|
from pydantic import Field
|
|
from ..utils import BaseRequest
|
|
from torch import Tensor
|
|
import base64
|
|
import io
|
|
from PIL import Image
|
|
import numpy as np
|
|
|
|
|
|
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 GeminiFlashImageRequest(BaseRequest):
|
|
"""
|
|
Request builder for Gemini Flash Image generateContent endpoint.
|
|
Supports text-to-image and image-edit (text+image to image).
|
|
"""
|
|
|
|
API_PATH_TEMPLATE = "/v1beta/models/{model}:generateContent"
|
|
|
|
prompt: str = Field(..., description="Text prompt")
|
|
model: str = Field(default="gemini-3.1-flash-image", description="Gemini Flash Image model")
|
|
image: Optional[Tensor] = Field(default=None, description="Optional input image for edit")
|
|
mime_type: str = Field(default="image/png", description="MIME type for inline image data")
|
|
|
|
def __init__(self, prompt: str, model: str = "gemini-3.1-flash-image", image: Optional[Tensor] = None, mime_type: str = "image/png", **kwargs):
|
|
super().__init__(**kwargs)
|
|
self.prompt = prompt
|
|
self.model = model
|
|
self.API_PATH = self.API_PATH_TEMPLATE.format(model=model)
|
|
self.image = image
|
|
self.mime_type = mime_type
|
|
|
|
def build_payload(self) -> Dict[str, Any]:
|
|
parts: List[Dict[str, Any]] = []
|
|
if self.prompt:
|
|
parts.append({"text": self.prompt})
|
|
|
|
if isinstance(self.image, Tensor):
|
|
parts.append({
|
|
"inlineData": _tensor_to_base64(self.image, self.mime_type)
|
|
})
|
|
|
|
payload = {
|
|
"contents": [
|
|
{
|
|
"role": "user",
|
|
"parts": parts,
|
|
}
|
|
],
|
|
"generationConfig": {
|
|
"responseModalities": ["TEXT", "IMAGE"],
|
|
},
|
|
}
|
|
return self._remove_empty_fields(payload)
|
|
|
|
def field_required(self):
|
|
return ["prompt"]
|
|
|
|
def field_order(self):
|
|
return ["model", "prompt", "image", "mime_type"]
|