287 lines
12 KiB
Python
287 lines
12 KiB
Python
"""
|
|
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)
|
|
|