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

#!/usr/bin/env python3
"""Generate multiple images from a caption file 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(caption_file: str) -> int:
    """Count existing image variations for a caption file.

    Args:
        caption_file: Path to the caption text file

    Returns:
        Number of existing generated images matching the caption name
    """
    caption_name = os.path.splitext(os.path.basename(caption_file))[0]
    caption_dir = os.path.dirname(caption_file) or "."

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

    return count


def generate_multiple(caption_file: str, count: int = 2, server: str = "127.0.0.1:8188", dry_run: bool = False) -> tuple[list[str], int]:
    """Generate multiple images with different random seeds.

    Args:
        caption_file: Path to the caption text file
        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(caption_file):
        print(f"Error: Caption file not found: {caption_file}")
        sys.exit(1)

    # Validate caption file content
    with open(caption_file, "r") as f:
        caption = f.read().strip()

    if not caption:
        print(f"Error: Caption file is empty: {caption_file}")
        sys.exit(1)

    if dry_run:
        print(f"  Caption preview: {caption[:60]}{'...' if len(caption) > 60 else ''}")
        return [], count_existing_variations(caption_file)

    # Count existing variations
    existing = count_existing_variations(caption_file)
    print(f"Found {existing} existing variations")
    print(f"Requested total: {count}")

    # Calculate how many more to generate
    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_dirs = []
    script_dir = os.path.dirname(os.path.abspath(__file__))
    generator_script = os.path.join(script_dir, "generate_from_caption.py")

    for i in range(needed):
        # Generate random seed
        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}")

        # Run the generator script
        result = subprocess.run(
            [sys.executable, generator_script, caption_file, str(seed), "--server", server],
            capture_output=False,
        )

        if result.returncode != 0:
            print(f"Warning: Generation {i + 1} of {needed} failed with exit code {result.returncode}")
        else:
            # Image is saved in current directory by generate_from_caption
            caption_name = os.path.splitext(os.path.basename(caption_file))[0]
            output_path = os.path.join(".", f"{caption_name}_{seed}_generated.png")
            output_dirs.append(output_path)

    return output_dirs, existing


def main():
    import argparse

    parser = argparse.ArgumentParser(
        description="Generate multiple images from caption with different seeds"
    )
    parser.add_argument("caption_file", help="Path to caption text file")
    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 caption files and server connection without generating",
    )

    args = parser.parse_args()

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

        # Check server
        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)

        # Validate caption
        print(f"\nValidating: {args.caption_file}")
        output_dirs, existing = generate_multiple(args.caption_file, args.count, args.server, dry_run=True)
        print(f"✓ Caption file is valid (found {existing} existing variations)")

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

    # Check server before starting
    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(args.caption_file, args.count, args.server)

    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()