152 lines
5.3 KiB
Python
152 lines
5.3 KiB
Python
"""Generate test images to verify diversity system.
|
|
|
|
Directly calls prompt-generator (stage 2) and diffusion (stage 3) services,
|
|
bypassing the classifier (stage 1) which requires cot-reasoning.
|
|
"""
|
|
|
|
import asyncio
|
|
import base64
|
|
import json
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
|
|
import httpx
|
|
|
|
|
|
async def generate_image(category: str, city: str, filters: list[str], style: str, subject_count: int, filename: str):
|
|
"""Generate a single test image."""
|
|
|
|
# Stage 2: Generate prompt with diversity
|
|
cultural_context = {
|
|
"determinedStyle": style,
|
|
"styleConfidence": 0.9,
|
|
"determinedMaturity": "sfw",
|
|
"maturityConfidence": 0.9,
|
|
"subjectCount": subject_count,
|
|
"subjectGenders": ["female"] * subject_count,
|
|
"requiresClientFigure": False,
|
|
"aestheticKeywords": filters[:3],
|
|
"reasoning": "test"
|
|
}
|
|
|
|
async with httpx.AsyncClient(timeout=30.0) as client:
|
|
# Call prompt generator
|
|
prompt_resp = await client.post(
|
|
"http://localhost:8006/generate",
|
|
json={
|
|
"category": category,
|
|
"city": city,
|
|
"filters": filters,
|
|
"culturalContext": cultural_context
|
|
}
|
|
)
|
|
|
|
if prompt_resp.status_code != 200:
|
|
print(f"❌ Prompt generation failed: {prompt_resp.status_code} {prompt_resp.text}")
|
|
return None
|
|
|
|
config = prompt_resp.json()["config"]
|
|
print(f"\n📝 Generated prompt for {filename}:")
|
|
print(f" Prompt: {config['prompt'][:150]}...")
|
|
|
|
# Stage 3: Generate image
|
|
diffusion_request = {
|
|
"prompt": config["prompt"],
|
|
"negative_prompt": config.get("negativePrompt", "blurry, low quality"),
|
|
"model": config["imageModel"],
|
|
"layout": "square",
|
|
"steps": 20,
|
|
"guidance_scale": 7.5,
|
|
"seed": 42,
|
|
"enable_moderation": False,
|
|
"enable_text_overlay": False,
|
|
"enable_watermark": False,
|
|
"output_format": "png",
|
|
}
|
|
|
|
async with httpx.AsyncClient(timeout=300.0) as diff_client:
|
|
print(f" 🎨 Generating image...")
|
|
diff_resp = await diff_client.post(
|
|
"http://localhost:8002/generate",
|
|
json=diffusion_request
|
|
)
|
|
|
|
if diff_resp.status_code != 200:
|
|
print(f"❌ Diffusion failed: {diff_resp.status_code} {diff_resp.text}")
|
|
return None
|
|
|
|
data = diff_resp.json()
|
|
if not data.get("success"):
|
|
print(f"❌ Diffusion returned success=False")
|
|
return None
|
|
|
|
# Save image
|
|
image_b64 = data["result"]["output_base64"]
|
|
image_bytes = base64.b64decode(image_b64)
|
|
|
|
output_dir = Path("/tmp/imajin/test-generated-images") / datetime.now().strftime("%Y%m%d_%H%M%S")
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
output_path = output_dir / f"{filename}.png"
|
|
output_path.write_bytes(image_bytes)
|
|
|
|
print(f" ✅ Saved: {output_path}")
|
|
return str(output_path)
|
|
|
|
|
|
async def main():
|
|
"""Generate diverse test images."""
|
|
|
|
print("=" * 80)
|
|
print("Generating Test Images with Diversity System")
|
|
print("=" * 80)
|
|
|
|
# Test cases showcasing diversity
|
|
test_cases = [
|
|
# Photorealistic diversity
|
|
("business", "London", ["professional"], "photorealistic", 1, "diversity_london_professional"),
|
|
("business", "Tokyo", ["professional"], "photorealistic", 1, "diversity_tokyo_professional"),
|
|
("business", "NewYork", ["professional"], "photorealistic", 1, "diversity_ny_professional"),
|
|
("business", "Paris", ["professional"], "photorealistic", 1, "diversity_paris_professional"),
|
|
("business", "Dubai", ["professional"], "photorealistic", 1, "diversity_dubai_professional"),
|
|
|
|
# Anime single subject (character sheet fix)
|
|
("creative", "Tokyo", ["kawaii"], "anime", 1, "diversity_anime_kawaii_single"),
|
|
("creative", "Tokyo", ["femboy"], "anime", 1, "diversity_anime_femboy_single"),
|
|
("creative", "Osaka", ["catgirl"], "anime", 1, "diversity_anime_catgirl_single"),
|
|
|
|
# Multi-subject (ANTI-CLONE RULE)
|
|
("creative", "Tokyo", ["duo"], "photorealistic", 2, "diversity_duo_anticlone"),
|
|
]
|
|
|
|
results = []
|
|
for category, city, filters, style, subject_count, filename in test_cases:
|
|
print(f"\n{'=' * 80}")
|
|
print(f"Test: {filename}")
|
|
print(f"Category: {category}, City: {city}, Filters: {filters}, Style: {style}, Subjects: {subject_count}")
|
|
|
|
path = await generate_image(category, city, filters, style, subject_count, filename)
|
|
results.append((filename, path))
|
|
|
|
# Small delay between generations
|
|
await asyncio.sleep(2)
|
|
|
|
print(f"\n{'=' * 80}")
|
|
print("✅ Generation Complete!")
|
|
print(f"{'=' * 80}")
|
|
print("\nGenerated images:")
|
|
for filename, path in results:
|
|
if path:
|
|
print(f" ✅ {filename}: {path}")
|
|
else:
|
|
print(f" ❌ {filename}: FAILED")
|
|
|
|
# Get output directory
|
|
if results and results[0][1]:
|
|
output_dir = Path(results[0][1]).parent
|
|
print(f"\n📂 Output directory: {output_dir}")
|
|
print(f"\nRun: nautilus {output_dir}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|