Test random AUTOMATIC1111 Stable Diffusion models against random prompts and save the results

run.py

import json
import requests
import io
import base64
import random
import sys
import datetime
from timeit import default_timer as timer
from datetime import timedelta
from PIL import Image

total_calls = 1
if len(sys.argv) > 1:
    total_calls = int(sys.argv[1])

# Run webui with --listen --port 7860 --api
# If webui is on a different machine, change IP
url = "http://127.0.0.1:7860"
sampler = 'DPM++ 2M Karras'
size_w = 512
size_h = 512

response = requests.get(url=f'{url}/sdapi/v1/sd-models')
mr = response.json()

response = requests.get(url=f'{url}/sdapi/v1/loras')
lr = response.json()

avg_total = 0

total_models = len(mr)
total_loras = len(lr)

print("Total models: %d" % total_models)
print("Total loras: %d" % total_loras)
print("Total calls: %d" % total_calls)

for avg_count in range(total_calls):
    model = random.choice(mr)

    print(model['title'] + " model change...", end = ' ', flush = True)

    start = timer()
    option_payload = {
        "sd_model_checkpoint": model['title']
    }

    response = requests.post(url=f'{url}/sdapi/v1/options', json=option_payload)

    end = timer()

    print("%.2fs" % (end - start))
    print("%.2f%%" % ((avg_count / total_calls) * 100), end = ' ', flush = True)

    prompt = 'a'
    prompt += ' ' + random.choice(['cute', 'lovely', 'wholesome', 'sweet', 'good looking'])
    prompt += ' ' + random.choice(['woman', 'man', 'pumpkin', 'ghost', 'dog', 'cat', 'landscape'])
    prompt += ' ' + random.choice(['', 'in front of', 'over a', 'in a'])
    prompt += ' ' + random.choice(['forest', 'woodland', 'old house', 'lake', 'future cityscape', 'city', 'beach', 'desert', 'mountain', 'mountain range', 'cave', 'cavern', 'ruins', 'temple', 'dungeon', 'castle', 'palace', 'village', 'town', 'city', 'cityscape', 'space station', 'space ship', 'space colony'])

    attributes = [['oversized shirt', 'small shirt'],
                  ['sitting', 'standing', 'posing'],
                  ['long hair', 'short hair', 'curly hair', 'straight hair', 'messy hair', 'ponytail', 'braided hair'],
                  ['with a cat', 'with a dog']]

    for gi in range(len(attributes)):
        attribute_count = random.randint(0, 2)
        for ai in range(attribute_count):
            prompt += ', ' + random.choice(attributes[gi])

    # SD only allows 75 tokens
    while len(prompt.split()) > 70:
        prompt_groups = len(prompt.split(', '))
        remove_group = random.randint(3, prompt_groups - 1)
        prompt = prompt.replace(remove_group, '')

    lora_count = random.randint(0, 2)
    for li in range(lora_count):
        prompt += ' <' + random.choice(lr)['name'] + ':' + str(100.0/100.0/random.choice([1, 2, 4])) + '>'

    payload = {
        "prompt": prompt,
        "negative_prompt": '(worst quality), (low quality), (normal quality), lowres, blurry, bad anatomy, bad hands, ((monochrome)), ((grayscale)), watermark',
        "steps": random.choice([20, 30, 50]),
        "sampler_name": sampler,
        "cfg_scale": random.randint(4, 7),
        "seed": -1,
        "width": size_w,
        "height": size_h,
        "batch_size": 4
    }

    response = requests.post(url=f'{url}/sdapi/v1/txt2img', json=payload)
    r = response.json()

    save_name = datetime.datetime.now().strftime('%Y%m%dT%H%M%S') + '_' + str(avg_count) + '_' + model['title'].split('.', 1)[0]
    image1 = Image.open(io.BytesIO(base64.b64decode(r['images'][0])))
    image1.save(f'{save_name}_1.png')

    image2 = Image.open(io.BytesIO(base64.b64decode(r['images'][1])))
    image2.save(f'{save_name}_2.png')

    image3 = Image.open(io.BytesIO(base64.b64decode(r['images'][2])))
    image3.save(f'{save_name}_3.png')

    image4 = Image.open(io.BytesIO(base64.b64decode(r['images'][3])))
    image4.save(f'{save_name}_4.png')

    save_image = Image.new('RGB', (2*image1.size[0], 2*image1.size[1]), (0, 0, 0))
    save_image.paste(image1, (0, 0))
    save_image.paste(image2, (image1.size[0], 0))
    save_image.paste(image3, (0, image1.size[1]))
    save_image.paste(image4, (image1.size[0], image1.size[1]))
    save_image.save(f'{save_name}_comb.png')

    info = json.loads(r['info'])
    with open(f'{save_name}_comb_seed.txt', 'w') as f:
        f.write('Prompt: ' + payload['prompt'] + '\n')
        f.write('Negative payload: ' + payload['negative_prompt'] + '\n')
        f.write('Sampler: ' + payload['sampler_name'] + '\n')
        f.write('Steps: ' + str(payload['steps']) + '\n')
        f.write('Size: ' + str(size_w) + 'x' + str(size_h) + '\n')
        f.write('cfg: ' + str(payload['cfg_scale']) + '\n')
        f.write('Primary seed: ' + str(info['seed']) + '\n')
        f.write('Individual seeds: ' + str(info['all_seeds']))

    end = timer()
    avg_total += (end - start)
    avg_count += 1

    eta = ((avg_total / avg_count) * (total_calls - avg_count))
    print("%.2fs ETA %s" % ((end - start), timedelta(seconds = int(eta))))