ai-game-2/asset-work/generate_multiple_combo.py

#!/usr/bin/env python3
"""Generate multiple combo images from a kq4-decompile visual image with different seeds."""

import random
import subprocess
import sys
import os
import urllib.request


def check_server(server_address: str = "127.0.0.1:8188", timeout: int = 5) -> bool:
    """Check if ComfyUI server is running and accessible."""
    try:
        req = urllib.request.Request(
            f"http://{server_address}/system_stats",
            method="GET",
        )
        with urllib.request.urlopen(req, timeout=timeout) as response:
            return response.status == 200
    except Exception:
        return False


def count_existing_variations(image_path: str, caption_num: int = 1) -> int:
    """Count existing image variations for a visual image.

    Args:
        image_path: Path to the kq4-decompile visual image

    Returns:
        Number of existing generated images matching the room name
    """
    script_dir = os.path.dirname(os.path.abspath(image_path)) or "."
    asset_work_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)))
    combo_outputs_dir = os.path.join(asset_work_dir, "combo_outputs")

    if not os.path.exists(combo_outputs_dir):
        return 0

    image_basename = os.path.basename(image_path)
    room_name = (
        image_basename.replace("pic_", "").replace("_visual.png", "").replace("_", "-")
    )

    room_output_dir = os.path.join(combo_outputs_dir, room_name)
    if not os.path.exists(room_output_dir):
        return 0

    count = 0
    for item in os.listdir(room_output_dir):
        if item.startswith(f"{room_name}_caption_{caption_num}_") and item.endswith("_generated.png"):
            count += 1

    return count


def generate_multiple_combo(
    image_path: str,
    caption_text: str,
    count: int = 2,
    server: str = "127.0.0.1:8188",
    dry_run: bool = False,
    caption_num: int = 1,
) -> tuple[list[str], int]:
    """Generate multiple combo images with different random seeds.

    Args:
        image_path: Path to kq4-decompile visual image
        caption_text: Direct caption text for the scene
        count: Total number of variations to have (default: 2)
        server: ComfyUI server address
        dry_run: If True, validate only without generating

    Returns:
        Tuple of (list of created image paths, number of existing variations)
    """
    if not os.path.exists(image_path):
        print(f"Error: Image file not found: {image_path}")
        sys.exit(1)

    image_basename = os.path.basename(image_path)
    room_name = (
        image_basename.replace("pic_", "").replace("_visual.png", "").replace("_", "-")
    )

    if dry_run:
        print(f"  Image: {image_basename}")
        print(
            f"  Caption: {caption_text[:50]}..."
            if len(caption_text) > 50
            else f"  Caption: {caption_text}"
        )
        return [], count_existing_variations(image_path, caption_num)

    existing = count_existing_variations(image_path, caption_num)
    print(f"Found {existing} existing variations")
    print(f"Requested total: {count}")

    needed = count - existing

    if needed <= 0:
        print(f"\nAlready have {existing} variations (requested {count})")
        print("No new images needed.")
        return [], existing

    print(f"Will generate {needed} new variation(s)")

    output_paths = []
    script_dir = os.path.dirname(os.path.abspath(__file__))
    generator_script = os.path.join(script_dir, "generate_from_caption_combo.py")

    for i in range(needed):
        seed = random.randint(1, 2**32 - 1)
        print(f"\n{'=' * 60}")
        print(f"Generation {i + 1} of {needed} (seed: {seed})")
        print(f"Progress: {existing + i + 1} of {count} total")
        print(f"{'=' * 60}")

        cmd = [
            sys.executable,
            generator_script,
            image_path,
            caption_text,
            str(seed),
            "--server",
            server,
        ]

        result = subprocess.run(cmd, capture_output=False)

        if result.returncode != 0:
            print(
                f"Warning: Generation {i + 1} of {needed} failed with exit code {result.returncode}"
            )
        else:
            asset_work_dir = script_dir
            combo_outputs_dir = os.path.join(asset_work_dir, "combo_outputs")
            output_path = os.path.join(
                combo_outputs_dir, room_name, f"{room_name}_caption_{caption_num}_{seed}_generated.png"
            )
            output_paths.append(output_path)

    return output_paths, existing


def main():
    import argparse

    parser = argparse.ArgumentParser(
        description="Generate multiple combo images from kq4-decompile visual with different seeds"
    )
    parser.add_argument(
        "image_path",
        help="Path to kq4-decompile visual image (e.g., kq4-sierra-decompile/rooms/kq4-057-witch-cave/pic_057_visual.png)",
    )
    parser.add_argument(
        "caption_text",
        help="Direct caption text for the scene (not AI-generated)",
    )
    parser.add_argument(
        "--count",
        "-n",
        type=int,
        default=2,
        help="Number of variations to generate (default: 2)",
    )
    parser.add_argument(
        "--server",
        default="127.0.0.1:8188",
        help="ComfyUI server address (default: 127.0.0.1:8188)",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Test mode: validate image file and server connection without generating",
    )
    parser.add_argument(
        "--caption-num",
        type=int,
        default=1,
        help="Caption number for output filename (default: 1)",
    )

    args = parser.parse_args()

    if args.dry_run:
        print("\n" + "=" * 60)
        print("DRY RUN MODE - Validating without generating")
        print("=" * 60)

        print(f"\nChecking ComfyUI server at {args.server}...")
        if check_server(args.server):
            print("✓ Server is accessible")
        else:
            print(f"✗ Server is NOT accessible at {args.server}")
            print("  Please ensure ComfyUI is running before generating images.")
            sys.exit(1)

        print(f"\nValidating: {args.image_path}")
        output_paths, existing = generate_multiple_combo(
            args.image_path,
            args.caption_text,
            args.count,
            args.server,
            dry_run=True,
            caption_num=args.caption_num,
        )
        print(f"✓ Image file is valid (found {existing} existing variations)")

        print("\n" + "=" * 60)
        print("✓ Dry run successful! All checks passed.")
        print("=" * 60)
        sys.exit(0)

    print(f"\nChecking ComfyUI server at {args.server}...")
    if not check_server(args.server):
        print(f"Error: ComfyUI server is not running at {args.server}")
        print("Please start ComfyUI first or check the server address.")
        print(f"\nTo test without generating, use: --dry-run")
        sys.exit(1)
    print("✓ Server is running\n")

    output_paths, existing = generate_multiple_combo(
        args.image_path,
        args.caption_text,
        args.count,
        args.server,
        dry_run=False,
        caption_num=args.caption_num,
    )

    print(f"\n{'=' * 60}")
    print("All generations complete!")
    print(f"{'=' * 60}")
    if output_paths:
        print(f"\nCreated {len(output_paths)} new image(s):")
        for output_path in output_paths:
            print(f"  - {output_path}")
    print(f"\nTotal variations now: {existing + len(output_paths)}")


if __name__ == "__main__":
    main()