import streamlit as st import numpy as np import random import pandas as pd import glob import csv from PIL import Image from datasets import load_dataset, Dataset, load_from_disk from huggingface_hub import login import os import datasets import requests from bs4 import BeautifulSoup 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([1, 1, 2, 0.5]) 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) with selecters[1]: order = st.selectbox('Order', ['Ascending', 'Descending'], index=1 if sort_by == 'clip_score' or sort_by == 'model_download_count' else 0) if order == 'Ascending': order = True else: order = False items = items.sort_values(by=[sort_by], ascending=order).reset_index(drop=True) with selecters[2]: info = st.multiselect('Show Info', ['model_download_count', 'clip_score', 'avg_rank', 'model_name', 'model_id', 'modelVersion_name', 'modelVersion_id'], default=sort_by) col_num = st.slider('Number of columns', min_value=1, max_value=9, value=4, step=1, key='col_num') with selecters[3]: filter = st.selectbox('Filter', ['All', 'Checked', 'Unchecked']) if filter == 'Checked': items = items[items['checked'] == True].reset_index(drop=True) elif filter == 'Unchecked': items = items[items['checked'] == False].reset_index(drop=True) 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()