""" Test script to generate images for existing articles Tests image generation on project 23: first 2 T1 articles and first 3 T2 articles """ import sys from pathlib import Path sys.path.insert(0, str(Path(__file__).parent.parent)) from src.database.session import db_manager from src.database.repositories import ( ProjectRepository, GeneratedContentRepository, SiteDeploymentRepository ) from src.generation.service import ContentGenerator from src.generation.ai_client import AIClient, PromptManager from src.generation.image_generator import ImageGenerator, truncate_title, slugify from src.generation.image_injection import insert_hero_after_h1, insert_content_images_after_h2s, generate_alt_text from src.generation.image_upload import upload_image_to_storage from src.deployment.bunny_storage import BunnyStorageClient from src.core.config import get_config import click import random from pathlib import Path def test_image_generation(project_id: int): """Test image generation on existing articles""" # Create output directory for test images output_dir = Path("test_images") output_dir.mkdir(exist_ok=True) click.echo(f"Test images will be saved to: {output_dir.absolute()}\n") session = db_manager.get_session() try: # Get repositories project_repo = ProjectRepository(session) content_repo = GeneratedContentRepository(session) site_repo = SiteDeploymentRepository(session) # Get project project = project_repo.get_by_id(project_id) if not project: click.echo(f"Project {project_id} not found") return click.echo(f"\n{'='*60}") click.echo(f"Testing Image Generation for Project {project_id}") click.echo(f"Project: {project.name}") click.echo(f"Main Keyword: {project.main_keyword}") click.echo(f"{'='*60}\n") # Get articles t1_articles = content_repo.get_by_project_and_tier(project_id, "tier1", require_site=False) t2_articles = content_repo.get_by_project_and_tier(project_id, "tier2", require_site=False) click.echo(f"Found {len(t1_articles)} T1 articles, using first 2") click.echo(f"Found {len(t2_articles)} T2 articles, using first 3\n") # Initialize AI client and image generator import os from dotenv import load_dotenv load_dotenv() api_key = os.getenv("OPENROUTER_API_KEY") if not api_key: click.echo("Error: OPENROUTER_API_KEY not set in environment", err=True) return fal_api_key = os.getenv("FAL_API_KEY") if not fal_api_key: click.echo("\n[WARN] FAL_API_KEY not set - image generation will fail") click.echo(" Set FAL_API_KEY in your .env file to test image generation\n") ai_client = AIClient( api_key=api_key, model=os.getenv("AI_MODEL", "gpt-4o-mini") ) prompt_manager = PromptManager() image_generator = ImageGenerator( ai_client=ai_client, prompt_manager=prompt_manager, project_repo=project_repo ) storage_client = BunnyStorageClient() # Test T1 articles (first 2) click.echo(f"\n{'='*60}") click.echo("T1 ARTICLES") click.echo(f"{'='*60}\n") for i, article in enumerate(t1_articles[:2], 1): click.echo(f"\n--- T1 Article {i}: {article.title[:60]}... ---") if not article.site_deployment_id: click.echo(" [WARN] No site assigned, skipping image upload") site = None else: site = site_repo.get_by_id(article.site_deployment_id) if not site: click.echo(" [WARN] Site not found, skipping image upload") site = None # Generate theme prompt (if not exists) click.echo("\n1. Theme Prompt:") if project.image_theme_prompt: click.echo(f" (Using cached): {project.image_theme_prompt}") else: click.echo(" Generating theme prompt...") theme = image_generator.get_theme_prompt(project_id) click.echo(f" Generated: {theme}") # Generate hero image click.echo("\n2. Hero Image:") try: # Show the prompt that will be used theme = image_generator.get_theme_prompt(project_id) click.echo(f" Prompt: {theme}") click.echo(f" Title (will be overlaid): {article.title}") hero_image = image_generator.generate_hero_image( project_id=project_id, title=article.title, width=1280, height=720 ) if hero_image: click.echo(f" [OK] Generated ({len(hero_image):,} bytes)") # Save to local file main_keyword_slug = slugify(project.main_keyword) local_file = output_dir / f"hero-t1-{main_keyword_slug}-{i}.jpg" local_file.write_bytes(hero_image) click.echo(f" [OK] Saved to: {local_file}") if site: file_path = f"images/{main_keyword_slug}.jpg" hero_url = upload_image_to_storage(storage_client, site, hero_image, file_path) if hero_url: click.echo(f" [OK] Uploaded: {hero_url}") else: click.echo(" [FAIL] Upload failed") else: click.echo(" (Skipped upload - no site)") else: click.echo(" [FAIL] Generation failed") except Exception as e: click.echo(f" [ERROR] {str(e)[:200]}") # Generate content images (1-3 for T1) click.echo("\n3. Content Images:") num_content_images = random.randint(1, 3) click.echo(f" Generating {num_content_images} content image(s)...") entities = project.entities or [] related_searches = project.related_searches or [] if not entities or not related_searches: click.echo(" [WARN] No entities/related_searches, skipping") else: for j in range(num_content_images): entity = random.choice(entities) related_search = random.choice(related_searches) click.echo(f"\n Image {j+1}/{num_content_images}:") click.echo(f" Entity: {entity}") click.echo(f" Related Search: {related_search}") try: # Show the prompt that will be used theme = image_generator.get_theme_prompt(project_id) content_prompt = f"{theme} Focus on {entity} and {related_search}, professional illustration style." click.echo(f" Prompt: {content_prompt}") content_image = image_generator.generate_content_image( project_id=project_id, entity=entity, related_search=related_search, width=512, height=512 ) if content_image: click.echo(f" [OK] Generated ({len(content_image):,} bytes)") # Save to local file main_keyword_slug = slugify(project.main_keyword) entity_slug = slugify(entity) related_slug = slugify(related_search) local_file = output_dir / f"content-{main_keyword_slug}-{i}-{j+1}-{entity_slug}-{related_slug}.jpg" local_file.write_bytes(content_image) click.echo(f" [OK] Saved to: {local_file}") if site: file_path = f"images/{main_keyword_slug}-{entity_slug}-{related_slug}.jpg" img_url = upload_image_to_storage(storage_client, site, content_image, file_path) if img_url: click.echo(f" [OK] Uploaded: {img_url}") else: click.echo(" [FAIL] Upload failed") else: click.echo(" (Skipped upload - no site)") else: click.echo(" [FAIL] Generation failed") except Exception as e: click.echo(f" [ERROR] {str(e)[:200]}") # Test T2 articles (first 3) click.echo(f"\n\n{'='*60}") click.echo("T2 ARTICLES") click.echo(f"{'='*60}\n") for i, article in enumerate(t2_articles[:3], 1): click.echo(f"\n--- T2 Article {i}: {article.title[:60]}... ---") if not article.site_deployment_id: click.echo(" [WARN] No site assigned, skipping image upload") site = None else: site = site_repo.get_by_id(article.site_deployment_id) if not site: click.echo(" [WARN] Site not found, skipping image upload") site = None # Generate hero image only (T2 doesn't get content images by default) click.echo("\n1. Hero Image:") try: # Show the prompt that will be used theme = image_generator.get_theme_prompt(project_id) click.echo(f" Prompt: {theme}") click.echo(f" Title (will be overlaid): {article.title}") hero_image = image_generator.generate_hero_image( project_id=project_id, title=article.title, width=1280, height=720 ) if hero_image: click.echo(f" [OK] Generated ({len(hero_image):,} bytes)") # Save to local file main_keyword_slug = slugify(project.main_keyword) local_file = output_dir / f"hero-t2-{main_keyword_slug}-{i}.jpg" local_file.write_bytes(hero_image) click.echo(f" [OK] Saved to: {local_file}") if site: file_path = f"images/{main_keyword_slug}.jpg" hero_url = upload_image_to_storage(storage_client, site, hero_image, file_path) if hero_url: click.echo(f" [OK] Uploaded: {hero_url}") else: click.echo(" [FAIL] Upload failed") else: click.echo(" (Skipped upload - no site)") else: click.echo(" [FAIL] Generation failed") except Exception as e: click.echo(f" [ERROR] {str(e)[:200]}") click.echo("\n2. Content Images:") click.echo(" (Skipped - T2 articles don't get content images by default)") click.echo(f"\n\n{'='*60}") click.echo("TEST COMPLETE") click.echo(f"{'='*60}\n") except Exception as e: click.echo(f"Error: {e}", err=True) import traceback traceback.print_exc() finally: session.close() if __name__ == "__main__": test_image_generation(23)