app.py

import gradio as gr
import json
import numpy as np
import pandas as pd
from datasets import load_from_disk
from itertools import chain
import operator

pd.options.plotting.backend = "plotly"


TITLE = "Diffusion Professions Cluster Explorer"

professions_dset = load_from_disk("professions")
professions_df = professions_dset.to_pandas()


clusters_dicts = dict(
    (num_cl, json.load(open(f"clusters/professions_to_clusters_{num_cl}.json")))
    for num_cl in [12, 24, 48]
)

cluster_summaries_by_size = json.load(open("clusters/cluster_summaries_by_size.json"))

prompts = pd.read_csv("promptsadjectives.csv")
professions = list(sorted([p.lower() for p in prompts["Occupation-Noun"].tolist()]))
models = {
    "All": "All Models",
    "SD_14": "Stable Diffusion 1.4",
    "SD_2": "Stable Diffusion 2",
    "DallE": "Dall-E 2",
}

df_models = {
    "All Models": "All",
    "Stable Diffusion 1.4": "SD_14",
    "Stable Diffusion 2": "SD_2",
    "Dall-E 2": "DallE",
}


def describe_cluster(num_clusters, block="label"):
    cl_dict = clusters_dicts[num_clusters]
    labels_values = sorted(cl_dict.items(), key=operator.itemgetter(1))
    labels_values.reverse()
    total = float(sum(cl_dict.values()))
    lv_prcnt = list(
        (item[0], round(item[1] * 100 / total, 0)) for item in labels_values
    )
    top_label = lv_prcnt[0][0]
    description_string = (
        "<span>The most represented %s is <b>%s</b>, making up about <b>%d%%</b> of the cluster.</span>"
        % (to_string(block), to_string(top_label), lv_prcnt[0][1])
    )
    description_string += "<p>This is followed by: "
    for lv in lv_prcnt[1:]:
        description_string += "<BR/><b>%s:</b> %d%%" % (to_string(lv[0]), lv[1])
    description_string += "</p>"
    return description_string


def make_profession_plot(num_clusters, prof_name):
    sorted_cl_scores = [
        (k, v)
        for k, v in sorted(
            clusters_dicts[num_clusters]["All"][prof_name][
                "cluster_proportions"
            ].items(),
            key=lambda x: x[1],
            reverse=True,
        )
        if v > 0
    ]
    pre_pandas = dict(
        [
            (
                models[mod_name],
                dict(
                    (
                        f"Cluster {k}",
                        clusters_dicts[num_clusters][mod_name][prof_name][
                            "cluster_proportions"
                        ][k],
                    )
                    for k, _ in sorted_cl_scores
                ),
            )
            for mod_name in models
        ]
    )
    df = pd.DataFrame.from_dict(pre_pandas)
    prof_plot = df.plot(kind="bar", barmode="group")
    cl_summary_text = f"Profession ``{prof_name}'':\n"
    for cl_id, _ in sorted_cl_scores:
        cl_summary_text += f"- {cluster_summaries_by_size[str(num_clusters)][int(cl_id)].replace(' gender terms', '').replace('; ethnicity terms:', ',')} \n"
    return (
        prof_plot,
        gr.update(
            choices=[k for k, _ in sorted_cl_scores], value=sorted_cl_scores[0][0]
        ),
        gr.update(value=cl_summary_text),
    )


