421 lines
18 KiB
Python
421 lines
18 KiB
Python
# FSLGeminiGenerateImage.py
|
|
import os, time, base64, random
|
|
from io import BytesIO
|
|
from typing import Any, Dict, List, Optional, Tuple
|
|
|
|
import numpy as np
|
|
import requests
|
|
|
|
try:
|
|
import torch
|
|
except Exception:
|
|
torch = None
|
|
|
|
try:
|
|
from PIL import Image, ImageOps, ImageFilter
|
|
except Exception:
|
|
Image = None
|
|
|
|
|
|
# -------- constants and helpers --------
|
|
|
|
# FINAL ASPECT RATIOS: Includes common screen and photography ratios
|
|
ASPECT_CHOICES = ("1:1", "16:9", "9:16", "4:3", "3:4", "3:2", "2:3")
|
|
QUALITY_CHOICES = ("Standard", "High")
|
|
|
|
def _tensor_like_to_numpy(img_like):
|
|
"""Accept torch or np tensor/array and return uint8 ndarray in HxWxC or HxW."""
|
|
if torch is not None and isinstance(img_like, torch.Tensor):
|
|
t = img_like.detach().cpu().float()
|
|
if t.dim() == 4:
|
|
t = t[0]
|
|
if t.dim() == 3 and t.shape[0] in (1, 3, 4) and t.shape[-1] not in (1, 3, 4):
|
|
t = t.permute(1, 2, 0) # CHW -> HWC
|
|
arr = t.numpy()
|
|
elif isinstance(img_like, np.ndarray):
|
|
arr = img_like
|
|
else:
|
|
raise TypeError(f"Expected torch.Tensor or np.ndarray, got {type(img_like)}")
|
|
|
|
# scale to 0..255
|
|
if arr.dtype != np.uint8 or arr.max() <= 1.5:
|
|
arr = np.clip(arr * (255.0 if arr.max() <= 1.5 else 1.0), 0, 255).astype(np.uint8)
|
|
return arr
|
|
|
|
|
|
def _to_pil(img_like) -> Image.Image:
|
|
arr = _tensor_like_to_numpy(img_like)
|
|
if Image is None:
|
|
raise RuntimeError("Pillow is required.")
|
|
if arr.ndim == 2:
|
|
return Image.fromarray(arr, "L")
|
|
c = arr.shape[2] if arr.ndim == 3 else 1
|
|
if c == 1:
|
|
return Image.fromarray(arr[:, :, 0], "L")
|
|
if c == 3:
|
|
return Image.fromarray(arr, "RGB")
|
|
if c == 4:
|
|
return Image.fromarray(arr, "RGBA")
|
|
return Image.fromarray(arr[:, :, :3], "RGB")
|
|
|
|
|
|
def _tensor_to_png_b64(img_like) -> str:
|
|
pil = _to_pil(img_like)
|
|
buf = BytesIO()
|
|
pil.save(buf, format="PNG")
|
|
return base64.b64encode(buf.getvalue()).decode("utf-8")
|
|
|
|
|
|
def _prepare_mask_png_b64(
|
|
mask_like=None,
|
|
mask_image_like=None,
|
|
target_size=None,
|
|
invert: bool = False,
|
|
feather_radius: int = 0,
|
|
strength: float = 1.0,
|
|
) -> Optional[str]:
|
|
"""
|
|
Build an 8-bit 'L' mask PNG b64.
|
|
- White (255) = editable
|
|
- Black (0) = preserved
|
|
Optional transforms: invert, feather (GaussianBlur), strength (0..1 linear scale).
|
|
"""
|
|
if mask_like is None and mask_image_like is None:
|
|
return None
|
|
src = mask_like if mask_like is not None else mask_image_like
|
|
m = _to_pil(src).convert("L")
|
|
if target_size is not None and m.size != target_size:
|
|
m = m.resize(target_size, resample=Image.NEAREST)
|
|
if invert:
|
|
m = ImageOps.invert(m)
|
|
if feather_radius and feather_radius > 0:
|
|
m = m.filter(ImageFilter.GaussianBlur(radius=int(feather_radius)))
|
|
if strength is not None:
|
|
s = max(0.0, min(1.0, float(strength)))
|
|
if s < 1.0:
|
|
# linear scale towards 0 (stronger protection)
|
|
m = m.point(lambda px: int(px * s))
|
|
buf = BytesIO()
|
|
m.save(buf, format="PNG")
|
|
return base64.b64encode(buf.getvalue()).decode("utf-8")
|
|
|
|
|
|
def _extract_base64_image(resp_json: Dict[str, Any]) -> Optional[str]:
|
|
def _walk(n):
|
|
if isinstance(n, dict):
|
|
for k, v in n.items():
|
|
yield k, v
|
|
yield from _walk(v)
|
|
elif isinstance(n, list):
|
|
for i in n:
|
|
yield from _walk(i)
|
|
for k, v in _walk(resp_json):
|
|
if k in ("inline_data", "inlineData") and isinstance(v, dict):
|
|
data = v.get("data")
|
|
mime = v.get("mime_type") or v.get("mimeType") or ""
|
|
if data and ("image/" in mime or mime in ("", None)):
|
|
return data
|
|
return None
|
|
|
|
|
|
def _post_with_backoff(url: str, headers: Dict[str, str], params: Dict[str, Any], payload: Dict[str, Any]):
|
|
delay = 1.0
|
|
last = None
|
|
for _ in range(4):
|
|
r = requests.post(url, headers=headers, params=params, json=payload, timeout=180)
|
|
last = r
|
|
if r.status_code != 429:
|
|
return r
|
|
time.sleep(delay)
|
|
delay *= 2
|
|
return last
|
|
|
|
|
|
MODEL_CHOICES = (
|
|
"gemini-3-pro-image-preview",
|
|
"gemini-2.5-flash-image",
|
|
)
|
|
|
|
|
|
class FSLGeminiGenerateImage:
|
|
"""Gemini image generator with inpainting support, quality, and aspect ratio controls via prompt injection.
|
|
Base selection: auto / init_image / images.
|
|
Mask inputs: MASK or IMAGE, with invert/feather/strength options.
|
|
"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls) -> Dict[str, Any]:
|
|
tooltip_wh = "These values are used for metadata only. Actual image size is determined by 'aspect_ratio' and API limits (~1K max)."
|
|
return {
|
|
"required": {
|
|
"prompt": ("STRING", {"multiline": True, "default": ""}),
|
|
"model": (MODEL_CHOICES, {"default": "gemini-3-pro-image-preview"}),
|
|
"width": ("INT", {"default": 1024, "min": 64, "max": 2048, "step": 64, "tooltip": tooltip_wh}),
|
|
"height": ("INT", {"default": 1024, "min": 64, "max": 2048, "step": 64, "tooltip": tooltip_wh}),
|
|
"aspect_ratio": (ASPECT_CHOICES, {"default": "1:1", "tooltip": "Request this aspect ratio via prompt injection."}),
|
|
"quality": (QUALITY_CHOICES, {"default": "Standard", "tooltip": "Standard requests 1K-style resolution. High requests 4K-style resolution via prompt injection."}),
|
|
"temperature": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.1}),
|
|
"seed": ("INT", {"default": -1, "min": -1, "max": 0x7FFFFFFF}),
|
|
"enable_safety": ("BOOLEAN", {"default": False, "tooltip": "If true, enables default API safety filters. If false, attempts to disable safety filters (BLOCK_NONE)."}),
|
|
"use_init_image": ("BOOLEAN", {"default": True}),
|
|
"use_images": ("BOOLEAN", {"default": False}),
|
|
"use_mask": ("BOOLEAN", {"default": False}),
|
|
"base_for_inpaint": (("auto", "init_image", "images"), {"default": "auto"}),
|
|
"output_prompt": ("BOOLEAN", {"default": False}),
|
|
},
|
|
"optional": {
|
|
"init_image": ("IMAGE",),
|
|
"images": ("IMAGE",),
|
|
"mask": ("MASK",), # true mask
|
|
"mask_image": ("IMAGE",), # image as mask
|
|
"invert_mask": ("BOOLEAN", {"default": False}),
|
|
"feather_radius": ("INT", {"default": 0, "min": 0, "max": 64, "step": 1}),
|
|
"mask_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.05}),
|
|
"api_key": ("STRING", {"default": "", "multiline": False, "placeholder": "Leave blank to use GEMINI_API_KEY env var"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "STRING", "STRING", "INT", "INT", "FLOAT", "INT", "STRING", "DICT")
|
|
RETURN_NAMES = ("image", "prompt", "model", "width", "height", "temperature", "seed", "seed_mode", "metadata")
|
|
FUNCTION = "generate"
|
|
CATEGORY = "FSL"
|
|
|
|
# ---- internals ----
|
|
def _resolve_api_key(self, k: str) -> str:
|
|
k = (k or "").strip() or os.environ.get("GEMINI_API_KEY", "").strip()
|
|
if not k:
|
|
raise RuntimeError("Gemini API key missing. Provide it in the node or set GEMINI_API_KEY.")
|
|
return k
|
|
|
|
def _append_tensor_as_part(self, img_like, out: List[Dict[str, Any]]):
|
|
out.append({"inline_data": {"mime_type": "image/png", "data": _tensor_to_png_b64(img_like)}})
|
|
|
|
def _iter_imgs(self, imgs):
|
|
if imgs is None:
|
|
return []
|
|
if isinstance(imgs, (list, tuple)):
|
|
return list(imgs)
|
|
if torch is not None and isinstance(imgs, torch.Tensor):
|
|
if imgs.dim() == 4:
|
|
return [imgs[i] for i in range(imgs.shape[0])]
|
|
return [imgs]
|
|
if isinstance(imgs, np.ndarray):
|
|
return [imgs]
|
|
return []
|
|
|
|
def _collect_parts_with_base_choice(
|
|
self,
|
|
base_choice: str,
|
|
use_init_image: bool, init_image,
|
|
use_images: bool, images,
|
|
use_mask: bool, mask, mask_image,
|
|
invert_mask: bool, feather_radius: int, mask_strength: float,
|
|
) -> Tuple[List[Dict[str, Any]], Optional[Tuple[int, int]], str]:
|
|
"""
|
|
Returns (parts, base_size, base_source)
|
|
base_source in {"init_image","images","none"}
|
|
"""
|
|
parts: List[Dict[str, Any]] = []
|
|
base_size = None
|
|
base_source = "none"
|
|
|
|
imgs_list = self._iter_imgs(images)
|
|
|
|
if base_choice == "init_image" and use_init_image and init_image is not None:
|
|
# base = init_image
|
|
if isinstance(init_image, (list, tuple)) and len(init_image) > 0:
|
|
pil = _to_pil(init_image[0]); base_size = pil.size; self._append_tensor_as_part(init_image[0], parts)
|
|
elif torch is not None and isinstance(init_image, torch.Tensor):
|
|
t = init_image[0] if (init_image.dim() == 4 and init_image.shape[0] > 0) else init_image
|
|
pil = _to_pil(t); base_size = pil.size; self._append_tensor_as_part(t, parts)
|
|
elif isinstance(init_image, np.ndarray):
|
|
pil = _to_pil(init_image); base_size = pil.size; self._append_tensor_as_part(init_image, parts)
|
|
base_source = "init_image"
|
|
|
|
# mask after base
|
|
if use_mask and (mask is not None or mask_image is not None):
|
|
b64 = _prepare_mask_png_b64(mask, mask_image, base_size, invert_mask, feather_radius, mask_strength)
|
|
if b64:
|
|
parts.append({"inline_data": {"mime_type": "image/png", "data": b64}})
|
|
|
|
# remaining refs
|
|
if use_images and imgs_list:
|
|
for t in imgs_list:
|
|
self._append_tensor_as_part(t, parts)
|
|
return parts, base_size, base_source
|
|
|
|
if base_choice == "images" and use_images and imgs_list:
|
|
# promote first image as base
|
|
first = imgs_list[0]
|
|
pil = _to_pil(first); base_size = pil.size; self._append_tensor_as_part(first, parts)
|
|
base_source = "images"
|
|
|
|
# mask after base
|
|
if use_mask and (mask is not None or mask_image is not None):
|
|
b64 = _prepare_mask_png_b64(mask, mask_image, base_size, invert_mask, feather_radius, mask_strength)
|
|
if b64:
|
|
# Keep the structure consistent
|
|
parts.append({"inline_data": {"mime_type": "image/png", "data": b64}} )
|
|
|
|
# rest as refs
|
|
for t in imgs_list[1:]:
|
|
self._append_tensor_as_part(t, parts)
|
|
return parts, base_size, base_source
|
|
|
|
# no valid base selected — still push refs if present
|
|
if use_images and imgs_list:
|
|
for t in imgs_list:
|
|
if base_size is None:
|
|
base_size = _to_pil(t).size
|
|
self._append_tensor_as_part(t, parts)
|
|
|
|
# If mask without a base, we can still send it (model may ignore)
|
|
if use_mask and (mask is not None or mask_image is not None):
|
|
b64 = _prepare_mask_png_b64(mask, mask_image, base_size, invert_mask, feather_radius, mask_strength)
|
|
if b64:
|
|
parts.append({"inline_data": {"mime_type": "image/png", "data": b64}})
|
|
|
|
return parts, base_size, base_source
|
|
|
|
# ---- main ----
|
|
def generate(
|
|
self,
|
|
prompt: str, model: str, width: int, height: int, aspect_ratio: str, quality: str, temperature: float, seed: int,
|
|
enable_safety: bool,
|
|
use_init_image: bool, use_images: bool, use_mask: bool, base_for_inpaint: str,
|
|
output_prompt: bool,
|
|
init_image: Optional["torch.Tensor"] = None, images: Optional["torch.Tensor"] = None,
|
|
mask: Optional["torch.Tensor"] = None, mask_image: Optional["torch.Tensor"] = None,
|
|
invert_mask: bool = False, feather_radius: int = 0, mask_strength: float = 1.0,
|
|
api_key: str = "",
|
|
):
|
|
seed_mode = "fixed"
|
|
if seed == -1:
|
|
seed = random.randint(0, 0x7FFFFFFF)
|
|
seed_mode = "randomize"
|
|
|
|
api_key = self._resolve_api_key(api_key)
|
|
endpoint = f"https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent"
|
|
headers = {"Content-Type": "application/json"}
|
|
|
|
# Resolve effective base with AUTO fallback
|
|
requested = base_for_inpaint
|
|
have_init = bool(use_init_image and (init_image is not None))
|
|
have_images = bool(use_images and (images is not None))
|
|
|
|
if requested == "auto":
|
|
base_choice = "init_image" if have_init else ("images" if have_images else "none")
|
|
else:
|
|
if requested == "init_image" and not have_init:
|
|
base_choice = "images" if have_images else "none"
|
|
elif requested == "images" and not have_images:
|
|
base_choice = "init_image" if have_init else "none"
|
|
else:
|
|
base_choice = requested
|
|
|
|
# --- Quality/Aspect Ratio Injection Logic ---
|
|
if quality == "High":
|
|
res_text = "high-fidelity, 4K resolution image"
|
|
else:
|
|
res_text = "detailed, 1K resolution image" # Standard
|
|
|
|
# The new, stronger injection statement, which requests the desired size/aspect ratio
|
|
injection_text = (
|
|
f"\n\n***FOR IMAGE GENERATION:*** Generate this as a **{res_text}** in a **{aspect_ratio}** aspect ratio."
|
|
)
|
|
|
|
final_prompt = prompt.strip() + injection_text
|
|
|
|
# Build parts
|
|
parts: List[Dict[str, Any]] = []
|
|
is_inpainting = use_mask and (mask is not None or mask_image is not None)
|
|
|
|
if is_inpainting:
|
|
parts.append({"text":
|
|
"INPAINT MODE: You are given a base image and a mask image. "
|
|
"Modify ONLY pixels where the mask is white; keep black (unmasked) pixels exactly the same as the base image. "
|
|
"Respect edges and composition."
|
|
})
|
|
|
|
# Append the final prompt (with injection) as the main text part
|
|
parts.append({"text": final_prompt})
|
|
|
|
|
|
ref_parts, base_size, resolved_base = self._collect_parts_with_base_choice(
|
|
base_choice=base_choice,
|
|
use_init_image=use_init_image, init_image=init_image,
|
|
use_images=use_images, images=images,
|
|
use_mask=use_mask, mask=mask, mask_image=mask_image,
|
|
invert_mask=invert_mask, feather_radius=feather_radius, mask_strength=mask_strength,
|
|
)
|
|
parts.extend(ref_parts)
|
|
|
|
if is_inpainting and resolved_base == "none":
|
|
raise RuntimeError("Inpaint requires a base image. Enable and connect either init_image or images.")
|
|
|
|
payload = {
|
|
"contents": [{"role": "user", "parts": parts}],
|
|
"generationConfig": {"temperature": float(temperature), "seed": int(seed)},
|
|
}
|
|
|
|
# Conditionally disable safety settings
|
|
if not enable_safety:
|
|
payload["safetySettings"] = [
|
|
{"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"},
|
|
{"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE"},
|
|
{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_NONE"},
|
|
{"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_NONE"},
|
|
]
|
|
|
|
resp = _post_with_backoff(endpoint, headers=headers, params={"key": api_key}, payload=payload)
|
|
|
|
# Parse response
|
|
try:
|
|
resp_json = resp.json()
|
|
except Exception as e:
|
|
raise RuntimeError(f"Gemini response was not JSON: {e}")
|
|
|
|
if isinstance(resp_json, dict) and "error" in resp_json:
|
|
err = resp_json["error"]
|
|
raise RuntimeError(f"Gemini error: {err.get('status')} - {err.get('message')}")
|
|
|
|
b64_img = _extract_base64_image(resp_json)
|
|
if not b64_img:
|
|
finish_reason = resp_json.get('candidates', [{}])[0].get('finishReason', 'UNKNOWN')
|
|
raise RuntimeError(f"Failed to find inline image in Gemini response. Finish Reason: {finish_reason}. Submitted Prompt: {final_prompt[:250]}... First 1k of JSON: {str(resp_json)[:1000]}")
|
|
|
|
|
|
try:
|
|
img_bytes = base64.b64decode(b64_img)
|
|
if Image is None:
|
|
raise RuntimeError("Pillow is required to decode returned image.")
|
|
pil = Image.open(BytesIO(img_bytes)).convert("RGBA")
|
|
except Exception as e:
|
|
raise RuntimeError(f"Failed to decode inline image: {e}")
|
|
|
|
np_img = np.array(pil).astype(np.float32) / 255.0
|
|
if torch is None:
|
|
raise RuntimeError("torch is required for IMAGE tensor output.")
|
|
image_tensor = torch.from_numpy(np_img)[None, ...]
|
|
|
|
out_prompt = final_prompt if output_prompt else ""
|
|
metadata = {
|
|
"prompt": prompt, "model": model, "width": width, "height": height, "aspect_ratio": aspect_ratio, "quality": quality,
|
|
"temperature": temperature, "seed": seed, "seed_mode": seed_mode,
|
|
"enable_safety": enable_safety,
|
|
"use_init_image": bool(use_init_image and init_image is not None),
|
|
"use_images": bool(use_images and images is not None),
|
|
"use_mask": bool(use_mask and (mask is not None or mask_image is not None)),
|
|
"mask_source": "MASK" if mask is not None else ("IMAGE" if mask_image is not None else "none"),
|
|
"requested_base_for_inpaint": base_for_inpaint,
|
|
"effective_base_for_inpaint": base_choice,
|
|
"resolved_base_source": resolved_base,
|
|
"invert_mask": bool(invert_mask),
|
|
"feather_radius": int(feather_radius),
|
|
"mask_strength": float(mask_strength),
|
|
}
|
|
return (image_tensor, out_prompt, model, width, height, temperature, seed, seed_mode, metadata)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {"FSLGeminiGenerateImage": FSLGeminiGenerateImage}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {"FSLGeminiGenerateImage": "FSL Gemini Generate Image"} |