wip

2023-12-01 10:15:35 +08:00
parent 973987dfbb
commit 9a9eb8abfd
55 changed files with 2596 additions and 1251 deletions
--- a/web_app/src/lib/api.ts
+++ b/web_app/src/lib/api.ts
@@ -1,18 +1,20 @@
-import { PluginName } from "@/lib/types"
-import { ControlNetMethodMap, Rect, Settings } from "@/lib/store"
-import { dataURItoBlob, loadImage, srcToFile } from "@/lib/utils"
+import { ModelInfo, Rect } from "@/lib/types"
+import { Settings } from "@/lib/states"
+import { dataURItoBlob, srcToFile } from "@/lib/utils"
+import axios from "axios"

 export const API_ENDPOINT = import.meta.env.VITE_BACKEND
  ? import.meta.env.VITE_BACKEND
  : ""

+const api = axios.create({
+  baseURL: API_ENDPOINT,
+})
+
 export default async function inpaint(
  imageFile: File,
  settings: Settings,
  croperRect: Rect,
-  prompt?: string,
-  negativePrompt?: string,
-  seed?: number,
  maskBase64?: string,
  customMask?: File,
  paintByExampleImage?: File
@@ -26,38 +28,29 @@ export default async function inpaint(
    fd.append("mask", customMask)
  }

-  const hdSettings = settings.hdSettings[settings.model]
  fd.append("ldmSteps", settings.ldmSteps.toString())
  fd.append("ldmSampler", settings.ldmSampler.toString())
  fd.append("zitsWireframe", settings.zitsWireframe.toString())
-  fd.append("hdStrategy", hdSettings.hdStrategy)
-  fd.append("hdStrategyCropMargin", hdSettings.hdStrategyCropMargin.toString())
-  fd.append(
-    "hdStrategyCropTrigerSize",
-    hdSettings.hdStrategyCropTrigerSize.toString()
-  )
-  fd.append(
-    "hdStrategyResizeLimit",
-    hdSettings.hdStrategyResizeLimit.toString()
-  )
+  fd.append("hdStrategy", "Crop")
+  fd.append("hdStrategyCropMargin", "128")
+  fd.append("hdStrategyCropTrigerSize", "640")
+  fd.append("hdStrategyResizeLimit", "2048")

-  fd.append("prompt", prompt === undefined ? "" : prompt)
-  fd.append(
-    "negativePrompt",
-    negativePrompt === undefined ? "" : negativePrompt
-  )
+  fd.append("prompt", settings.prompt)
+  fd.append("negativePrompt", settings.negativePrompt)
  fd.append("croperX", croperRect.x.toString())
  fd.append("croperY", croperRect.y.toString())
  fd.append("croperHeight", croperRect.height.toString())
  fd.append("croperWidth", croperRect.width.toString())
-  fd.append("useCroper", settings.showCroper ? "true" : "false")
+  // fd.append("useCroper", settings.showCroper ? "true" : "false")
+  fd.append("useCroper", "false")

  fd.append("sdMaskBlur", settings.sdMaskBlur.toString())
  fd.append("sdStrength", settings.sdStrength.toString())
  fd.append("sdSteps", settings.sdSteps.toString())
  fd.append("sdGuidanceScale", settings.sdGuidanceScale.toString())
  fd.append("sdSampler", settings.sdSampler.toString())
-  fd.append("sdSeed", seed ? seed.toString() : "-1")
+  fd.append("sdSeed", settings.seed.toString())
  fd.append("sdMatchHistograms", settings.sdMatchHistograms ? "true" : "false")
  fd.append("sdScale", (settings.sdScale / 100).toString())

@@ -69,7 +62,7 @@ export default async function inpaint(
    "paintByExampleGuidanceScale",
    settings.paintByExampleGuidanceScale.toString()
  )
-  fd.append("paintByExampleSeed", seed ? seed.toString() : "-1")
+  fd.append("paintByExampleSeed", settings.seed.toString())
  fd.append(
    "paintByExampleMaskBlur",
    settings.paintByExampleMaskBlur.toString()
@@ -94,10 +87,7 @@ export default async function inpaint(
    "controlnet_conditioning_scale",
    settings.controlnetConditioningScale.toString()
  )
-  fd.append(
-    "controlnet_method",
-    ControlNetMethodMap[settings.controlnetMethod.toString()]
-  )
+  fd.append("controlnet_method", settings.controlnetMethod.toString())

  try {
    const res = await fetch(`${API_ENDPOINT}/inpaint`, {
@@ -137,6 +127,10 @@ export function currentModel() {
  })
 }

+export function fetchModelInfos(): Promise<ModelInfo[]> {
+  return api.get("/models").then((response) => response.data)
+}
+
 export function isDesktop() {
  return fetch(`${API_ENDPOINT}/is_desktop`, {
    method: "GET",
--- a/web_app/src/lib/states.ts
+++ b/web_app/src/lib/states.ts
@@ -1,14 +1,17 @@
-import { create, StoreApi, UseBoundStore } from "zustand"
+import { create } from "zustand"
 import { persist } from "zustand/middleware"
 import { immer } from "zustand/middleware/immer"
-import { SortBy, SortOrder } from "./types"
+import { CV2Flag, LDMSampler, ModelInfo, SortBy, SortOrder } from "./types"
 import { DEFAULT_BRUSH_SIZE } from "./const"
+import { SDSampler } from "./store"

 type FileManagerState = {
  sortBy: SortBy
  sortOrder: SortOrder
  layout: "rows" | "masonry"
  searchText: string
+  inputDirectory: string
+  outputDirectory: string
 }

 type CropperState = {
@@ -18,81 +21,256 @@ type CropperState = {
  height: number
 }

+export type Settings = {
+  model: ModelInfo
+  enableDownloadMask: boolean
+  enableManualInpainting: boolean
+  enableUploadMask: boolean
+  showCroper: boolean
+
+  // For LDM
+  ldmSteps: number
+  ldmSampler: LDMSampler
+
+  // For ZITS
+  zitsWireframe: boolean
+
+  // For OpenCV2
+  cv2Radius: number
+  cv2Flag: CV2Flag
+
+  // For Diffusion moel
+  prompt: string
+  negativePrompt: string
+  seed: number
+  seedFixed: boolean
+
+  // For SD
+  sdMaskBlur: number
+  sdStrength: number
+  sdSteps: number
+  sdGuidanceScale: number
+  sdSampler: SDSampler
+  sdMatchHistograms: boolean
+  sdScale: number
+
+  // Paint by Example
+  paintByExampleSteps: number
+  paintByExampleGuidanceScale: number
+  paintByExampleMaskBlur: number
+  paintByExampleMatchHistograms: boolean
+
+  // InstructPix2Pix
+  p2pSteps: number
+  p2pImageGuidanceScale: number
+  p2pGuidanceScale: number
+
+  // ControlNet
+  controlnetConditioningScale: number
+  controlnetMethod: string
+}
+
+type ServerConfig = {
+  plugins: string[]
+  availableControlNet: Record<string, string[]>
+  enableFileManager: boolean
+  enableAutoSaving: boolean
+}
+
+type InteractiveSegState = {
+  isInteractiveSeg: boolean
+  isInteractiveSegRunning: boolean
+  clicks: number[][]
+}
+
 type AppState = {
  file: File | null
+  customMask: File | null
  imageHeight: number
  imageWidth: number
  brushSize: number
  brushSizeScale: number
-
  isInpainting: boolean
-  isInteractiveSeg: boolean // 是否正处于 sam 状态
-  isInteractiveSegRunning: boolean
-  interactiveSegClicks: number[][]
-
-  prompt: string
+  isPluginRunning: boolean

+  interactiveSegState: InteractiveSegState
  fileManagerState: FileManagerState
  cropperState: CropperState
+  serverConfig: ServerConfig
+
+  settings: Settings
 }

 type AppAction = {
  setFile: (file: File) => void
+  setCustomFile: (file: File) => void
  setIsInpainting: (newValue: boolean) => void
+  setIsPluginRunning: (newValue: boolean) => void
  setBrushSize: (newValue: number) => void
  setImageSize: (width: number, height: number) => void
-  setPrompt: (newValue: string) => void
-
-  setFileManagerSortBy: (newValue: SortBy) => void
-  setFileManagerSortOrder: (newValue: SortOrder) => void
-  setFileManagerLayout: (
-    newValue: AppState["fileManagerState"]["layout"]
-  ) => void
-  setFileManagerSearchText: (newValue: string) => void

  setCropperX: (newValue: number) => void
  setCropperY: (newValue: number) => void
  setCropperWidth: (newValue: number) => void
  setCropperHeight: (newValue: number) => void
+
+  setServerConfig: (newValue: ServerConfig) => void
+  setSeed: (newValue: number) => void
+  updateSettings: (newSettings: Partial<Settings>) => void
+  updateFileManagerState: (newState: Partial<FileManagerState>) => void
+  updateInteractiveSegState: (newState: Partial<InteractiveSegState>) => void
+  resetInteractiveSegState: () => void
+  shouldShowPromptInput: () => boolean
+}
+
+const defaultValues: AppState = {
+  file: null,
+  customMask: null,
+  imageHeight: 0,
+  imageWidth: 0,
+  brushSize: DEFAULT_BRUSH_SIZE,
+  brushSizeScale: 1,
+  isInpainting: false,
+  isPluginRunning: false,
+
+  interactiveSegState: {
+    isInteractiveSeg: false,
+    isInteractiveSegRunning: false,
+    clicks: [],
+  },
+
+  cropperState: {
+    x: 0,
+    y: 0,
+    width: 512,
+    height: 512,
+  },
+  fileManagerState: {
+    sortBy: SortBy.CTIME,
+    sortOrder: SortOrder.DESCENDING,
+    layout: "masonry",
+    searchText: "",
+    inputDirectory: "",
+    outputDirectory: "",
+  },
+  serverConfig: {
+    plugins: [],
+    availableControlNet: { SD: [], SD2: [], SDXL: [] },
+    enableFileManager: false,
+    enableAutoSaving: false,
+  },
+  settings: {
+    model: {
+      name: "lama",
+      path: "lama",
+      model_type: "inpaint",
+      support_controlnet: false,
+      support_freeu: false,
+      support_lcm_lora: false,
+      is_single_file_diffusers: false,
+    },
+    showCroper: false,
+    enableDownloadMask: false,
+    enableManualInpainting: false,
+    enableUploadMask: false,
+    ldmSteps: 30,
+    ldmSampler: LDMSampler.ddim,
+    zitsWireframe: true,
+    cv2Radius: 5,
+    cv2Flag: CV2Flag.INPAINT_NS,
+    prompt: "",
+    negativePrompt: "",
+    seed: 42,
+    seedFixed: false,
+    sdMaskBlur: 5,
+    sdStrength: 1.0,
+    sdSteps: 50,
+    sdGuidanceScale: 7.5,
+    sdSampler: SDSampler.uni_pc,
+    sdMatchHistograms: false,
+    sdScale: 100,
+    paintByExampleSteps: 50,
+    paintByExampleGuidanceScale: 7.5,
+    paintByExampleMaskBlur: 5,
+    paintByExampleMatchHistograms: false,
+    p2pSteps: 50,
+    p2pImageGuidanceScale: 1.5,
+    p2pGuidanceScale: 7.5,
+    controlnetConditioningScale: 0.4,
+    controlnetMethod: "lllyasviel/control_v11p_sd15_canny",
+  },
 }

 export const useStore = create<AppState & AppAction>()(
  immer(
    persist(
      (set, get) => ({
-        file: null,
-        imageHeight: 0,
-        imageWidth: 0,
-        brushSize: DEFAULT_BRUSH_SIZE,
-        brushSizeScale: 1,
-        isInpainting: false,
-        isInteractiveSeg: false,
-        isInteractiveSegRunning: false,
-        interactiveSegClicks: [],
-        prompt: "",
-        cropperState: {
-          x: 0,
-          y: 0,
-          width: 0,
-          height: 0,
+        ...defaultValues,
+
+        shouldShowPromptInput: (): boolean => {
+          const model_type = get().settings.model.model_type
+          return ["diffusers_sd"].includes(model_type)
        },
-        fileManagerState: {
-          sortBy: SortBy.CTIME,
-          sortOrder: SortOrder.DESCENDING,
-          layout: "masonry",
-          searchText: "",
+
+        setServerConfig: (newValue: ServerConfig) => {
+          set((state: AppState) => {
+            state.serverConfig = newValue
+          })
        },
+
+        updateSettings: (newSettings: Partial<Settings>) => {
+          set((state: AppState) => {
+            state.settings = {
+              ...state.settings,
+              ...newSettings,
+            }
+          })
+        },
+
+        updateFileManagerState: (newState: Partial<FileManagerState>) => {
+          set((state: AppState) => {
+            state.fileManagerState = {
+              ...state.fileManagerState,
+              ...newState,
+            }
+          })
+        },
+
+        updateInteractiveSegState: (newState: Partial<InteractiveSegState>) => {
+          set((state: AppState) => {
+            state.interactiveSegState = {
+              ...state.interactiveSegState,
+              ...newState,
+            }
+          })
+        },
+        resetInteractiveSegState: () => {
+          set((state: AppState) => {
+            state.interactiveSegState = defaultValues.interactiveSegState
+          })
+        },
+
        setIsInpainting: (newValue: boolean) =>
          set((state: AppState) => {
            state.isInpainting = newValue
          }),

+        setIsPluginRunning: (newValue: boolean) =>
+          set((state: AppState) => {
+            state.isPluginRunning = newValue
+          }),
+
        setFile: (file: File) =>
          set((state: AppState) => {
            // TODO: 清空各种状态
            state.file = file
          }),

+        setCustomFile: (file: File) =>
+          set((state: AppState) => {
+            state.customMask = file
+          }),
+
        setBrushSize: (newValue: number) =>
          set((state: AppState) => {
            state.brushSize = newValue
@@ -107,11 +285,6 @@ export const useStore = create<AppState & AppAction>()(
          })
        },

-        setPrompt: (newValue: string) =>
-          set((state: AppState) => {
-            state.prompt = newValue
-          }),
-
        setCropperX: (newValue: number) =>
          set((state: AppState) => {
            state.cropperState.x = newValue
@@ -132,32 +305,18 @@ export const useStore = create<AppState & AppAction>()(
            state.cropperState.height = newValue
          }),

-        setFileManagerSortBy: (newValue: SortBy) =>
+        setSeed: (newValue: number) =>
          set((state: AppState) => {
-            state.fileManagerState.sortBy = newValue
-          }),
-
-        setFileManagerSortOrder: (newValue: SortOrder) =>
-          set((state: AppState) => {
-            state.fileManagerState.sortOrder = newValue
-          }),
-
-        setFileManagerLayout: (newValue: "rows" | "masonry") =>
-          set((state: AppState) => {
-            state.fileManagerState.layout = newValue
-          }),
-
-        setFileManagerSearchText: (newValue: string) =>
-          set((state: AppState) => {
-            state.fileManagerState.searchText = newValue
+            state.settings.seed = newValue
          }),
      }),
      {
        name: "ZUSTAND_STATE", // name of the item in the storage (must be unique)
+        version: 0,
        partialize: (state) =>
          Object.fromEntries(
            Object.entries(state).filter(([key]) =>
-              ["fileManagerState", "prompt"].includes(key)
+              ["fileManagerState", "prompt", "settings"].includes(key)
            )
          ),
      }
--- a/web_app/src/lib/store.ts
+++ b/web_app/src/lib/store.ts
@@ -1,6 +1,6 @@
 import { atom, selector } from "recoil"
 import _ from "lodash"
-import { CV2Flag, HDStrategy, LDMSampler, ModelsHDSettings } from "./types"
+import { CV2Flag, LDMSampler } from "./types"

 export enum AIModel {
  LAMA = "lama",
@@ -320,7 +320,6 @@ export interface Settings {
  graduallyInpainting: boolean
  runInpaintingManually: boolean
  model: AIModel
-  hdSettings: ModelsHDSettings

  // For LDM
  ldmSteps: number
@@ -363,107 +362,6 @@ export interface Settings {
  controlnetMethod: string
 }

-const defaultHDSettings: ModelsHDSettings = {
-  [AIModel.LAMA]: {
-    hdStrategy: HDStrategy.CROP,
-    hdStrategyResizeLimit: 2048,
-    hdStrategyCropTrigerSize: 800,
-    hdStrategyCropMargin: 196,
-    enabled: true,
-  },
-  [AIModel.LDM]: {
-    hdStrategy: HDStrategy.CROP,
-    hdStrategyResizeLimit: 1080,
-    hdStrategyCropTrigerSize: 1080,
-    hdStrategyCropMargin: 128,
-    enabled: true,
-  },
-  [AIModel.ZITS]: {
-    hdStrategy: HDStrategy.CROP,
-    hdStrategyResizeLimit: 1024,
-    hdStrategyCropTrigerSize: 1024,
-    hdStrategyCropMargin: 128,
-    enabled: true,
-  },
-  [AIModel.MAT]: {
-    hdStrategy: HDStrategy.CROP,
-    hdStrategyResizeLimit: 1024,
-    hdStrategyCropTrigerSize: 512,
-    hdStrategyCropMargin: 128,
-    enabled: true,
-  },
-  [AIModel.FCF]: {
-    hdStrategy: HDStrategy.CROP,
-    hdStrategyResizeLimit: 512,
-    hdStrategyCropTrigerSize: 512,
-    hdStrategyCropMargin: 128,
-    enabled: false,
-  },
-  [AIModel.SD15]: {
-    hdStrategy: HDStrategy.ORIGINAL,
-    hdStrategyResizeLimit: 768,
-    hdStrategyCropTrigerSize: 512,
-    hdStrategyCropMargin: 128,
-    enabled: false,
-  },
-  [AIModel.ANYTHING4]: {
-    hdStrategy: HDStrategy.ORIGINAL,
-    hdStrategyResizeLimit: 768,
-    hdStrategyCropTrigerSize: 512,
-    hdStrategyCropMargin: 128,
-    enabled: false,
-  },
-  [AIModel.REALISTIC_VISION_1_4]: {
-    hdStrategy: HDStrategy.ORIGINAL,
-    hdStrategyResizeLimit: 768,
-    hdStrategyCropTrigerSize: 512,
-    hdStrategyCropMargin: 128,
-    enabled: false,
-  },
-  [AIModel.SD2]: {
-    hdStrategy: HDStrategy.ORIGINAL,
-    hdStrategyResizeLimit: 768,
-    hdStrategyCropTrigerSize: 512,
-    hdStrategyCropMargin: 128,
-    enabled: false,
-  },
-  [AIModel.PAINT_BY_EXAMPLE]: {
-    hdStrategy: HDStrategy.ORIGINAL,
-    hdStrategyResizeLimit: 768,
-    hdStrategyCropTrigerSize: 512,
-    hdStrategyCropMargin: 128,
-    enabled: false,
-  },
-  [AIModel.PIX2PIX]: {
-    hdStrategy: HDStrategy.ORIGINAL,
-    hdStrategyResizeLimit: 768,
-    hdStrategyCropTrigerSize: 512,
-    hdStrategyCropMargin: 128,
-    enabled: false,
-  },
-  [AIModel.Mange]: {
-    hdStrategy: HDStrategy.CROP,
-    hdStrategyResizeLimit: 1280,
-    hdStrategyCropTrigerSize: 1024,
-    hdStrategyCropMargin: 196,
-    enabled: true,
-  },
-  [AIModel.CV2]: {
-    hdStrategy: HDStrategy.RESIZE,
-    hdStrategyResizeLimit: 1080,
-    hdStrategyCropTrigerSize: 512,
-    hdStrategyCropMargin: 128,
-    enabled: true,
-  },
-  [AIModel.KANDINSKY22]: {
-    hdStrategy: HDStrategy.ORIGINAL,
-    hdStrategyResizeLimit: 768,
-    hdStrategyCropTrigerSize: 512,
-    hdStrategyCropMargin: 128,
-    enabled: false,
-  },
-}
-
 export enum SDSampler {
  ddim = "ddim",
  pndm = "pndm",
@@ -487,7 +385,6 @@ export const settingStateDefault: Settings = {
  graduallyInpainting: true,
  runInpaintingManually: false,
  model: AIModel.LAMA,
-  hdSettings: defaultHDSettings,

  ldmSteps: 25,
  ldmSampler: LDMSampler.plms,
@@ -588,24 +485,6 @@ export const seedState = selector({
  },
 })

-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 },
-    })
-  },
-})
-
 export const isSDState = selector({
  key: "isSD",
  get: ({ get }) => {
--- a/web_app/src/lib/types.ts
+++ b/web_app/src/lib/types.ts
@@ -1,3 +1,19 @@
+export interface ModelInfo {
+  name: string
+  path: string
+  model_type:
+    | "inpaint"
+    | "diffusers_sd"
+    | "diffusers_sdxl"
+    | "diffusers_sd_inpaint"
+    | "diffusers_sdxl_inpaint"
+    | "diffusers_other"
+  support_controlnet: boolean
+  support_freeu: boolean
+  support_lcm_lora: boolean
+  is_single_file_diffusers: boolean
+}
+
 export enum PluginName {
  RemoveBG = "RemoveBG",
  AnimeSeg = "AnimeSeg",
@@ -18,12 +34,6 @@ export enum SortOrder {
  ASCENDING = "asc",
 }

-export enum HDStrategy {
-  ORIGINAL = "Original",
-  RESIZE = "Resize",
-  CROP = "Crop",
-}
-
 export enum LDMSampler {
  ddim = "ddim",
  plms = "plms",
@@ -34,12 +44,9 @@ export enum CV2Flag {
  INPAINT_TELEA = "INPAINT_TELEA",
 }

-export interface HDSettings {
-  hdStrategy: HDStrategy
-  hdStrategyResizeLimit: number
-  hdStrategyCropTrigerSize: number
-  hdStrategyCropMargin: number
-  enabled: boolean
+export interface Rect {
+  x: number
+  y: number
+  width: number
+  height: number
 }
-
-export type ModelsHDSettings = { [key in AIModel]: HDSettings }