From c54950c05c6a669b5fadbde22711b9232c1470c1 Mon Sep 17 00:00:00 2001
From: Qing <cwq1913@gmail.com>
Date: Tue, 15 Nov 2022 21:53:16 +0800
Subject: [PATCH] add mask rerun button

---
 lama_cleaner/app/package.json                 |  1 +
 .../app/src/components/Editor/Editor.tsx      |  2 +-
 .../app/src/components/Header/Header.tsx      | 85 +++++++++++++------
 lama_cleaner/app/src/store/Atoms.tsx          |  5 ++
 lama_cleaner/app/yarn.lock                    |  5 ++
 lama_cleaner/server.py                        |  2 +-
 6 files changed, 70 insertions(+), 30 deletions(-)

diff --git a/lama_cleaner/app/package.json b/lama_cleaner/app/package.json
index 47c148e..92d7905 100644
--- a/lama_cleaner/app/package.json
+++ b/lama_cleaner/app/package.json
@@ -6,6 +6,7 @@
   "dependencies": {
     "@heroicons/react": "^1.0.4",
     "@radix-ui/react-dialog": "0.1.8-rc.25",
+    "@radix-ui/react-icons": "^1.1.1",
     "@radix-ui/react-popover": "^1.0.0",
     "@radix-ui/react-select": "0.1.2-rc.27",
     "@radix-ui/react-switch": "^0.1.5",
diff --git a/lama_cleaner/app/src/components/Editor/Editor.tsx b/lama_cleaner/app/src/components/Editor/Editor.tsx
index 22da920..003313b 100644
--- a/lama_cleaner/app/src/components/Editor/Editor.tsx
+++ b/lama_cleaner/app/src/components/Editor/Editor.tsx
@@ -1229,7 +1229,7 @@ export default function Editor() {
             onClick={download}
           />
 
-          {settings.runInpaintingManually && (
+          {settings.runInpaintingManually && !isSD && (
             <Button
               toolTip="Run Inpainting"
               tooltipPosition="top"
diff --git a/lama_cleaner/app/src/components/Header/Header.tsx b/lama_cleaner/app/src/components/Header/Header.tsx
index 3b9ae56..2b7c91d 100644
--- a/lama_cleaner/app/src/components/Header/Header.tsx
+++ b/lama_cleaner/app/src/components/Header/Header.tsx
@@ -1,7 +1,13 @@
 import { ArrowLeftIcon, UploadIcon } from '@heroicons/react/outline'
+import { PlayIcon } from '@radix-ui/react-icons'
 import React, { useState } from 'react'
 import { useRecoilState, useRecoilValue } from 'recoil'
-import { fileState, isInpaintingState, isSDState } from '../../store/Atoms'
+import {
+  fileState,
+  isInpaintingState,
+  isSDState,
+  maskState,
+} from '../../store/Atoms'
 import Button from '../shared/Button'
 import Shortcuts from '../Shortcuts/Shortcuts'
 import { ThemeChanger } from './ThemeChanger'
@@ -13,6 +19,7 @@ import emitter, { EVENT_CUSTOM_MASK } from '../../event'
 const Header = () => {
   const isInpainting = useRecoilValue(isInpaintingState)
   const [file, setFile] = useRecoilState(fileState)
+  const [mask, setMask] = useRecoilState(maskState)
   const [uploadElemId] = useState(`file-upload-${Math.random().toString()}`)
   const [maskUploadElemId] = useState(`mask-upload-${Math.random().toString()}`)
   const isSD = useRecoilValue(isSDState)
@@ -29,7 +36,11 @@ const Header = () => {
           }}
         >
           <label htmlFor={uploadElemId}>
-            <Button icon={<UploadIcon />} style={{ border: 0 }}>
+            <Button
+              icon={<UploadIcon />}
+              style={{ border: 0 }}
+              disabled={isInpainting}
+            >
               <input
                 style={{ display: 'none' }}
                 id={uploadElemId}
@@ -47,33 +58,51 @@ const Header = () => {
             </Button>
           </label>
 
-          <label
-            htmlFor={maskUploadElemId}
-            style={{ visibility: file ? 'visible' : 'hidden' }}
+          <div
+            style={{
+              visibility: file ? 'visible' : 'hidden',
+              display: 'flex',
+              justifyContent: 'center',
+              alignItems: 'center',
+            }}
           >
-            <Button style={{ border: 0 }} disabled={isInpainting}>
-              <input
-                style={{ display: 'none' }}
-                id={maskUploadElemId}
-                name={maskUploadElemId}
-                type="file"
-                onClick={e => {
-                  const element = e.target as HTMLInputElement
-                  element.value = ''
-                }}
-                onChange={ev => {
-                  const newFile = ev.currentTarget.files?.[0]
-                  if (newFile) {
-                    // TODO: check mask size
-                    console.info('Send custom mask')
-                    emitter.emit(EVENT_CUSTOM_MASK, { mask: newFile })
-                  }
-                }}
-                accept="image/png, image/jpeg"
-              />
-              Mask
-            </Button>
-          </label>
+            <label htmlFor={maskUploadElemId}>
+              <Button style={{ border: 0 }} disabled={isInpainting}>
+                <input
+                  style={{ display: 'none' }}
+                  id={maskUploadElemId}
+                  name={maskUploadElemId}
+                  type="file"
+                  onClick={e => {
+                    const element = e.target as HTMLInputElement
+                    element.value = ''
+                  }}
+                  onChange={ev => {
+                    const newFile = ev.currentTarget.files?.[0]
+                    if (newFile) {
+                      // TODO: check mask size
+                      console.info('Send custom mask')
+                      emitter.emit(EVENT_CUSTOM_MASK, { mask: newFile })
+                      setMask(newFile)
+                    }
+                  }}
+                  accept="image/png, image/jpeg"
+                />
+                Mask
+              </Button>
+            </label>
+            <Button
+              style={{
+                visibility: mask ? 'visible' : 'hidden',
+              }}
+              icon={<PlayIcon />}
+              onClick={() => {
+                if (mask) {
+                  emitter.emit(EVENT_CUSTOM_MASK, { mask })
+                }
+              }}
+            />
+          </div>
         </div>
 
         {isSD && file ? <PromptInput /> : <></>}
diff --git a/lama_cleaner/app/src/store/Atoms.tsx b/lama_cleaner/app/src/store/Atoms.tsx
index c3aefc2..25002c6 100644
--- a/lama_cleaner/app/src/store/Atoms.tsx
+++ b/lama_cleaner/app/src/store/Atoms.tsx
@@ -18,6 +18,11 @@ export const fileState = atom<File | undefined>({
   default: undefined,
 })
 
+export const maskState = atom<File | undefined>({
+  key: 'maskState',
+  default: undefined,
+})
+
 export interface Rect {
   x: number
   y: number
diff --git a/lama_cleaner/app/yarn.lock b/lama_cleaner/app/yarn.lock
index 1cba09b..2b3e340 100644
--- a/lama_cleaner/app/yarn.lock
+++ b/lama_cleaner/app/yarn.lock
@@ -1762,6 +1762,11 @@
     "@radix-ui/react-primitive" "1.0.0"
     "@radix-ui/react-use-callback-ref" "1.0.0"
 
+"@radix-ui/react-icons@^1.1.1":
+  version "1.1.1"
+  resolved "https://registry.npmmirror.com/@radix-ui/react-icons/-/react-icons-1.1.1.tgz#38d2aa548035dd3b799c169bd17177b1cec3152b"
+  integrity sha512-xc3wQC59rsFylVbSusQCrrM+6695ppF730Q6yqzhRdqDcRNWIm2R6ngpzBoSOQMcwnq4p805F+Gr7xo4fmtN1A==
+
 "@radix-ui/react-id@0.1.5":
   version "0.1.5"
   resolved "https://registry.npmmirror.com/@radix-ui/react-id/-/react-id-0.1.5.tgz#010d311bedd5a2884c1e9bb6aaaa4e6cc1d1d3b8"
diff --git a/lama_cleaner/server.py b/lama_cleaner/server.py
index f7014b0..6211154 100644
--- a/lama_cleaner/server.py
+++ b/lama_cleaner/server.py
@@ -57,7 +57,7 @@ BUILD_DIR = os.environ.get("LAMA_CLEANER_BUILD_DIR", "app/build")
 
 class NoFlaskwebgui(logging.Filter):
     def filter(self, record):
-        return "GET /flaskwebgui-keep-server-alive" not in record.getMessage()
+        return "flaskwebgui-keep-server-alive" not in record.getMessage()
 
 
 logging.getLogger("werkzeug").addFilter(NoFlaskwebgui())