forge/backend/scripts/test_nano_banana_simple.py

76 lines
2.5 KiB
Python

import asyncio
import os
import sys
import logging
from dotenv import load_dotenv
# Add backend to path
sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
load_dotenv()
# Mock settings
class Settings:
google_api_key = os.getenv("GOOGLE_API_KEY")
import app.services.image_generator as img_gen
img_gen.settings = Settings()
# Configure logging
logging.basicConfig(level=logging.INFO)
async def test_generation():
prompt = "A cyberpunk street scene"
# 1. First generate an image (or use a dummy one if we had one, but let's generate)
print("--- Testing Generation (Text to Image) ---")
input_data = {
"prompt": prompt,
"model": "gemini-2.0-flash-exp" # Try 2.0 first
}
# We can't easily call _generate_nano_banana without mocking httpx or having a real key.
# Assuming the user has a real key since they are getting results.
try:
image_data, filename = await img_gen._generate_nano_banana(input_data)
if image_data:
print(f"✓ Generated base image: {len(image_data)} bytes")
# 2. Now try to 'Edit' it
print("\n--- Testing Edit (Image + Text to Image) ---")
edit_prompt = "Make it daytime"
# Test Gemini 3 Variants
variants = [
"gemini-3.0-pro-image-preview",
"gemini-3.0-pro-image",
"gemini-3.0-flash-image",
"gemini-3.0-flash-image-preview"
]
for model_name in variants:
print(f"\n--- Testing Model: {model_name} ---")
input_data_test = {
"prompt": edit_prompt,
"model": model_name
}
try:
# We pass 'image/png' because our fake input_image is implied to be handled
img, fname = await img_gen._generate_nano_banana(input_data_test, image_data, "image/png")
if img:
print(f"✓ Success with {model_name}: {len(img)} bytes")
break # Stop on first success just to know valid one
else:
print(f"✗ Failed with {model_name} (No data)")
except Exception as e:
print(f"✗ Error with {model_name}: {e}")
else:
print("✗ Initial generation failed")
except Exception as e:
print(f"✗ Error: {e}")
if __name__ == "__main__":
asyncio.run(test_generation())