new images
This commit is contained in:
243
asset-work/generate_multiple_combo.py
Executable file
243
asset-work/generate_multiple_combo.py
Executable file
@@ -0,0 +1,243 @@
|
||||
#!/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()
|
||||
Reference in New Issue
Block a user