novafarma/scripts/test_hybrid_comfyui.py

#!/usr/bin/env python3
"""
ComfyUI Test - Dark Hand-Drawn 2D Stylized Indie Style
Test batch generation with approved hybrid style formula
"""
import os
import json
import time
import requests
from pathlib import Path

# Try both common ComfyUI ports
COMFYUI_PORTS = [8188, 8000]
COMFYUI_URL = None
OUTPUT_DIR = "/Users/davidkotnik/repos/novafarma/style_test_samples/comfyui_tests"

# APPROVED HYBRID STYLE FORMULA
STYLE_PREFIX = """dark hand-drawn 2D stylized indie game art, 
bold thick black hand-drawn outlines, 
smooth vector rendering, 
cartoon-style exaggerated proportions but dark mature atmosphere NOT Disney, 
mature 90s cartoon aesthetic,"""

STYLE_SUFFIX = """stylized character NOT realistic, 
mature indie game art, 
clean white background"""

# Test prompts for different categories
TEST_ASSETS = [
    {
        "name": "test_npc_farmer",
        "prompt": f"{STYLE_PREFIX} Full body NPC character sprite, wasteland farmer with weathered face, brown vest over torn shirt, holding pitchfork tool, standing proud pose, gritty muted color palette with earthy browns, dusty greens, ash grays, {STYLE_SUFFIX}"
    },
    {
        "name": "test_zombie_basic",
        "prompt": f"{STYLE_PREFIX} Full body monster sprite, shambling zombie with grey-green decaying skin, torn clothes, undead appearance, standing menacing pose, gritty dark color palette with death grays, sickly greens, blood rust, cartoon-style exaggerated horror features but mature NOT cute, gruesome but stylized undead, {STYLE_SUFFIX}"
    },
    {
        "name": "test_building_shack",
        "prompt": f"{STYLE_PREFIX} Environment building sprite, old wooden shack with warped perspective and crooked angles, wonky roofline, weathered planks, broken window, Don't Starve meets Stardew Valley aesthetic, gritty muted color palette with weathered wood browns, rusty nails, faded grays, cartoon-style warped proportions but dark mature atmosphere, post-apocalyptic decay visible, mature indie game environment asset, {STYLE_SUFFIX}"
    },
    {
        "name": "test_animal_pig",
        "prompt": f"{STYLE_PREFIX} Full body animal sprite, farm pig with pink skin, standing side view, slightly wonky posture, gritty color palette with muted pinks, earthy browns for dirt, cartoon-style proportions but not cute Disney style, mature indie farm aesthetic, {STYLE_SUFFIX}"
    },
    {
        "name": "test_crop_carrot",
        "prompt": f"{STYLE_PREFIX} Crop plant sprite, carrot plant with orange carrots visible in soil, green leafy tops, slightly warped stems with exaggerated curves, gritty color palette with vibrant oranges for carrots against muted greens for leaves, cartoon-style exaggerated plant shapes but not cute, mature indie farm aesthetic, stylized vegetation, game farming asset, {STYLE_SUFFIX}"
    }
]

def find_comfyui_server():
    """Find running ComfyUI instance"""
    global COMFYUI_URL
    
    for port in COMFYUI_PORTS:
        url = f"http://127.0.0.1:{port}"
        try:
            r = requests.get(f"{url}/system_stats", timeout=2)
            if r.status_code == 200:
                version = r.json().get("system", {}).get("comfyui_version", "unknown")
                print(f"✅ Found ComfyUI v{version} on port {port}")
                COMFYUI_URL = url
                return True
        except:
            continue
    
    print(f"❌ ComfyUI not found on ports: {COMFYUI_PORTS}")
    print("   Start ComfyUI with: open /Applications/ComfyUI.app")
    return False

