Add keyboard shortcut support

app.py

@@ -21,6 +21,8 @@ import gradio as gr
 import base64
 import skimage
 import skimage.measure
+import yaml
+import json
 
 try:
     abspath = os.path.abspath(__file__)
@@ -59,6 +61,10 @@ if device != "cuda":
 
     autocast = contextlib.nullcontext
 
+with open("config.yaml", "r") as yaml_in:
+    yaml_object = yaml.safe_load(yaml_in)
+    config_json = json.dumps(yaml_object)
+
 
 def load_html():
     body, canvaspy = "", ""
@@ -480,8 +486,8 @@ def run_outpaint(
         height=height,
     )
     base64_str_lst = []
-    if use_correction!="disabled" and enable_img2img:
-        use_correction="border_mode"
+    if use_correction != "disabled" and enable_img2img:
+        use_correction = "border_mode"
     for image in images:
         image = correction_func.run(pil.resize(image.size), image, mode=use_correction)
         resized_img = image.resize((width, height), resample=SAMPLING_MODE,)
@@ -613,12 +619,8 @@ with blocks as demo:
                 type="value",
             )
             postprocess_check = gr.Radio(
-                label="Photometric Correction Mode", 
-                choices=[
-                    "disabled",
-                    "mask_mode",
-                    "border_mode",
-                ],
+                label="Photometric Correction Mode",
+                choices=["disabled", "mask_mode", "border_mode",],
                 value="disabled",
                 type="value",
             )
@@ -662,6 +664,11 @@ with blocks as demo:
     <img src='https://not.exist' onerror='{xss_js}'>""",
         visible=False,
     )
+    xss_html_setup_shortcut = gr.HTML(
+        value=f"""
+    <img src='https://not.exist' onerror='let app=document.querySelector("gradio-app");app=app.shadowRoot??app;let frame=app.querySelector("#sdinfframe").contentWindow;frame.postMessage(["shortcut", "f{config_json}"], "*");'>""",
+        visible=False,
+    )
     # sd pipeline parameters
     sd_img2img = gr.Checkbox(label="Enable Img2Img", value=False, visible=False)
     sd_resize = gr.Checkbox(label="Resize small input", value=True, visible=False)

index.html

@@ -105,7 +105,8 @@ from js import (
     update_scale,
     adjust_selection,
     update_count,
-    enable_result_lst
+    enable_result_lst,
+    setup_shortcut,
 )
 
 
@@ -230,6 +231,8 @@ async def upload_func(event):
     #base.buffer[yo:yo+h,xo:xo+w,-1]=arr[:,:,-1]
     base.draw_buffer()
 
+async def setup_shortcut_func(event):
+    setup_shortcut(event.data[1])
   
 
 document.querySelector("#export").addEventListener("click",create_proxy(export_func))
@@ -357,6 +360,8 @@ async def message_func(event):
         await eraser_size_func(event)
     elif event.data[0]=="resize_selection":
         await resize_selection_func(event)
+    elif event.data[0]=="shortcut":
+        await setup_shortcut_func(event)
     
 window.addEventListener("message",create_proxy(message_func))
 

js/toolbar.js

@@ -431,4 +431,27 @@ function end_overlay()
     document.querySelector("#overlay_container").style.opacity = 1.0;
     document.querySelector("#overlay_container").style.pointerEvents="none";
 }
-document.querySelector("#container").addEventListener("wheel",(e)=>{e.preventDefault()})
+document.querySelector("#container").addEventListener("wheel",(e)=>{e.preventDefault()})
+window.setup_shortcut=function(json)
+{
+    var config=JSON.parse(json);
+    var key_map={};
+    Object.keys(config.shortcut).forEach(k=>{
+        key_map[config.shortcut[k]]=k;
+    })
+    document.addEventListener("keydown",(e)=>{
+        if(e.target.tagName!="INPUT")
+        {
+            let key=e.key;
+            if(e.ctrlKey)
+            {
+                key="Ctrl+"+e.key;
+                e.preventDefault();
+            }
+            if(key in key_map)
+            {
+                w2ui.toolbar.click(key_map[key]);
+            }
+        }
+    })
+}

postprocess.py

@@ -73,6 +73,8 @@ class PhotometricCorrection:
         output_arr=np.array(inpainted_image)
         mask=input_arr[:,:,-1]
         mask=255-mask
+        if mask.sum()<1 and mode=="mask_mode":
+            mode=""
         if mode=="mask_mode":
             mask = skimage.measure.block_reduce(mask, (8, 8), np.max)
             mask = mask.repeat(8, axis=0).repeat(8, axis=1)