| from pathlib import Path
|
| from PIL import Image
|
| from PIL.PngImagePlugin import PngImageFile, PngInfo
|
| import json
|
| from pprint import pprint
|
| import argparse
|
| from rich.console import Console
|
| from rich.progress import Progress
|
| from rich_argparse import RichHelpFormatter
|
|
|
|
|
| def parse_a111(params, verbose=False):
|
|
|
| params = params.split("\n")
|
|
|
| prompt = params[0].strip()
|
| neg = params[1].split(":")[1].strip()
|
|
|
| settings = {}
|
| try:
|
| settings = {
|
| s.split(":")[0].strip(): s.split(":")[1].strip()
|
| for s in params[2].split(",")
|
| }
|
|
|
| except IndexError:
|
| settings = {"raw": params[2].strip()}
|
|
|
| if verbose:
|
| print(f"PROMPT: {prompt}")
|
| print(f"NEG: {neg}")
|
| print("SETTINGS:")
|
| pprint(settings, indent=4)
|
|
|
| return {"prompt": prompt, "negative": neg, "settings": settings}
|
|
|
|
|
| import glob
|
|
|
| if __name__ == "__main__":
|
| parser = argparse.ArgumentParser(
|
| description="Crude metadata extractor from A111 pngs",
|
| formatter_class=RichHelpFormatter
|
| )
|
| parser.add_argument("inputs", nargs="*", help="Input image files")
|
| parser.add_argument("--output", help="Output JSON file")
|
| parser.add_argument("-v", "--verbose", action="store_true", help="Verbose mode")
|
| parser.add_argument(
|
| "--glob", help="Enable glob pattern matching", metavar="PATTERN"
|
| )
|
|
|
| args = parser.parse_args()
|
|
|
|
|
| if not args.glob and not args.inputs:
|
| parser.error("Either --glob flag or inputs must be provided.")
|
| if args.glob:
|
| glob_pattern = args.glob
|
| try:
|
| pattern_path = str(Path(glob_pattern).expanduser().resolve())
|
|
|
| if not any(glob.glob(pattern_path)):
|
| raise ValueError(f"No files found for glob pattern: {glob_pattern}")
|
| except Exception as e:
|
| console = Console()
|
| console.print(
|
| f"[bold red]Error: Invalid glob pattern '{glob_pattern}': {e}[/bold red]"
|
| )
|
|
|
| exit(1)
|
| else:
|
| glob_pattern = None
|
|
|
| input_files = []
|
|
|
| if glob_pattern:
|
| input_files = list(glob.glob(str(Path(glob_pattern).expanduser().resolve())))
|
| else:
|
| input_files = [Path(p) for p in args.inputs]
|
|
|
| console = Console()
|
| console.print("Input Files:", style="bold", end=" ")
|
| console.print(f"{len(input_files):03d} files", style="cyan")
|
|
|
|
|
| console.print("\nOutput File:", style="bold", end=" ")
|
| console.print(f"{Path(args.output).resolve().absolute()}", style="cyan")
|
|
|
| with Progress(console=console, auto_refresh=True) as progress:
|
|
|
| unique_info = {}
|
| last = None
|
|
|
| task = progress.add_task("[cyan]Extracting meta...", total=len(input_files) + 1)
|
| for p in input_files:
|
| im = Image.open(p)
|
| parsed = parse_a111(im.info["parameters"], args.verbose)
|
|
|
| if parsed != last:
|
| unique_info[Path(p).stem] = parsed
|
|
|
| last = parsed
|
| progress.update(task, advance=1)
|
| progress.refresh()
|
|
|
| unique_info = json.dumps(unique_info, indent=4)
|
| with open(args.output, "w") as f:
|
| f.write(unique_info)
|
| progress.update(task, advance=1)
|
| progress.refresh()
|
|
|
| console.print("\nProcessing completed!", style="bold green")
|
|
|
