fix safety check

app.py

@@ -195,6 +195,7 @@ class GalleryApp:
         return items, info, col_num
 
     def selection_panel_2(self, items):
+
         selecters = st.columns([1, 4])
 
         # select sort type
@@ -219,17 +220,17 @@ class GalleryApp:
                 # add custom weights
                 sub_selecters = st.columns([1, 1, 1, 1])
 
-                if 'default_weights' not in st.session_state:
-                    st.session_state.default_weights = [1.0, 0.8, 0.2]
+                if 'score_weights' not in st.session_state:
+                    st.session_state.score_weights = [1.0, 0.8, 0.2]
 
                 with sub_selecters[0]:
-                    clip_weight = st.number_input('Clip Score Weight', min_value=-100.0, max_value=100.0, value=st.session_state.default_weights[0], step=0.1, help='the weight for normalized clip score')
+                    clip_weight = st.number_input('Clip Score Weight', min_value=-100.0, max_value=100.0, value=st.session_state.score_weights[0], step=0.1, help='the weight for normalized clip score')
                 with sub_selecters[1]:
-                    rank_weight = st.number_input('Distinctiveness Weight', min_value=-100.0, max_value=100.0, value=st.session_state.default_weights[1], step=0.1, help='the weight for average rank')
+                    rank_weight = st.number_input('Distinctiveness Weight', min_value=-100.0, max_value=100.0, value=st.session_state.score_weights[1], step=0.1, help='the weight for average rank')
                 with sub_selecters[2]:
-                    pop_weight = st.number_input('Popularity Weight', min_value=-100.0, max_value=100.0, value=st.session_state.default_weights[2], step=0.1, help='the weight for normalized popularity score')
+                    pop_weight = st.number_input('Popularity Weight', min_value=-100.0, max_value=100.0, value=st.session_state.score_weights[2], step=0.1, help='the weight for normalized popularity score')
 
-                st.session_state.default_weights = [clip_weight, rank_weight, pop_weight]
+                st.session_state.score_weights = [clip_weight, rank_weight, pop_weight]
 
                 items.loc[:, 'weighted_score_sum'] = round(items['norm_clip'] * clip_weight + items['avg_rank'] * rank_weight + items[
                     'norm_pop'] * pop_weight, 4)
@@ -241,17 +242,6 @@ class GalleryApp:
                 dist_threshold = st.number_input('Distinctiveness Threshold', min_value=0.0, max_value=1.0, value=0.84, step=0.01, help='Only show models with distinctiveness score lower than this threshold, set 1.0 to show all images')
                 items = items[items['avg_rank'] < dist_threshold].reset_index(drop=True)
 
-                # filter = st.selectbox('Filter', ['Safe', 'All', 'Unsafe'])
-                # print('filter', filter)
-                # # initialize unsafe_modelVersion_ids
-                # if filter == 'Safe':
-                #     # return unchecked items
-                #     items = items[items['checked'] == False].reset_index(drop=True)
-                #
-                # elif filter == 'Unsafe':
-                #     # return checked items
-                #     items = items[items['checked'] == True].reset_index(drop=True)
-
         # draw a distribution histogram
         if sort_type == 'Scores':
             try:
@@ -369,7 +359,7 @@ class GalleryApp:
 
         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.')
+            safety_check = st.checkbox('I understand that this prompt may contain unsafe content. Show these images anyway.', key=f'{prompt_id}')
 
         if safety_check:
             items, info, col_num = self.selection_panel_2(items)
@@ -392,6 +382,7 @@ class GalleryApp:
                 with buttons_space:
                     st.form_submit_button('Refresh', on_click=gallery_space.empty, use_container_width=True, type='primary')
 
+
     def reset_current_prompt(self, prompt_id):
         # reset current prompt
         self.promptBook.loc[self.promptBook['prompt_id'] == prompt_id, 'checked'] = False