each model has its own hdsettings

lama_cleaner/app/src/store/Atoms.tsx

@@ -1,8 +1,14 @@
-import { atom } from 'recoil'
+import { atom, selector } from 'recoil'
+import _ from 'lodash'
 import { HDStrategy, LDMSampler } from '../components/Settings/HDSettingBlock'
-import { AIModel } from '../components/Settings/ModelSettingBlock'
 import { ToastState } from '../components/shared/Toast'
 
+export enum AIModel {
+  LAMA = 'lama',
+  LDM = 'ldm',
+  ZITS = 'zits',
+}
+
 export const fileState = atom<File | undefined>({
   key: 'fileState',
   default: undefined,
@@ -30,18 +36,22 @@ export const shortcutsState = atom<boolean>({
   default: false,
 })
 
+export interface HDSettings {
+  hdStrategy: HDStrategy
+  hdStrategyResizeLimit: number
+  hdStrategyCropTrigerSize: number
+  hdStrategyCropMargin: number
+}
+
+type ModelsHDSettings = { [key in AIModel]: HDSettings }
+
 export interface Settings {
   show: boolean
   downloadMask: boolean
   graduallyInpainting: boolean
   runInpaintingManually: boolean
   model: AIModel
-
-  // For LaMa
-  hdStrategy: HDStrategy
-  hdStrategyResizeLimit: number
-  hdStrategyCropTrigerSize: number
-  hdStrategyCropMargin: number
+  hdSettings: ModelsHDSettings
 
   // For LDM
   ldmSteps: number
@@ -51,22 +61,39 @@ export interface Settings {
   zitsWireframe: boolean
 }
 
-export const settingStateDefault = {
+const defaultHDSettings: ModelsHDSettings = {
+  [AIModel.LAMA]: {
+    hdStrategy: HDStrategy.RESIZE,
+    hdStrategyResizeLimit: 2048,
+    hdStrategyCropTrigerSize: 2048,
+    hdStrategyCropMargin: 128,
+  },
+  [AIModel.LDM]: {
+    hdStrategy: HDStrategy.CROP,
+    hdStrategyResizeLimit: 1080,
+    hdStrategyCropTrigerSize: 1080,
+    hdStrategyCropMargin: 128,
+  },
+  [AIModel.ZITS]: {
+    hdStrategy: HDStrategy.CROP,
+    hdStrategyResizeLimit: 1024,
+    hdStrategyCropTrigerSize: 1024,
+    hdStrategyCropMargin: 128,
+  },
+}
+
+export const settingStateDefault: Settings = {
   show: false,
   downloadMask: false,
   graduallyInpainting: true,
   runInpaintingManually: false,
   model: AIModel.LAMA,
+  hdSettings: defaultHDSettings,
 
-  ldmSteps: 50,
+  ldmSteps: 25,
   ldmSampler: LDMSampler.plms,
 
   zitsWireframe: true,
-
-  hdStrategy: HDStrategy.RESIZE,
-  hdStrategyResizeLimit: 2048,
-  hdStrategyCropTrigerSize: 2048,
-  hdStrategyCropMargin: 128,
 }
 
 const localStorageEffect =
@@ -76,20 +103,44 @@ const localStorageEffect =
     if (savedValue != null) {
       const storageSettings = JSON.parse(savedValue)
       storageSettings.show = false
-      setSelf({ ...settingStateDefault, ...storageSettings })
+
+      const restored = _.merge(
+        _.cloneDeep(settingStateDefault),
+        storageSettings
+      )
+      setSelf(restored)
     }
 
-    onSet((newValue: Settings, _: string, isReset: boolean) =>
+    onSet((newValue: Settings, val: string, isReset: boolean) =>
       isReset
         ? localStorage.removeItem(key)
         : localStorage.setItem(key, JSON.stringify(newValue))
     )
   }
 
+const ROOT_STATE_KEY = 'settingsState2'
 // Each atom can reference an array of these atom effect functions which are called in priority order when the atom is initialized
 // https://recoiljs.org/docs/guides/atom-effects/#local-storage-persistence
 export const settingState = atom<Settings>({
-  key: 'settingsState',
+  key: ROOT_STATE_KEY,
   default: settingStateDefault,
-  effects: [localStorageEffect('settingsState')],
+  effects: [localStorageEffect(ROOT_STATE_KEY)],
+})
+
+export const hdSettingsState = selector({
+  key: 'hdSettings',
+  get: ({ get }) => {
+    const settings = get(settingState)
+    return settings.hdSettings[settings.model]
+  },
+  set: ({ get, set }, newValue: any) => {
+    const settings = get(settingState)
+    const hdSettings = settings.hdSettings[settings.model]
+    const newHDSettings = { ...hdSettings, ...newValue }
+
+    set(settingState, {
+      ...settings,
+      hdSettings: { ...settings.hdSettings, [settings.model]: newHDSettings },
+    })
+  },
 })