add iput image

app.py

@@ -63,6 +63,7 @@ def move_image(source: Union[str, Image.Image],
             del st.session_state['unique_colors']
 
     st.session_state[dest] = source_image
+    st.session_state['dest'] = source_image
     if rerun:
         st.experimental_rerun()
 
@@ -149,7 +150,8 @@ def make_output_image():
     st.write("#### Output image")
     st.image(output_image, width=512)
     if st.button("Move to input image"):
-        move_image('output_image', 'initial_image', remove_state=True, rerun=True)
+        move_image('output_image', 'initial_image', remove_state=False, rerun=False)
+        move_image('output_image', 'input_image', remove_state=True, rerun=True)
 
 def make_editing_canvas(canvas_color, brush, _reset_state, generation_mode, paint_mode):
     st.write("#### Input image")