refactor

app.py

@@ -154,30 +154,24 @@ def make_editing_canvas(canvas_color, brush, _reset_state, generation_mode, pain
         if 'seg' not in st.session_state:
             with st.spinner(text="Preparing image segmentation"):
                 image = get_image()
-                print("Preparing image segmentation")
                 real_seg = np.array(segment_image(Image.fromarray(image)))
                 st.session_state['seg'] = real_seg
-                # get unique RGB colors from this segmentation map
-                st.image(Image.fromarray(real_seg), width=512)
-                print(real_seg.shape)
-                unique_colors = np.unique(real_seg.reshape(-1, real_seg.shape[2]), axis=0)
-                unique_colors = [tuple(color) for color in unique_colors]
-                print(unique_colors)
-                st.session_state['unique_colors'] = unique_colors
-                # multiselect to choose multiple colors
-                chosen_colors = st.multiselect(
-                    label="Choose colors", options=unique_colors, default=None, key='unique_colors', format_func=map_colors_rgb
-                )
-                print(chosen_colors)
+
+        if 'unique_colors' not in st.session_state:
+            unique_colors = np.unique(real_seg.reshape(-1, real_seg.shape[2]), axis=0)
+            unique_colors = [tuple(color) for color in unique_colors]
+            st.session_state['unique_colors'] = unique_colors
+
+        chosen_colors = st.multiselect(
+            label="Choose colors", options=unique_colors, default=None, key='unique_colors', format_func=map_colors_rgb
+        )
 
         if st.button("generate image", key='generate_button'):
             image = get_image()
-            print("Preparing image segmentation")
             real_seg = segment_image(Image.fromarray(image))
             mask, seg = preprocess_seg_mask(canvas, real_seg)
 
             with st.spinner(text="Generating image"):
-                print("Making image")
                 result_image = make_image_controlnet(image=image,
                                                         mask_image=mask,
                                                         controlnet_conditioning_image=seg,