File size: 6,423 Bytes
2cd207c
 
 
 
 
 
 
 
 
 
dbb7d2e
2cd207c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
dbb7d2e
 
 
 
 
 
2cd207c
 
 
 
 
 
 
99a4efe
2cd207c
 
 
 
 
 
 
 
 
99a4efe
2cd207c
 
99a4efe
2cd207c
 
 
 
 
99a4efe
 
 
 
2cd207c
 
 
99a4efe
2cd207c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5acd4d2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2cd207c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
dbb7d2e
 
2cd207c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
import streamlit as st
import numpy as np
import random
import pandas as pd
import glob
import csv
from PIL import Image
from datasets import load_dataset, Dataset, load_from_disk
from huggingface_hub import login
import os
import datasets


class GalleryApp:
    def __init__(self, promptBook):
        self.promptBook = promptBook
        st.set_page_config(layout="wide")

    def gallery(self, items, col_num, info):
        cols = st.columns(col_num)
        # # sort items by brisque score
        # items = items.sort_values(by=['brisque'], ascending=True).reset_index(drop=True)
        for idx in range(len(items)):
            with cols[idx % col_num]:
                image = st.session_state.images[items.iloc[idx]['row_idx'].item()]['image']
                st.image(image,
                         use_column_width=True,
                )
                with st.expander('Similarity Info'):
                    tab1, tab2 = st.tabs(['Most Similar', 'Least Similar'])
                    with tab1:
                        st.image(image, use_column_width=True)
                    with tab2:
                        st.image(image, use_column_width=True)
                for key in info:
                    st.write(f"**{key}**: {items.iloc[idx][key]}")

    def app(self):
        st.title('Model Coffer Gallery')
        st.write('This is a gallery of images generated by the models in the Model Coffer')

        # metadata, images = st.columns([1, 3])
        # with images:
            # prompt_tags = self.promptBook['tag'].unique()
            # # sort tags by alphabetical order
            # prompt_tags = np.sort(prompt_tags)
            #
            # selecters = st.columns(3)
            # with selecters[0]:
            #     tag = st.selectbox('Select a tag', prompt_tags)

        with st.sidebar:
            prompt_tags = self.promptBook['tag'].unique()
            # sort tags by alphabetical order
            prompt_tags = np.sort(prompt_tags)[::-1]

            tag = st.selectbox('Select a tag', prompt_tags)

            items = self.promptBook[self.promptBook['tag'] == tag].reset_index(drop=True)

            original_prompts = np.sort(items['prompt'].unique())[::-1]

            # remove the first four items in the prompt, which are mostly the same
            prompts = [', '.join(x.split(', ')[4:]) for x in original_prompts]
            prompt = st.selectbox('Select prompt', prompts)

            idx = prompts.index(prompt)
            prompt_full = ', '.join(original_prompts[idx].split(', ')[:4]) + ', ' + prompt

            prompt_id = items[items['prompt'] == prompt_full]['prompt_id'].unique()[0]
            items = items[items['prompt_id'] == prompt_id].reset_index(drop=True)

            st.write('**Prompt ID**')
            st.caption(f"{prompt_id}")
            st.write('**Prompt**')
            st.caption(f"{items['prompt'][0]}")
            st.write('**Negative Prompt**')
            st.caption(f"{items['negativePrompt'][0]}")
            st.write('**Sampler**')
            st.caption(f"{items['sampler'][0]}")
            st.write('**cfgScale**')
            st.caption(f"{items['cfgScale'][0]}")
            st.write('**Size**')
            st.caption(f"width: {items['size'][0].split('x')[0]}, height: {items['size'][0].split('x')[1]}")
            st.write('**Seed**')
            st.caption(f"{items['seed'][0]}")

        # with images:
        selecters = st.columns([1, 1, 2])

        with selecters[0]:
            sort_by = st.selectbox('Sort by', items.columns[11: -1])

        with selecters[1]:
            order = st.selectbox('Order', ['Ascending', 'Descending'], index=1 if sort_by == 'clip_score' or sort_by == 'model_download_count' else 0)
            if order == 'Ascending':
                order = True
            else:
                order = False

        items = items.sort_values(by=[sort_by], ascending=order).reset_index(drop=True)
        with selecters[2]:
            info = st.multiselect('Show Info',
                                  ['brisque_score', 'clip_score', 'model_download_count', 'model_name', 'model_id',
                                   'modelVersion_name', 'modelVersion_id'],
                                  default=sort_by)

        col_num = st.slider('Number of columns', min_value=1, max_value=9, value=4, step=1, key='col_num')
        self.gallery(items, col_num, info)


if __name__ == '__main__':
    login(token=os.environ.get("HF_TOKEN"))

    if 'roster' not in st.session_state:
        print('loading roster')
        # st.session_state.roster = pd.DataFrame(load_dataset('NYUSHPRP/ModelCofferRoster', split='train'))
        st.session_state.roster = pd.DataFrame(load_from_disk(os.path.join(os.getcwd(), 'data', 'roster')))
        st.session_state.roster = st.session_state.roster[['model_id', 'model_name', 'modelVersion_id', 'modelVersion_name',
                                          'model_download_count']].drop_duplicates().reset_index(drop=True)
    # add model download count from roster to promptbook dataframe
    if 'promptBook' not in st.session_state:
        print('loading promptBook')

        st.session_state.promptBook = pd.DataFrame(load_dataset('NYUSHPRP/ModelCofferMetadata', split='train'))
        st.session_state.images = load_from_disk(os.path.join(os.getcwd(), 'data', 'promptbook'))['train']
        print(st.session_state.images)
        print('images loaded')
        # st.session_state.promptBook = pd.DataFrame(load_dataset('NYUSHPRP/ModelCofferPromptBook', split='train'))
        st.session_state.promptBook = st.session_state.promptBook.merge(st.session_state.roster[['model_id', 'model_name', 'modelVersion_id', 'modelVersion_name', 'model_download_count']], on=['model_id', 'modelVersion_id'], how='left')

        # add column to record current row index
        st.session_state.promptBook['row_idx'] = st.session_state.promptBook.index
        print('promptBook loaded')
    # print(st.session_state.promptBook)

    check_roster_error = False
    if check_roster_error:
        # print all rows with the same model_id and modelVersion_id but different model_download_count in roster
        print(st.session_state.roster[st.session_state.roster.duplicated(subset=['model_id', 'modelVersion_id'], keep=False)].sort_values(by=['model_id', 'modelVersion_id']))

    app = GalleryApp(promptBook=st.session_state.promptBook)
    app.app()