imajin/tooling/scripts/generate_original_tests.py

"""Generate the original test cases including femboy+latex Tokyo.

This bypasses Stage 1 (classifier) which requires cot-reasoning service.
We'll manually provide the cultural context that Stage 1 would have generated.
"""

import asyncio
import base64
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."""

    # Cultural context that classifier would have generated
    cultural_context = {
        "determinedStyle": style,
        "styleConfidence": 0.9,
        "determinedMaturity": "suggestive" if "latex" in filters or "femboy" in filters else "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:
        # Stage 2: 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📝 Prompt: {config['prompt'][:200]}...")

        # Stage 3: Diffusion
        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}")
                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 original test cases from test_full_pipeline.py."""

    print("=" * 80)
    print("Generating Original Test Cases (with Diversity System)")
    print("=" * 80)

    # Original test cases from test_full_pipeline.py parametrize
    test_cases = [
        # Anime filters (should now be single subjects, not character sheets)
        ("escorts", "Tokyo", ["femboy", "latex"], "anime", 1, "new_femboy_latex_tokyo"),
        ("escorts", "Tokyo", ["femboy"], "anime", 1, "new_femboy"),
        ("escorts", "Tokyo", ["kawaii"], "anime", 1, "new_kawaii"),
        ("escorts", "Tokyo", ["catgirl"], "anime", 1, "new_catgirl"),
        ("escorts", "Tokyo", ["egirl"], "anime", 1, "new_egirl"),
        ("escorts", "Tokyo", ["vtuber"], "anime", 1, "new_vtuber"),
        ("escorts", "Tokyo", ["cosplay"], "anime", 1, "new_cosplay"),
        ("escorts", "Tokyo", ["gyaru"], "anime", 1, "new_gyaru"),

        # Photorealistic filters (should now have diversity)
        ("escorts", "Tokyo", ["professional_escort"], "photorealistic", 1, "new_professional_escort"),
        ("escorts", "Tokyo", ["businesswoman"], "photorealistic", 1, "new_businesswoman"),
        ("escorts", "Tokyo", ["lawyer"], "photorealistic", 1, "new_lawyer"),

        # Empty filters default
        ("escorts", "New York", [], "photorealistic", 1, "new_empty_filters_default"),
    ]

    results = []
    for category, city, filters, style, subject_count, filename in test_cases:
        print(f"\n{'=' * 80}")
        print(f"Test: {filename}")
        print(f"Filters: {filters}, Style: {style}")

        path = await generate_image(category, city, filters, style, subject_count, filename)
        results.append((filename, path))

        await asyncio.sleep(1)

    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")

    if results and results[0][1]:
        output_dir = Path(results[0][1]).parent
        print(f"\n📂 Output directory: {output_dir}")
        print(f"\nCompare with original: /tmp/imajin/test-generated-images/20260113_070715/")


if __name__ == "__main__":
    asyncio.run(main())