#!/usr/bin/env python3
"""
Test Character Consistency - SEED method
Generate Kai walking frames with CONSISTENT SEED
"""

import vertexai
from vertexai.preview.vision_models import ImageGenerationModel
from pathlib import Path

# Initialize
vertexai.init(project="gen-lang-client-0428644398", location="us-central1")

print("="*60)
print("🎮 KAI CHARACTER CONSISTENCY TEST - SEED METHOD")
print("="*60)

# Base character description
base_character = """
Young male farmer character named Kai. Green spiky hair, determined expression, wearing dark blue jean jacket over white shirt, backpack, ripped blue jeans, and brown boots. Cartoon vector art style with bold black outlines (1.4px), flat vibrant colors, clean playful aesthetic.
"""

# Animation frames with SAME SEED
frames = [
    {
        "name": "kai_seed_idle",
        "prompt": f"{base_character} Standing idle pose, front facing view. Solid bright green background (#00FF00), centered character."
    },
    {
        "name": "kai_seed_walk1",
        "prompt": f"{base_character} Walking animation, left leg forward step, right arm swinging forward. Solid bright green background (#00FF00), centered character."
    },
    {
        "name": "kai_seed_walk2",
        "prompt": f"{base_character} Walking animation, neutral standing pose between steps. Solid bright green background (#00FF00), centered character."
    },
    {
        "name": "kai_seed_walk3",
        "prompt": f"{base_character} Walking animation, right leg forward step, left arm swinging forward. Solid bright green background (#00FF00), centered character."
    }
]

model = ImageGenerationModel.from_pretrained("imagegeneration@006")

SEED = 12345  # ← ISTO število za vse frame = boljša konsistentnost!

# Generate frames
for i, frame in enumerate(frames, 1):
    print(f"\n🎨 Generating Frame {i}/4: {frame['name']}")
    print(f"📝 Using SEED: {SEED}")
    
    try:
        response = model.generate_images(
            prompt=frame['prompt'],
            number_of_images=1,
            seed=SEED,  # ← ISTI SEED!
            aspect_ratio="1:1",
            guidance_scale=15,  # Higher = stricter prompt following
            safety_filter_level="block_some",
            person_generation="allow_adult"
        )
        
        output_path = Path(f"test_character/{frame['name']}.png")
        output_path.parent.mkdir(exist_ok=True)
        
        response.images[0].save(location=str(output_path))
        
        print(f"✅ Saved: {output_path}")
        
    except Exception as e:
        print(f"❌ Error: {e}")

print("\n" + "="*60)
print("✅ TEST COMPLETE!")
print("📁 Check: test_character/ folder")
print("🔍 Compare frames for consistency!")
print("="*60)