From eee91e7b82aab80722764c74410840b52bf14491 Mon Sep 17 00:00:00 2001
From: Qing <cwq1913@gmail.com>
Date: Wed, 13 Jul 2022 09:04:28 +0800
Subject: [PATCH] frontend: add ZITS config

---
 lama_cleaner/app/src/adapters/inpainting.ts   |  1 +
 .../components/Settings/ModelSettingBlock.tsx | 36 +++++++++++++++++--
 lama_cleaner/app/src/store/Atoms.tsx          |  7 ++++
 lama_cleaner/server.py                        |  1 +
 4 files changed, 43 insertions(+), 2 deletions(-)

diff --git a/lama_cleaner/app/src/adapters/inpainting.ts b/lama_cleaner/app/src/adapters/inpainting.ts
index 9c9e459..f7fd3ef 100644
--- a/lama_cleaner/app/src/adapters/inpainting.ts
+++ b/lama_cleaner/app/src/adapters/inpainting.ts
@@ -17,6 +17,7 @@ export default async function inpaint(
 
   fd.append('ldmSteps', settings.ldmSteps.toString())
   fd.append('ldmSampler', settings.ldmSampler.toString())
+  fd.append('zitsWireframe', settings.zitsWireframe.toString())
   fd.append('hdStrategy', settings.hdStrategy)
   fd.append('hdStrategyCropMargin', settings.hdStrategyCropMargin.toString())
   fd.append(
diff --git a/lama_cleaner/app/src/components/Settings/ModelSettingBlock.tsx b/lama_cleaner/app/src/components/Settings/ModelSettingBlock.tsx
index 2d7a1f9..02fdfb2 100644
--- a/lama_cleaner/app/src/components/Settings/ModelSettingBlock.tsx
+++ b/lama_cleaner/app/src/components/Settings/ModelSettingBlock.tsx
@@ -2,6 +2,7 @@ import React, { ReactNode } from 'react'
 import { useRecoilState } from 'recoil'
 import { settingState } from '../../store/Atoms'
 import Selector from '../shared/Selector'
+import { Switch, SwitchThumb } from '../shared/Switch'
 import { LDMSampler } from './HDSettingBlock'
 import NumberInputSetting from './NumberInputSetting'
 import SettingBlock from './SettingBlock'
@@ -9,6 +10,7 @@ import SettingBlock from './SettingBlock'
 export enum AIModel {
   LAMA = 'lama',
   LDM = 'ldm',
+  ZITS = 'ZITS',
 }
 
 function ModelSettingBlock() {
@@ -39,7 +41,7 @@ function ModelSettingBlock() {
           target="_blank"
           rel="noreferrer noopener"
         >
-          {name}
+          Paper: {name}
         </a>
 
         <a
@@ -48,7 +50,7 @@ function ModelSettingBlock() {
           target="_blank"
           rel="noreferrer noopener"
         >
-          {githubUrl}
+          Offical Repository: {githubUrl}
         </a>
       </div>
     )
@@ -89,6 +91,34 @@ function ModelSettingBlock() {
     )
   }
 
+  const renderZITSModelDesc = () => {
+    return (
+      <div>
+        {renderModelDesc(
+          'Incremental Transformer Structure Enhanced Image Inpainting with Masking Positional Encoding',
+          'https://arxiv.org/abs/2203.00867',
+          'https://github.com/DQiaole/ZITS_inpainting'
+        )}
+        <SettingBlock
+          className="sub-setting-block"
+          title="Wireframe"
+          input={
+            <Switch
+              checked={setting.zitsWireframe}
+              onCheckedChange={checked => {
+                setSettingState(old => {
+                  return { ...old, zitsWireframe: checked }
+                })
+              }}
+            >
+              <SwitchThumb />
+            </Switch>
+          }
+        />
+      </div>
+    )
+  }
+
   const renderOptionDesc = (): ReactNode => {
     switch (setting.model) {
       case AIModel.LAMA:
@@ -99,6 +129,8 @@ function ModelSettingBlock() {
         )
       case AIModel.LDM:
         return renderLDMModelDesc()
+      case AIModel.ZITS:
+        return renderZITSModelDesc()
       default:
         return <></>
     }
diff --git a/lama_cleaner/app/src/store/Atoms.tsx b/lama_cleaner/app/src/store/Atoms.tsx
index 89ab854..7e8f374 100644
--- a/lama_cleaner/app/src/store/Atoms.tsx
+++ b/lama_cleaner/app/src/store/Atoms.tsx
@@ -46,6 +46,9 @@ export interface Settings {
   // For LDM
   ldmSteps: number
   ldmSampler: LDMSampler
+
+  // For ZITS
+  zitsWireframe: boolean
 }
 
 export const settingStateDefault = {
@@ -54,8 +57,12 @@ export const settingStateDefault = {
   graduallyInpainting: true,
   runInpaintingManually: false,
   model: AIModel.LAMA,
+
   ldmSteps: 50,
   ldmSampler: LDMSampler.plms,
+
+  zitsWireframe: true,
+
   hdStrategy: HDStrategy.RESIZE,
   hdStrategyResizeLimit: 2048,
   hdStrategyCropTrigerSize: 2048,
diff --git a/lama_cleaner/server.py b/lama_cleaner/server.py
index 79b8a0b..769692f 100644
--- a/lama_cleaner/server.py
+++ b/lama_cleaner/server.py
@@ -95,6 +95,7 @@ def process():
         ldm_steps=form["ldmSteps"],
         ldm_sampler=form["ldmSampler"],
         hd_strategy=form["hdStrategy"],
+        zits_wireframe=form["zitsWireframe"],
         hd_strategy_crop_margin=form["hdStrategyCropMargin"],
         hd_strategy_crop_trigger_size=form["hdStrategyCropTrigerSize"],
         hd_strategy_resize_limit=form["hdStrategyResizeLimit"],