def make_profession_table(num_clusters, prof_names, mod_name, max_cols=8):
    professions_list_clusters = [
        (
            prof_name,
            clusters_dicts[num_clusters][df_models[mod_name]][prof_name][
                "cluster_proportions"
            ],
        )
        for prof_name in prof_names
    ]
    totals = sorted(
        [
            (
                k,
                sum(
                    prof_clusters[str(k)]
                    for _, prof_clusters in professions_list_clusters
                ),
            )
            for k in range(num_clusters)
        ],
        key=lambda x: x[1],
        reverse=True,
    )[:max_cols]
    prof_list_pre_pandas = [
        dict(
            [
                ("Profession", prof_name),
                (
                    "Entropy",
                    clusters_dicts[num_clusters][df_models[mod_name]][prof_name][
                        "entropy"
                    ],
                ),
                (
                    "Labor Women",
                    clusters_dicts[num_clusters][df_models[mod_name]][prof_name][
                        "labor_fm"
                    ][0],
                ),
                ("", ""),
            ]
            + [(f"Cluster {k}", prof_clusters[str(k)]) for k, v in totals if v > 0]
        )
        for prof_name, prof_clusters in professions_list_clusters
    ]
    clusters_df = pd.DataFrame.from_dict(prof_list_pre_pandas)
    cl_summary_text = ""
    for cl_id, _ in totals[:max_cols]:
        cl_summary_text += f"- {cluster_summaries_by_size[str(num_clusters)][cl_id].replace(' gender terms', '').replace('; ethnicity terms:', ',')} \n"
    return (
        [c[0] for c in totals],
        (
            clusters_df.style.background_gradient(
                axis=None, vmin=0, vmax=100, cmap="YlGnBu"
            )
            .format(precision=1)
            .to_html()
        ),
        gr.update(value=cl_summary_text),
    )


def get_image(model, fname, score):
    return (
        professions_dset.select(
            professions_df[
                (professions_df["image_path"] == fname)
                & (professions_df["model"] == model)
            ].index
        )["image"][0],
        " ".join(fname.split("/")[0].split("_")[4:])
        + f" | {score:.2f}"
        + f" | {models[model]}",
    )


def show_examplars(num_clusters, prof_name, cl_id, confidence_threshold=0.6):
    # only show images where the similarity to the centroid is > confidence_threshold
    examplars_dict = clusters_dicts[num_clusters]["All"][prof_name][
        "cluster_examplars"
    ][str(cl_id)]
    l = [
        tuple(img)
        for img in examplars_dict["close"]
        + examplars_dict["mid"][:2]
        + examplars_dict["far"]
    ]
    l = [
        img
        for i, img in enumerate(l)
        if img[0] > confidence_threshold and img not in l[:i]
    ]
    return (
        [get_image(model, fname, score) for score, model, fname in l],
        gr.update(
            label=f"Generations for profession ''{prof_name}'' assigned to cluster {cl_id} of {num_clusters}"
        ),
    )


