diff --git a/app/data/action/generate_image.py b/app/data/action/generate_image.py index 2832bf10..91d65015 100644 --- a/app/data/action/generate_image.py +++ b/app/data/action/generate_image.py @@ -3,11 +3,12 @@ @action( name="generate_image", - description="""Generates an image using Google's Nano Banana Pro (Gemini 3 Pro Image) model. -- State-of-the-art image generation with 1K, 2K, or 4K resolution support -- Excellent text rendering for infographics, menus, diagrams -- Uses GOOGLE_API_KEY environment variable (same as Gemini LLM provider) -- If API key is not set, returns an error with setup instructions + description="""Generates an image using either OpenAI's Images 2.0 (gpt-image-2) or Google's Nano Banana 2 (gemini-3.1-flash-image-preview) model. +- Automatically selects the provider based on which API key(s) are configured +- If only one API key is set, that provider is used automatically +- If both keys are configured, asks the user which provider to use and remembers the choice +- If no API keys are configured, returns an error with setup instructions +- Supports 1K, 2K, or 4K resolution and multiple aspect ratios - TIP: When generating multiple images for the same project or related work, use 'reference_images' parameter with previously generated images to maintain consistent style across all outputs""", default=True, mode="CLI", @@ -55,7 +56,12 @@ "safety_filter_level": { "type": "string", "example": "block_medium_and_above", - "description": "Safety filter level. Options: 'block_none', 'block_only_high', 'block_medium_and_above', 'block_low_and_above'. Default: 'block_medium_and_above'.", + "description": "Safety filter level (Gemini only). Options: 'block_none', 'block_only_high', 'block_medium_and_above', 'block_low_and_above'. Default: 'block_medium_and_above'. Ignored when using OpenAI.", + }, + "provider_preference": { + "type": "string", + "example": "openai", + "description": "Which provider to use: 'openai' (Images 2.0 / gpt-image-2) or 'gemini' (Nano Banana 2 / gemini-3.1-flash-image-preview). Only needed when both API keys are configured and no saved preference exists. Providing this saves it as the default for future calls.", }, }, output_schema={ @@ -81,7 +87,7 @@ "description": "Status message or error message.", }, }, - requirement=["google-genai", "Pillow"], + requirement=["google-genai", "openai", "Pillow"], test_payload={ "prompt": "A cute cartoon cat sitting on a rainbow", "resolution": "1K", @@ -92,7 +98,7 @@ ) def generate_image(input_data: dict) -> dict: """ - Generates an image using Google's Nano Banana Pro (Gemini 3 Pro Image) model. + Generates an image using OpenAI's Images 2.0 (gpt-image-2) or Google's Nano Banana 2 (gemini-3.1-flash-image-preview). """ import os import sys @@ -112,19 +118,47 @@ def generate_image(input_data: dict) -> dict: "message": "Image generated successfully (simulated mode).", } - # Pre-flight validation: check API key is configured - from app.config import get_api_key + # Determine which provider to use based on available API keys and user preference + from app.config import get_api_key, get_settings, save_settings + + openai_key = get_api_key("openai") + gemini_key = get_api_key("gemini") - api_key = get_api_key("gemini") - if not api_key: + if not openai_key and not gemini_key: return { "status": "error", "image_paths": [], "prompt_used": "", "resolution": "", - "message": "Gemini API key is not configured. Tell the user the Google Gemini API key is required for image generation, and ask if they need help setting it up.", + "message": ( + "No image generation API key is configured. " + "Tell the user they need either an OpenAI API key (for Images 2.0 / gpt-image-2) " + "or a Google Gemini API key (for Nano Banana 2 / gemini-3.1-flash-image-preview), " + "and ask if they need help setting one up." + ), } + provider_preference = input_data.get("provider_preference", "").strip().lower() + _cfg = get_settings() + saved_provider = _cfg.get("image_generation", {}).get("preferred_provider", "") + + if openai_key and not gemini_key: + provider = "openai" + elif gemini_key and not openai_key: + provider = "gemini" + else: + # Both keys present + if provider_preference in ("openai", "gemini"): + provider = provider_preference + _cfg.setdefault("image_generation", {})["preferred_provider"] = provider + save_settings(_cfg) + elif saved_provider in ("openai", "gemini"): + provider = saved_provider + else: + provider = "gemini" + + api_key = openai_key if provider == "openai" else gemini_key + # Validate required input prompt = input_data.get("prompt", "").strip() if not prompt: @@ -273,6 +307,7 @@ def _ensure_package(pkg_name): try: _ensure_package("google-genai") + _ensure_package("openai") _ensure_package("Pillow") except Exception as e: return { @@ -284,38 +319,43 @@ def _ensure_package(pkg_name): } try: - from google import genai - from google.genai import types from PIL import Image import io import base64 - client = genai.Client(api_key=api_key) - - # Prepare reference images if provided - image_parts = [] - for ref_path in reference_images: - if os.path.exists(ref_path): - try: - with open(ref_path, "rb") as f: - image_data = f.read() - ext = os.path.splitext(ref_path)[1].lower() - mime_map = { - ".png": "image/png", - ".jpg": "image/jpeg", - ".jpeg": "image/jpeg", - ".gif": "image/gif", - ".webp": "image/webp", - } - mime_type = mime_map.get(ext, "image/png") - image_parts.append( - types.Part.from_bytes(data=image_data, mime_type=mime_type) - ) - except Exception: - pass # Skip invalid reference images + image_paths = [] + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") - # Build the prompt with generation instructions - generation_prompt = f"""Generate an image based on the following description: + if provider == "gemini": + from google import genai + from google.genai import types + + client = genai.Client(api_key=api_key) + + # Prepare reference images if provided + image_parts = [] + for ref_path in reference_images: + if os.path.exists(ref_path): + try: + with open(ref_path, "rb") as f: + image_data = f.read() + ext = os.path.splitext(ref_path)[1].lower() + mime_map = { + ".png": "image/png", + ".jpg": "image/jpeg", + ".jpeg": "image/jpeg", + ".gif": "image/gif", + ".webp": "image/webp", + } + mime_type = mime_map.get(ext, "image/png") + image_parts.append( + types.Part.from_bytes(data=image_data, mime_type=mime_type) + ) + except Exception: + pass # Skip invalid reference images + + # Build the prompt with generation instructions + generation_prompt = f"""Generate an image based on the following description: {prompt} @@ -324,89 +364,161 @@ def _ensure_package(pkg_name): - Aspect ratio: {aspect_ratio} - Number of variations: {number_of_images}""" - if negative_prompt: - generation_prompt += f"\n- Avoid: {negative_prompt}" - - content_parts = list(image_parts) - content_parts.append(generation_prompt) - - # Safety settings - safety_settings = None - if safety_filter_level != "block_none": - harm_block_threshold = { - "block_only_high": "BLOCK_ONLY_HIGH", - "block_medium_and_above": "BLOCK_MEDIUM_AND_ABOVE", - "block_low_and_above": "BLOCK_LOW_AND_ABOVE", - }.get(safety_filter_level, "BLOCK_MEDIUM_AND_ABOVE") - - safety_settings = [ - types.SafetySetting(category=category, threshold=harm_block_threshold) - for category in ( - "HARM_CATEGORY_HARASSMENT", - "HARM_CATEGORY_HATE_SPEECH", - "HARM_CATEGORY_SEXUALLY_EXPLICIT", - "HARM_CATEGORY_DANGEROUS_CONTENT", + if negative_prompt: + generation_prompt += f"\n- Avoid: {negative_prompt}" + + content_parts = list(image_parts) + content_parts.append(generation_prompt) + + # Safety settings + safety_settings = None + if safety_filter_level != "block_none": + harm_block_threshold = { + "block_only_high": "BLOCK_ONLY_HIGH", + "block_medium_and_above": "BLOCK_MEDIUM_AND_ABOVE", + "block_low_and_above": "BLOCK_LOW_AND_ABOVE", + }.get(safety_filter_level, "BLOCK_MEDIUM_AND_ABOVE") + + safety_settings = [ + types.SafetySetting(category=category, threshold=harm_block_threshold) + for category in ( + "HARM_CATEGORY_HARASSMENT", + "HARM_CATEGORY_HATE_SPEECH", + "HARM_CATEGORY_SEXUALLY_EXPLICIT", + "HARM_CATEGORY_DANGEROUS_CONTENT", + ) + ] + + generate_config = types.GenerateContentConfig( + candidate_count=1, + response_modalities=["TEXT", "IMAGE"], + image_config=types.ImageConfig(image_size=resolution), + safety_settings=safety_settings, + ) + + response = client.models.generate_content( + model="gemini-3.1-flash-image-preview", + contents=content_parts, + config=generate_config, + ) + + images_found = _extract_images_from_response(response) + + if not images_found: + block_reason = _get_block_reason(response) + if block_reason: + return { + "status": "error", + "image_paths": [], + "prompt_used": prompt, + "resolution": resolution, + "message": f"Image generation was blocked by safety filters: {block_reason}. Try modifying your prompt or adjusting safety_filter_level.", + } + return { + "status": "error", + "image_paths": [], + "prompt_used": prompt, + "resolution": resolution, + "message": "No images were generated. The model did not produce image output for this prompt. Try rephrasing your prompt or check if your API key has access to image generation.", + } + + for i, img_data in enumerate(images_found[:number_of_images]): + save_path = _build_save_path( + output_path, timestamp, i, number_of_images, len(images_found) ) - ] + parent_dir = os.path.dirname(os.path.abspath(save_path)) + if parent_dir: + os.makedirs(parent_dir, exist_ok=True) + pil_image = _to_pil_image(img_data, Image, io, base64) + pil_image.save(save_path, "PNG") + image_paths.append(save_path) - generate_config = types.GenerateContentConfig( - candidate_count=1, - response_modalities=["TEXT", "IMAGE"], - image_config=types.ImageConfig(image_size=resolution), - safety_settings=safety_settings, - ) + model_label = "Nano Banana 2" - response = client.models.generate_content( - model="gemini-3-pro-image-preview", - contents=content_parts, - config=generate_config, - ) + elif provider == "openai": + from openai import OpenAI - # Extract images from response - image_paths = [] - timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + if safety_filter_level != "block_medium_and_above": + warnings.append( + "safety_filter_level is not supported by OpenAI and has been ignored." + ) + + # Map aspect_ratio to OpenAI size string + openai_size_map = { + "1:1": "1024x1024", + "16:9": "1536x1024", + "4:3": "1536x1024", + "9:16": "1024x1536", + "3:4": "1024x1536", + } + openai_size = openai_size_map.get(aspect_ratio, "1024x1024") - # Process response to find generated images - images_found = _extract_images_from_response(response) + # Map resolution to OpenAI quality + openai_quality_map = {"1K": "medium", "2K": "high", "4K": "high"} + openai_quality = openai_quality_map.get(resolution, "medium") - if not images_found: - # Check if response was blocked by safety filters - block_reason = _get_block_reason(response) - if block_reason: + # Build prompt (OpenAI has no negative_prompt param — append to prompt) + full_prompt = prompt + if negative_prompt: + full_prompt += f"\n\nAvoid: {negative_prompt}" + + client = OpenAI(api_key=api_key) + + valid_ref_paths = [p for p in reference_images if os.path.exists(p)] + if valid_ref_paths: + image_files = [open(p, "rb") for p in valid_ref_paths] + try: + response = client.images.edit( + model="gpt-image-2", + image=image_files, + prompt=full_prompt, + n=number_of_images, + size=openai_size, + ) + finally: + for f in image_files: + f.close() + else: + response = client.images.generate( + model="gpt-image-2", + prompt=full_prompt, + n=number_of_images, + size=openai_size, + quality=openai_quality, + ) + + import urllib.request as _urllib_request + images_found = [] + for item in response.data: + if item.b64_json: + images_found.append(base64.b64decode(item.b64_json)) + elif item.url: + with _urllib_request.urlopen(item.url) as _r: + images_found.append(_r.read()) + + if not images_found: return { "status": "error", "image_paths": [], "prompt_used": prompt, "resolution": resolution, - "message": f"Image generation was blocked by safety filters: {block_reason}. Try modifying your prompt or adjusting safety_filter_level.", + "message": "No images were generated. The model did not produce image output for this prompt. Try rephrasing your prompt.", } - return { - "status": "error", - "image_paths": [], - "prompt_used": prompt, - "resolution": resolution, - "message": "No images were generated. The model did not produce image output for this prompt. Try rephrasing your prompt or check if your API key has access to image generation.", - } - # Save each generated image - for i, img_data in enumerate(images_found[:number_of_images]): - save_path = _build_save_path( - output_path, timestamp, i, number_of_images, len(images_found) - ) - - # Ensure parent directory exists - parent_dir = os.path.dirname(os.path.abspath(save_path)) - if parent_dir: - os.makedirs(parent_dir, exist_ok=True) + for i, img_bytes in enumerate(images_found[:number_of_images]): + save_path = _build_save_path( + output_path, timestamp, i, number_of_images, len(images_found) + ) + parent_dir = os.path.dirname(os.path.abspath(save_path)) + if parent_dir: + os.makedirs(parent_dir, exist_ok=True) + pil_image = Image.open(io.BytesIO(img_bytes)) + pil_image.save(save_path, "PNG") + image_paths.append(save_path) - # Save the image - pil_image = _to_pil_image(img_data, Image, io, base64) - pil_image.save(save_path, "PNG") - image_paths.append(save_path) + model_label = "Images 2.0" - message = ( - f"Successfully generated {len(image_paths)} image(s) using Nano Banana Pro." - ) + message = f"Successfully generated {len(image_paths)} image(s) using {model_label}." if warnings: message += " Warnings: " + " ".join(warnings) @@ -421,17 +533,26 @@ def _ensure_package(pkg_name): except Exception as e: error_message = str(e) - # Provide more helpful error messages - if "quota" in error_message.lower() or "rate" in error_message.lower(): - error_message = f"API rate limit or quota exceeded: {error_message}" - elif "invalid" in error_message.lower() and "key" in error_message.lower(): - error_message = f"Invalid API key: {error_message}. Please verify your GOOGLE_API_KEY is correct." - elif "permission" in error_message.lower() or "access" in error_message.lower(): - error_message = f"API access denied: {error_message}. Ensure your API key has access to Nano Banana Pro model." - elif "safety" in error_message.lower() or "blocked" in error_message.lower(): - error_message = f"Content blocked by safety filters: {error_message}. Try modifying your prompt." - elif "not found" in error_message.lower() or "404" in error_message: - error_message = f"Model not available: {error_message}. The gemini-3-pro-image-preview model may not be accessible with your API key. Try using Google AI Studio to verify access." + if provider == "gemini": + if "quota" in error_message.lower() or "rate" in error_message.lower(): + error_message = f"Gemini API rate limit or quota exceeded: {error_message}" + elif "invalid" in error_message.lower() and "key" in error_message.lower(): + error_message = f"Invalid Gemini API key: {error_message}. Please verify your Google API key is correct." + elif "permission" in error_message.lower() or "access" in error_message.lower(): + error_message = f"Gemini API access denied: {error_message}. Ensure your API key has access to the Nano Banana 2 model." + elif "safety" in error_message.lower() or "blocked" in error_message.lower(): + error_message = f"Content blocked by Gemini safety filters: {error_message}. Try modifying your prompt." + elif "not found" in error_message.lower() or "404" in error_message: + error_message = f"Gemini model not available: {error_message}. The gemini-3.1-flash-image-preview model may not be accessible with your API key. Try using Google AI Studio to verify access." + else: + if "billing" in error_message.lower() or "insufficient_quota" in error_message.lower() or "rate" in error_message.lower(): + error_message = f"OpenAI API rate limit or quota exceeded: {error_message}" + elif "invalid_api_key" in error_message.lower() or ("invalid" in error_message.lower() and "key" in error_message.lower()): + error_message = f"Invalid OpenAI API key: {error_message}. Please verify your OpenAI API key is correct." + elif "content_policy" in error_message.lower() or "safety" in error_message.lower() or "blocked" in error_message.lower(): + error_message = f"Content blocked by OpenAI safety policy: {error_message}. Try modifying your prompt." + elif "not found" in error_message.lower() or "404" in error_message: + error_message = f"OpenAI model not available: {error_message}. The gpt-image-2 model may not be accessible with your API key." return { "status": "error",