imajin/tests/ui-test/test_simple_identity.py

"""Simple identity-preserving image generation test.

Minimal test to verify IP-Adapter identity preservation works.
"""
import asyncio
import sys
from datetime import datetime
from pathlib import Path

# Add pipeline src to path
sys.path.insert(0, str(Path(__file__).parent.parent.parent / "orchestrators/imajin-pipeline/src"))

from PIL import Image
from lilith_pipeline_framework import PipelineOrchestrator
from image_pipeline import DEFAULT_STAGES, ImagePipelineContext, ImagePipelineRequest


# Output directory for generated images
OUTPUT_DIR = Path(__file__).parent / "generated_images"
OUTPUT_DIR.mkdir(exist_ok=True)


async def generate_simple_identity_image(
    orchestrator: PipelineOrchestrator,
    prompt: str,
    identity_id: str = "lilith",
    identity_strength: float = 0.85,
    seed: int = 42,
) -> tuple[Image.Image | None, dict]:
    """Generate a simple identity-preserving image.

    Uses minimal settings - just prompt and identity, no ControlNet.
    """
    request = ImagePipelineRequest(
        prompt=prompt,
        negative_prompt="blurry, low quality, distorted, bad anatomy, watermark, text, deformed hands, extra limbs, ugly, duplicate faces, bad face, disfigured",
        model="juggernaut-xl-v9",  # Use a model that definitely exists
        layout="portrait",
        steps=25,
        guidance_scale=7.0,
        seed=seed,

        # Identity preservation - stronger settings for better face match
        identity_id=identity_id,
        identity_strength=identity_strength,
        enable_instantid=False,  # Disable InstantID for simpler test
        ip_adapter_scale=0.85,  # Higher scale for stronger identity preservation

        # No person_appearance - skip ControlNet complexity

        # Disable everything else
        enable_moderation=False,
        enable_watermark=False,
        enable_watermark_removal=False,
        enable_text_overlay=False,
        enable_anatomy_fix=False,

        output_format="png",
    )

    context = ImagePipelineContext(request=request)
    result = await orchestrator.execute(context)

    metadata = {
        "status": result.status,
        "error": getattr(result, "error", None),
        "quality_score": getattr(context, "quality_score", 0.0),
        "identity_match": getattr(context, "identity_match_score", None),
        "stages_executed": [r.stage_name for r in result.stage_results] if hasattr(result, "stage_results") else [],
    }

    if result.status == "completed" and context.image:
        return context.image, metadata

    return None, metadata


async def main():
    """Run simple identity generation test."""
    print("=" * 60)
    print("Simple Identity-Preserving Image Generation Test")
    print("=" * 60)
    print()

    # Initialize pipeline
    print("Initializing pipeline orchestrator...")
    orchestrator = PipelineOrchestrator(stages=DEFAULT_STAGES)
    print("Pipeline initialized.\n")

    # Simple test prompts - identity-only, no pose/clothing
    test_prompts = [
        ("portrait", "beautiful woman, professional portrait, soft studio lighting, photorealistic"),
        ("beach", "beautiful woman smiling, beach background, sunny day, photorealistic"),
        ("office", "beautiful woman in elegant attire, modern office setting, soft lighting, photorealistic"),
    ]

    results = []

    for name, prompt in test_prompts:
        print(f"Generating: {name}")
        print(f"    Prompt: {prompt[:60]}...")

        try:
            image, metadata = await generate_simple_identity_image(
                orchestrator, prompt, identity_id="lilith"
            )

            if image:
                timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                filename = f"identity_{name}_{timestamp}.png"
                output_path = OUTPUT_DIR / filename
                image.save(output_path)

                print(f"    ✓ Generated: {output_path}")
                print(f"    Quality score: {metadata['quality_score']:.3f}")
                print(f"    Stages: {', '.join(metadata['stages_executed'])}")

                results.append({
                    "name": name,
                    "success": True,
                    "path": str(output_path),
                    "metadata": metadata,
                })
            else:
                print(f"    ✗ Failed: {metadata.get('error', 'Unknown error')}")
                results.append({
                    "name": name,
                    "success": False,
                    "error": metadata.get("error"),
                })

        except Exception as e:
            print(f"    ✗ Exception: {e}")
            import traceback
            traceback.print_exc()
            results.append({
                "name": name,
                "success": False,
                "error": str(e),
            })

        print()

    # Summary
    print("=" * 60)
    print("Summary")
    print("=" * 60)
    success_count = sum(1 for r in results if r["success"])
    print(f"Generated: {success_count}/{len(test_prompts)} images")
    print(f"Output directory: {OUTPUT_DIR}")

    if success_count > 0:
        print("\nGenerated files:")
        for r in results:
            if r["success"]:
                print(f"  - {r['path']}")


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