add edit function

Home.py

@@ -43,6 +43,7 @@ def logout():
     st.session_state.pop('selected_dict', None)
     st.session_state.pop('score_weights', None)
     st.session_state.pop('gallery_state', None)
+    st.session_state.pop('edit_state', None)
     st.session_state.pop('progress', None)
     st.session_state.pop('gallery_focus', None)
     st.session_state.pop('assigned_rank_mode', None)

pages/Gallery.py

@@ -300,7 +300,7 @@ class GalleryApp:
         # save tag to session state on change
         tag = st.radio('Select a tag', prompt_tags, index=tag_focus_idx, horizontal=True, key='tag', label_visibility='collapsed')
 
-        print('current state: ', st.session_state.gallery_state)
+        # print('current state: ', st.session_state.gallery_state)
 
         if st.session_state.gallery_state == 'graph':
 
@@ -364,13 +364,12 @@ class GalleryApp:
                     if has_selection:
                         checkout = st.button('Check out selections', use_container_width=True, type='primary')
                         if checkout:
-                            print('checkout')
                             # 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'
-                            print(st.session_state.gallery_state)
+                            # print(st.session_state.gallery_state)
                             st.experimental_rerun()
                     else:
                         st.write(':orange[👇 **Select images you like below**]')
@@ -389,25 +388,6 @@ class GalleryApp:
 
             self.checkout_mode(tag, items)
 
-
-            # items = items[items['modelVersion_id'].isin(st.session_state.selected_dict[prompt_id])].reset_index(
-            #     drop=True)
-            # self.gallery_mode(prompt_id, items)
-            #
-            # with subset_selector[-1]:
-            #     state_operations = st.columns([1, 1])
-            #     with state_operations[0]:
-            #         back = st.button('Back to 🖼️', use_container_width=True)
-            #         if back:
-            #             st.session_state.gallery_state[prompt_id] = 'graph'
-            #             st.experimental_rerun()
-            #
-            #     with state_operations[1]:
-            #         forward = st.button('Check out', use_container_width=True, type='primary', on_click=self.submit_actions, args=('Continue', prompt_id))
-            #         if forward:
-            #             switch_page('ranking')
-
-
     def graph_mode(self, prompt_id, items):
         graph_cols = st.columns([3, 1])
         # prompt = st.chat_input(f"Selected model version ids: {str(st.session_state.selected_dict.get(prompt_id, []))}",
@@ -466,10 +446,6 @@ class GalleryApp:
                     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)
 
-                    # for info in infos:
-                    #     st.write(f"**{info}**:")
-                    #     st.write(item[info])
-
             else:
                 st.info('Please click on an image to show')
 
@@ -525,31 +501,60 @@ class GalleryApp:
                 default_expand = True if st.session_state.gallery_focus['prompt'] == prompt else False
                 with st.expander(f'**{prompt}**', expanded=default_expand):
                     # st.caption('select info to show')
-                    checkout_panel = st.columns([6, 1, 1])
+                    checkout_panel = st.columns([5, 3])
                     with checkout_panel[0]:
                         pass
                         info = st.multiselect('Show Info',
                                               ['model_name', 'model_id', 'modelVersion_name', 'modelVersion_id',
                                                'weighted_score_sum', 'model_download_count', 'clip_score', 'mcos_score',
                                                'nsfw_score', 'norm_nsfw'],
-                                              label_visibility='collapsed', key=f'info_{prompt_id}', placeholder='Select info to show')
-                    with checkout_panel[1]:
-                        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.experimental_rerun()
-                    with checkout_panel[2]:
-                        proceed = st.button('Proceed ➡️', key=f'checkout_proceed_{prompt_id}', use_container_width=True,
-                                            type='primary')
-                        if proceed:
-                            st.session_state.gallery_focus['tag'] = tag
-                            st.session_state.gallery_focus['prompt'] = prompt
-                            switch_page('ranking')
-
-                    self.gallery_standard(items[items['prompt_id'] == prompt_id].reset_index(drop=True), 4, info, show_checkbox=False)
+                                              label_visibility='collapsed', key=f'info_{prompt_id}', placeholder='Select what infos 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.experimental_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.experimental_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]))
+
+                                    st.session_state.edit_state = False
+                                    st.experimental_rerun()
+
+                        with checkout_buttons[-1]:
+                            proceed = st.button('Proceed ➡️', key=f'checkout_proceed_{prompt_id}', use_container_width=True,
+                                                type='primary')
+                            if proceed:
+                                st.session_state.gallery_focus['tag'] = tag
+                                st.session_state.gallery_focus['prompt'] = prompt
+                                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')