Spaces:
Running
Running
| import streamlit as st | |
| import numpy as np | |
| import random | |
| import pandas as pd | |
| import glob | |
| import csv | |
| from PIL import Image | |
| import datasets | |
| from datasets import load_dataset, Dataset, load_from_disk | |
| from huggingface_hub import login | |
| import os | |
| import requests | |
| SCORE_NAME_MAPPING = {'clip': 'clip_score', 'rank': 'avg_rank', 'pop': 'model_download_count'} | |
| class GalleryApp: | |
| def __init__(self, promptBook): | |
| self.promptBook = promptBook | |
| st.set_page_config(layout="wide") | |
| def gallery_masonry(self, items, col_num, info): | |
| cols = st.columns(col_num) | |
| # # sort items by brisque score | |
| # items = items.sort_values(by=['brisque'], ascending=True).reset_index(drop=True) | |
| for idx in range(len(items)): | |
| with cols[idx % col_num]: | |
| image = st.session_state.images[items.iloc[idx]['row_idx'].item()]['image'] | |
| st.image(image, | |
| use_column_width=True, | |
| ) | |
| # with st.expander('Similarity Info'): | |
| # tab1, tab2 = st.tabs(['Most Similar', 'Least Similar']) | |
| # with tab1: | |
| # st.image(image, use_column_width=True) | |
| # with tab2: | |
| # st.image(image, use_column_width=True) | |
| # show checkbox | |
| self.promptBook.loc[items.iloc[idx]['row_idx'].item(), 'checked'] = st.checkbox( | |
| 'Select', value=self.promptBook.loc[items.iloc[idx]['row_idx'].item(), 'checked'], | |
| key=f'select_{idx}') | |
| for key in info: | |
| st.write(f"**{key}**: {items.iloc[idx][key]}") | |
| def gallery_standard(self, items, col_num, info): | |
| rows = len(items) // col_num + 1 | |
| containers = [st.container() for _ in range(rows*2)] | |
| for idx in range(0, len(items), col_num): | |
| # assign one container for each row | |
| row_idx = (idx // col_num) * 2 | |
| 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 = st.session_state.images[items.iloc[idx+j]['row_idx'].item()]['image'] | |
| # image = list(st.session_state.images.skip(items.iloc[idx+j]['row_idx'].item()).take(1))[0]['image'] | |
| st.image(image, | |
| use_column_width=True, | |
| ) | |
| # show checkbox | |
| self.promptBook.loc[items.iloc[idx+j]['row_idx'].item(), 'checked'] = st.checkbox('Select', value=self.promptBook.loc[items.iloc[idx+j]['row_idx'].item(), 'checked'], key=f'select_{idx+j}') | |
| # show selected info | |
| for key in info: | |
| st.write(f"**{key}**: {items.iloc[idx+j][key]}") | |
| # st.write(row_idx/2, idx+j, rows) | |
| # extra_info = st.checkbox('Extra Info', key=f'extra_info_{idx+j}') | |
| # if extra_info: | |
| # with containers[row_idx+1]: | |
| # st.image(image, use_column_width=True) | |
| def app(self): | |
| st.title('Model Coffer Gallery') | |
| st.write('This is a gallery of images generated by the models in the Model Coffer') | |
| with st.sidebar: | |
| prompt_tags = self.promptBook['tag'].unique() | |
| # sort tags by alphabetical order | |
| prompt_tags = np.sort(prompt_tags)[::-1] | |
| tag = st.selectbox('Select a tag', prompt_tags) | |
| items = self.promptBook[self.promptBook['tag'] == tag].reset_index(drop=True) | |
| original_prompts = np.sort(items['prompt'].unique())[::-1] | |
| # remove the first four items in the prompt, which are mostly the same | |
| if tag != 'abstract': | |
| prompts = [', '.join(x.split(', ')[4:]) for x in original_prompts] | |
| prompt = st.selectbox('Select prompt', prompts) | |
| idx = prompts.index(prompt) | |
| prompt_full = ', '.join(original_prompts[idx].split(', ')[:4]) + ', ' + prompt | |
| else: | |
| prompt_full = st.selectbox('Select prompt', original_prompts) | |
| prompt_id = items[items['prompt'] == prompt_full]['prompt_id'].unique()[0] | |
| items = items[items['prompt_id'] == prompt_id].reset_index(drop=True) | |
| st.write('**Prompt ID**') | |
| st.caption(f"{prompt_id}") | |
| st.write('**Prompt**') | |
| st.caption(f"{items['prompt'][0]}") | |
| st.write('**Negative Prompt**') | |
| st.caption(f"{items['negativePrompt'][0]}") | |
| st.write('**Sampler**') | |
| st.caption(f"{items['sampler'][0]}") | |
| st.write('**cfgScale**') | |
| st.caption(f"{items['cfgScale'][0]}") | |
| st.write('**Size**') | |
| st.caption(f"width: {items['size'][0].split('x')[0]}, height: {items['size'][0].split('x')[1]}") | |
| st.write('**Seed**') | |
| st.caption(f"{items['seed'][0]}") | |
| # # for tag as civitai, add civitai reference | |
| # if tag == 'civitai': | |
| # st.write('**Reference**') | |
| # | |
| # res = requests.get(f'https://civitai.com/images', params={'post_id': prompt_id}) | |
| # st.write(res) | |
| # image_url = res.json()['items'][0]['url'] | |
| # st.image(image_url, use_column_width=True) | |
| # with images: | |
| # selecters = st.columns([2, 1, 2, 0.5]) | |
| selecters = st.columns([4, 1, 1]) | |
| with selecters[0]: | |
| # # sort_by = st.selectbox('Sort by', items.columns[11: -1]) | |
| # sort_by = st.selectbox('Sort by', ['model_download_count', 'clip_score', 'avg_rank', 'model_name', 'model_id', | |
| # 'modelVersion_name', 'modelVersion_id']) | |
| print(items.columns) | |
| types = st.columns([1, 3]) | |
| with types[0]: | |
| sort_type = st.selectbox('Sort by', ['IDs and Names', 'Scores']) | |
| with types[1]: | |
| if sort_type == 'IDs and Names': | |
| sort_by = st.selectbox('Sort by', ['model_name', 'model_id', 'modelVersion_name', 'modelVersion_id'], label_visibility='hidden') | |
| elif sort_type == 'Scores': | |
| sort_by = st.multiselect('Sort by', ['clip_score', 'avg_rank', 'popularity'], label_visibility='hidden', default=['clip_score', 'avg_rank', 'popularity']) | |
| # process sort_by to map to the column name | |
| if len(sort_by) == 3: | |
| sort_by = 'clip+rank+pop' | |
| elif len(sort_by) == 2: | |
| if 'clip_score' in sort_by and 'avg_rank' in sort_by: | |
| sort_by = 'clip+rank' | |
| elif 'clip_score' in sort_by and 'popularity' in sort_by: | |
| sort_by = 'clip+pop' | |
| elif 'avg_rank' in sort_by and 'popularity' in sort_by: | |
| sort_by = 'rank+pop' | |
| elif len(sort_by) == 1: | |
| if 'popularity' in sort_by: | |
| sort_by = 'model_download_count' | |
| else: | |
| sort_by = sort_by[0] | |
| print(sort_by) | |
| with selecters[1]: | |
| order = st.selectbox('Order', ['Ascending', 'Descending'], index=1 if sort_type == 'Scores' else 0) | |
| if order == 'Ascending': | |
| order = True | |
| else: | |
| order = False | |
| items = items.sort_values(by=[sort_by], ascending=order).reset_index(drop=True) | |
| with selecters[2]: | |
| filter = st.selectbox('Filter', ['All', 'Checked', 'Unchecked']) | |
| if filter == 'Checked': | |
| items = items[items['checked'] is True].reset_index(drop=True) | |
| elif filter == 'Unchecked': | |
| items = items[items['checked'] is False].reset_index(drop=True) | |
| info = st.multiselect('Show Info', | |
| ['model_download_count', 'clip_score', 'avg_rank', 'model_name', 'model_id', | |
| 'modelVersion_name', 'modelVersion_id', 'clip+rank', 'clip+pop', 'rank+pop', 'clip+rank+pop'], | |
| default=sort_by) | |
| print('info', info) | |
| # add one annotation | |
| mentioned_scores = [] | |
| for i in info: | |
| if '+' in i: | |
| mentioned = i.split('+') | |
| for m in mentioned: | |
| if SCORE_NAME_MAPPING[m] not in mentioned_scores: | |
| mentioned_scores.append(SCORE_NAME_MAPPING[m]) | |
| if len(mentioned_scores) > 0: | |
| st.write(f"**Note: ** The scores {mentioned_scores} are normalized to [0, 1] for each score type, and then added together. The higher the score, the better the model.") | |
| col_num = st.slider('Number of columns', min_value=1, max_value=9, value=4, step=1, key='col_num') | |
| with st.form(key=f'{prompt_id}', clear_on_submit=False): | |
| buttons = st.columns([1, 1, 1]) | |
| with buttons[0]: | |
| submit = st.form_submit_button('Save selections', on_click=self.save_checked, use_container_width=True, type='primary') | |
| with buttons[1]: | |
| submit = st.form_submit_button('Reset current prompt', on_click=self.reset_current_prompt, kwargs={'prompt_id': prompt_id} , use_container_width=True) | |
| with buttons[2]: | |
| submit = st.form_submit_button('Reset all selections', on_click=self.reset_all, use_container_width=True) | |
| self.gallery_standard(items, col_num, info) | |
| def reset_current_prompt(self, prompt_id): | |
| # reset current prompt | |
| self.promptBook.loc[self.promptBook['prompt_id'] == prompt_id, 'checked'] = False | |
| self.save_checked() | |
| def reset_all(self): | |
| # reset all | |
| self.promptBook.loc[:, 'checked'] = False | |
| self.save_checked() | |
| def save_checked(self): | |
| # save checked images to huggingface dataset | |
| dataset = load_dataset('NYUSHPRP/ModelCofferMetadata', split='train') | |
| # get checked images | |
| checked_info = self.promptBook['checked'] | |
| # print('checked_info: ', checked_info) | |
| # for d in checked_info: | |
| # if d is True: | |
| # print('checked') | |
| if 'checked' in dataset.column_names: | |
| dataset = dataset.remove_columns('checked') | |
| dataset = dataset.add_column('checked', checked_info) | |
| # print('metadata dataset: ', dataset) | |
| dataset.push_to_hub('NYUSHPRP/ModelCofferMetadata', split='train') | |
| if __name__ == '__main__': | |
| login(token=os.environ.get("HF_TOKEN")) | |
| if 'roster' not in st.session_state: | |
| print('loading roster') | |
| # st.session_state.roster = pd.DataFrame(load_dataset('NYUSHPRP/ModelCofferRoster', split='train')) | |
| st.session_state.roster = pd.DataFrame(load_from_disk(os.path.join(os.getcwd(), 'data', 'roster'))) | |
| st.session_state.roster = st.session_state.roster[['model_id', 'model_name', 'modelVersion_id', 'modelVersion_name', | |
| 'model_download_count']].drop_duplicates().reset_index(drop=True) | |
| # add model download count from roster to promptbook dataframe | |
| if 'promptBook' not in st.session_state: | |
| print('loading promptBook') | |
| st.session_state.promptBook = pd.DataFrame(load_dataset('NYUSHPRP/ModelCofferMetadata', split='train')) | |
| # add 'checked' column to promptBook if not exist | |
| if 'checked' not in st.session_state.promptBook.columns: | |
| st.session_state.promptBook.loc[:, 'checked'] = False | |
| st.session_state.images = load_from_disk(os.path.join(os.getcwd(), 'data', 'promptbook')) | |
| # st.session_state.images = load_dataset('NYUSHPRP/ModelCofferPromptBook', split='train', streaming=True) | |
| print(st.session_state.images) | |
| print('images loaded') | |
| # st.session_state.promptBook = pd.DataFrame(load_dataset('NYUSHPRP/ModelCofferPromptBook', split='train')) | |
| st.session_state.promptBook = st.session_state.promptBook.merge(st.session_state.roster[['model_id', 'model_name', 'modelVersion_id', 'modelVersion_name', 'model_download_count']], on=['model_id', 'modelVersion_id'], how='left') | |
| # add column to record current row index | |
| st.session_state.promptBook['row_idx'] = st.session_state.promptBook.index | |
| print('promptBook loaded') | |
| # print(st.session_state.promptBook) | |
| check_roster_error = False | |
| if check_roster_error: | |
| # print all rows with the same model_id and modelVersion_id but different model_download_count in roster | |
| print(st.session_state.roster[st.session_state.roster.duplicated(subset=['model_id', 'modelVersion_id'], keep=False)].sort_values(by=['model_id', 'modelVersion_id'])) | |
| app = GalleryApp(promptBook=st.session_state.promptBook) | |
| app.app() | |