254 lines
8.8 KiB
Python
254 lines
8.8 KiB
Python
from google import genai
|
|
from google.genai import types
|
|
import os
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
import io
|
|
from datetime import datetime
|
|
|
|
class GeminiImageGenerator:
|
|
"""
|
|
A ComfyUI node for generating images using Google Gemini API
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"user_prompt": ("STRING", {
|
|
"multiline": True,
|
|
"default": "Generate an image"
|
|
}),
|
|
"model_name": (["gemini-3-pro-image-preview", "gemini-2.5-flash-image-preview"], {
|
|
"default": "gemini-3-pro-image-preview"
|
|
}),
|
|
"aspect_ratio": (["1:1", "16:9", "9:16", "4:3", "3:4", "21:9", "3:2", "2:3"], {
|
|
"default": "1:1"
|
|
}),
|
|
"image_resolution": (["1K", "2K", "4K"], {
|
|
"default": "2K"
|
|
}),
|
|
"seed": ("INT", {
|
|
"default": 0,
|
|
"min": 0,
|
|
"max": 0xffffffffffffffff
|
|
}),
|
|
},
|
|
"optional": {
|
|
"system_prompt": ("STRING", {
|
|
"multiline": True,
|
|
"default": ""
|
|
}),
|
|
"image_1": ("IMAGE",),
|
|
"image_2": ("IMAGE",),
|
|
"image_3": ("IMAGE",),
|
|
"image_4": ("IMAGE",),
|
|
"image_5": ("IMAGE",),
|
|
"image_6": ("IMAGE",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "STRING")
|
|
RETURN_NAMES = ("image", "info_text")
|
|
FUNCTION = "generate_image"
|
|
CATEGORY = "Gemini"
|
|
|
|
def tensor_to_pil(self, tensor):
|
|
"""Convert ComfyUI tensor to PIL Image"""
|
|
# ComfyUI uses [B, H, W, C] format with values in [0, 1]
|
|
if tensor.dim() == 4:
|
|
tensor = tensor[0] # Take first batch
|
|
|
|
# Convert to numpy and scale to 0-255
|
|
np_image = (tensor.cpu().numpy() * 255).astype(np.uint8)
|
|
return Image.fromarray(np_image)
|
|
|
|
def pil_to_tensor(self, pil_image):
|
|
"""Convert PIL Image to ComfyUI tensor"""
|
|
np_image = np.array(pil_image).astype(np.float32) / 255.0
|
|
|
|
# Ensure 3 channels (RGB)
|
|
if len(np_image.shape) == 2: # Grayscale
|
|
np_image = np.stack([np_image] * 3, axis=-1)
|
|
elif np_image.shape[2] == 4: # RGBA
|
|
np_image = np_image[:, :, :3]
|
|
|
|
# Add batch dimension [B, H, W, C]
|
|
tensor = torch.from_numpy(np_image)[None,]
|
|
return tensor
|
|
|
|
|
|
def generate_image(self, user_prompt, model_name, aspect_ratio, image_resolution, seed,
|
|
system_prompt="",
|
|
image_1=None, image_2=None, image_3=None,
|
|
image_4=None, image_5=None, image_6=None):
|
|
|
|
start_time = datetime.now()
|
|
info_dict = {
|
|
"user_prompt": user_prompt,
|
|
"model_name": model_name,
|
|
"aspect_ratio": aspect_ratio,
|
|
"image_resolution": image_resolution,
|
|
"start_time": start_time.strftime("%Y-%m-%d %H:%M:%S"),
|
|
"success": False,
|
|
"error": None
|
|
}
|
|
|
|
try:
|
|
# Initialize Gemini client
|
|
client = genai.Client(
|
|
vertexai=True,
|
|
api_key=os.environ.get("GOOGLE_CLOUD_API_KEY"),
|
|
)
|
|
|
|
# Prepare content parts
|
|
content_parts = []
|
|
|
|
# Add input images if provided
|
|
for idx, img_tensor in enumerate([image_1, image_2, image_3, image_4, image_5, image_6], 1):
|
|
if img_tensor is not None:
|
|
# Convert tensor to PIL Image
|
|
pil_image = self.tensor_to_pil(img_tensor)
|
|
|
|
# Convert PIL Image to bytes
|
|
img_byte_arr = io.BytesIO()
|
|
pil_image.save(img_byte_arr, format='PNG')
|
|
img_bytes = img_byte_arr.getvalue()
|
|
|
|
# Add to content parts
|
|
content_parts.append(types.Part.from_bytes(
|
|
data=img_bytes,
|
|
mime_type="image/png"
|
|
))
|
|
|
|
# Build prompt text
|
|
prompt_text = ""
|
|
if system_prompt:
|
|
prompt_text += f"System: {system_prompt}\n\n"
|
|
prompt_text += user_prompt
|
|
|
|
# Add text prompt
|
|
content_parts.append(types.Part.from_text(text=prompt_text))
|
|
|
|
# Create content
|
|
contents = [
|
|
types.Content(
|
|
role="user",
|
|
parts=content_parts
|
|
)
|
|
]
|
|
|
|
# Use resolution directly (1K, 2K, 4K)
|
|
image_size = image_resolution
|
|
|
|
# Configure generation
|
|
generate_content_config = types.GenerateContentConfig(
|
|
temperature=1,
|
|
top_p=0.95,
|
|
max_output_tokens=32768,
|
|
response_modalities=["TEXT", "IMAGE"],
|
|
safety_settings=[
|
|
types.SafetySetting(
|
|
category="HARM_CATEGORY_HATE_SPEECH",
|
|
threshold="OFF"
|
|
),
|
|
types.SafetySetting(
|
|
category="HARM_CATEGORY_DANGEROUS_CONTENT",
|
|
threshold="OFF"
|
|
),
|
|
types.SafetySetting(
|
|
category="HARM_CATEGORY_SEXUALLY_EXPLICIT",
|
|
threshold="OFF"
|
|
),
|
|
types.SafetySetting(
|
|
category="HARM_CATEGORY_HARASSMENT",
|
|
threshold="OFF"
|
|
)
|
|
],
|
|
image_config=types.ImageConfig(
|
|
aspect_ratio=aspect_ratio,
|
|
image_size=image_size,
|
|
output_mime_type="image/png",
|
|
),
|
|
)
|
|
|
|
# Generate image
|
|
generated_image_data = None
|
|
response_text = []
|
|
|
|
for chunk in client.models.generate_content_stream(
|
|
model=model_name,
|
|
contents=contents,
|
|
config=generate_content_config,
|
|
):
|
|
if chunk.candidates:
|
|
for candidate in chunk.candidates:
|
|
if candidate.content and candidate.content.parts:
|
|
for part in candidate.content.parts:
|
|
# Collect text
|
|
if hasattr(part, 'text') and part.text:
|
|
response_text.append(part.text)
|
|
|
|
# Get image data
|
|
if hasattr(part, 'inline_data') and part.inline_data:
|
|
generated_image_data = part.inline_data.data
|
|
|
|
# Process generated image
|
|
if generated_image_data:
|
|
pil_image = Image.open(io.BytesIO(generated_image_data))
|
|
output_tensor = self.pil_to_tensor(pil_image)
|
|
|
|
info_dict["success"] = True
|
|
end_time = datetime.now()
|
|
info_dict["generation_time"] = f"{(end_time - start_time).total_seconds():.2f} seconds"
|
|
|
|
# Format info text
|
|
info_text = f"""=== Generation Info ===
|
|
User Prompt: {user_prompt}
|
|
Model Name: {model_name}
|
|
Aspect Ratio: {aspect_ratio}
|
|
Image Resolution: {image_resolution}
|
|
Start Time: {info_dict["start_time"]}
|
|
Generation Time: {info_dict["generation_time"]}
|
|
Status: ✅ Success
|
|
Seed: {seed}
|
|
"""
|
|
if response_text:
|
|
info_text += f"\nResponse Text:\n{''.join(response_text)}"
|
|
|
|
return (output_tensor, info_text)
|
|
else:
|
|
raise Exception("No image generated in response")
|
|
|
|
except Exception as e:
|
|
info_dict["success"] = False
|
|
info_dict["error"] = str(e)
|
|
|
|
# Format error info text
|
|
info_text = f"""=== Generation Info ===
|
|
User Prompt: {user_prompt}
|
|
Model Name: {model_name}
|
|
Aspect Ratio: {aspect_ratio}
|
|
Image Resolution: {image_resolution}
|
|
Start Time: {info_dict["start_time"]}
|
|
Status: ❌ Failed
|
|
Error: {str(e)}
|
|
Seed: {seed}
|
|
"""
|
|
|
|
# Return a black image as placeholder
|
|
black_image = torch.zeros((1, 512, 512, 3))
|
|
return (black_image, info_text)
|
|
|
|
|
|
# ComfyUI node registration
|
|
NODE_CLASS_MAPPINGS = {
|
|
"GeminiImageGenerator": GeminiImageGenerator
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"GeminiImageGenerator": "Gemini Image Generator"
|
|
}
|
|
|