Files
ai-game-2/asset-work/generate_all_captions.py
2026-02-25 17:41:12 -08:00

103 lines
3.4 KiB
Python
Executable File

#!/usr/bin/env python3
"""Generate multiple images from all caption files recursively."""
import subprocess
import sys
import os
import re
from pathlib import Path
def main():
script_dir = Path(__file__).parent.resolve()
generator = script_dir / "generate_multiple.py"
dry_run = "--dry-run" in sys.argv
if not generator.exists():
print(f"Error: generate_multiple.py not found in {script_dir}")
sys.exit(1)
print("Finding all caption files...")
caption_files = list(Path(".").rglob("caption_*.txt"))
total = len(caption_files)
if total == 0:
print("No caption files found!")
sys.exit(0)
print(f"Found {total} caption file(s)")
if dry_run:
print("\nDRY RUN MODE - Validating all caption files")
counter = 0
failed = 0
total_existing = 0
total_created = 0
total_needed = 0
# Default variations per caption file (matches generate_multiple.py default)
variations_per_file = 1
for caption_file in caption_files:
counter += 1
cmd = ["python3", str(generator), str(caption_file), "--count", str(variations_per_file)]
if dry_run:
cmd.append("--dry-run")
try:
result = subprocess.run(cmd, capture_output=True, text=True, check=True)
# Parse output to count existing and created images
# Look for "Found X existing variations" (also works with "found {existing} existing variations" in dry-run mode)
existing_match = re.search(r'[Ff]ound (\d+) existing variations?', result.stdout)
created_match = re.search(r'Created (\d+) new output directories?', result.stdout)
if existing_match:
existing = int(existing_match.group(1))
total_existing += existing
# Calculate how many would be needed (max 0, can't be negative)
needed = max(0, variations_per_file - existing)
total_needed += needed
if created_match:
total_created += int(created_match.group(1))
# Progress indicator
status = "" if result.returncode == 0 else ""
print(f"[{counter}/{total}] {status} {caption_file}")
except subprocess.CalledProcessError as e:
failed += 1
print(f"[{counter}/{total}] ✗ {caption_file} - FAILED")
print()
print("=" * 50)
if dry_run:
print("DRY RUN COMPLETE")
print("=" * 50)
print(f"Caption files validated: {total}")
print(f"Images already existing: {total_existing}")
print(f"Images would be created: {total_needed}")
print(f"Total when complete: {total_existing + total_needed}")
if failed == 0:
print(f"✓ All caption files validated successfully!")
else:
print(f"{failed} of {total} caption file(s) failed validation")
sys.exit(1)
else:
print("GENERATION COMPLETE")
print("=" * 50)
print(f"Caption files processed: {total}")
print(f"Images already existing: {total_existing}")
print(f"Images newly created: {total_created}")
print(f"Total images now: {total_existing + total_created}")
if failed > 0:
print(f"Failed: {failed}")
if __name__ == "__main__":
main()