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( # """ # # """, # # unsafe_allow_html=True, # )