def queue_comfy_prompt(prompt_text: str, output_name: str) -> dict:
    """Send prompt to ComfyUI"""
    workflow = {
        "3": {
            "inputs": {
                "seed": int(time.time() * 1000),
                "steps": 30,
                "cfg": 7.5,
                "sampler_name": "euler",
                "scheduler": "normal",
                "denoise": 1,
                "model": ["4", 0],
                "positive": ["6", 0],
                "negative": ["7", 0],
                "latent_image": ["5", 0]
            },
            "class_type": "KSampler"
        },
        "4": {
            "inputs": {
                "ckpt_name": "sd_xl_base_1.0.safetensors"
            },
            "class_type": "CheckpointLoaderSimple"
        },
        "5": {
            "inputs": {
                "width": 512,
                "height": 512,
                "batch_size": 1
            },
            "class_type": "EmptyLatentImage"
        },
        "6": {
            "inputs": {
                "text": prompt_text,
                "clip": ["4", 1]
            },
            "class_type": "CLIPTextEncode"
        },
        "7": {
            "inputs": {
                "text": "blurry, low quality, pixelated, voxel, 3D render, realistic photo, photorealistic, photography, Disney cute style, bright colors, clean cartoon",
                "clip": ["4", 1]
            },
            "class_type": "CLIPTextEncode"
        },
        "8": {
            "inputs": {
                "samples": ["3", 0],
                "vae": ["4", 2]
            },
            "class_type": "VAEDecode"
        },
        "9": {
            "inputs": {
                "filename_prefix": output_name,
                "images": ["8", 0]
            },
            "class_type": "SaveImage"
        }
    }
    
    try:
        response = requests.post(
            f"{COMFYUI_URL}/prompt",
            json={"prompt": workflow, "client_id": f"hybrid_{int(time.time())}"}
        )
        return response.json()
    except Exception as e:
        print(f"❌ Queue error: {e}")
        return None

def wait_for_completion(prompt_id: str, timeout: int = 300) -> bool:
    """Wait for generation to complete"""
    start = time.time()
    while time.time() - start < timeout:
        try:
            response = requests.get(f"{COMFYUI_URL}/history/{prompt_id}")
            history = response.json()
            
            if prompt_id in history:
                status = history[prompt_id].get("status", {})
                if status.get("completed", False):
                    return True
                    
        except Exception as e:
            pass
        
        time.sleep(3)
    
    return False

def download_image(prompt_id: str, output_path: Path) -> bool:
    """Download generated image"""
    try:
        response = requests.get(f"{COMFYUI_URL}/history/{prompt_id}")
        history = response.json()
        
        if prompt_id not in history:
            return False
            
        outputs = history[prompt_id].get("outputs", {})
        for node_id, node_output in outputs.items():
            if "images" in node_output:
                for img in node_output["images"]:
                    filename = img["filename"]
                    subfolder = img.get("subfolder", "")
                    
                    img_url = f"{COMFYUI_URL}/view"
                    params = {
                        "filename": filename,
                        "subfolder": subfolder,
                        "type": "output"
                    }
                    
                    img_response = requests.get(img_url, params=params)
                    if img_response.status_code == 200:
                        output_path.parent.mkdir(parents=True, exist_ok=True)
                        with open(output_path, 'wb') as f:
                            f.write(img_response.content)
                        return True
        
        return False
        
    except Exception as e:
        print(f"❌ Download error: {e}")
        return False

# MAIN
print("=" * 80)
print("🎨 COMFYUI HYBRID STYLE TEST")
print("=" * 80)
print()
print("Testing 'Dark Hand-Drawn 2D Stylized Indie' style")
print()

# Find server
if not find_comfyui_server():
    exit(1)

print()
print(f"📊 Will generate {len(TEST_ASSETS)} test assets:")
for asset in TEST_ASSETS:
    print(f"   - {asset['name']}")

print()
input("Press ENTER to start generation... ")
print()

# Generate each asset
success_count = 0
for i, asset in enumerate(TEST_ASSETS, 1):
    print(f"\n[{i}/{len(TEST_ASSETS)}] 🎨 Generating: {asset['name']}")
    print(f"   Prompt: {asset['prompt'][:100]}...")
    
    # Queue
    result = queue_comfy_prompt(asset['prompt'], asset['name'])
    if not result or "prompt_id" not in result:
        print(f"   ❌ Failed to queue!")
        continue
    
    prompt_id = result["prompt_id"]
    print(f"   ⏳ Queued: {prompt_id}")
    
    # Wait
    print(f"   ⏳ Waiting for completion...")
    if wait_for_completion(prompt_id):
        output_path = Path(OUTPUT_DIR) / f"{asset['name']}.png"
        if download_image(prompt_id, output_path):
            print(f"   ✅ SUCCESS! Saved to: {output_path}")
            success_count += 1
        else:
            print(f"   ❌ Failed to download!")
    else:
        print(f"   ⏱️ Timeout!")

print()
print("=" * 80)
print(f"✅ Complete! {success_count}/{len(TEST_ASSETS)} assets generated")
print(f"📁 Check: {OUTPUT_DIR}")
print("=" * 80)