Update app.py

app.py

@@ -30,6 +30,11 @@ def divisible_by_8(image):
     return resized_image
 
 
+def restore_version(index, versions):
+    final_dict = {'background': versions[index], 'layers': None, 'composite': versions[index]}
+    return final_dict
+
+
 def generate(image_editor, prompt, neg_prompt, versions):
     print('versions')
     print(versions)
@@ -89,5 +94,6 @@ with gr.Blocks() as demo:
 
     version_gallery.select(get_select_index, None, selected)
     generate_button.click(fn=generate, inputs=[sketch_pad,prompt, neg_prompt, version_gallery], outputs=[sketch_pad, out1, out2, out3, version_gallery])
+    restore_button.click(fn=restore_version, inputs=[selected, version_gallery], outputs=sketch_pad)
 
 demo.launch()