from datasets import load_dataset from collections import Counter, defaultdict from random import sample, shuffle import datasets from pandas import DataFrame from huggingface_hub import list_datasets import os import gradio as gr import secrets parti_prompt_results = [] ORG = "diffusers-parti-prompts" SUBMISSIONS = { "sd-v1-5": load_dataset(os.path.join(ORG, "sd-v1-5"))["train"], "sd-v2-1": load_dataset(os.path.join(ORG, "sd-v2.1"))["train"], "if-v1-0": load_dataset(os.path.join(ORG, "karlo-v1"))["train"], "karlo": load_dataset(os.path.join(ORG, "if-v-1.0"))["train"], # "Kadinsky": } LINKS = { "sd-v1-5": "https://huggingface.co/runwayml/stable-diffusion-v1-5", "sd-v2-1": "https://huggingface.co/stabilityai/stable-diffusion-2-1", "if-v1-0": "https://huggingface.co/DeepFloyd/IF-I-XL-v1.0", "karlo": "https://huggingface.co/kakaobrain/karlo-v1-alpha", } NUM_QUESTIONS = 10 MODEL_KEYS = "-".join(SUBMISSIONS.keys()) SUBMISSION_ORG = f"results-{MODEL_KEYS}" PROMPT_FORMAT = "Pick the picture that best matches the prompt: **{}**" submission_names = list(SUBMISSIONS.keys()) num_images = len(SUBMISSIONS[submission_names[0]]) def load_submissions(): all_datasets = list_datasets(author=SUBMISSION_ORG) relevant_ids = [d.id for d in all_datasets] submitted_ids = [] for _id in relevant_ids: ds = load_dataset(_id)["train"] submitted_ids += ds["id"] submitted_ids = Counter(submitted_ids) return submitted_ids SUBMITTED_IDS = load_submissions() def generate_random_hash(length=8): """ Generates a random hash of specified length. Args: length (int): The length of the hash to generate. Returns: str: A random hash of specified length. """ if length % 2 != 0: raise ValueError("Length should be an even number.") num_bytes = length // 2 random_bytes = secrets.token_bytes(num_bytes) random_hash = secrets.token_hex(num_bytes) return random_hash def refresh(row_number, dataframe): if row_number == NUM_QUESTIONS: submitted_ids = load_submissions() return start(submitted_ids) else: return dataframe def start(): ids = {id: 0 for id in range(num_images)} ids = {**ids, **SUBMITTED_IDS} # sort by count ids = sorted(ids.items(), key=lambda x: x[1]) freq_ids = defaultdict(list) for k, v in ids: freq_ids[v].append(k) # shuffle in-between categories for k, v_list in freq_ids.items(): shuffle(v_list) freq_ids[v] = v_list shuffled_ids = sum(list(freq_ids.values()), []) # get lowest count ids id_candidates = shuffled_ids[: (10 * NUM_QUESTIONS)] # get random `NUM_QUESTIONS` ids to check image_ids = sample(id_candidates, k=NUM_QUESTIONS) images = {} for i in range(NUM_QUESTIONS): order = list(range(len(SUBMISSIONS))) shuffle(order) id = image_ids[i] row = SUBMISSIONS[submission_names[0]][id] images[i] = { "prompt": row["Prompt"], "result": "", "id": id, "Challenge": row["Challenge"], "Category": row["Category"], "Note": row["Note"], } for n, m in enumerate(order): images[i][f"choice_{n}"] = m images_frame = DataFrame.from_dict(images, orient="index") return images_frame def process(dataframe, row_number=0): if row_number == NUM_QUESTIONS: return None, "", "" image_id = dataframe.iloc[row_number]["id"] choices = [ submission_names[dataframe.iloc[row_number][f"choice_{i}"]] for i in range(len(SUBMISSIONS)) ] images = [SUBMISSIONS[c][int(image_id)]["images"] for c in choices] prompt = SUBMISSIONS[choices[0]][int(image_id)]["Prompt"] prompt = PROMPT_FORMAT.format(prompt) counter = f"{row_number + 1}/{NUM_QUESTIONS}" return images, prompt, counter def write_result(user_choice, row_number, dataframe): if row_number == NUM_QUESTIONS: return row_number, dataframe user_choice = int(user_choice) chosen_model = submission_names[dataframe.iloc[row_number][f"choice_{user_choice}"]] dataframe.loc[row_number, "result"] = chosen_model return row_number + 1, dataframe def get_index(evt: gr.SelectData) -> int: return evt.index def change_view(row_number, dataframe): if row_number == NUM_QUESTIONS: favorite_model = dataframe["result"].value_counts().idxmax() dataset = datasets.Dataset.from_pandas(dataframe) dataset = dataset.remove_columns(set(dataset.column_names) - set(["id", "result"])) hash = generate_random_hash() repo_id = os.path.join(SUBMISSION_ORG, hash) dataset.push_to_hub(repo_id, token=os.getenv("HF_TOKEN")) return { intro_view: gr.update(visible=True), result_view: gr.update(visible=True), gallery_view: gr.update(visible=False), result: f"You are of type: [**{favorite_model}**]({LINKS[favorite_model]}) 🔥", } else: return { intro_view: gr.update(visible=False), result_view: gr.update(visible=False), gallery_view: gr.update(visible=True), result: "", } TITLE = "# Community Parti Prompts - Who is your open-source genAI model?" DESCRIPTION = """ *This is an interactive game in which you click through pre-generated images from SD-v1-5, SD-v2.1, Karlo, and IF using [Parti Prompts](https://huggingface.co/datasets/nateraw/parti-prompts) prompts.* \n *You choices will go into the public community [genAI leaderboard](TODO).* """ EXPLANATION = """\n\n ## How it works 📖 \n\n 1. Click on 'Start' 2. A prompt and 4 different images are displayed 3. Select your favorite image 4. Click on 'Select' 5. After 10 rounds your favorite diffusion model is displayed """ GALLERY_COLUMN_NUM = len(SUBMISSIONS) with gr.Blocks() as demo: gr.Markdown(TITLE) gr.Markdown(DESCRIPTION) with gr.Column(visible=True) as intro_view: gr.Markdown(EXPLANATION) start_button = gr.Button("Start").style(full_width=False) headers = ["prompt", "result", "id", "Challenge", "Category", "Note"] + [ f"choice_{i}" for i in range(len(SUBMISSIONS)) ] datatype = ["str", "str", "number", "str", "str", "str"] + len(SUBMISSIONS) * [ "number" ] with gr.Column(visible=False): row_number = gr.Number( label="Current row selection index", value=0, precision=0, interactive=False, ) # Create Data Frame with gr.Column(visible=False) as result_view: result = gr.Markdown("") dataframe = gr.Dataframe( headers=headers, datatype=datatype, row_count=NUM_QUESTIONS, col_count=(6 + len(SUBMISSIONS), "fixed"), interactive=False, ) gr.Markdown("Click on start to play again!") with gr.Column(visible=False) as gallery_view: counter = gr.Markdown(f" ### 1/{NUM_QUESTIONS}") prompt = gr.Markdown(PROMPT_FORMAT.format("")) gallery = gr.Gallery( label="All images", show_label=False, elem_id="gallery" ).style(columns=GALLERY_COLUMN_NUM, object_fit="contain") next_button = gr.Button("Select").style(full_width=False) with gr.Column(visible=False): selected_image = gr.Number(label="Selected index", value=-1, precision=0) start_button.click( fn=start, inputs=[], outputs=dataframe, show_progress=True ).then( fn=lambda x: 0 if x == NUM_QUESTIONS else x, inputs=[row_number], outputs=[row_number], ).then( fn=change_view, inputs=[row_number, dataframe], outputs=[intro_view, result_view, gallery_view, result] ).then( fn=process, inputs=[dataframe], outputs=[gallery, prompt, counter] ) gallery.select( fn=get_index, outputs=selected_image, queue=False, ) next_button.click( fn=write_result, inputs=[selected_image, row_number, dataframe], outputs=[row_number, dataframe], ).then( fn=change_view, inputs=[row_number, dataframe], outputs=[intro_view, result_view, gallery_view, result] ).then( fn=process, inputs=[dataframe, row_number], outputs=[gallery, prompt, counter] ).then( fn=lambda x: 0 if x == NUM_QUESTIONS else x, inputs=[row_number], outputs=[row_number], ).then( fn=refresh, inputs=[row_number, dataframe], outputs=[dataframe], ) demo.launch()