#!/usr/bin/env python3 """ Genesis Imagen Generator ======================== Integrates Imagen 4.0 for image generation capabilities. Pricing: - Imagen 4: $0.04/image - Imagen 4 Ultra: $0.06/image - Imagen 4 Fast: $0.02/image Usage: from imagen_generator import ImagenAgent agent = ImagenAgent() result = agent.generate("A futuristic dashboard") """ import os import sys import json import base64 from datetime import datetime from pathlib import Path from typing import Optional, Dict, Any, List # Add genesis-system to path sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) from budget_manager import BudgetManager # Try the new google.genai SDK first, fall back to deprecated SDK try: from google import genai from google.genai import types NEW_SDK = True print("[OK] Using google.genai SDK (recommended)") except ImportError: import google.generativeai as genai NEW_SDK = False print("[WARN] Using deprecated google.generativeai SDK") # Load config CONFIG_PATH = os.path.join(os.path.dirname(__file__), "genesis_config.json") with open(CONFIG_PATH) as f: CONFIG = json.load(f) # Pricing constants IMAGEN_PRICING = { "imagen-4.0-generate-001": 0.04, "imagen-4.0-ultra-generate-001": 0.06, "imagen-4.0-fast-generate-001": 0.02 } # Output directory OUTPUT_DIR = Path(os.path.dirname(__file__)) / "generated_images" OUTPUT_DIR.mkdir(exist_ok=True) class ImagenAgent: """ Agent for generating images using Imagen 4.0. Supports multiple model variants and aspect ratios. """ def __init__(self, model: str = "imagen-4.0-generate-001", budget_limit: float = 10.0): self.model_name = model self.budget = BudgetManager(daily_limit=budget_limit) self.cost_per_image = IMAGEN_PRICING.get(model, 0.04) # Configure API api_key = CONFIG["gemini"]["api_key"] if NEW_SDK: self.client = genai.Client(api_key=api_key) else: genai.configure(api_key=api_key) self.client = None print(f"[OK] Imagen Agent initialized") print(f"[OK] Model: {model}") print(f"[OK] Cost: ${self.cost_per_image}/image") def generate(self, prompt: str, aspect_ratio: str = "1:1", num_images: int = 1, save_to_disk: bool = True) -> Dict[str, Any]: """ Generate image(s) from a text prompt. Args: prompt: Description of the image to generate aspect_ratio: "1:1", "16:9", "9:16", "4:3", "3:4" num_images: Number of images (1-4) save_to_disk: Whether to save images locally Returns: Dict with 'images', 'paths', 'cost', 'timestamp' """ # Budget check total_cost = self.cost_per_image * num_images if not self.budget.is_within_budget(): return { "images": [], "paths": [], "cost": 0, "error": "budget_exceeded" } try: if NEW_SDK: # Use new google.genai SDK response = self.client.models.generate_images( model=self.model_name, prompt=prompt, config=types.GenerateImagesConfig( number_of_images=num_images, aspect_ratio=aspect_ratio, person_generation="ALLOW_ALL" ) ) images_data = [] for img in response.generated_images: images_data.append(img.image.image_bytes) else: # Fallback: Use deprecated SDK with Gemini model # Note: Direct Imagen generation may not be available model = genai.GenerativeModel('gemini-2.0-flash') response = model.generate_content( f"Generate a detailed description for an image: {prompt}" ) # Return description instead of actual image return { "images": [], "description": response.text, "paths": [], "cost": 0.0001, # Minimal cost for text "note": "Image generation requires google.genai SDK", "timestamp": datetime.now().isoformat() } # Log cost self.budget.log_cost(total_cost, self.model_name) # Save to disk if requested paths = [] if save_to_disk and images_data: for i, img_bytes in enumerate(images_data): timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") filename = f"imagen_{timestamp}_{i}.png" filepath = OUTPUT_DIR / filename with open(filepath, "wb") as f: f.write(img_bytes) paths.append(str(filepath)) print(f"[SAVED] {filepath}") return { "images": [base64.b64encode(img).decode() for img in images_data], "paths": paths, "cost": total_cost, "prompt": prompt, "model": self.model_name, "timestamp": datetime.now().isoformat() } except Exception as e: return { "images": [], "paths": [], "cost": 0, "error": str(e), "timestamp": datetime.now().isoformat() } def generate_variations(self, prompt: str, variations: List[str], aspect_ratio: str = "1:1") -> List[Dict]: """ Generate multiple image variations based on a base prompt. Args: prompt: Base prompt variations: List of variation modifiers aspect_ratio: Aspect ratio for all images Returns: List of generation results """ results = [] for var in variations: full_prompt = f"{prompt}, {var}" result = self.generate(full_prompt, aspect_ratio=aspect_ratio) result["variation"] = var results.append(result) return results def get_stats(self) -> Dict[str, Any]: """Get agent statistics.""" return { "model": self.model_name, "cost_per_image": self.cost_per_image, "budget_spent": self.budget.get_current_spend(), "budget_limit": self.budget.daily_limit, "within_budget": self.budget.is_within_budget(), "output_dir": str(OUTPUT_DIR) } # CLI for testing if __name__ == "__main__": import argparse parser = argparse.ArgumentParser(description="Imagen Generator") parser.add_argument("--generate", "-g", type=str, help="Generate image from prompt") parser.add_argument("--model", "-m", type=str, default="imagen-4.0-generate-001", choices=list(IMAGEN_PRICING.keys()), help="Model variant") parser.add_argument("--aspect", "-a", type=str, default="1:1", help="Aspect ratio") parser.add_argument("--test", action="store_true", help="Run self-test") args = parser.parse_args() agent = ImagenAgent(model=args.model) if args.test: print("\n[TEST] Imagen Generator Self-Test") print(f"[OK] Agent initialized: {agent.get_stats()}") print(f"[OK] Output directory: {OUTPUT_DIR}") print(f"[OK] SDK: {'google.genai (new)' if NEW_SDK else 'google.generativeai (deprecated)'}") print("\n[DONE] Self-test complete") elif args.generate: result = agent.generate(args.generate, aspect_ratio=args.aspect) if result.get("error"): print(f"[ERROR] {result['error']}") elif result.get("paths"): print(f"[OK] Generated {len(result['paths'])} image(s)") for path in result["paths"]: print(f" - {path}") print(f"[COST] ${result['cost']:.4f}") else: print(f"[NOTE] {result.get('note', 'No image generated')}") if result.get("description"): print(f"[DESC] {result['description'][:200]}...") else: print("Usage: python imagen_generator.py --generate 'your prompt'") print(" python imagen_generator.py --test")