novafarma/scripts/test_generation.py

#!/usr/bin/env python3
"""
🧪 TEST GENERATION - Generira 5 testnih slik
Z avtomatskim Git commitom
"""

import os
import sys
import json
import time
import uuid
import requests
import subprocess
from pathlib import Path
from datetime import datetime

try:
    from rembg import remove
    from PIL import Image
except ImportError:
    print("📦 Installing rembg...")
    os.system("pip3 install --user rembg pillow onnxruntime")
    from rembg import remove
    from PIL import Image

# Configuration
COMFYUI_URL = "http://127.0.0.1:8000"
OUTPUT_DIR = Path("/Users/davidkotnik/repos/novafarma/assets/images")
REPO_DIR = Path("/Users/davidkotnik/repos/novafarma")

STYLE = "2D indie game sprite, cartoon vector style, clean smooth lines, full body visible, isolated on pure white background #FFFFFF"
NEGATIVE = "pixel art, pixelated, 3d, realistic, photo, blurry, watermark, text, green background"

# 🧪 TEST ASSETS - Par hitrih testnih slik
TEST_ASSETS = [
    ("npcs", "npc_elder", "wise tribal elder NPC, walking stick, long white beard, robes"),
    ("npcs", "npc_child", "survivor orphan child NPC, oversized clothes, teddy bear"),
    ("zivali", "goat_brown", "brown goat farm animal, small horns, playful"),
    ("items", "tool_hammer", "hammer building tool, wooden handle"),
    ("environment", "flower_red", "red flowers patch, blooming"),
]


def log(msg):
    ts = datetime.now().strftime("%H:%M:%S")
    print(f"[{ts}] {msg}")
    sys.stdout.flush()


def git_commit(file_path, category, name):
    """Auto-commit generated asset"""
    try:
        rel_path = file_path.relative_to(REPO_DIR)
        
        # Git add
        subprocess.run(
            ["git", "add", str(rel_path)],
            cwd=REPO_DIR,
            check=True,
            capture_output=True
        )
        
        # Git commit
        commit_msg = f"🎨 Generated asset: {category}/{name}"
        subprocess.run(
            ["git", "commit", "-m", commit_msg],
            cwd=REPO_DIR,
            check=True,
            capture_output=True
        )
        
        log(f"   ✅ Git committed: {rel_path}")
        return True
    except subprocess.CalledProcessError as e:
        log(f"   ⚠️  Git commit failed (možno že committed): {e}")
        return False


def create_workflow(prompt_text, output_name, size=512):
    seed = int(time.time() * 1000) % 2147483647
    full_prompt = f"{STYLE}, {prompt_text}"
    
    return {
        "1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "dreamshaper_8.safetensors"}},
        "2": {"class_type": "EmptyLatentImage", "inputs": {"width": size, "height": size, "batch_size": 1}},
        "3": {"class_type": "CLIPTextEncode", "inputs": {"text": full_prompt, "clip": ["1", 1]}},
        "4": {"class_type": "CLIPTextEncode", "inputs": {"text": NEGATIVE, "clip": ["1", 1]}},
        "5": {
            "class_type": "KSampler",
            "inputs": {
                "seed": seed, "steps": 30, "cfg": 7.5,
                "sampler_name": "euler_ancestral", "scheduler": "karras",
                "denoise": 1.0, "model": ["1", 0],
                "positive": ["3", 0], "negative": ["4", 0], "latent_image": ["2", 0]
            }
        },
        "6": {"class_type": "VAEDecode", "inputs": {"samples": ["5", 0], "vae": ["1", 2]}},
        "7": {"class_type": "SaveImage", "inputs": {"filename_prefix": output_name, "images": ["6", 0]}}
    }


def queue_prompt(workflow):
    try:
        r = requests.post(
            f"{COMFYUI_URL}/prompt",
            json={"prompt": workflow, "client_id": f"test_{uuid.uuid4().hex[:8]}"},
            timeout=10
        )
        return r.json().get("prompt_id")
    except Exception as e:
        log(f"❌ Queue error: {e}")
        return None


def wait_completion(prompt_id, timeout=120):
    start = time.time()
    while time.time() - start < timeout:
        try:
            r = requests.get(f"{COMFYUI_URL}/history/{prompt_id}", timeout=5)
            data = r.json()
            if prompt_id in data and data[prompt_id].get("outputs"):
                return True
        except:
            pass
        time.sleep(2)
    return False


def download_and_process(prompt_id, output_path):
    try:
        h = requests.get(f"{COMFYUI_URL}/history/{prompt_id}").json()
        for out in h.get(prompt_id, {}).get("outputs", {}).values():
            for img in out.get("images", []):
                r = requests.get(f"{COMFYUI_URL}/view", params={
                    "filename": img["filename"],
                    "subfolder": img.get("subfolder", ""),
                    "type": "output"
                })
                if r.status_code == 200:
                    output_path.parent.mkdir(parents=True, exist_ok=True)
                    
                    # Remove background
                    from io import BytesIO
                    src_img = Image.open(BytesIO(r.content))
                    transparent_img = remove(src_img)
                    transparent_img.save(str(output_path), "PNG")
                    return True
        return False
    except Exception as e:
        log(f"❌ Download error: {e}")
        return False


def main():
    log("=" * 60)
    log("🧪 TEST GENERATION - 5 Testnih Slik")
    log("   Z avtomatskim Git commitom")
    log("=" * 60)
    
    # Check ComfyUI
    try:
        r = requests.get(f"{COMFYUI_URL}/system_stats", timeout=5)
        v = r.json()["system"]["comfyui_version"]
        log(f"✅ ComfyUI v{v} running")
    except:
        log("❌ ComfyUI not running!")
        return
    
    log(f"\n🎯 Generating {len(TEST_ASSETS)} test assets...\n")
    
    success, fail = 0, 0
    start_time = time.time()
    
    for i, (cat, name, prompt) in enumerate(TEST_ASSETS, 1):
        path = OUTPUT_DIR / cat / f"{name}.png"
        
        # Skip if exists
        if path.exists():
            log(f"[{i}/{len(TEST_ASSETS)}] ⏭️  {name} - already exists")
            continue
        
        log(f"[{i}/{len(TEST_ASSETS)}] 🎨 Generating {name}...")
        
        wf = create_workflow(prompt, name, 512)
        pid = queue_prompt(wf)
        
        if pid and wait_completion(pid) and download_and_process(pid, path):
            log(f"   ✅ Image saved (transparent)")
            
            # Auto-commit
            git_commit(path, cat, name)
            
            success += 1
        else:
            log(f"   ❌ FAILED")
            fail += 1
    
    elapsed = time.time() - start_time
    log("\n" + "=" * 60)
    log("🧪 TEST COMPLETE!")
    log(f"   ✅ Success: {success}")
    log(f"   ❌ Failed: {fail}")
    log(f"   ⏱️  Time: {elapsed:.1f}s ({elapsed/60:.1f} min)")
    log(f"   📁 Output: {OUTPUT_DIR}")
    log("=" * 60)


if __name__ == "__main__":
    main()