beta version of new gallery view

Archive/Gallery_beta0913.py

@@ -0,0 +1,643 @@
+import os
+import requests
+
+import altair as alt
+import numpy as np
+import pandas as pd
+import streamlit as st
+import streamlit.components.v1 as components
+
+from bs4 import BeautifulSoup
+from datasets import load_dataset, Dataset, load_from_disk
+from huggingface_hub import login
+from streamlit_agraph import agraph, Node, Edge, Config
+from streamlit_extras.switch_page_button import switch_page
+from sklearn.svm import LinearSVC
+
+SCORE_NAME_MAPPING = {'clip': 'clip_score', 'rank': 'msq_score', 'pop': 'model_download_count'}
+
+
+class GalleryApp:
+    def __init__(self, promptBook, images_ds):
+        self.promptBook = promptBook
+        self.images_ds = images_ds
+
+    def gallery_standard(self, items, col_num, info):
+        rows = len(items) // col_num + 1
+        containers = [st.container() for _ in range(rows)]
+        for idx in range(0, len(items), col_num):
+            row_idx = idx // col_num
+            with containers[row_idx]:
+                cols = st.columns(col_num)
+                for j in range(col_num):
+                    if idx + j < len(items):
+                        with cols[j]:
+                            # show image
+                            # image = self.images_ds[items.iloc[idx + j]['row_idx'].item()]['image']
+                            image = f"https://modelcofferbucket.s3-accelerate.amazonaws.com/{items.iloc[idx + j]['image_id']}.png"
+                            st.image(image, use_column_width=True)
+
+                            # handel checkbox information
+                            prompt_id = items.iloc[idx + j]['prompt_id']
+                            modelVersion_id = items.iloc[idx + j]['modelVersion_id']
+
+                            check_init = True if modelVersion_id in st.session_state.selected_dict.get(prompt_id, []) else False
+
+                            # st.write("Position: ", idx + j)
+
+                            # show checkbox
+                            st.checkbox('Select', key=f'select_{prompt_id}_{modelVersion_id}', value=check_init)
+
+                            # show selected info
+                            for key in info:
+                                st.write(f"**{key}**: {items.iloc[idx + j][key]}")
+
+    def gallery_graph(self, items):
+        items = load_tsne_coordinates(items)
+
+        # sort items to be popularity from low to high, so that most popular ones will be on the top
+        items = items.sort_values(by=['model_download_count'], ascending=True).reset_index(drop=True)
+
+        scale = 50
+        items.loc[:, 'x'] = items['x'] * scale
+        items.loc[:, 'y'] = items['y'] * scale
+
+        nodes = []
+        edges = []
+
+        for idx in items.index:
+            # if items.loc[idx, 'modelVersion_id'] in st.session_state.selected_dict.get(items.loc[idx, 'prompt_id'], 0):
+            #     opacity = 0.2
+            # else:
+            #     opacity = 1.0
+
+            nodes.append(Node(id=items.loc[idx, 'image_id'],
+                              # label=str(items.loc[idx, 'model_name']),
+                              title=f"model name: {items.loc[idx, 'model_name']}\nmodelVersion name: {items.loc[idx, 'modelVersion_name']}\nclip score: {items.loc[idx, 'clip_score']}\nmcos score: {items.loc[idx, 'mcos_score']}\npopularity: {items.loc[idx, 'model_download_count']}",
+                              size=20,
+                              shape='image',
+                              image=f"https://modelcofferbucket.s3-accelerate.amazonaws.com/{items.loc[idx, 'image_id']}.png",
+                              x=items.loc[idx, 'x'].item(),
+                              y=items.loc[idx, 'y'].item(),
+                              # fixed=True,
+                              color={'background': '#E0E0E1', 'border': '#ffffff', 'highlight': {'border': '#F04542'}},
+                              # opacity=opacity,
+                              shadow={'enabled': True, 'color': 'rgba(0,0,0,0.4)', 'size': 10, 'x': 1, 'y': 1},
+                              borderWidth=2,
+                              shapeProperties={'useBorderWithImage': True},
+                              )
+                         )
+
+        config = Config(width='100%',
+                        height='600',
+                        directed=True,
+                        physics=False,
+                        hierarchical=False,
+                        interaction={'navigationButtons': True, 'dragNodes': False, 'multiselect': False},
+                        # **kwargs
+                        )
+
+        return agraph(nodes=nodes,
+                      edges=edges,
+                      config=config,
+                      )
+
+    def selection_panel(self, items):
+        # temperal function
+
+        selecters = st.columns([1, 4])
+
+        if 'score_weights' not in st.session_state:
+            st.session_state.score_weights = [1.0, 0.8, 0.2, 0.8]
+
+        # select sort type
+        with selecters[0]:
+            sort_type = st.selectbox('Sort by', ['Scores', 'IDs and Names'])
+            if sort_type == 'Scores':
+                sort_by = 'weighted_score_sum'
+
+        # select other options
+        with selecters[1]:
+            if sort_type == 'IDs and Names':
+                sub_selecters = st.columns([3, 1])
+                # select sort by
+                with sub_selecters[0]:
+                    sort_by = st.selectbox('Sort by',
+                                           ['model_name', 'model_id', 'modelVersion_name', 'modelVersion_id', 'norm_nsfw'],
+                                           label_visibility='hidden')
+
+                continue_idx = 1
+
+            else:
+                # add custom weights
+                sub_selecters = st.columns([1, 1, 1, 1])
+
+                with sub_selecters[0]:
+                    clip_weight = st.number_input('Clip Score Weight', min_value=-100.0, max_value=100.0, value=1.0, step=0.1, help='the weight for normalized clip score')
+                with sub_selecters[1]:
+                    mcos_weight = st.number_input('Dissimilarity Weight', min_value=-100.0, max_value=100.0, value=0.8, step=0.1, help='the weight for m(eam) s(imilarity) q(antile) score for measuring distinctiveness')
+                with sub_selecters[2]:
+                    pop_weight = st.number_input('Popularity Weight', min_value=-100.0, max_value=100.0, value=0.2, step=0.1, help='the weight for normalized popularity score')
+
+                items.loc[:, 'weighted_score_sum'] = round(items[f'norm_clip'] * clip_weight + items[f'norm_mcos'] * mcos_weight + items[
+                    'norm_pop'] * pop_weight, 4)
+
+                continue_idx = 3
+
+                # save latest weights
+                st.session_state.score_weights[0] = round(clip_weight, 2)
+                st.session_state.score_weights[1] = round(mcos_weight, 2)
+                st.session_state.score_weights[2] = round(pop_weight, 2)
+
+            # select threshold
+            with sub_selecters[continue_idx]:
+                nsfw_threshold = st.number_input('NSFW Score Threshold', min_value=0.0, max_value=1.0, value=0.8, step=0.01, help='Only show models with nsfw score lower than this threshold, set 1.0 to show all images')
+                items = items[items['norm_nsfw'] <= nsfw_threshold].reset_index(drop=True)
+
+            # save latest threshold
+            st.session_state.score_weights[3] = nsfw_threshold
+
+        # draw a distribution histogram
+        if sort_type == 'Scores':
+            try:
+                with st.expander('Show score distribution histogram and select score range'):
+                    st.write('**Score distribution histogram**')
+                    chart_space = st.container()
+                    # st.write('Select the range of scores to show')
+                    hist_data = pd.DataFrame(items[sort_by])
+                    mini = hist_data[sort_by].min().item()
+                    mini = mini//0.1 * 0.1
+                    maxi = hist_data[sort_by].max().item()
+                    maxi = maxi//0.1 * 0.1 + 0.1
+                    st.write('**Select the range of scores to show**')
+                    r = st.slider('Select the range of scores to show', min_value=mini, max_value=maxi, value=(mini, maxi), step=0.05, label_visibility='collapsed')
+                    with chart_space:
+                        st.altair_chart(altair_histogram(hist_data, sort_by, r[0], r[1]), use_container_width=True)
+                    # event_dict = altair_component(altair_chart=altair_histogram(hist_data, sort_by))
+                    # r = event_dict.get(sort_by)
+                    if r:
+                        items = items[(items[sort_by] >= r[0]) & (items[sort_by] <= r[1])].reset_index(drop=True)
+                        # st.write(r)
+            except:
+                pass
+
+        display_options = st.columns([1, 4])
+
+        with display_options[0]:
+            # select order
+            order = st.selectbox('Order', ['Ascending', 'Descending'], index=1 if sort_type == 'Scores' else 0)
+            if order == 'Ascending':
+                order = True
+            else:
+                order = False
+
+        with display_options[1]:
+
+            # select info to show
+            info = st.multiselect('Show Info',
+                                  ['model_name', 'model_id', 'modelVersion_name', 'modelVersion_id',
+                                   'weighted_score_sum', 'model_download_count', 'clip_score', 'mcos_score',
+                                   'nsfw_score', 'norm_nsfw'],
+                                  default=sort_by)
+
+        # apply sorting to dataframe
+        items = items.sort_values(by=[sort_by], ascending=order).reset_index(drop=True)
+
+        # select number of columns
+        col_num = st.slider('Number of columns', min_value=1, max_value=9, value=4, step=1, key='col_num')
+
+        return items, info, col_num
+
+    def sidebar(self):
+        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, index=5)
+
+            items = self.promptBook[self.promptBook['tag'] == tag].reset_index(drop=True)
+
+            prompts = np.sort(items['prompt'].unique())[::1]
+
+            selected_prompt = st.selectbox('Select prompt', prompts, index=3)
+
+            mode = st.radio('Select a mode', ['Gallery', 'Graph'], horizontal=True, index=1)
+
+            items = items[items['prompt'] == selected_prompt].reset_index(drop=True)
+            prompt_id = items['prompt_id'].unique()[0]
+            note = items['note'].unique()[0]
+
+            # show source
+            if isinstance(note, str):
+                if note.isdigit():
+                    st.caption(f"`Source: civitai`")
+                else:
+                    st.caption(f"`Source: {note}`")
+            else:
+                st.caption("`Source: Parti-prompts`")
+
+            # show image metadata
+            image_metadatas = ['prompt', 'negativePrompt', 'sampler', 'cfgScale', 'size', 'seed']
+            for key in image_metadatas:
+                label = ' '.join(key.split('_')).capitalize()
+                st.write(f"**{label}**")
+                if items[key][0] == ' ':
+                    st.write('`None`')
+                else:
+                    st.caption(f"{items[key][0]}")
+
+            # for note as civitai image id, add civitai reference
+            if isinstance(note, str) and note.isdigit():
+                try:
+                    st.write(f'**[Civitai Reference](https://civitai.com/images/{note})**')
+                    res = requests.get(f'https://civitai.com/images/{note}')
+                    # st.write(res.text)
+                    soup = BeautifulSoup(res.text, 'html.parser')
+                    image_section = soup.find('div', {'class': 'mantine-12rlksp'})
+                    image_url = image_section.find('img')['src']
+                    st.image(image_url, use_column_width=True)
+                except:
+                    pass
+
+        return prompt_tags, tag, prompt_id, items, mode
+
+    def app(self):
+        st.title('Model Visualization and Retrieval')
+        st.write('This is a gallery of images generated by the models')
+
+        prompt_tags, tag, prompt_id, items, mode = self.sidebar()
+        # items, info, col_num = self.selection_panel(items)
+
+        # subset = st.radio('Select a subset', ['All', 'Selected Only'], index=0, horizontal=True)
+        # try:
+        #     if subset == 'Selected Only':
+        #         items = items[items['modelVersion_id'].isin(st.session_state.selected_dict[prompt_id])].reset_index(drop=True)
+        # except:
+        #     pass
+
+        # add safety check for some prompts
+        safety_check = True
+        unsafe_prompts = {}
+        # initialize unsafe prompts
+        for prompt_tag in prompt_tags:
+            unsafe_prompts[prompt_tag] = []
+        # manually add unsafe prompts
+        unsafe_prompts['world knowledge'] = [83]
+        unsafe_prompts['abstract'] = [1, 3]
+
+        if int(prompt_id.item()) in unsafe_prompts[tag]:
+            st.warning('This prompt may contain unsafe content. They might be offensive, depressing, or sexual.')
+            safety_check = st.checkbox('I understand that this prompt may contain unsafe content. Show these images anyway.', key=f'safety_{prompt_id}')
+
+        if safety_check:
+            if mode == 'Gallery':
+                self.gallery_mode(prompt_id, items)
+            elif mode == 'Graph':
+                self.graph_mode(prompt_id, items)
+
+
+    def graph_mode(self, prompt_id, items):
+        graph_cols = st.columns([3, 1])
+        prompt = st.chat_input(f"Selected model version ids: {str(st.session_state.selected_dict.get(prompt_id, []))}",
+                               disabled=False, key=f'{prompt_id}')
+        if prompt:
+            switch_page("ranking")
+
+        with graph_cols[0]:
+            graph_space = st.empty()
+
+            with graph_space.container():
+                return_value = self.gallery_graph(items)
+
+        with graph_cols[1]:
+            if return_value:
+                with st.form(key=f'{prompt_id}'):
+                    image_url = f"https://modelcofferbucket.s3-accelerate.amazonaws.com/{return_value}.png"
+
+                    st.image(image_url)
+
+                    item = items[items['image_id'] == return_value].reset_index(drop=True).iloc[0]
+                    modelVersion_id = item['modelVersion_id']
+
+                    # handle selection
+                    if 'selected_dict' in st.session_state:
+                        if item['prompt_id'] not in st.session_state.selected_dict:
+                            st.session_state.selected_dict[item['prompt_id']] = []
+
+                        if modelVersion_id in st.session_state.selected_dict[item['prompt_id']]:
+                            checked = True
+                        else:
+                            checked = False
+
+                    if checked:
+                        # deselect = st.button('Deselect', key=f'select_{item["prompt_id"]}_{item["modelVersion_id"]}', use_container_width=True)
+                        deselect = st.form_submit_button('Deselect', use_container_width=True)
+                        if deselect:
+                            st.session_state.selected_dict[item['prompt_id']].remove(item['modelVersion_id'])
+                            self.remove_ranking_states(item['prompt_id'])
+                            st.experimental_rerun()
+
+                    else:
+                        # select = st.button('Select', key=f'select_{item["prompt_id"]}_{item["modelVersion_id"]}', use_container_width=True, type='primary')
+                        select = st.form_submit_button('Select', use_container_width=True, type='primary')
+                        if select:
+                            st.session_state.selected_dict[item['prompt_id']].append(item['modelVersion_id'])
+                            self.remove_ranking_states(item['prompt_id'])
+                            st.experimental_rerun()
+
+                    # st.write(item)
+                    infos = ['model_name', 'modelVersion_name', 'model_download_count', 'clip_score', 'mcos_score',
+                             'nsfw_score']
+
+                    infos_df = item[infos]
+                    # rename columns
+                    infos_df = infos_df.rename(index={'model_name': 'Model', 'modelVersion_name': 'Version', 'model_download_count': 'Downloads', 'clip_score': 'Clip Score', 'mcos_score': 'mcos Score', 'nsfw_score': 'NSFW Score'})
+                    st.table(infos_df)
+
+                    # for info in infos:
+                    #     st.write(f"**{info}**:")
+                    #     st.write(item[info])
+
+            else:
+                st.info('Please click on an image to show')
+
+
+    def gallery_mode(self, prompt_id, items):
+        items, info, col_num = self.selection_panel(items)
+
+        if 'selected_dict' in st.session_state:
+            # st.write('checked: ', str(st.session_state.selected_dict.get(prompt_id, [])))
+            dynamic_weight_options = ['Grid Search', 'SVM', 'Greedy']
+            dynamic_weight_panel = st.columns(len(dynamic_weight_options))
+
+            if len(st.session_state.selected_dict.get(prompt_id, [])) > 0:
+                btn_disable = False
+            else:
+                btn_disable = True
+
+            for i in range(len(dynamic_weight_options)):
+                method = dynamic_weight_options[i]
+                with dynamic_weight_panel[i]:
+                    btn = st.button(method, use_container_width=True, disabled=btn_disable, on_click=self.dynamic_weight, args=(prompt_id, items, method))
+
+        prompt = st.chat_input(f"Selected model version ids: {str(st.session_state.selected_dict.get(prompt_id, []))}", disabled=False, key=f'{prompt_id}')
+        if prompt:
+            switch_page("ranking")
+
+        with st.form(key=f'{prompt_id}'):
+            # buttons = st.columns([1, 1, 1])
+            buttons_space = st.columns([1, 1, 1, 1])
+            gallery_space = st.empty()
+
+            with buttons_space[0]:
+                continue_btn = st.form_submit_button('Confirm Selection', use_container_width=True, type='primary')
+                if continue_btn:
+                    self.submit_actions('Continue', prompt_id)
+
+            with buttons_space[1]:
+                select_btn = st.form_submit_button('Select All', use_container_width=True)
+                if select_btn:
+                    self.submit_actions('Select', prompt_id)
+
+            with buttons_space[2]:
+                deselect_btn = st.form_submit_button('Deselect All', use_container_width=True)
+                if deselect_btn:
+                    self.submit_actions('Deselect', prompt_id)
+
+            with buttons_space[3]:
+                refresh_btn = st.form_submit_button('Refresh', on_click=gallery_space.empty, use_container_width=True)
+
+            with gallery_space.container():
+                with st.spinner('Loading images...'):
+                    self.gallery_standard(items, col_num, info)
+
+                st.info("Don't forget to scroll back to top and click the 'Confirm Selection' button to save your selection!!!")
+
+
+
+    def submit_actions(self, status, prompt_id):
+        # remove counter from session state
+        # st.session_state.pop('counter', None)
+        self.remove_ranking_states('prompt_id')
+        if status == 'Select':
+            modelVersions = self.promptBook[self.promptBook['prompt_id'] == prompt_id]['modelVersion_id'].unique()
+            st.session_state.selected_dict[prompt_id] = modelVersions.tolist()
+            print(st.session_state.selected_dict, 'select')
+            st.experimental_rerun()
+        elif status == 'Deselect':
+            st.session_state.selected_dict[prompt_id] = []
+            print(st.session_state.selected_dict, 'deselect')
+            st.experimental_rerun()
+            # self.promptBook.loc[self.promptBook['prompt_id'] == prompt_id, 'checked'] = False
+        elif status == 'Continue':
+            st.session_state.selected_dict[prompt_id] = []
+            for key in st.session_state:
+                keys = key.split('_')
+                if keys[0] == 'select' and keys[1] == str(prompt_id):
+                    if st.session_state[key]:
+                        st.session_state.selected_dict[prompt_id].append(int(keys[2]))
+            # switch_page("ranking")
+            print(st.session_state.selected_dict, 'continue')
+            st.experimental_rerun()
+
+    def dynamic_weight(self, prompt_id, items, method='Grid Search'):
+        selected = items[
+            items['modelVersion_id'].isin(st.session_state.selected_dict[prompt_id])].reset_index(drop=True)
+        optimal_weight = [0, 0, 0]
+
+        if method == 'Grid Search':
+            # grid search method
+            top_ranking = len(items) * len(selected)
+
+            for clip_weight in np.arange(-1, 1, 0.1):
+                for mcos_weight in np.arange(-1, 1, 0.1):
+                    for pop_weight in np.arange(-1, 1, 0.1):
+
+                        weight_all = clip_weight*items[f'norm_clip'] + mcos_weight*items[f'norm_mcos'] + pop_weight*items['norm_pop']
+                        weight_all_sorted = weight_all.sort_values(ascending=False).reset_index(drop=True)
+                        # print('weight_all_sorted:', weight_all_sorted)
+                        weight_selected = clip_weight*selected[f'norm_clip'] + mcos_weight*selected[f'norm_mcos'] + pop_weight*selected['norm_pop']
+
+                        # get the index of values of weight_selected in weight_all_sorted
+                        rankings = []
+                        for weight in weight_selected:
+                            rankings.append(weight_all_sorted.index[weight_all_sorted == weight].tolist()[0])
+                        if sum(rankings) <= top_ranking:
+                            top_ranking = sum(rankings)
+                            print('current top ranking:', top_ranking, rankings)
+                            optimal_weight = [clip_weight, mcos_weight, pop_weight]
+            print('optimal weight:', optimal_weight)
+
+        elif method == 'SVM':
+            # svm method
+            print('start svm method')
+            # get residual dataframe that contains models not selected
+            residual = items[~items['modelVersion_id'].isin(selected['modelVersion_id'])].reset_index(drop=True)
+            residual = residual[['norm_clip_crop', 'norm_mcos_crop', 'norm_pop']]
+            residual = residual.to_numpy()
+            selected = selected[['norm_clip_crop', 'norm_mcos_crop', 'norm_pop']]
+            selected = selected.to_numpy()
+
+            y = np.concatenate((np.full((len(selected), 1), -1), np.full((len(residual), 1), 1)), axis=0).ravel()
+            X = np.concatenate((selected, residual), axis=0)
+
+            # fit svm model, and get parameters for the hyperplane
+            clf = LinearSVC(random_state=0, C=1.0, fit_intercept=False, dual='auto')
+            clf.fit(X, y)
+            optimal_weight = clf.coef_[0].tolist()
+            print('optimal weight:', optimal_weight)
+            pass
+
+        elif method == 'Greedy':
+            for idx in selected.index:
+                # find which score is the highest, clip, mcos, or pop
+                clip_score = selected.loc[idx, 'norm_clip_crop']
+                mcos_score = selected.loc[idx, 'norm_mcos_crop']
+                pop_score = selected.loc[idx, 'norm_pop']
+                if clip_score >= mcos_score and clip_score >= pop_score:
+                    optimal_weight[0] += 1
+                elif mcos_score >= clip_score and mcos_score >= pop_score:
+                    optimal_weight[1] += 1
+                elif pop_score >= clip_score and pop_score >= mcos_score:
+                    optimal_weight[2] += 1
+
+            # normalize optimal_weight
+            optimal_weight = [round(weight/len(selected), 2) for weight in optimal_weight]
+            print('optimal weight:', optimal_weight)
+            print('optimal weight:', optimal_weight)
+
+        st.session_state.score_weights[0: 3] = optimal_weight
+
+
+    def remove_ranking_states(self, prompt_id):
+        # for drag sort
+        try:
+            st.session_state.counter[prompt_id] = 0
+            st.session_state.ranking[prompt_id] = {}
+            print('remove ranking states')
+        except:
+            print('no sort ranking states to remove')
+
+        # for battles
+        try:
+            st.session_state.pointer[prompt_id] = {'left': 0, 'right': 1}
+            print('remove battles states')
+        except:
+            print('no battles states to remove')
+
+        # for page progress
+        try:
+            st.session_state.progress[prompt_id] = 'ranking'
+            print('reset page progress states')
+        except:
+            print('no page progress states to be reset')
+
+
+# hist_data = pd.DataFrame(np.random.normal(42, 10, (200, 1)), columns=["x"])
+@st.cache_resource
+def altair_histogram(hist_data, sort_by, mini, maxi):
+    brushed = alt.selection_interval(encodings=['x'], name="brushed")
+
+    chart = (
+        alt.Chart(hist_data)
+        .mark_bar(opacity=0.7, cornerRadius=2)
+        .encode(alt.X(f"{sort_by}:Q", bin=alt.Bin(maxbins=25)), y="count()")
+        # .add_selection(brushed)
+        # .properties(width=800, height=300)
+    )
+
+    # Create a transparent rectangle for highlighting the range
+    highlight = (
+        alt.Chart(pd.DataFrame({'x1': [mini], 'x2': [maxi]}))
+        .mark_rect(opacity=0.3)
+        .encode(x='x1', x2='x2')
+        # .properties(width=800, height=300)
+    )
+
+    # Layer the chart and the highlight rectangle
+    layered_chart = alt.layer(chart, highlight)
+
+    return layered_chart
+
+
+@st.cache_data
+def load_hf_dataset():
+    # login to huggingface
+    login(token=os.environ.get("HF_TOKEN"))
+
+    # load from huggingface
+    roster = pd.DataFrame(load_dataset('MAPS-research/GEMRec-Roster', split='train'))
+    promptBook = pd.DataFrame(load_dataset('MAPS-research/GEMRec-Metadata', split='train'))
+    # images_ds = load_from_disk(os.path.join(os.getcwd(), 'data', 'promptbook'))
+    images_ds = None  # set to None for now since we use s3 bucket to store images
+
+    # # process dataset
+    # roster = roster[['model_id', 'model_name', 'modelVersion_id', 'modelVersion_name',
+    #                                                    'model_download_count']].drop_duplicates().reset_index(drop=True)
+
+    # add 'custom_score_weights' column to promptBook if not exist
+    if 'weighted_score_sum' not in promptBook.columns:
+        promptBook.loc[:, 'weighted_score_sum'] = 0
+
+    # merge roster and promptbook
+    promptBook = promptBook.merge(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
+    promptBook.loc[:, 'row_idx'] = promptBook.index
+
+    # apply a nsfw filter
+    promptBook = promptBook[promptBook['nsfw_score'] <= 0.84].reset_index(drop=True)
+
+    # add a column that adds up 'norm_clip', 'norm_mcos', and 'norm_pop'
+    score_weights = [1.0, 0.8, 0.2]
+    promptBook.loc[:, 'total_score'] = round(promptBook['norm_clip'] * score_weights[0] + promptBook['norm_mcos'] * score_weights[1] + promptBook['norm_pop'] * score_weights[2], 4)
+
+    return roster, promptBook, images_ds
+
+@st.cache_data
+def load_tsne_coordinates(items):
+    # load tsne coordinates
+    tsne_df = pd.read_parquet('./data/feats_tsne.parquet')
+
+    # print(tsne_df['modelVersion_id'].dtype)
+
+    print('before merge:', items)
+    items = items.merge(tsne_df, on=['modelVersion_id', 'prompt_id'], how='left')
+    print('after merge:', items)
+    return items
+
+
+if __name__ == "__main__":
+    st.set_page_config(page_title="Model Coffer Gallery", page_icon="🖼️", layout="wide")
+
+    if 'user_id' not in st.session_state:
+        st.warning('Please log in first.')
+        home_btn = st.button('Go to Home Page')
+        if home_btn:
+            switch_page("home")
+    else:
+        # st.write('You have already logged in as ' + st.session_state.user_id[0])
+        roster, promptBook, images_ds = load_hf_dataset()
+        # print(promptBook.columns)
+
+        # initialize selected_dict
+        if 'selected_dict' not in st.session_state:
+            st.session_state['selected_dict'] = {}
+
+        app = GalleryApp(promptBook=promptBook, images_ds=images_ds)
+        app.app()
+
+        # components.html(
+        #     """
+        #     <script>
+        #       var iframe = window.parent.document.querySelector('[title="streamlit_agraph.agraph"]');
+        #       console.log(iframe);
+        #       var targetElement = iframe.contentDocument.querySelector('div.vis-network div.vis-navigation div.vis-button.vis-zoomExtends');
+        #       console.log(targetElement);
+        #       targetElement.style.background-image = "url(https://www.flaticon.com/free-icon-font/menu-burger_3917215?related_id=3917215#)";
+        #     </script>
+        #     """,
+        #     # unsafe_allow_html=True,
+        # )

Archive/agraphTest.py

@@ -1,170 +0,0 @@
-import os
-
-import streamlit as st
-import torch
-import pandas as pd
-import numpy as np
-
-from datasets import load_dataset, Dataset, load_from_disk
-from huggingface_hub import login
-from streamlit_agraph import agraph, Node, Edge, Config
-from sklearn.manifold import TSNE
-
-
-@st.cache_data
-def load_hf_dataset():
-    # login to huggingface
-    login(token=os.environ.get("HF_TOKEN"))
-
-    # load from huggingface
-    roster = pd.DataFrame(load_dataset('MAPS-research/GEMRec-Roster', split='train'))
-    promptBook = pd.DataFrame(load_dataset('MAPS-research/GEMRec-Metadata', split='train'))
-
-    # process dataset
-    roster = roster[['model_id', 'model_name', 'modelVersion_id', 'modelVersion_name',
-                                                       'model_download_count']].drop_duplicates().reset_index(drop=True)
-
-    # add 'custom_score_weights' column to promptBook if not exist
-    if 'weighted_score_sum' not in promptBook.columns:
-        promptBook.loc[:, 'weighted_score_sum'] = 0
-
-    # merge roster and promptbook
-    promptBook = promptBook.merge(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
-    promptBook.loc[:, 'row_idx'] = promptBook.index
-
-    return roster, promptBook
-
-
-@st.cache_data
-def calc_tsne(prompt_id):
-    print('==> loading feats')
-    feats = {}
-    for pt in os.listdir('../data/feats'):
-        if pt.split('.')[-1] == 'pt' and pt.split('.')[0].isdigit():
-            feats[pt.split('.')[0]] = torch.load(os.path.join('../data/feats', pt))
-
-    print('==> applying t-SNE')
-    # apply t-SNE to entries in each feat in feats to get 2D coordinates
-    tsne = TSNE(n_components=2, random_state=0)
-    # for k, v in tqdm(feats.items()):
-    #     feats[k]['tsne'] = tsne.fit_transform(v['all'].numpy())
-    # prompt_id = '90'
-    feats[prompt_id]['tsne'] = tsne.fit_transform(feats[prompt_id]['all'].numpy())
-
-    feats_df = pd.DataFrame(feats[prompt_id]['tsne'], columns=['x', 'y'])
-    feats_df['prompt_id'] = prompt_id
-
-    keys = []
-    for k in feats[prompt_id].keys():
-        if k != 'all' and k != 'tsne':
-            keys.append(int(k.item()))
-
-    feats_df['modelVersion_id'] = keys
-
-
-    return feats_df
-
-    # print(feats[prompt_id]['tsne'])
-
-
-if __name__ == '__main__':
-    st.set_page_config(layout="wide")
-
-    # load dataset
-    roster, promptBook = load_hf_dataset()
-    # prompt_id = '20'
-
-    with st.sidebar:
-        st.write('## Select Prompt')
-        prompts = promptBook['prompt_id'].unique().tolist()
-        # sort prompts by prompt_id
-        prompts.sort()
-        prompt_id = st.selectbox('Select Prompt', prompts, index=0)
-        physics = st.checkbox('Enable Physics')
-
-    feats_df = calc_tsne(str(prompt_id))
-
-    # keys = []
-    # for k in feats[prompt_id].keys():
-    #     if k != 'all' and k != 'tsne':
-    #         keys.append(int(k.item()))
-
-    # print(keys)
-
-    data = []
-    for idx in feats_df.index:
-        modelVersion_id = feats_df.loc[idx, 'modelVersion_id']
-        image_id = promptBook[(promptBook['modelVersion_id'] == modelVersion_id) & (
-                    promptBook['prompt_id'] == int(prompt_id))].reset_index(drop=True).loc[0, 'image_id']
-        image_url = f"https://modelcofferbucket.s3-accelerate.amazonaws.com/{image_id}.png"
-        scale = 50
-        data.append((feats_df.loc[idx, 'x'] * scale, feats_df.loc[idx, 'y'] * scale, image_url))
-
-        image_size = promptBook[(promptBook['image_id'] == image_id)].reset_index(drop=True).loc[0, 'size'].split('x')
-
-    nodes = []
-    edges = []
-
-    for d in data:
-        nodes.append( Node(id=d[2],
-                           # label=str(items.loc[idx, 'model_name']),
-                           size=20,
-                           shape="image",
-                           image=d[2],
-                           x=[d[0]],
-                           y=[d[1]],
-                           fixed=False if physics else True,
-                           color={'background': '#00000', 'border': '#ffffff'},
-                           shadow={'enabled': True, 'color': 'rgba(0,0,0,0.4)', 'size': 10, 'x': 1, 'y': 1},
-                           # borderWidth=1,
-                           # shapeProperties={'useBorderWithImage': True},
-                           )
-                    )
-
-
-    # nodes.append( Node(id="Spiderman",
-    #                    label="Peter Parker",
-    #                    size=25,
-    #                    shape="circularImage",
-    #                    image="http://marvel-force-chart.surge.sh/marvel_force_chart_img/top_spiderman.png")
-    #             ) # includes **kwargs
-    # nodes.append( Node(id="Captain_Marvel",
-    #                    label="Carol Danvers",
-    #                    fixed=True,
-    #                    size=25,
-    #                    shape="circularImage",
-    #                    image="http://marvel-force-chart.surge.sh/marvel_force_chart_img/top_captainmarvel.png")
-    #             )
-    # edges.append( Edge(source="Captain_Marvel",
-    #                    label="friend_of",
-    #                    target="Spiderman",
-    #                    length=200,
-    #                    # **kwargs
-    #                    )
-    #             )
-    #
-    config = Config(width='100%',
-                    height=800,
-                    directed=True,
-                    physics=physics,
-                    hierarchical=False,
-                    # **kwargs
-                    )
-
-    cols = st.columns([3, 1], gap='large')
-
-    with cols[0]:
-            return_value = agraph(nodes=nodes,
-                                  edges=edges,
-                                  config=config)
-
-    # st.write(return_value)
-
-    with cols[1]:
-        try:
-            st.image(return_value, use_column_width=True)
-        except:
-            st.write('No image selected')

Archive/bokehTest.py

@@ -1,182 +0,0 @@
-import os
-
-import streamlit as st
-import torch
-import pandas as pd
-import numpy as np
-import requests
-
-from bokeh.plotting import figure, show
-from bokeh.models import HoverTool, ColumnDataSource, CustomJSHover
-from bokeh.embed import file_html
-from bokeh.resources import CDN  # Import CDN here
-from datasets import load_dataset, Dataset, load_from_disk
-from huggingface_hub import login
-from sklearn.manifold import TSNE
-from tqdm import tqdm
-
-
-@st.cache_data
-def load_hf_dataset():
-    # login to huggingface
-    login(token=os.environ.get("HF_TOKEN"))
-
-    # load from huggingface
-    roster = pd.DataFrame(load_dataset('MAPS-research/GEMRec-Roster', split='train'))
-    promptBook = pd.DataFrame(load_dataset('MAPS-research/GEMRec-Metadata', split='train'))
-
-    # process dataset
-    roster = roster[['model_id', 'model_name', 'modelVersion_id', 'modelVersion_name',
-                                                       'model_download_count']].drop_duplicates().reset_index(drop=True)
-
-    # add 'custom_score_weights' column to promptBook if not exist
-    if 'weighted_score_sum' not in promptBook.columns:
-        promptBook.loc[:, 'weighted_score_sum'] = 0
-
-    # merge roster and promptbook
-    promptBook = promptBook.merge(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
-    promptBook.loc[:, 'row_idx'] = promptBook.index
-
-    return roster, promptBook
-
-def show_with_bokeh(data, streamlit=False):
-    # Extract x, y coordinates and image URLs
-    x_coords, y_coords, image_urls = zip(*data)
-
-    # Create a ColumnDataSource
-    source = ColumnDataSource(data=dict(x=x_coords, y=y_coords, image=image_urls))
-
-    # Create a figure
-    p = figure(width=800, height=600)
-
-    # Add scatter plot
-    scatter = p.scatter(x='x', y='y', size=20, source=source)
-
-    # Define hover tool
-    hover = HoverTool()
-    # hover.tooltips = """
-    #         <div>
-    #             <iframe src="@image" width="512" height="512"></iframe>
-    #         </div>
-    #     """
-    # hover.formatters = {'@image': CustomJSHover(code="""
-    #         const index = cb_data.index;
-    #         const url = cb_data.source.data['image'][index];
-    #         return '<iframe src="' + url + '" width="512" height="512"></iframe>';
-    #     """)}
-
-    hover.tooltips = """
-            <div>
-                <img src="@image" style='object-fit: contain'; height=100%">
-            </div>
-        """
-    hover.formatters = {'@image': CustomJSHover(code="""
-            const index = cb_data.index;
-            const url = cb_data.source.data['image'][index];
-            return '<img src="' + url + '">';
-        """)}
-
-    p.add_tools(hover)
-
-    # Generate HTML with the plot
-    html = file_html(p, CDN, "Interactive Scatter Plot with Hover Images")
-
-    # Save the HTML file or show it
-    # with open("scatter_plot_with_hover_images.html", "w") as f:
-    #     f.write(html)
-
-    if streamlit:
-        st.bokeh_chart(p, use_container_width=True)
-    else:
-        show(p)
-
-
-def show_with_bokeh_2(data, image_size=[40, 40], streamlit=False):
-    # Extract x, y coordinates and image URLs
-    x_coords, y_coords, image_urls = zip(*data)
-
-    # Create a ColumnDataSource
-    source = ColumnDataSource(data=dict(x=x_coords, y=y_coords, image=image_urls))
-
-    # Create a figure
-    p = figure(width=800, height=600, aspect_ratio=1.0)
-
-    # Add image glyphs
-    # image_size = 40  # Adjust this size as needed
-    scale = 0.1
-    image_size = [int(image_size[0])*scale, int(image_size[1])*scale]
-    print(image_size)
-    p.image_url(url='image', x='x', y='y', source=source, w=image_size[0], h=image_size[1], anchor="center")
-
-    # Define hover tool
-    hover = HoverTool()
-    hover.tooltips = """
-            <div>
-                <img src="@image" style='object-fit: contain'; height=100%'">
-            </div>
-        """
-    p.add_tools(hover)
-
-    # Generate HTML with the plot
-    html = file_html(p, CDN, "Scatter Plot with Images")
-
-    # Save the HTML file or show it
-    # with open("scatter_plot_with_images.html", "w") as f:
-    #     f.write(html)
-
-    if streamlit:
-        st.bokeh_chart(p, use_container_width=True)
-    else:
-        show(p)
-
-
-if __name__ == '__main__':
-    # load dataset
-    roster, promptBook = load_hf_dataset()
-
-    print('==> loading feats')
-    feats = {}
-    for pt in os.listdir('../data/feats'):
-        if pt.split('.')[-1] == 'pt' and pt.split('.')[0].isdigit():
-            feats[pt.split('.')[0]] = torch.load(os.path.join('../data/feats', pt))
-
-    print('==> applying t-SNE')
-    # apply t-SNE to entries in each feat in feats to get 2D coordinates
-    tsne = TSNE(n_components=2, random_state=0)
-    # for k, v in tqdm(feats.items()):
-    #     feats[k]['tsne'] = tsne.fit_transform(v['all'].numpy())
-    prompt_id = '49'
-    feats[prompt_id]['tsne'] = tsne.fit_transform(feats[prompt_id]['all'].numpy())
-
-    print(feats[prompt_id]['tsne'])
-
-    keys = []
-    for k in feats[prompt_id].keys():
-        if k != 'all' and k != 'tsne':
-            keys.append(int(k.item()))
-
-    print(keys)
-
-    data = []
-    for idx in range(len(keys)):
-        modelVersion_id = keys[idx]
-        image_id = promptBook[(promptBook['modelVersion_id'] == modelVersion_id) & (promptBook['prompt_id'] == int(prompt_id))].reset_index(drop=True).loc[0, 'image_id']
-        image_url = f"https://modelcofferbucket.s3-accelerate.amazonaws.com/{image_id}.png"
-        scale = 50
-        data.append((feats[prompt_id]['tsne'][idx][0]*scale, feats[prompt_id]['tsne'][idx][1]*scale, image_url))
-
-        image_size = promptBook[(promptBook['image_id'] == image_id)].reset_index(drop=True).loc[0, 'size'].split('x')
-
-    # # Sample data: (x, y) coordinates and corresponding image URLs
-    # data = [
-    #     (2, 5, "https://www.crunchyroll.com/imgsrv/display/thumbnail/480x720/catalog/crunchyroll/669dae5dbea3d93bb5f1012078501976.jpeg"),
-    #     (4, 8, "https://i.pinimg.com/originals/40/6d/38/406d38957bc4fd12f34c5dfa3d73b86d.jpg"),
-    #     (7, 3, "https://i.pinimg.com/550x/76/27/d2/7627d227adc6fb5fb6662ebfb9d82d7e.jpg"),
-    #     # Add more data points and image URLs
-    # ]
-
-    # show_with_bokeh(data, streamlit=True)
-    show_with_bokeh_2(data, image_size=image_size, streamlit=True)

Archive/optimization.py

@@ -1,37 +0,0 @@
-import numpy as np
-from scipy.optimize import minimize, differential_evolution
-
-
-# Define the function y = x_1*w_1 + x_2*w_2 + x_3*w_3
-def objective_function(w_indices):
-    x_1 = x_1_values[int(w_indices[0])]
-    x_2 = x_2_values[int(w_indices[1])]
-    x_3 = x_3_values[int(w_indices[2])]
-    return - (x_1 * w_indices[3] + x_2 * w_indices[4] + x_3 * w_indices[5])  # Use w_indices to get w_1, w_2, w_3
-
-
-if __name__ == '__main__':
-    # Given sets of discrete values for x_1, x_2, and x_3
-    x_1_values = [1, 2, 3, 5, 6]
-    x_2_values = [0, 5, 7, 2, 1]
-    x_3_values = [3, 7, 4, 5, 2]
-
-    # Perform differential evolution optimization with integer variables
-    # bounds = [(0, len(x_1_values) - 2), (0, len(x_2_values) - 1), (0, len(x_3_values) - 1), (-1, 1), (-1, 1), (-1, 1)]
-    bounds = [(3, 4), (3, 4), (3, 4), (-1, 1), (-1, 1), (-1, 1)]
-    result = differential_evolution(objective_function, bounds)
-
-    # Get the optimal indices of x_1, x_2, and x_3
-    x_1_index, x_2_index, x_3_index, w_1_opt, w_2_opt, w_3_opt = result.x
-
-    # Calculate the peak point (x_1, x_2, x_3) corresponding to the optimal indices
-    x_1_peak = x_1_values[int(x_1_index)]
-    x_2_peak = x_2_values[int(x_2_index)]
-    x_3_peak = x_3_values[int(x_3_index)]
-
-    # Print the results
-    print("Optimal w_1:", w_1_opt)
-    print("Optimal w_2:", w_2_opt)
-    print("Optimal w_3:", w_3_opt)
-    print("Peak Point (x_1, x_2, x_3):", (x_1_peak, x_2_peak, x_3_peak))
-    print("Maximum Value of y:", -result.fun)  # Use negative sign as we previously used to maximize

Archive/optimization2.py

@@ -1,40 +0,0 @@
-import numpy as np
-from scipy.optimize import minimize
-
-if __name__ == '__main__':
-
-    # Given subset of m values for x_1, x_2, and x_3
-    x1_subset = [2, 3, 4]
-    x2_subset = [0, 1]
-    x3_subset = [5, 6, 7]
-
-    # Full set of possible values for x_1, x_2, and x_3
-    x1_full = [1, 2, 3, 4, 5]
-    x2_full = [0, 1, 2, 3, 4, 5]
-    x3_full = [3, 5, 7]
-
-    # Define the objective function for quantile-based ranking
-    def objective_function(w):
-        y_subset = [x1 * w[0] + x2 * w[1] + x3 * w[2] for x1, x2, x3 in zip(x1_subset, x2_subset, x3_subset)]
-        y_full_set = [x1 * w[0] + x2 * w[1] + x3 * w[2] for x1 in x1_full for x2 in x2_full for x3 in x3_full]
-
-        # Calculate the 90th percentile of y values for the full set
-        y_full_set_90th_percentile = np.percentile(y_full_set, 90)
-
-        # Maximize the difference between the 90th percentile of the subset and the 90th percentile of the full set
-        return - min(y_subset) + y_full_set_90th_percentile
-
-
-    # Bounds for w_1, w_2, and w_3 (-1 to 1)
-    bounds = [(-1, 1), (-1, 1), (-1, 1)]
-
-    # Perform bounded optimization to find the values of w_1, w_2, and w_3 that maximize the objective function
-    result = minimize(objective_function, np.zeros(3), method='TNC', bounds=bounds)
-
-    # Get the optimal values of w_1, w_2, and w_3
-    w_1_opt, w_2_opt, w_3_opt = result.x
-
-    # Print the results
-    print("Optimal w_1:", w_1_opt)
-    print("Optimal w_2:", w_2_opt)
-    print("Optimal w_3:", w_3_opt)

pages/Gallery.py

@@ -157,29 +157,29 @@ class GalleryApp:
             # save latest threshold
             st.session_state.score_weights[3] = nsfw_threshold
 
-        # draw a distribution histogram
-        if sort_type == 'Scores':
-            try:
-                with st.expander('Show score distribution histogram and select score range'):
-                    st.write('**Score distribution histogram**')
-                    chart_space = st.container()
-                    # st.write('Select the range of scores to show')
-                    hist_data = pd.DataFrame(items[sort_by])
-                    mini = hist_data[sort_by].min().item()
-                    mini = mini//0.1 * 0.1
-                    maxi = hist_data[sort_by].max().item()
-                    maxi = maxi//0.1 * 0.1 + 0.1
-                    st.write('**Select the range of scores to show**')
-                    r = st.slider('Select the range of scores to show', min_value=mini, max_value=maxi, value=(mini, maxi), step=0.05, label_visibility='collapsed')
-                    with chart_space:
-                        st.altair_chart(altair_histogram(hist_data, sort_by, r[0], r[1]), use_container_width=True)
-                    # event_dict = altair_component(altair_chart=altair_histogram(hist_data, sort_by))
-                    # r = event_dict.get(sort_by)
-                    if r:
-                        items = items[(items[sort_by] >= r[0]) & (items[sort_by] <= r[1])].reset_index(drop=True)
-                        # st.write(r)
-            except:
-                pass
+        # # draw a distribution histogram
+        # if sort_type == 'Scores':
+        #     try:
+        #         with st.expander('Show score distribution histogram and select score range'):
+        #             st.write('**Score distribution histogram**')
+        #             chart_space = st.container()
+        #             # st.write('Select the range of scores to show')
+        #             hist_data = pd.DataFrame(items[sort_by])
+        #             mini = hist_data[sort_by].min().item()
+        #             mini = mini//0.1 * 0.1
+        #             maxi = hist_data[sort_by].max().item()
+        #             maxi = maxi//0.1 * 0.1 + 0.1
+        #             st.write('**Select the range of scores to show**')
+        #             r = st.slider('Select the range of scores to show', min_value=mini, max_value=maxi, value=(mini, maxi), step=0.05, label_visibility='collapsed')
+        #             with chart_space:
+        #                 st.altair_chart(altair_histogram(hist_data, sort_by, r[0], r[1]), use_container_width=True)
+        #             # event_dict = altair_component(altair_chart=altair_histogram(hist_data, sort_by))
+        #             # r = event_dict.get(sort_by)
+        #             if r:
+        #                 items = items[(items[sort_by] >= r[0]) & (items[sort_by] <= r[1])].reset_index(drop=True)
+        #                 # st.write(r)
+        #     except:
+        #         pass
 
         display_options = st.columns([1, 4])
 
@@ -208,25 +208,24 @@ class GalleryApp:
 
         return items, info, col_num
 
-    def sidebar(self):
+    def sidebar(self, items, prompt_id, note):
         with st.sidebar:
-            prompt_tags = self.promptBook['tag'].unique()
-            # sort tags by alphabetical order
-            prompt_tags = np.sort(prompt_tags)[::1]
+            # 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, index=5)
+            #
+            # items = self.promptBook[self.promptBook['tag'] == tag].reset_index(drop=True)
+            #
+            # prompts = np.sort(items['prompt'].unique())[::1]
+            #
+            # selected_prompt = st.selectbox('Select prompt', prompts, index=3)
 
-            tag = st.selectbox('Select a tag', prompt_tags, index=5)
+            # mode = st.radio('Select a mode', ['Gallery', 'Graph'], horizontal=True, index=1)
 
-            items = self.promptBook[self.promptBook['tag'] == tag].reset_index(drop=True)
+            # items = items[items['prompt'] == selected_prompt].reset_index(drop=True)
 
-            prompts = np.sort(items['prompt'].unique())[::1]
-
-            selected_prompt = st.selectbox('Select prompt', prompts, index=3)
-
-            mode = st.radio('Select a mode', ['Gallery', 'Graph'], horizontal=True, index=1)
-
-            items = items[items['prompt'] == selected_prompt].reset_index(drop=True)
-            prompt_id = items['prompt_id'].unique()[0]
-            note = items['note'].unique()[0]
 
             # show source
             if isinstance(note, str):
@@ -260,49 +259,66 @@ class GalleryApp:
                 except:
                     pass
 
-        return prompt_tags, tag, prompt_id, items, mode
+        # return prompt_tags, tag, prompt_id, items
 
-    def app(self):
-        st.title('Model Visualization and Retrieval')
-        st.write('This is a gallery of images generated by the models')
-
-        prompt_tags, tag, prompt_id, items, mode = self.sidebar()
-        # items, info, col_num = self.selection_panel(items)
-
-        # subset = st.radio('Select a subset', ['All', 'Selected Only'], index=0, horizontal=True)
-        # try:
-        #     if subset == 'Selected Only':
-        #         items = items[items['modelVersion_id'].isin(st.session_state.selected_dict[prompt_id])].reset_index(drop=True)
-        # except:
-        #     pass
-
-        # add safety check for some prompts
-        safety_check = True
-        unsafe_prompts = {}
-        # initialize unsafe prompts
-        for prompt_tag in prompt_tags:
-            unsafe_prompts[prompt_tag] = []
-        # manually add unsafe prompts
-        unsafe_prompts['world knowledge'] = [83]
-        unsafe_prompts['abstract'] = [1, 3]
-
-        if int(prompt_id.item()) in unsafe_prompts[tag]:
-            st.warning('This prompt may contain unsafe content. They might be offensive, depressing, or sexual.')
-            safety_check = st.checkbox('I understand that this prompt may contain unsafe content. Show these images anyway.', key=f'safety_{prompt_id}')
-
-        if safety_check:
-            if mode == 'Gallery':
-                self.gallery_mode(prompt_id, items)
-            elif mode == 'Graph':
-                self.graph_mode(prompt_id, items)
 
+    def app(self):
+        # st.title('Model Visualization and Retrieval')
+        # st.write('This is a gallery of images generated by the models')
+
+        # build the tabular view
+        prompt_tags = self.promptBook['tag'].unique()
+        # sort tags by alphabetical order
+        prompt_tags = np.sort(prompt_tags)[::1].tolist()
+
+        tabs = st.tabs(prompt_tags)
+        with st.spinner('Loading...'):
+            for i in range(len(prompt_tags)):
+                with tabs[i]:
+                    tag = prompt_tags[i]
+                    items = self.promptBook[self.promptBook['tag'] == tag].reset_index(drop=True)
+
+                    prompts = np.sort(items['prompt'].unique())[::1]
+
+                    subset_selector = st.columns([3, 1])
+                    with subset_selector[0]:
+                        selected_prompt = st.selectbox('Select prompt', prompts, index=3)
+                    with subset_selector[1]:
+                        subset = st.selectbox('Select a subset', ['All', 'Selected Only'], index=0, key=f'subset_{selected_prompt}')
+                    items = items[items['prompt'] == selected_prompt].reset_index(drop=True)
+                    prompt_id = items['prompt_id'].unique()[0]
+                    note = items['note'].unique()[0]
+
+                    # add safety check for some prompts
+                    safety_check = True
+                    unsafe_prompts = {}
+                    # initialize unsafe prompts
+                    for prompt_tag in prompt_tags:
+                        unsafe_prompts[prompt_tag] = []
+                    # manually add unsafe prompts
+                    unsafe_prompts['world knowledge'] = [83]
+                    unsafe_prompts['abstract'] = [1, 3]
+
+                    if int(prompt_id.item()) in unsafe_prompts[tag]:
+                        st.warning('This prompt may contain unsafe content. They might be offensive, depressing, or sexual.')
+                        safety_check = st.checkbox('I understand that this prompt may contain unsafe content. Show these images anyway.', key=f'safety_{prompt_id}')
+
+                    if safety_check:
+
+                        # if subset == 'Selected Only' and 'selected_dict' in st.session_state:
+                        #     items = items[items['modelVersion_id'].isin(st.session_state.selected_dict[prompt_id])].reset_index(drop=True)
+                        #     self.gallery_mode(prompt_id, items)
+                        # else:
+                        self.graph_mode(prompt_id, items)
+
+        self.sidebar(items, prompt_id, note)
 
     def graph_mode(self, prompt_id, items):
         graph_cols = st.columns([3, 1])
-        prompt = st.chat_input(f"Selected model version ids: {str(st.session_state.selected_dict.get(prompt_id, []))}",
-                               disabled=False, key=f'{prompt_id}')
-        if prompt:
-            switch_page("ranking")
+        # prompt = st.chat_input(f"Selected model version ids: {str(st.session_state.selected_dict.get(prompt_id, []))}",
+        #                        disabled=False, key=f'{prompt_id}')
+        # if prompt:
+        #     switch_page("ranking")
 
         with graph_cols[0]:
             graph_space = st.empty()
@@ -366,20 +382,20 @@ class GalleryApp:
     def gallery_mode(self, prompt_id, items):
         items, info, col_num = self.selection_panel(items)
 
-        if 'selected_dict' in st.session_state:
-            # st.write('checked: ', str(st.session_state.selected_dict.get(prompt_id, [])))
-            dynamic_weight_options = ['Grid Search', 'SVM', 'Greedy']
-            dynamic_weight_panel = st.columns(len(dynamic_weight_options))
-
-            if len(st.session_state.selected_dict.get(prompt_id, [])) > 0:
-                btn_disable = False
-            else:
-                btn_disable = True
-
-            for i in range(len(dynamic_weight_options)):
-                method = dynamic_weight_options[i]
-                with dynamic_weight_panel[i]:
-                    btn = st.button(method, use_container_width=True, disabled=btn_disable, on_click=self.dynamic_weight, args=(prompt_id, items, method))
+        # if 'selected_dict' in st.session_state:
+        #     # st.write('checked: ', str(st.session_state.selected_dict.get(prompt_id, [])))
+        #     dynamic_weight_options = ['Grid Search', 'SVM', 'Greedy']
+        #     dynamic_weight_panel = st.columns(len(dynamic_weight_options))
+        #
+        #     if len(st.session_state.selected_dict.get(prompt_id, [])) > 0:
+        #         btn_disable = False
+        #     else:
+        #         btn_disable = True
+        #
+        #     for i in range(len(dynamic_weight_options)):
+        #         method = dynamic_weight_options[i]
+        #         with dynamic_weight_panel[i]:
+        #             btn = st.button(method, use_container_width=True, disabled=btn_disable, on_click=self.dynamic_weight, args=(prompt_id, items, method))
 
         prompt = st.chat_input(f"Selected model version ids: {str(st.session_state.selected_dict.get(prompt_id, []))}", disabled=False, key=f'{prompt_id}')
         if prompt:
@@ -387,32 +403,28 @@ class GalleryApp:
 
         with st.form(key=f'{prompt_id}'):
             # buttons = st.columns([1, 1, 1])
-            buttons_space = st.columns([1, 1, 1, 1])
+            buttons_space = st.columns([1, 1, 1])
             gallery_space = st.empty()
 
             with buttons_space[0]:
-                continue_btn = st.form_submit_button('Confirm Selection', use_container_width=True, type='primary')
+                continue_btn = st.form_submit_button('Proceed selections to ranking', use_container_width=True, type='primary')
                 if continue_btn:
-                    self.submit_actions('Continue', prompt_id)
+                    # self.submit_actions('Continue', prompt_id)
+                    switch_page("ranking")
 
             with buttons_space[1]:
-                select_btn = st.form_submit_button('Select All', use_container_width=True)
-                if select_btn:
-                    self.submit_actions('Select', prompt_id)
-
-            with buttons_space[2]:
                 deselect_btn = st.form_submit_button('Deselect All', use_container_width=True)
                 if deselect_btn:
                     self.submit_actions('Deselect', prompt_id)
 
-            with buttons_space[3]:
+            with buttons_space[2]:
                 refresh_btn = st.form_submit_button('Refresh', on_click=gallery_space.empty, use_container_width=True)
 
             with gallery_space.container():
                 with st.spinner('Loading images...'):
                     self.gallery_standard(items, col_num, info)
 
-                st.info("Don't forget to scroll back to top and click the 'Confirm Selection' button to save your selection!!!")
+                # st.info("Don't forget to scroll back to top and click the 'Confirm Selection' button to save your selection!!!")