add new custom weighting mode

app.py

@@ -10,6 +10,8 @@ from datasets import load_dataset, Dataset, load_from_disk
 from huggingface_hub import login
 import os
 import requests
+from bs4 import BeautifulSoup
+import re
 
 SCORE_NAME_MAPPING = {'clip': 'clip_score', 'rank': 'avg_rank', 'pop': 'model_download_count'}
 
@@ -57,7 +59,7 @@ class GalleryApp:
                         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,
                             )
@@ -75,75 +77,22 @@ class GalleryApp:
                             #     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])
+    def selection_panel(self, items):
         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')
+                    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'])
+                    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:
@@ -172,18 +121,24 @@ class GalleryApp:
         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)
+            filter = st.selectbox('Filter', ['Safe', 'All', 'Unsafe'])
+            print('filter', filter)
+            # initialize unsafe_modelVersion_ids
+            if filter == 'Safe':
+                # return checked items
+                items = items[items['checked'] == False].reset_index(drop=True)
+
+            elif filter == 'Unsafe':
+                # return unchecked items
+                items = items[items['checked'] == True].reset_index(drop=True)
+                print(items)
 
         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'],
+                               '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:
@@ -193,20 +148,173 @@ class GalleryApp:
                     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.")
+            st.info(
+                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')
+
+        return items, info, col_num
+
+
+    def selection_panel_2(self, items):
+        selecters = st.columns([1, 5])
+
+        with selecters[0]:
+            sort_type = st.selectbox('Sort by', ['IDs and Names', 'Scores'])
+            if sort_type == 'Scores':
+                sort_by = 'weighted_score_sum'
+
+        with selecters[1]:
+            if sort_type == 'IDs and Names':
+                sub_selecters = st.columns([3, 1, 1])
+                with sub_selecters[0]:
+                    sort_by = st.selectbox('Sort by',
+                                           ['model_name', 'model_id', 'modelVersion_name', 'modelVersion_id'],
+                                           label_visibility='hidden')
+
+                continue_idx = 1
+
+            else:
+                sub_selecters = st.columns([1, 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)
+                with sub_selecters[1]:
+                    rank_weight = st.number_input('Rank Score Weight', min_value=-100.0, max_value=100.0, value=1.0, step=0.1)
+                with sub_selecters[2]:
+                    pop_weight = st.number_input('Popularity Weight', min_value=-100.0, max_value=100.0, value=1.0, step=0.1)
+
+                items.loc[:, 'weighted_score_sum'] = items['norm_clip'] * clip_weight + items['avg_rank'] * rank_weight + items[
+                    'norm_pop'] * pop_weight
+
+                continue_idx = 3
+
+            with sub_selecters[continue_idx]:
+                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 sub_selecters[continue_idx+1]:
+                filter = st.selectbox('Filter', ['Safe', 'All', 'Unsafe'])
+                print('filter', filter)
+                # initialize unsafe_modelVersion_ids
+                if filter == 'Safe':
+                    # return checked items
+                    items = items[items['checked'] == False].reset_index(drop=True)
+
+                elif filter == 'Unsafe':
+                    # return unchecked items
+                    items = items[items['checked'] == True].reset_index(drop=True)
+                    print(items)
+
+        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', 'weighted_score_sum'],
+                              default=sort_by)
 
+        # 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.info(
+                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)
+        return items, info, col_num
+
+    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)
+
+            # show image metadata
+            image_metadatas = ['prompt_id', '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 tag as civitai, add civitai reference
+            if tag == 'civitai':
+                try:
+                    st.write('**Civitai Reference**')
+                    res = requests.get(f'https://civitai.com/images/{prompt_id.item()}')
+                    # 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
+
+
+        # 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['civitai'] = [375790, 366222, 295008, 256477]
+        unsafe_prompts['people'] = [53]
+        unsafe_prompts['art'] = [23]
+        unsafe_prompts['abstract'] = [10, 12]
+
+        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.')
+
+        if safety_check:
+            items, info, col_num = self.selection_panel_2(items)
+            # self.gallery_standard(items, col_num, info)
+
+            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
@@ -223,10 +331,6 @@ class GalleryApp:
         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')
@@ -254,6 +358,10 @@ if __name__ == '__main__':
         if 'checked' not in st.session_state.promptBook.columns:
             st.session_state.promptBook.loc[:, 'checked'] = False
 
+        # add 'custom_score_weights' column to promptBook if not exist
+        if 'weighted_score_sum' not in st.session_state.promptBook.columns:
+            st.session_state.promptBook.loc[:, 'weighted_score_sum'] = 0
+
         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)