OpenGenAI/open-parti-prompts
21
1from datasets import load_dataset2from collections import Counter, defaultdict3from random import sample, shuffle4from collections import Counter5import datasets6from pandas import DataFrame7from huggingface_hub import list_datasets8import os9import gradio as gr10 11import secrets12 13 14parti_prompt_results = []15ORG = "diffusers-parti-prompts"16SUBMISSIONS = {17 "kand2": load_dataset(os.path.join(ORG, "kandinsky-2-2"))["train"],18 "sdxl": load_dataset(os.path.join(ORG, "sdxl-1.0-refiner"))["train"],19 "wuerst": load_dataset(os.path.join(ORG, "wuerstchen"))["train"],20 "karlo": load_dataset(os.path.join(ORG, "karlo-v1"))["train"],21}22 23LINKS = {24 "kand2": "https://huggingface.co/kandinsky-community/kandinsky-2-2-decoder",25 "sdxl": "https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0",26 "wuerst": "https://huggingface.co/warp-ai/wuerstchen",27 "karlo": "https://huggingface.co/kakaobrain/karlo-v1-alpha",28}29KANDINSKY = """30"## The creative one 🎨! 3132\n You mostly resonate with **Kandinsky 2.2** released by AI Forever.33\n Kandinsky 2.2 has a similar architecture to DALLE-2 and works extremely well for artistic, colorful generations.34\n Check out your soulmate [here](https://huggingface.co/kandinsky-community/kandinsky-2-2-decoder).35"""36SDXL_RESULT = """37## The powerful one ⚡! 3839\n You mostly resonate with **Stable Diffusion XL** released by Stability AI.40\n Stable Diffusion XL consists of a two diffusion models that are chained together, a base model and a refiner model. Together, the system contains roughly 5 billion parameters.41\n It's the latest open-source release of Stable Diffusion and allows to render stunning images of much larger sizes than Stable Diffusion v1.42Try it out [here](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0).43"""44WUERSTCHEN = """45## The innovative one ⚗️ !4647\n You mostly resonate with **Wuerstchen** released by the WARP team.48\n Wuerstchen is a three stage diffusion model that proposed a very novel, innovative model architecture.49\n Wuerstchen is able to generate very large images (up to 1024x2048) in just a few seconds.50\n The model has an amazing image quality vs. speed trade-off.51\n Check out your new best friend [here](https://huggingface.co/warp-ai/wuerstchen).52"""53KARLO = """54## The precise one 🎯!5556\n You mostly resonate with **Karlo** released by KakaoBrain.57\n Karlo is based on the same architecture as DALLE-2 and has been trained on the [well curated COYO dataset](https://huggingface.co/datasets/kakaobrain/coyo-700m).58\n Play around with it [here]("https://huggingface.co/kakaobrain/karlo-v1-alpha").59"""60 61RESULT = {62 "kand2": KANDINSKY,63 "wuerst": WUERSTCHEN,64 "sdxl": SDXL_RESULT,65 "karlo": KARLO,66}67NUM_QUESTIONS = 1068MODEL_KEYS = "-".join(SUBMISSIONS.keys())69SUBMISSION_ORG = f"result-{MODEL_KEYS}"70PROMPT_FORMAT = " Select the image that best matches the prompt and click on 'Submit'. Remember that if multiple images match the prompt equally well, select them all. If no image matches the prompt, no image shall be selected."71 72submission_names = list(SUBMISSIONS.keys())73num_images = len(SUBMISSIONS[submission_names[0]])74 75 76def load_submissions():77 all_datasets = list_datasets(author=SUBMISSION_ORG)78 relevant_ids = [d.id for d in all_datasets]79 80 submitted_ids = []81 for _id in relevant_ids:82 ds = load_dataset(_id)["train"]83 submitted_ids += ds["id"]84 85 submitted_ids = Counter(submitted_ids)86 return submitted_ids87 88 89SUBMITTED_IDS = load_submissions()90 91 92def generate_random_hash(length=8):93 """94 Generates a random hash of specified length.95 96 Args:97 length (int): The length of the hash to generate.98 99 Returns:100 str: A random hash of specified length.101 """102 if length % 2 != 0:103 raise ValueError("Length should be an even number.")104 105 num_bytes = length // 2106 random_bytes = secrets.token_bytes(num_bytes)107 random_hash = secrets.token_hex(num_bytes)108 109 return random_hash110 111 112def refresh(row_number, dataframe):113 if row_number == NUM_QUESTIONS:114 submitted_ids = load_submissions()115 return start(submitted_ids)116 else:117 return dataframe118 119def start():120 ids = {id: 0 for id in range(num_images)}121 ids = {**ids, **SUBMITTED_IDS}122 123 # sort by count124 ids = sorted(ids.items(), key=lambda x: x[1])125 freq_ids = defaultdict(list)126 for k, v in ids:127 freq_ids[v].append(k)128 129 # shuffle in-between categories130 for k, v_list in freq_ids.items():131 shuffle(v_list)132 freq_ids[k] = v_list133 134 shuffled_ids = sum(list(freq_ids.values()), [])135 136 # get lowest count ids137 id_candidates = shuffled_ids[: (10 * NUM_QUESTIONS)]138 139 # get random `NUM_QUESTIONS` ids to check140 image_ids = sample(id_candidates, k=NUM_QUESTIONS)141 images = {}142 143 for i in range(NUM_QUESTIONS):144 order = list(range(len(SUBMISSIONS)))145 shuffle(order)146 147 id = image_ids[i]148 row = SUBMISSIONS[submission_names[0]][id]149 images[i] = {150 "prompt": row["Prompt"],151 "result": "",152 "id": id,153 "Challenge": row["Challenge"],154 "Category": row["Category"],155 "Note": row["Note"],156 }157 for n, m in enumerate(order):158 images[i][f"choice_{n}"] = m159 160 images_frame = DataFrame.from_dict(images, orient="index")161 return images_frame162 163 164def process(dataframe, row_number=0):165 if row_number == NUM_QUESTIONS:166 nones = len(RESULT) * [None]167 falses = len(RESULT) * [False]168 return *nones, *falses, "", ""169 170 image_id = dataframe.iloc[row_number]["id"]171 choices = [172 submission_names[dataframe.iloc[row_number][f"choice_{i}"]]173 for i in range(len(SUBMISSIONS))174 ]175 images = [SUBMISSIONS[c][int(image_id)]["images"] for c in choices]176 177 prompt = SUBMISSIONS[choices[0]][int(image_id)]["Prompt"]178 prompt = f'# "{prompt}"'179 counter = f"***{row_number + 1}/{NUM_QUESTIONS} {PROMPT_FORMAT}***"180 image_buttons = len(images) * [False]181 182 return *images, *image_buttons, prompt, counter183 184 185def write_result(user_choice, row_number, dataframe):186 if row_number == NUM_QUESTIONS:187 return row_number, dataframe188 189 user_choices = []190 for i, b in enumerate(str(user_choice)):191 if bool(int(b)):192 user_choices.append(i)193 194 chosen_models = []195 for user_choice in user_choices:196 chosen_models.append(submission_names[dataframe.iloc[row_number][f"choice_{user_choice}"]])197 198 print(chosen_models)199 dataframe.loc[row_number, "result"] = ",".join(chosen_models)200 return row_number + 1, dataframe201 202 203def get_index(evt: gr.SelectData) -> int:204 return evt.index205 206 207def change_view(row_number, dataframe):208 if row_number == NUM_QUESTIONS:209 210 results = sum([x.split(",") for x in dataframe["result"].values], [])211 results = [r for r in results if len(r) > 0]212 favorite_model = Counter(results).most_common(1)[0][0]213 214 dataset = datasets.Dataset.from_pandas(dataframe)215 dataset = dataset.remove_columns(set(dataset.column_names) - set(["id", "result"]))216 hash = generate_random_hash()217 repo_id = os.path.join(SUBMISSION_ORG, hash)218 219 dataset.push_to_hub(repo_id, token=os.getenv("HF_TOKEN"))220 return {221 intro_view: gr.update(visible=False),222 result_view: gr.update(visible=True),223 gallery_view: gr.update(visible=False),224 start_view: gr.update(visible=True),225 result: RESULT[favorite_model],226 }227 else:228 return {229 intro_view: gr.update(visible=False),230 result_view: gr.update(visible=False),231 gallery_view: gr.update(visible=True),232 start_view: gr.update(visible=False),233 result: "",234 }235 236 237TITLE = "# What AI model is best for you? 👩⚕️"238 239DESCRIPTION = """240***How it works*** 📖 \n\n241- Upon clicking start, you are shown image descriptions alongside four AI generated images.242\n- Select the image that best matches the prompt. If multiple images match the prompt equally well, select all images. If no image matches the prompt, leave all images unchecked.243\n- Answer **10** questions to find out what AI generator most resonates with you. 244\n- Your submissions contribute to [**Open Parti Prompts Leaderboard**](https://huggingface.co/spaces/OpenGenAI/parti-prompts-leaderboard) ❤️.245\n\n246"""247 248NOTE = """\n\n\n\n249The prompts you are shown originate from the [Parti Prompts](https://huggingface.co/datasets/nateraw/parti-prompts) dataset.250Parti Prompts is designed to test text-to-image AI models on 1600+ prompts of varying difficulty and categories.251The images you are shown have been pre-generated with 4 state-of-the-art open-sourced text-to-image models.252You answers will be used to contribute to the official [**Open Parti Prompts Leaderboard**](https://huggingface.co/spaces/OpenGenAI/parti-prompts-leaderboard).253Every couple months, the generated images will be updated with possibly improved models. The current models and code that was used to generate the images can be verified here:\n254- [kandinsky-2-2](https://huggingface.co/kandinsky-community/kandinsky-2-2-decoder) \n255- [wuerstchen](https://huggingface.co/warp-ai/wuerstchen) \n256- [sdxl-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0) \n257- [karlo](https://huggingface.co/datasets/diffusers-parti-prompts/karlo-v1) \n258"""259 260GALLERY_COLUMN_NUM = len(SUBMISSIONS)261 262with gr.Blocks() as demo:263 gr.Markdown(TITLE)264 with gr.Column(visible=True) as intro_view:265 gr.Markdown(DESCRIPTION)266 267 headers = ["prompt", "result", "id", "Challenge", "Category", "Note"] + [268 f"choice_{i}" for i in range(len(SUBMISSIONS))269 ]270 datatype = ["str", "str", "number", "str", "str", "str"] + len(SUBMISSIONS) * [271 "number"272 ]273 274 with gr.Column(visible=False):275 row_number = gr.Number(276 label="Current row selection index",277 value=0,278 precision=0,279 interactive=False,280 )281 282 # Create Data Frame283 with gr.Column(visible=False) as result_view:284 result = gr.Markdown("")285 dataframe = gr.Dataframe(286 headers=headers,287 datatype=datatype,288 row_count=NUM_QUESTIONS,289 col_count=(6 + len(SUBMISSIONS), "fixed"),290 interactive=False,291 )292 gr.Markdown("Click on start to play again!")293 294 with gr.Column(visible=True) as start_view:295 start_button = gr.Button("Start").style(full_width=True)296 gr.Markdown(NOTE)297 298 with gr.Column(visible=False):299 selected_image = gr.Textbox(label="Selected indexes")300 301 with gr.Column(visible=False) as gallery_view:302 with gr.Row():303 counter = gr.Markdown(f"***1/{NUM_QUESTIONS} {PROMPT_FORMAT}***")304 with gr.Row():305 prompt = gr.Markdown("")306 with gr.Blocks():307 with gr.Row():308 with gr.Column(min_width=200) as c1:309 image_1 = gr.Image(interactive=False)310 image_1_button = gr.Checkbox(False, label="Image 1").style(full_width=True)311 with gr.Column(min_width=200) as c2:312 image_2 = gr.Image(interactive=False)313 image_2_button = gr.Checkbox(False, label="Image 2").style(full_width=True)314 with gr.Column(min_width=200) as c3:315 image_3 = gr.Image(interactive=False)316 image_3_button = gr.Checkbox(False, label="Image 3").style(full_width=True)317 with gr.Column(min_width=200) as c4:318 image_4 = gr.Image(interactive=False)319 image_4_button = gr.Checkbox(False, label="Image 4").style(full_width=True) 320 with gr.Row():321 submit_button = gr.Button("Submit").style(full_width=True) 322 323 start_button.click(324 fn=start,325 inputs=[],326 outputs=dataframe,327 show_progress=True328 ).then(329 fn=lambda x: 0 if x == NUM_QUESTIONS else x,330 inputs=[row_number],331 outputs=[row_number],332 ).then(333 fn=change_view,334 inputs=[row_number, dataframe],335 outputs=[intro_view, result_view, gallery_view, start_view, result],336 ).then(337 fn=process, 338 inputs=[dataframe],339 outputs=[image_1, image_2, image_3, image_4, image_1_button, image_2_button, image_3_button, image_4_button, prompt, counter]340 )341 342 def integerize(x1, x2, x3, x4):343 number = f"{int(x1)}{int(x2)}{int(x3)}{int(x4)}"344 return number345 346 submit_button.click(347 fn=integerize,348 inputs=[image_1_button, image_2_button, image_3_button, image_4_button],349 outputs=[selected_image],350 ).then(351 fn=write_result,352 inputs=[selected_image, row_number, dataframe],353 outputs=[row_number, dataframe],354 ).then(355 fn=change_view,356 inputs=[row_number, dataframe],357 outputs=[intro_view, result_view, gallery_view, start_view, result]358 ).then(359 fn=process,360 inputs=[dataframe, row_number],361 outputs=[image_1, image_2, image_3, image_4, image_1_button, image_2_button, image_3_button, image_4_button, prompt, counter],362 ).then(363 fn=lambda x: 0 if x == NUM_QUESTIONS else x,364 inputs=[row_number],365 outputs=[row_number],366 ).then(367 fn=refresh,368 inputs=[row_number, dataframe],369 outputs=[dataframe],370 )371 372demo.launch()373