with gr.Blocks(title=TITLE) as demo:
    gr.Markdown(
        """
        # Identity Biases in Diffusion Models: Professions
        
        This tool helps you explore the different clusters that we discovered in the images generated by 3 text-to-image models: Dall-E 2, Stable Diffusion v.1.4 and v.2.
        This work was done in the scope of the [Stable Bias Project](https://huggingface.co/spaces/society-ethics/StableBias).
        """
    )
    gr.HTML(
        """<span style="color:red" font-size:smaller>⚠️ DISCLAIMER: the images displayed by this tool were generated by text-to-image systems and may depict offensive stereotypes or contain explicit content.</span>"""
    )
    with gr.Tab("Professions Overview"):
        gr.Markdown(
            "Select one or more professions and models from the dropdowns on the left to see which clusters are most representative for this combination. Try choosing different numbers of clusters to see if the results change, and then go to the 'Profession Focus' tab to go more in-depth into these results."
        )
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("Select the parameters here:")
                num_clusters = gr.Radio(
                    [12, 24, 48],
                    value=12,
                    label="How many clusters do you want to use to represent identities?",
                )
                model_choices = gr.Dropdown(
                    [
                        "All Models",
                        "Stable Diffusion 1.4",
                        "Stable Diffusion 2",
                        "Dall-E 2",
                    ],
                    value="All Models",
                    label="Which models do you want to compare?",
                    interactive=True,
                )
                profession_choices_overview = gr.Dropdown(
                    professions,
                    value=["CEO", "director", "social assistant", "social worker"],
                    label="Which professions do you want to compare?",
                    multiselect=True,
                    interactive=True,
                )
            with gr.Column(scale=3):
                with gr.Row():
                    table = gr.HTML(
                        label="Profession assignment per cluster", wrap=True
                    )
                with gr.Row():
                    # clusters = gr.Dataframe(type="array", visible=False, col_count=1)
                    clusters = gr.Textbox(label="clusters", visible=False)
                    gr.Markdown(
                        """
                        ##### What do the clusters mean?
                        Below is a summary of the identity cluster compositions.
                        For more details, see the [companion demo](https://huggingface.co/spaces/society-ethics/DiffusionFaceClustering):
                        """
                    )
                with gr.Row():
                    with gr.Accordion(label="Cluster summaries", open=True):
                        cluster_descriptions_table = gr.Text(
                            "TODO", label="Cluster summaries", show_label=False
                        )
    with gr.Tab("Profession Focus"):
        with gr.Row():
            with gr.Column():
                gr.Markdown(
                    "Select profession to visualize and see which clusters and identity groups are most represented in the profession, as well as some examples of generated images below."
                )
                profession_choice_focus = gr.Dropdown(
                    choices=professions,
                    value="scientist",
                    label="Select profession:",
                )
                num_clusters_focus = gr.Radio(
                    [12, 24, 48],
                    value=12,
                    label="How many clusters do you want to use to represent identities?",
                )
            with gr.Column():
                plot = gr.Plot(
                    label=f"Makeup of the cluster assignments for profession {profession_choice_focus}"
                )
        with gr.Row():
            with gr.Column():
                gr.Markdown(
                    """
                    ##### What do the clusters mean?
                    Below is a summary of the identity cluster compositions.
                    For more details, see the [companion demo](https://huggingface.co/spaces/society-ethics/DiffusionFaceClustering):
                    """
                )
                with gr.Accordion(label="Cluster summaries", open=True):
                    cluster_descriptions = gr.Text(
                        "TODO", label="Cluster summaries", show_label=False
                    )
            with gr.Column():
                gr.Markdown(
                    """
                    ##### What's in the clusters?
                    You can show examples of profession images assigned to each identity cluster by selecting one here:
                    """
                )
                with gr.Accordion(label="Cluster selection", open=True):
                    cluster_id_focus = gr.Dropdown(
                        choices=[i for i in range(num_clusters_focus.value)],
                        value=0,
                        label="Select cluster to visualize:",
                    )
        with gr.Row():
            examplars_plot = gr.Gallery(
                label="Profession images assigned to the selected cluster."
            ).style(grid=4, height="auto", container=True)
    demo.load(
        make_profession_table,
        [num_clusters, profession_choices_overview, model_choices],
        [clusters, table, cluster_descriptions_table],
        queue=False,
    )
    demo.load(
        make_profession_plot,
        [num_clusters_focus, profession_choice_focus],
        [plot, cluster_id_focus, cluster_descriptions],
        queue=False,
    )
    demo.load(
        show_examplars,
        [
            num_clusters_focus,
            profession_choice_focus,
            cluster_id_focus,
        ],
        [examplars_plot, examplars_plot],
        queue=False,
    )
    for var in [num_clusters, model_choices, profession_choices_overview]:
        var.change(
            make_profession_table,
            [num_clusters, profession_choices_overview, model_choices],
            [clusters, table, cluster_descriptions_table],
            queue=False,
        )
    for var in [num_clusters_focus, profession_choice_focus]:
        var.change(
            make_profession_plot,
            [num_clusters_focus, profession_choice_focus],
            [plot, cluster_id_focus, cluster_descriptions],
            queue=False,
        )
    for var in [num_clusters_focus, profession_choice_focus, cluster_id_focus]:
        var.change(
            show_examplars,
            [
                num_clusters_focus,
                profession_choice_focus,
                cluster_id_focus,
            ],
            [examplars_plot, examplars_plot],
            queue=False,
        )


if __name__ == "__main__":
    demo.queue().launch(debug=True)