Update app.py

app.py

@@ -13,7 +13,7 @@ import time
 import imghdr
 from loguru import logger
 
-from share_btn import community_icon_html, loading_icon_html, share_js
+from share_btn import community_icon_html, loading_icon_html, download_icon_html, share_js
 
 from lama_cleaner.model_manager import ModelManager
 from lama_cleaner.schema import Config
@@ -192,19 +192,26 @@ 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), gr.update(visible=True), gr.update(visible=True)
+    return output, gr.update(visible=True), gr.update(visible=True)
 
-
-css = '''
-    .container {max-width: 100%;margin: auto;padding-top: 1.5rem}
+css = '''    
+    .container {max-width: 1150px; margin: auto;padding-top: 1.5rem}
     #begin-btn {color: blue; font-size:20px;}
     #work-container {min-width: min(160px, 100%) !important;flex-grow: 0 !important}
-    #image_output{margin: 0 auto; text-align: center;width:640px}
-    #erase-container{margin: 0 auto; text-align: center;width:150px;border-width:5px;border-color:#2c9748}
+   
+    #op-container{margin: 0 auto; text-align: center;width:fit-content;min-width: min(150px, 100%);flex-grow: 0; flex-wrap: nowrap;}
+    #erase-btn-container{margin: 0 auto; text-align: center;width:150px;border-width:5px;border-color:#2c9748}
+    #erase-btn {padding:0;}
     #enhancer-checkbox{width:520px}
     #enhancer-tip{width:450px}
     #enhancer-tip-div{text-align: left}
-    #prompt-container{margin: 0 auto; text-align: center;width:fit-content;min-width: min(150px, 100%);flex-grow: 0; flex-wrap: nowrap;}
+
+    #image_output{margin: 0 auto; text-align: center;width:640px}
+
+    #download-container{margin: 0 auto; text-align: center;width:fit-content; min-width: min(150px, 100%);flex-grow: 0; flex-wrap: nowrap;}
+    #download-btn-container{margin: 0 auto; text-align: center;width:100px;border-width:1px;border-color:#2c9748}
+    #download-btn {padding:0;}
+
     #image_upload .touch-none{display: flex}
     @keyframes spin {
         from {
@@ -239,6 +246,7 @@ set_page_elements = """async() => {
             return false; 
         }
     }
+
     var gradioEl = document.querySelector('body > gradio-app').shadowRoot;
     if (!gradioEl) {
         gradioEl = document.querySelector('body > gradio-app');
@@ -249,16 +257,21 @@ set_page_elements = """async() => {
     const image_output = gradioEl.querySelectorAll('#image_output')[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];
+
     if (isMobile()) {
         const group1_width = group1.offsetWidth;
+
         image_upload.setAttribute('style', 'width:' + (group1_width - 13*2) + 'px; min-height:none;');
         data_image.setAttribute('style', 'width: ' + (group1_width - 14*2) + 'px;min-height:none;');
         data_image_div.setAttribute('style', 'width: ' + (group1_width - 14*2) + 'px;min-height:none;');
         image_output.setAttribute('style', 'width: ' + (group1_width - 13*2) + 'px;min-height:none;'); 
+
         const enhancer = gradioEl.querySelectorAll('#enhancer-checkbox')[0];  
         enhancer.style.display = "none";
+
         const share_btn_container = gradioEl.querySelectorAll('#share-btn-container')[0];  
         share_btn_container.setAttribute('style', 'width: 0px; height:0px;'); 
+
     } else {
         max_height = 800;
         image_upload.setAttribute('style', 'min-height: ' + max_height + 'px; overflow-x: overlay');
@@ -267,6 +280,45 @@ set_page_elements = """async() => {
     }
     group1.style.display = "none";
     group2.style.display = "block";
+
+}"""
+
+download_img = """async() => {
+    Date.prototype.Format = function (fmt) { 
+        var o = {
+            "M+": this.getMonth() + 1, 
+            "d+": this.getDate(), 
+            "h+": this.getHours(), 
+            "m+": this.getMinutes(),
+            "s+": this.getSeconds(),
+            "q+": Math.floor((this.getMonth() + 3) / 3), 
+            "S": this.getMilliseconds()
+        };
+        if (/(y+)/.test(fmt))
+            fmt = fmt.replace(RegExp.$1, (this.getFullYear() + "").substr(4 - RegExp.$1.length));
+        for (var k in o)
+            if (new RegExp("(" + k + ")").test(fmt)) fmt = fmt.replace(RegExp.$1, (RegExp.$1.length == 1) ? (o[k]) : (("00" + o[k]).substr(("" + o[k]).length)));
+                return fmt;
+    }   
+
+    var gradioEl = document.querySelector('body > gradio-app').shadowRoot;
+    if (!gradioEl) {
+        gradioEl = document.querySelector('body > gradio-app');
+    }  
+    const out_image = gradioEl.querySelectorAll('#image_output img')[0]; 
+    if (out_image) {
+        var x=new XMLHttpRequest();
+        x.open("GET", out_image.src, true);
+        x.responseType = 'blob';
+        x.onload = function(e){
+            var url = window.URL.createObjectURL(x.response)
+            var a = document.createElement('a');
+            a.href = url;
+            a.download = (new Date()).Format("yyyyMMdd_hhmmss");
+            a.click();
+        }
+        x.send();
+    }
 }"""
 
 image_blocks = gr.Blocks(css=css, title='Image Cleaner')
@@ -300,9 +352,9 @@ with image_blocks as demo:
                 with gr.Column():
                     image_input = gr.Image(source='upload', elem_id="image_upload",tool='sketch', type=f'{image_type}', 
                                      label="Upload", show_label=False).style(mobile_collapse=False)
-            with gr.Row(elem_id="prompt-container").style(mobile_collapse=False, equal_height=True):
-                with gr.Column(elem_id="erase-container"):
-                    btn_erase = gr.Button(value = "Erase(↓)",elem_id="erase_btn").style(
+            with gr.Row(elem_id="op-container").style(mobile_collapse=False, equal_height=True):
+                with gr.Column(elem_id="erase-btn-container"):
+                    erase_btn = gr.Button(value = "Erase(↓)",elem_id="erase-btn").style(
                         margin=True,
                         rounded=(True, True, True, True),
                         full_width=True,
@@ -313,12 +365,16 @@ with image_blocks as demo:
             with gr.Row(elem_id="output-container"):             
                 with gr.Column():
                     image_out = gr.Image(elem_id="image_output",label="Result", show_label=False)
-                    with gr.Group(elem_id="share-btn-container"):
-                        community_icon = gr.HTML(community_icon_html, visible=False)
-                        loading_icon = gr.HTML(loading_icon_html, visible=False)
-                        share_button = gr.Button("Share to community", elem_id="share-btn", visible=False)                    
+                    with gr.Group(elem_id="share-btn-container", visible=False) as share_btn_container:
+                        community_icon = gr.HTML(community_icon_html)
+                        loading_icon = gr.HTML(loading_icon_html)
+                        share_button = gr.Button("Share to community", elem_id="share-btn")  
+            with gr.Row(elem_id="download-container", visible=False) as download_container:             
+                with gr.Column(elem_id="download-btn-container"):
+                    download_button = gr.Button(elem_id="download-btn", value="Save(→)")
    
-            btn_erase.click(fn=predict, inputs=[image_input, img_enhancer], outputs=[image_out, community_icon, loading_icon, share_button])
+            erase_btn.click(fn=predict, inputs=[image_input, img_enhancer], outputs=[image_out, download_container, share_btn_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)