Update app.py

app.py

@@ -192,7 +192,7 @@ def predict(input, img_enhancer):
         image = np.array(image_pil)
         mask = np.array(mask_pil.convert("L"))
     output = model_process(image, mask, img_enhancer)
-    return output, gr.update(visible=True)
+    return output, [image_pil, output], gr.update(visible=True)
 
 css = '''    
     .container {max-width: 1150px; margin: auto;padding-top: 1.5rem}
@@ -241,7 +241,7 @@ css = '''
         display: none !important;
     }
 '''
-set_page_elements = """async() => {
+start_cleaner = """async() => {
     function isMobile() {
         try {
             document.createEvent("TouchEvent"); return true;
@@ -258,9 +258,11 @@ set_page_elements = """async() => {
     const group2 = gradioEl.querySelectorAll('#group_2')[0];    
     const image_upload = gradioEl.querySelectorAll('#image_upload')[0];
     const image_output = gradioEl.querySelectorAll('#image_output')[0];
+    const image_output_container = gradioEl.querySelectorAll('#image-output-container')[0];
     const data_image = gradioEl.querySelectorAll('#image_upload [data-testid="image"]')[0];
     const data_image_div = gradioEl.querySelectorAll('#image_upload [data-testid="image"] > div')[0];
 
+    image_output_container.setAttribute('style', 'width: 0px; height:0px; display:none;');
     if (isMobile()) {
         const group1_width = group1.offsetWidth;
 
@@ -276,6 +278,7 @@ set_page_elements = """async() => {
         share_btn_container.setAttribute('style', 'width: 0px; height:0px;'); 
         const share_btn_share_icon = gradioEl.querySelectorAll('#share-btn-share-icon')[0];  
         share_btn_share_icon.setAttribute('style', 'width: 0px; height:0px;'); 
+
     } else {
         max_height = 800;
         image_upload.setAttribute('style', 'min-height: ' + max_height + 'px; overflow-x: overlay');
@@ -289,7 +292,23 @@ set_page_elements = """async() => {
     }
     group1.style.display = "none";
     group2.style.display = "block";
-
+    
+    window['gradioEl'] = gradioEl;
+    window['doCheckAction'] = 0;
+    window['checkAction'] = function checkAction() {
+        try {
+            if (window['doCheckAction'] == 0) {
+                var gallery_items = window['gradioEl'].querySelectorAll('#gallery .gallery-item');
+                if (gallery_items && gallery_items.length == 2) {
+                    window.clearInterval(window['checkAction_interval']);
+                    window['doCheckAction'] = 1;
+                    gallery_items[gallery_items.length-1].click();
+                }
+            }
+        } catch(e) {
+        }        
+    }
+    window['checkAction_interval'] = window.setInterval("window.checkAction()", 500);    
 }"""
 
 download_img = """async() => {
@@ -373,8 +392,12 @@ with image_blocks as demo:
                     enhancer_label = 'Enhanced image(processing is very slow, please check only for blurred images)'
                     img_enhancer = gr.Checkbox(label=enhancer_label).style(width=150) 
             with gr.Row(elem_id="output-container"):             
+                with gr.Column(elem_id="image-output-container"):
+                    image_out = gr.Image(elem_id="image_output",label="Result", show_label=False, visible=False)
                 with gr.Column():
-                    image_out = gr.Image(elem_id="image_output",label="Result", show_label=False)
+                    gallery = gr.Gallery(
+                            label="Generated images", show_label=False, elem_id="gallery"
+                        ).style(grid=[2], height="600px")
             with gr.Row(elem_id="download-container", visible=False) as download_container:             
                 with gr.Column(elem_id="download-btn-container") as download_btn_container:
                     download_button = gr.Button(elem_id="download-btn", value="Save(⏩)")  
@@ -384,10 +407,10 @@ with image_blocks as demo:
                         loading_icon = gr.HTML(loading_icon_html, elem_id="loading-icon", visible=True)
                         share_button = gr.Button("Share to community", elem_id="share-btn", visible=True)  
 
-            erase_btn.click(fn=predict, inputs=[image_input, img_enhancer], outputs=[image_out, download_container])
+            erase_btn.click(fn=predict, inputs=[image_input, img_enhancer], outputs=[image_out, gallery, download_container])
             download_button.click(None, [], [], _js=download_img)
             share_button.click(None, [], [], _js=share_js)
     
-        begin_button.click(fn=None, inputs=[], outputs=[group_1, group_2], _js=set_page_elements)
+        begin_button.click(fn=None, inputs=[], outputs=[group_1, group_2], _js=start_cleaner)
 
-image_blocks.launch(server_name='0.0.0.0')
+image_blocks.launch(server_name='0.0.0.0')