diffusers/tools

1#!/usr/bin/env python32import os3 4from collections import Counter5from datasets import load_dataset6from huggingface_hub import HfApi, list_datasets7 8 9api = HfApi(token=os.environ.get("HF_TOKEN", None))10def restart_space():11  api.restart_space(repo_id="OpenGenAI/parti-prompts-leaderboard")12 13parti_prompt_results = []14ORG = "diffusers-parti-prompts"15SUBMISSIONS = {16    "kand2": None,17    "sdxl": None,18    "wuerst": None,19    "karlo": None,20}21LINKS = {22    "kand2": "https://huggingface.co/kandinsky-community/kandinsky-2-2-decoder",23    "sdxl": "https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0",24    "wuerst": "https://huggingface.co/warp-ai/wuerstchen",25    "karlo": "https://huggingface.co/kakaobrain/karlo-v1-alpha",26}27MODEL_KEYS = "-".join(SUBMISSIONS.keys())28SUBMISSION_ORG = f"result-{MODEL_KEYS}"29 30submission_names = list(SUBMISSIONS.keys())31 32ORG = "diffusers-parti-prompts"33SUBMISSIONS = {34    "kand2": load_dataset(os.path.join(ORG, "kandinsky-2-2"))["train"],35    "sdxl": load_dataset(os.path.join(ORG, "sdxl-1.0-refiner"))["train"],36    "wuerst": load_dataset(os.path.join(ORG, "wuerstchen"))["train"],37    "karlo": load_dataset(os.path.join(ORG, "karlo-v1"))["train"],38}39ds = load_dataset("nateraw/parti-prompts")["train"]40 41parti_prompt_categories = ds["Category"]42parti_prompt_challenge = ds["Challenge"]43 44UPLOAD_ORG = "Almost-AGI-Diffusion"45 46def load_non_solved():47    all_datasets = list_datasets(author=SUBMISSION_ORG)48    relevant_ids = [d.id for d in all_datasets]49    50    all_non_solved_image_ids = []51 52    for _id in relevant_ids:53        try:54            ds = load_dataset(_id)["train"]55        except:56            continue57 58        for result, image_id in zip(ds["result"], ds["id"]):59            if result == "":60                all_non_solved_image_ids.append(image_id)61 62    all_non_solved_image_ids_dict = Counter(all_non_solved_image_ids)63    all_non_solved_image_ids = list(all_non_solved_image_ids_dict.keys())64    all_non_solved_image_votes = list(all_non_solved_image_ids_dict.values())65 66    return all_non_solved_image_ids, all_non_solved_image_votes67 68def main():69    non_solved_ids, upvotes = load_non_solved()70 71    for name, ds in SUBMISSIONS.items():72        ds_to_push = ds.select(non_solved_ids)73 74        votes_column = upvotes75 76        ds_to_push = ds_to_push.add_column("upvotes", votes_column)77        sorted_ds = ds_to_push.sort("upvotes", reverse=True)78 79        sorted_ds.push_to_hub(f"{UPLOAD_ORG}/{name}")80 81if __name__ == "__main__":82    main()83