SD-sort-images-by-UI

sd_meta_sorter.py

#
# Sort stable diffusion images by the UI they were generated in.
# 
import os
import json
import argparse
from PIL import Image
from tqdm import tqdm
from shutil import copyfile

parser = argparse.ArgumentParser(description="Sort images by UI into separate folders")
parser.add_argument("-i", "--input", default="./", help="Input folder to sort")
parser.add_argument("-o", "--output", default="./", help="Output folder for gens")
args = parser.parse_args()

target_auto = os.path.join(args.output, "auto1111")
target_comf = os.path.join(args.output, "comfyui")

if not os.path.isdir(args.output):
	parser.error("Output folder must exist!")
if not os.path.isdir(target_auto):
	os.mkdir(target_auto)
if not os.path.isdir(target_comf):
	os.mkdir(target_comf)

for name in tqdm(os.listdir(args.input)):
	path = os.path.join(args.input, name)
	if not os.path.isfile(path): continue
	if os.path.splitext(name)[1] != ".png": continue
	try: img = Image.open(path)
	except: continue
	if not img.text: continue
	if "parameters" in img.text.keys():
		target = os.path.join(target_auto, name)
	elif 'prompt' in img.text.keys() and 'workflow' in img.text.keys():
		target = os.path.join(target_comf, name)
	else: continue
	if os.path.isfile(target):
		tqdm.write(f"File aready exists! ignoring '{name}'")
	copyfile(path, target)

input("Press any key to continue.")