File size: 1,376 Bytes
e82dda8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 |
from datasets import load_dataset
from collections import Counter
from random import choices, shuffle
import os
parti_prompt_results = []
ORG = "diffusers-parti-prompts"
SUBMISSIONS = {
"Stable Diffusion 1-5": load_dataset(os.path.join(ORG, "sd-v1-5"))["train"],
# "Stable Diffusion 2-1": load_dataset(os.path.join(ORG, "sd-v2-1")),
# "IF-1-0": None,
# "Karlo": None,
# "Kadinsky":
}
NUM_QUESTIONS = 25
submission_names = list(SUBMISSIONS.keys())
num_images = len(SUBMISSIONS[submission_names[0]])
def start():
ids = {id: 0 for id in range(num_images)}
# submissions = load_dataset(os.path.join(ORG, "submissions"))
# submitted_ids = Counter(submissions["ids"])
submitted_ids = {}
ids = {**ids, **submitted_ids}
# sort by count
ids = sorted(ids)
# get lowest count ids
id_candidates = ids[:(8 * NUM_QUESTIONS)]
# get random `NUM_QUESTIONS` ids to check
image_ids = choices(id_candidates, k=NUM_QUESTIONS)
images = {}
for i in range(NUM_QUESTIONS):
order = list(range(len(SUBMISSIONS)))
shuffle(order)
id = image_ids[i]
images[i] = {
"prompt": SUBMISSIONS[submission_names[0]][id]["Prompt"],
"id": id,
"choices": {i: SUBMISSIONS[submission_names[i]][id]["images"] for i in order},
}
return images
hello = start()
|