From f7d7e89197be7868514cc0b2c0fc6f2b7816e3c3 Mon Sep 17 00:00:00 2001
From: Qing <cwq1913@gmail.com>
Date: Fri, 25 Nov 2022 08:55:10 +0800
Subject: [PATCH] fix is_disable_switch_model

---
 lama_cleaner/app/src/App.tsx                  | 20 ++++++++++++----
 .../components/Settings/ModelSettingBlock.tsx | 23 ++++++++-----------
 lama_cleaner/app/src/store/Atoms.tsx          | 14 +++++++++++
 lama_cleaner/server.py                        |  3 +++
 4 files changed, 42 insertions(+), 18 deletions(-)

diff --git a/lama_cleaner/app/src/App.tsx b/lama_cleaner/app/src/App.tsx
index 963f54e..a4c665d 100644
--- a/lama_cleaner/app/src/App.tsx
+++ b/lama_cleaner/app/src/App.tsx
@@ -5,11 +5,11 @@ import useInputImage from './hooks/useInputImage'
 import LandingPage from './components/LandingPage/LandingPage'
 import { themeState } from './components/Header/ThemeChanger'
 import Workspace from './components/Workspace'
-import { fileState, toastState } from './store/Atoms'
+import { fileState, isDisableModelSwitchState, toastState } from './store/Atoms'
 import { keepGUIAlive } from './utils'
 import Header from './components/Header/Header'
 import useHotKey from './hooks/useHotkey'
-import { isDesktop } from './adapters/inpainting'
+import { getIsDisableModelSwitch, isDesktop } from './adapters/inpainting'
 
 const SUPPORTED_FILE_TYPE = [
   'image/jpeg',
@@ -24,8 +24,9 @@ function App() {
   const [theme, setTheme] = useRecoilState(themeState)
   const [toastVal, setToastState] = useRecoilState(toastState)
   const userInputImage = useInputImage()
-  const [openPasteImageAlertDialog, setOpenPasteImageAlertDialog] =
-    useState(false)
+  const [isDisableModelSwitch, setIsDisableModelSwitch] = useRecoilState(
+    isDisableModelSwitchState
+  )
 
   // Set Input Image
   useEffect(() => {
@@ -43,6 +44,17 @@ function App() {
     fetchData()
   }, [])
 
+  useEffect(() => {
+    const fetchData = async () => {
+      const isDisable: string = await getIsDisableModelSwitch().then(res =>
+        res.text()
+      )
+      setIsDisableModelSwitch(isDisable === 'true')
+    }
+
+    fetchData()
+  }, [])
+
   // Dark Mode Hotkey
   useHotKey(
     'shift+d',
diff --git a/lama_cleaner/app/src/components/Settings/ModelSettingBlock.tsx b/lama_cleaner/app/src/components/Settings/ModelSettingBlock.tsx
index d8f7057..e6291a7 100644
--- a/lama_cleaner/app/src/components/Settings/ModelSettingBlock.tsx
+++ b/lama_cleaner/app/src/components/Settings/ModelSettingBlock.tsx
@@ -1,7 +1,13 @@
 import React, { ReactNode, useEffect, useState } from 'react'
-import { useRecoilState } from 'recoil'
+import { useRecoilState, useRecoilValue } from 'recoil'
 import { getIsDisableModelSwitch } from '../../adapters/inpainting'
-import { AIModel, CV2Flag, SDSampler, settingState } from '../../store/Atoms'
+import {
+  AIModel,
+  CV2Flag,
+  isDisableModelSwitchState,
+  SDSampler,
+  settingState,
+} from '../../store/Atoms'
 import Selector from '../shared/Selector'
 import { Switch, SwitchThumb } from '../shared/Switch'
 import Tooltip from '../shared/Tooltip'
@@ -11,18 +17,7 @@ import SettingBlock from './SettingBlock'
 
 function ModelSettingBlock() {
   const [setting, setSettingState] = useRecoilState(settingState)
-  const [isDisableModelSwitch, setIsDisableModelSwitch] = useState(false)
-
-  useEffect(() => {
-    const fetchData = async () => {
-      const isDisable: string = await getIsDisableModelSwitch().then(res =>
-        res.text()
-      )
-      setIsDisableModelSwitch(isDisable === 'true')
-    }
-
-    fetchData()
-  }, [])
+  const isDisableModelSwitch = useRecoilValue(isDisableModelSwitchState)
 
   const onModelChange = (value: AIModel) => {
     setSettingState(old => {
diff --git a/lama_cleaner/app/src/store/Atoms.tsx b/lama_cleaner/app/src/store/Atoms.tsx
index 54d1d78..c143a59 100644
--- a/lama_cleaner/app/src/store/Atoms.tsx
+++ b/lama_cleaner/app/src/store/Atoms.tsx
@@ -34,6 +34,7 @@ export interface Rect {
 interface AppState {
   disableShortCuts: boolean
   isInpainting: boolean
+  isDisableModelSwitch: boolean
 }
 
 export const appState = atom<AppState>({
@@ -41,6 +42,7 @@ export const appState = atom<AppState>({
   default: {
     disableShortCuts: false,
     isInpainting: false,
+    isDisableModelSwitch: false,
   },
 })
 
@@ -66,6 +68,18 @@ export const isInpaintingState = selector({
   },
 })
 
+export const isDisableModelSwitchState = selector({
+  key: 'isDisableModelSwitchState',
+  get: ({ get }) => {
+    const app = get(appState)
+    return app.isDisableModelSwitch
+  },
+  set: ({ get, set }, newValue: any) => {
+    const app = get(appState)
+    set(appState, { ...app, isDisableModelSwitch: newValue })
+  },
+})
+
 export const croperState = atom<Rect>({
   key: 'croperState',
   default: {
diff --git a/lama_cleaner/server.py b/lama_cleaner/server.py
index ef3a8b0..ffa2525 100644
--- a/lama_cleaner/server.py
+++ b/lama_cleaner/server.py
@@ -204,6 +204,9 @@ def get_is_desktop():
 
 @app.route("/model", methods=["POST"])
 def switch_model():
+    if is_disable_model_switch:
+        return "Switch model is disabled", 400
+
     new_name = request.form.get("name")
     if new_name == model.name:
         return "Same model", 200