Image masking support

index.html

@@ -30,10 +30,10 @@
             width: 95%;
         }
     }
-    #init_image_preview_container {
+    .image_preview_container {
         display: none;
     }
-    #init_image_clear {
+    .image_clear_btn {
         position: absolute;
         transform: translateX(-50%);
         background: black;
@@ -87,9 +87,17 @@
     <textarea id="prompt">a photograph of an astronaut riding a horse</textarea><br/>
 
     <label for="init_image"><b>Initial Image:</b> (optional) </label> <input id="init_image" name="init_image" type="file" /> </button><br/>
-    <div id="init_image_preview_container">
+    <div id="init_image_preview_container" class="image_preview_container">
         <img id="init_image_preview" src="" width="100" height="100" />
-        <button id="init_image_clear">X</button>
+        <button id="init_image_clear" class="image_clear_btn">X</button>
+    </div><br/>
+
+    <div id="mask_setting">
+        <label for="mask"><b>Image Mask:</b> (optional) </label> <input id="mask" name="mask" type="file" /> </button><br/>
+        <div id="mask_preview_container" class="image_preview_container">
+            <img id="mask_preview" src="" width="100" height="100" />
+            <button id="mask_clear" class="image_clear_btn">X</button>
+        </div>
     </div>
 
     <div id="configHeader"><b>Advanced settings:</b> [<a id="configToggleBtn" href="#">show</a>]</div>
@@ -133,6 +141,8 @@ let widthField = document.querySelector('#width')
 let heightField = document.querySelector('#height')
 let initImageSelector = document.querySelector("#init_image")
 let initImagePreview = document.querySelector("#init_image_preview")
+let maskImageSelector = document.querySelector("#mask")
+let maskImagePreview = document.querySelector("#mask_preview")
 let promptStrengthField = document.querySelector('#prompt_strength')
 let promptStrengthValueLabel = document.querySelector('#prompt_strength_value')
 
@@ -143,6 +153,10 @@ let initImagePreviewContainer = document.querySelector('#init_image_preview_cont
 let initImageClearBtn = document.querySelector('#init_image_clear')
 let promptStrengthContainer = document.querySelector('#prompt_strength_container')
 
+let maskSetting = document.querySelector('#mask_setting')
+let maskImagePreviewContainer = document.querySelector('#mask_preview_container')
+let maskImageClearBtn = document.querySelector('#mask_clear')
+
 let showConfigToggle = document.querySelector('#configToggleBtn')
 let configBox = document.querySelector('#config')
 let outputMsg = document.querySelector('#outputMsg')
@@ -229,6 +243,10 @@ async function makeImage() {
     if (initImagePreview.src.indexOf('data:image/png;base64') !== -1) {
         reqBody['init_image'] = initImagePreview.src
         reqBody['prompt_strength'] = promptStrengthField.value / 10
+
+        if (maskImagePreview.src.indexOf('data:image/png;base64') !== -1) {
+            reqBody['mask'] = maskImagePreview.src
+        }
     }
 
     let res = ''
@@ -321,6 +339,8 @@ async function makeImage() {
             initImagePreviewContainer.style.display = 'block'
             promptStrengthContainer.style.display = 'block'
 
+            maskSetting.style.display = 'block'
+
             randomSeedField.checked = false
             seedField.value = seed
             seedField.disabled = false
@@ -380,6 +400,7 @@ function showInitImagePreview() {
     if (initImageSelector.files.length === 0) {
         initImagePreviewContainer.style.display = 'none'
         promptStrengthContainer.style.display = 'none'
+        maskSetting.style.display = 'none'
         return
     }
 
@@ -391,6 +412,8 @@ function showInitImagePreview() {
         initImagePreview.src = reader.result
         initImagePreviewContainer.style.display = 'block'
         promptStrengthContainer.style.display = 'block'
+
+        maskSetting.style.display = 'block'
     })
 
     if (file) {
@@ -402,11 +425,46 @@ showInitImagePreview()
 
 initImageClearBtn.addEventListener('click', function() {
     initImageSelector.value = null
+    maskImageSelector.value = null
+
     initImagePreview.src = ''
+    maskImagePreview.src = ''
+
     initImagePreviewContainer.style.display = 'none'
+    maskImagePreviewContainer.style.display = 'none'
+
+    maskSetting.style.display = 'none'
+
     promptStrengthContainer.style.display = 'none'
 })
 
+function showMaskImagePreview() {
+    if (maskImageSelector.files.length === 0) {
+        maskImagePreviewContainer.style.display = 'none'
+        return
+    }
+
+    let reader = new FileReader()
+    let file = maskImageSelector.files[0]
+
+    reader.addEventListener('load', function() {
+        maskImagePreview.src = reader.result
+        maskImagePreviewContainer.style.display = 'block'
+    })
+
+    if (file) {
+        reader.readAsDataURL(file)
+    }
+}
+maskImageSelector.addEventListener('change', showMaskImagePreview)
+showMaskImagePreview()
+
+maskImageClearBtn.addEventListener('click', function() {
+    maskImageSelector.value = null
+    maskImagePreview.src = ''
+    maskImagePreviewContainer.style.display = 'none'
+})
+
 setInterval(healthCheck, HEALTH_PING_INTERVAL * 1000)
 </script>
 

main.py

@@ -13,6 +13,7 @@ app = FastAPI()
 class ImageRequest(BaseModel):
     prompt: str
     init_image: str = None # base64
+    mask: str = None # base64
     num_outputs: str = "1"
     num_inference_steps: str = "50"
     guidance_scale: str = "7.5"
@@ -51,6 +52,9 @@ async def image(req : ImageRequest):
         data['input']['init_image'] = req.init_image
         data['input']['prompt_strength'] = req.prompt_strength
 
+        if req.mask is not None:
+            data['input']['mask'] = req.mask
+
     if req.seed == "-1":
         del data['input']['seed']