124 lines
4.6 KiB
Python
124 lines
4.6 KiB
Python
import torch
|
|
import numpy as np
|
|
from PIL import Image
|
|
import io
|
|
import os
|
|
from google import genai
|
|
from google.genai import types
|
|
|
|
class GeminiImageGen:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"api_key": ("STRING", {"multiline": False, "default": "INSERT_API_KEY_HERE"}),
|
|
"prompt": ("STRING", {"multiline": True, "default": "A futuristic city with flying cars"}),
|
|
"model": (
|
|
[
|
|
"gemini-3-pro-image-preview",
|
|
"gemini-2.5-flash-image",
|
|
],
|
|
{"default": "gemini-3-pro-image-preview"},
|
|
),
|
|
"resolution": (
|
|
["1K", "2K", "4K"],
|
|
{"default": "1K"},
|
|
),
|
|
"aspect_ratio": (
|
|
[
|
|
"1:1",
|
|
"16:9",
|
|
"4:3",
|
|
"3:4",
|
|
"9:16",
|
|
"2:3",
|
|
"3:2",
|
|
"4:5",
|
|
"5:4",
|
|
"21:9",
|
|
],
|
|
{"default": "1:1"},
|
|
),
|
|
"seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
FUNCTION = "generate_image"
|
|
CATEGORY = "Gemini"
|
|
|
|
def generate_image(self, api_key, prompt, model, resolution, aspect_ratio, seed):
|
|
# Configure the client
|
|
if api_key == "INSERT_API_KEY_HERE":
|
|
api_key = os.environ.get("GOOGLE_API_KEY")
|
|
|
|
if not api_key:
|
|
raise ValueError("API Key is required. Please provide it in the node or set GOOGLE_API_KEY environment variable.")
|
|
|
|
client = genai.Client(api_key=api_key)
|
|
|
|
try:
|
|
# Configure proper image generation settings using ImageConfig
|
|
# This requires google-genai >= 1.51.0
|
|
config = types.GenerateContentConfig(
|
|
response_modalities=['TEXT', 'IMAGE'],
|
|
image_config=types.ImageConfig(
|
|
aspect_ratio=aspect_ratio,
|
|
image_size=resolution # Must be uppercase: "1K", "2K", "4K"
|
|
)
|
|
)
|
|
|
|
# Call the API
|
|
response = client.models.generate_content(
|
|
model=model,
|
|
contents=[prompt],
|
|
config=config
|
|
)
|
|
|
|
|
|
generated_image = None
|
|
|
|
# Extract image from response - the correct path is response.candidates[0].content.parts
|
|
if response.candidates and response.candidates[0].content and response.candidates[0].content.parts:
|
|
for part in response.candidates[0].content.parts:
|
|
if part.inline_data:
|
|
# Extract bytes directly from inline_data and create PIL Image
|
|
# as_image() returns incompatible Image type without .mode attribute
|
|
image_bytes = part.inline_data.data
|
|
generated_image = Image.open(io.BytesIO(image_bytes))
|
|
break
|
|
|
|
if generated_image is None:
|
|
raise RuntimeError("No image generated by Gemini API. The response did not contain inline image data.")
|
|
|
|
# Convert PIL Image to Tensor
|
|
if generated_image.mode != "RGB":
|
|
generated_image = generated_image.convert("RGB")
|
|
|
|
image = np.array(generated_image).astype(np.float32) / 255.0
|
|
image = torch.from_numpy(image)[None,]
|
|
|
|
return (image,)
|
|
|
|
except Exception as e:
|
|
error_msg = str(e)
|
|
if "429" in error_msg or "RESOURCE_EXHAUSTED" in error_msg:
|
|
raise RuntimeError(f"Gemini API Quota Exceeded (429). You are likely on the Free Tier and have exhausted your limit or this model is not available on the Free Tier. Please check your billing at Google AI Studio. Details: {error_msg}")
|
|
elif "404" in error_msg or "NOT_FOUND" in error_msg:
|
|
raise RuntimeError(f"Gemini Model Not Found (404). The model '{model}' might not be available in your region or for your API key. Details: {error_msg}")
|
|
else:
|
|
raise RuntimeError(f"Error generating image with Gemini API: {error_msg}")
|
|
|
|
# Node Mappings
|
|
NODE_CLASS_MAPPINGS = {
|
|
"GeminiImageGen": GeminiImageGen
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"GeminiImageGen": "Gemini 3 Pro Image Generator"
|
|
}
|