import itertools import json import os import requests import altair as alt import extra_streamlit_components as stx import random import numpy as np import pandas as pd import streamlit as st from bs4 import BeautifulSoup from datasets import load_dataset, Dataset, load_from_disk from datetime import datetime from huggingface_hub import login from streamlit_agraph import agraph, Node, Edge, Config from streamlit_extras.switch_page_button import switch_page from streamlit_extras.tags import tagger_component from streamlit_extras.no_default_selectbox import selectbox from sklearn.svm import LinearSVC from Home import connect_to_db class GalleryApp: def __init__(self, promptBook, images_ds): self.promptBook = promptBook self.images_ds = images_ds # init gallery state if 'gallery_state' not in st.session_state: st.session_state.gallery_state = 'graph' # initialize selected_dict if 'selected_dict' not in st.session_state: st.session_state['selected_dict'] = {} # clear up empty entries in seleted_dict for prompt_id in list(st.session_state.selected_dict.keys()): if len(st.session_state.selected_dict[prompt_id]) == 0: st.session_state.selected_dict.pop(prompt_id) if 'gallery_focus' not in st.session_state: st.session_state.gallery_focus = {'tag': None, 'prompt': None} def gallery_standard(self, items, col_num, info, show_checkbox=True): 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) if show_checkbox: # 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: 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']}", title=f"model name: {items.loc[idx, 'model_name']}", 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=3, borderWidthSelected=3, shapeProperties={'useBorderWithImage': True}, ) ) config = Config(width='100%', height='600', directed=True, physics=False, hierarchical=False, interaction={'navigationButtons': True, 'dragNodes': False, 'multiselect': False, 'hover': True}, # **kwargs ) return agraph(nodes=nodes, edges=edges, config=config, ) def sidebar(self, items, prompt_id, note): with st.sidebar: # 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 def text_coloring_add(self, tobe_colored:list, total_items, color_name='orange'): if color_name in ['orange', 'red', 'green', 'blue', 'violet', 'yellow']: colored = [f':{color_name}[{item}]' if item in tobe_colored else item for item in total_items] else: colored = [f'[{color_name}] {item}' if item in tobe_colored else item for item in total_items] return colored def text_coloring_remove(self, tobe_removed): if isinstance(tobe_removed, str): if tobe_removed.startswith(':'): tobe_removed = tobe_removed.split('[')[-1][:-1] elif tobe_removed.startswith('['): tobe_removed = tobe_removed.split(']')[-1][1:] return tobe_removed def app(self): # print(st.session_state.gallery_focus) st.write('### Prompt-Model Retrieval') with st.sidebar: tagger_component('**Gallery State:**', [st.session_state.gallery_state.title()], color_name=['orange']) # 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() # set focus tag and prompt index if exists if st.session_state.gallery_focus['tag'] is None: tag_focus_idx = 0 else: tag_focus_idx = prompt_tags.index(st.session_state.gallery_focus['tag']) # add coloring to tag based on selection tags_tobe_colored = self.promptBook[self.promptBook['prompt_id'].isin(st.session_state.selected_dict.keys())]['tag'].unique() # colored_prompt_tags = [f':orange[{tag}]' if tag in tags_tobe_colored else tag for tag in prompt_tags] colored_prompt_tags = self.text_coloring_add(tags_tobe_colored, prompt_tags, color_name='orange') # save tag to session state on change tag = st.radio('Select a tag', colored_prompt_tags, index=tag_focus_idx, horizontal=True, key='tag', label_visibility='collapsed') # remove coloring from tag tag = self.text_coloring_remove(tag) # print('tag: ', tag) # print('current state: ', st.session_state.gallery_state) if st.session_state.gallery_state == 'graph': items = self.promptBook[self.promptBook['tag'] == tag].reset_index(drop=True) prompts = np.sort(items['prompt'].unique())[::1].tolist() # print('prompts: ', prompts, 'tags: ', prompt_tags) # selt focus prompt index if exists if st.session_state.gallery_focus['prompt'] is None or tag != st.session_state.gallery_focus['tag']: prompt_focus_idx = 0 else: prompt_focus_idx = 1 + prompts.index(st.session_state.gallery_focus['prompt']) # st.caption('Select a prompt') subset_selector = st.columns([3, 1]) with subset_selector[0]: selector_bar = st.columns([1, 15]) with selector_bar[0]: shuffle = st.button('🎲', key='prompt_shuffle', on_click=self.random_gallery_focus, args=(prompt_tags,), use_container_width=True) with selector_bar[-1]: # add coloring to prompt based on selection prompts_tobe_colored = self.promptBook[self.promptBook['prompt_id'].isin(st.session_state.selected_dict.keys())]['prompt'].unique() colored_prompts = self.text_coloring_add(prompts_tobe_colored, prompts, color_name='✅') selected_prompt = selectbox('Select prompt', colored_prompts, key=f'prompt_{tag}', no_selection_label='---', label_visibility='collapsed', index=prompt_focus_idx) # remove coloring from prompt selected_prompt = self.text_coloring_remove(selected_prompt) # print('selected_prompt: ', selected_prompt) st.session_state.prompt_idx_last_time = prompts.index(selected_prompt) if selected_prompt else 0 if selected_prompt is None: # st.markdown(':orange[Please select a prompt above👆]') st.caption('Feel free to **navigate among tags and pages**! Your selection will be saved within one log-in session.') with subset_selector[-1]: st.button(':orange[👈 **Please select a prompt**]', disabled=True, use_container_width=True) else: 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 # load unsafe prompts unsafe_prompts = json.load(open('./data/unsafe_prompts.json', 'r')) for prompt_tag in prompt_tags: if prompt_tag not in unsafe_prompts: 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}') # print('current state: ', st.session_state.gallery_state) # # if st.session_state.gallery_state == 'graph': if safety_check: self.graph_mode(prompt_id, items) with subset_selector[-1]: has_selection = False try: if len(st.session_state.selected_dict.get(prompt_id, [])) > 0: has_selection = True except: pass if has_selection: checkout = st.button('Check out selections ➡️', use_container_width=True, type='primary', on_click=self.switch_to_checkout, args=(tag, selected_prompt)) else: st.button(':orange[👇 **Select images below**]', disabled=True, use_container_width=True) try: self.sidebar(items, prompt_id, note) except: pass elif st.session_state.gallery_state == 'check out': # select items under the current tag, while model_id in selected_dict keys with corresponding modelVersion_ids items = self.promptBook[self.promptBook['tag'] == tag].reset_index(drop=True) temp_items = pd.DataFrame() for prompt_id, selected_models in st.session_state.selected_dict.items(): temp_items = pd.concat([temp_items, items[items['modelVersion_id'].isin(selected_models) & (items['prompt_id'] == prompt_id)]], axis=0) items = temp_items.reset_index(drop=True) self.checkout_mode(tag, items) def switch_to_checkout(self, tag, selected_prompt): # add focus to session state st.session_state.gallery_focus['tag'] = tag st.session_state.gallery_focus['prompt'] = selected_prompt st.session_state.gallery_state = 'check out' def random_gallery_focus(self, tags): st.session_state.gallery_focus['tag'] = random.choice(tags) # st.session_state.gallery_focus['prompt'] = random.choice(prompts) prompts = self.promptBook[self.promptBook['tag'] == st.session_state.gallery_focus['tag']]['prompt'].unique() st.session_state.gallery_focus['prompt'] = random.choice(prompts) def graph_mode(self, prompt_id, items): graph_cols = st.columns([3, 1]) with graph_cols[0]: st.caption( 'Please **:red[click on and select]** as many images as you like! You will be able to compare them later in ranking stage.') 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 # get the latest record in database cursor = GALLERY_CONN.cursor() query = "SELECT * FROM gallery_clicks WHERE username = '{}' AND timestamp = '{}' AND prompt_id = '{}' AND modelVersion_id = {} ORDER BY clicktime DESC LIMIT 1".format( st.session_state.user_id[0], st.session_state.user_id[1], prompt_id, modelVersion_id) cursor.execute(query) record = cursor.fetchone() try: image_status = record['status'] except: image_status = None print('image_status: ', image_status) 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 'last_clicked' not in st.session_state or item['image_id'] != st.session_state.last_clicked: # print('last_clicked not in session state') # self.image_selection_control(item['tag'], item['prompt'], item['prompt_id'], modelVersion_id, 'select') # st.toast('Image selected.', icon='👍') # # st.session_state.last_clicked = item['image_id'] # if modelVersion_id in st.session_state.selected_dict[item['prompt_id']]: # checked = True # else: # checked = False if image_status == 'report': st.warning('You have reported this image') unreport = st.form_submit_button('Withdraw report', use_container_width=True, type='secondary', on_click=self.image_selection_control, args=(item['tag'], item['prompt'], item['prompt_id'], item['modelVersion_id'], 'deselect')) else: if image_status is None: self.image_selection_control(item['tag'], item['prompt'], item['prompt_id'], modelVersion_id, 'select') if image_status == 'select' or image_status == 'reselect' or image_status is None: # 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, on_click=self.image_selection_control, args=(item['tag'], item['prompt'], item['prompt_id'], item['modelVersion_id'], 'deselect')) elif image_status =='deselect': # select = st.button('Select', key=f'select_{item["prompt_id"]}_{item["modelVersion_id"]}', use_container_width=True, type='primary') reselect = st.form_submit_button('Reselect', use_container_width=True, type='primary', on_click=self.image_selection_control, args=(item['tag'], item['prompt'], item['prompt_id'], item['modelVersion_id'], 'reselect')) report = st.form_submit_button('⚠️Report', use_container_width=True, type='secondary', on_click=self.image_selection_control, args=( item['tag'], item['prompt'], item['prompt_id'], item['modelVersion_id'], 'report'), help='Report this image if it contains offensive, depressing, or sexual content.') if image_status == 'select' or image_status == 'reselect' or image_status is None: st.info( "Image selected. **Click 'Check out selections ➡️' on top to see all selected images**.") # 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) else: st.info('You can click on and select an image.') def image_selection_control(self, tag, prompt, prompt_id, modelVersion_id, operation:['select', 'reselect', 'deselect','report']): # self.remove_ranking_states(prompt_id) if operation == 'select' or operation == 'reselect': st.session_state.selected_dict[prompt_id].append(modelVersion_id) # add focus to session state st.session_state.gallery_focus['tag'] = tag st.session_state.gallery_focus['prompt'] = prompt elif operation == 'deselect': if modelVersion_id in st.session_state.selected_dict[prompt_id]: st.session_state.selected_dict[prompt_id].remove(modelVersion_id) elif operation == 'report': pass cursor = GALLERY_CONN.cursor() clicktime = datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S") query = "INSERT INTO gallery_clicks (username, timestamp, tag, prompt_id, modelVersion_id, clicktime, status) VALUES ('{}', '{}', '{}', '{}', {}, '{}', '{}')".format( st.session_state.user_id[0], st.session_state.user_id[1], tag, prompt_id, modelVersion_id, clicktime, operation) cursor.execute(query) GALLERY_CONN.commit() cursor.close() def checkout_mode(self, tag, items): # st.write(items) if len(items) > 0: prompt_ids = items['prompt_id'].unique() for i in range(len(prompt_ids)): prompt_id = prompt_ids[i] prompt = items[items['prompt_id'] == prompt_id]['prompt'].unique()[0] # default_expand = True if st.session_state.gallery_focus['prompt'] == prompt else False if tag == st.session_state.gallery_focus['tag'] and prompt == st.session_state.gallery_focus['prompt']: default_expand = True elif tag != st.session_state.gallery_focus['tag'] and i==0: default_expand = True else: default_expand = False with st.expander(f'**{prompt}**', expanded=default_expand): # st.caption('select info to show') checkout_panel = st.columns([5, 3]) with checkout_panel[0]: info = st.multiselect('Show Info', ['model_name', 'model_id', 'modelVersion_name', 'modelVersion_id', 'total_score', 'model_download_count', 'clip_score', 'mcos_score', 'norm_nsfw'], label_visibility='collapsed', key=f'info_{prompt_id}', placeholder='Select what info to show') with checkout_panel[-1]: checkout_buttons = st.columns([1, 1, 1]) with checkout_buttons[0]: back = st.button('Back to 🖼️', key=f'checkout_back_{prompt_id}', use_container_width=True) if back: st.session_state.gallery_focus['tag'] = tag st.session_state.gallery_focus['prompt'] = prompt print(st.session_state.gallery_focus) st.session_state.gallery_state = 'graph' st.rerun() with checkout_buttons[1]: # init edit state if 'edit_state' not in st.session_state: st.session_state.edit_state = False if not st.session_state.edit_state: edit = st.button('Edit', key=f'checkout_edit_{prompt_id}', use_container_width=True) if edit: st.session_state.edit_state = True st.rerun() else: done = st.button('Done', key=f'checkout_done_{prompt_id}', use_container_width=True) if done: st.session_state.selected_dict[prompt_id] = [] for key in st.session_state: # update selected_dict with edited selection 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])) self.image_selection_control(tag, prompt, prompt_id, int(keys[2]), 'select') # update database st.session_state.edit_state = False st.rerun() with checkout_buttons[-1]: proceed = st.button('Proceed ➡️', key=f'checkout_proceed_{prompt_id}', use_container_width=True, type='primary', disabled=st.session_state.edit_state) if proceed: self.remove_ranking_states(prompt_id) st.session_state.gallery_focus['tag'] = tag st.session_state.gallery_focus['prompt'] = prompt st.session_state.gallery_state = 'graph' print('selected_dict: ', st.session_state.selected_dict) # # save the user selection to database # cursor = GALLERY_CONN.cursor() # st.session_state.epoch['gallery'] += 1 # checkouttime = datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S") # # for modelVersion_id in st.session_state.selected_dict[prompt_id]: # for key, values in st.session_state.selected_dict.items(): # # print('key: ', key, 'values: ', values) # key_tag = self.promptBook[self.promptBook['prompt_id'] == key]['tag'].unique()[0] # for value in values: # query = "INSERT INTO gallery_selections (username, timestamp, tag, prompt_id, modelVersion_id, checkouttime, epoch) VALUES ('{}', '{}', '{}', '{}', {}, '{}', {})".format(st.session_state.user_id[0], st.session_state.user_id[1], key_tag, key, value, checkouttime, st.session_state.epoch['gallery']) # print(query) # cursor.execute(query) # GALLERY_CONN.commit() # cursor.close() # get the largest epoch number of this user and prompt cursor = GALLERY_CONN.cursor() db_table = 'battle_results' if st.session_state.assigned_rank_mode=='Battle' else 'sort_results' query = "SELECT MAX(epoch) FROM {} WHERE username = '{}' AND timestamp = '{}' AND prompt_id = {}".format(db_table, st.session_state.user_id[0], st.session_state.user_id[1], prompt_id) cursor.execute(query) max_epoch = cursor.fetchone()['MAX(epoch)'], # print('max epoch: ', max_epoch, type(max_epoch)) cursor.close() try: st.session_state.epoch['ranking'][prompt_id] = max_epoch[0] + 1 except TypeError: st.session_state.epoch['ranking'][prompt_id] = 1 # st.session_state.epoch['summary'][tag] = st.session_state.epoch['summary'].get(tag, 0) + 1 # st.session_state.epoch['summary']['overall'] += 1 print('epoch: ', st.session_state.epoch) switch_page('ranking') self.gallery_standard(items[items['prompt_id'] == prompt_id].reset_index(drop=True), 4, info, show_checkbox=st.session_state.edit_state) else: # with st.form(key=f'checkout_{tag}'): st.info('No selection under this tag') back = st.button('🖼️ Back to gallery and select something you like', key=f'checkout_{tag}', type='primary') if back: st.session_state.gallery_focus['tag'] = tag st.session_state.gallery_focus['prompt'] = None st.session_state.gallery_state = 'graph' st.rerun() 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') @st.cache_data def load_hf_dataset(show_NSFW=False): # 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 curation filter prompt_to_hide = json.load(open('./data/curation.json', 'r')) prompt_to_hide = list(itertools.chain.from_iterable(prompt_to_hide.values())) print('prompt to hide: ', prompt_to_hide) promptBook = promptBook[~promptBook['prompt_id'].isin(prompt_to_hide)].reset_index(drop=True) # apply a nsfw filter if not show_NSFW: promptBook = promptBook[promptBook['norm_nsfw'] <= 0.8].reset_index(drop=True) print('nsfw filter applied', len(promptBook)) # 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') items = items.merge(tsne_df, on=['modelVersion_id', 'prompt_id'], how='left') 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: GALLERY_CONN = connect_to_db() roster, promptBook, images_ds = load_hf_dataset(st.session_state.show_NSFW) app = GalleryApp(promptBook=promptBook, images_ds=images_ds) app.app() with open('./css/style.css') as f: st.markdown(f'', unsafe_allow_html=True)