make plugin work

2023-03-25 09:53:22 +08:00
parent 996a264797
commit 6e54f77ed6
16 changed files with 528 additions and 284 deletions
--- a/lama_cleaner/app/src/components/Editor/Editor.tsx
+++ b/lama_cleaner/app/src/components/Editor/Editor.tsx
@@ -18,10 +18,7 @@ import {
 } from 'react-zoom-pan-pinch'
 import { useRecoilState, useRecoilValue, useSetRecoilState } from 'recoil'
 import { useWindowSize, useKey, useKeyPressEvent } from 'react-use'
-import inpaint, {
-  downloadToOutput,
-  postInteractiveSeg,
-} from '../../adapters/inpainting'
+import inpaint, { downloadToOutput, runPlugin } from '../../adapters/inpainting'
 import Button from '../shared/Button'
 import Slider from './Slider'
 import SizeSelector from './SizeSelector'
@@ -50,6 +47,8 @@ import {
  isInteractiveSegRunningState,
  isInteractiveSegState,
  isPix2PixState,
+  isPluginRunningState,
+  isProcessingState,
  negativePropmtState,
  propmtState,
  runManuallyState,
@@ -69,6 +68,7 @@ import FileSelect from '../FileSelect/FileSelect'
 import InteractiveSeg from '../InteractiveSeg/InteractiveSeg'
 import InteractiveSegConfirmActions from '../InteractiveSeg/ConfirmActions'
 import InteractiveSegReplaceModal from '../InteractiveSeg/ReplaceModal'
+import { PluginName } from '../Plugins/Plugins'
 import MakeGIF from './MakeGIF'

 const TOOLBAR_SIZE = 200
@@ -118,13 +118,15 @@ export default function Editor() {
  const croperRect = useRecoilValue(croperState)
  const setToastState = useSetRecoilState(toastState)
  const [isInpainting, setIsInpainting] = useRecoilState(isInpaintingState)
+  const setIsPluginRunning = useSetRecoilState(isPluginRunningState)
+  const isProcessing = useRecoilValue(isProcessingState)
  const runMannually = useRecoilValue(runManuallyState)
  const isDiffusionModels = useRecoilValue(isDiffusionModelsState)
  const isPix2Pix = useRecoilValue(isPix2PixState)
  const [isInteractiveSeg, setIsInteractiveSeg] = useRecoilState(
    isInteractiveSegState
  )
-  const [isInteractiveSegRunning, setIsInteractiveSegRunning] = useRecoilState(
+  const setIsInteractiveSegRunning = useSetRecoilState(
    isInteractiveSegRunningState
  )

@@ -538,6 +540,77 @@ export default function Editor() {
    }
  }, [runInpainting])

+  const getCurrentRender = useCallback(async () => {
+    let targetFile = file
+    if (renders.length > 0) {
+      const lastRender = renders[renders.length - 1]
+      targetFile = await srcToFile(lastRender.currentSrc, file.name, file.type)
+    }
+    return targetFile
+  }, [file, renders])
+
+  useEffect(() => {
+    emitter.on(PluginName.InteractiveSeg, () => {
+      setIsInteractiveSeg(true)
+      if (interactiveSegMask !== null) {
+        setShowInteractiveSegModal(true)
+      }
+    })
+    return () => {
+      emitter.off(PluginName.InteractiveSeg)
+    }
+  })
+
+  const runRenderablePlugin = useCallback(
+    async (name: string) => {
+      if (isProcessing) {
+        return
+      }
+      try {
+        // TODO 要不要加 undoCurrentLine？？
+        setIsPluginRunning(true)
+        const targetFile = await getCurrentRender()
+        const res = await runPlugin(name, targetFile)
+        if (!res) {
+          throw new Error('Something went wrong on server side.')
+        }
+        const { blob } = res
+        const newRender = new Image()
+        await loadImage(newRender, blob)
+        const newRenders = [...renders, newRender]
+        setRenders(newRenders)
+      } catch (e: any) {
+        setToastState({
+          open: true,
+          desc: e.message ? e.message : e.toString(),
+          state: 'error',
+          duration: 3000,
+        })
+      } finally {
+        setIsPluginRunning(false)
+      }
+    },
+    [renders, setRenders, getCurrentRender, setIsPluginRunning, isProcessing]
+  )
+
+  useEffect(() => {
+    emitter.on(PluginName.RemoveBG, () => {
+      runRenderablePlugin(PluginName.RemoveBG)
+    })
+    return () => {
+      emitter.off(PluginName.RemoveBG)
+    }
+  }, [runRenderablePlugin])
+
+  useEffect(() => {
+    emitter.on(PluginName.RealESRGAN, () => {
+      runRenderablePlugin(PluginName.RealESRGAN)
+    })
+    return () => {
+      emitter.off(PluginName.RealESRGAN)
+    }
+  }, [runRenderablePlugin])
+
  const hadRunInpainting = () => {
    return renders.length !== 0
  }
@@ -759,13 +832,7 @@ export default function Editor() {
    }

    setIsInteractiveSegRunning(true)
-
-    let targetFile = file
-    if (renders.length > 0) {
-      const lastRender = renders[renders.length - 1]
-      targetFile = await srcToFile(lastRender.currentSrc, file.name, file.type)
-    }
-
+    const targetFile = await getCurrentRender()
    const prevMask = null
    // prev_mask seems to be not working better
    // if (tmpInteractiveSegMask !== null) {
@@ -777,7 +844,12 @@ export default function Editor() {
    // }

    try {
-      const res = await postInteractiveSeg(targetFile, prevMask, newClicks)
+      const res = await runPlugin(
+        PluginName.InteractiveSeg.toString(),
+        targetFile,
+        prevMask,
+        newClicks
+      )
      if (!res) {
        throw new Error('Something went wrong on server side.')
      }
@@ -990,10 +1062,7 @@ export default function Editor() {
  ])

  const disableUndo = () => {
-    if (isInteractiveSeg) {
-      return true
-    }
-    if (isInpainting) {
+    if (isProcessing) {
      return true
    }
    if (renders.length > 0) {
@@ -1074,10 +1143,7 @@ export default function Editor() {
  ])

  const disableRedo = () => {
-    if (isInteractiveSeg) {
-      return true
-    }
-    if (isInpainting) {
+    if (isProcessing) {
      return true
    }
    if (redoRenders.length > 0) {
@@ -1185,20 +1251,6 @@ export default function Editor() {
    return undefined
  }, [showBrush, isPanning])

-  useHotKey(
-    'i',
-    () => {
-      if (!isInteractiveSeg && isOriginalLoaded) {
-        setIsInteractiveSeg(true)
-        if (interactiveSegMask !== null) {
-          setShowInteractiveSegModal(true)
-        }
-      }
-    },
-    {},
-    [isInteractiveSeg, interactiveSegMask, isOriginalLoaded]
-  )
-
  // Standard Hotkeys for Brush Size
  useHotKey('[', () => {
    setBrushSize((currentBrushSize: number) => {
@@ -1370,11 +1422,7 @@ export default function Editor() {
        }}
      >
        <TransformComponent
-          contentClass={
-            isInpainting || isInteractiveSegRunning
-              ? 'editor-canvas-loading'
-              : ''
-          }
+          contentClass={isProcessing ? 'editor-canvas-loading' : ''}
          contentStyle={{
            visibility: initialCentered ? 'visible' : 'hidden',
          }}
@@ -1416,23 +1464,24 @@ export default function Editor() {
              }}
            >
              {showOriginal && (
-                <div
-                  className="editor-slider"
-                  style={{
-                    marginRight: `${sliderPos}%`,
-                  }}
-                />
+                <>
+                  <div
+                    className="editor-slider"
+                    style={{
+                      marginRight: `${sliderPos}%`,
+                    }}
+                  />
+                  <img
+                    className="original-image"
+                    src={original.src}
+                    alt="original"
+                    style={{
+                      width: `${original.naturalWidth}px`,
+                      height: `${original.naturalHeight}px`,
+                    }}
+                  />
+                </>
              )}
-
-              <img
-                className="original-image"
-                src={original.src}
-                alt="original"
-                style={{
-                  width: `${original.naturalWidth}px`,
-                  height: `${original.naturalHeight}px`,
-                }}
-              />
            </div>
          </div>

@@ -1467,6 +1516,7 @@ export default function Editor() {
      onMouseMove={onMouseMove}
      onMouseUp={onPointerUp}
    >
+      <MakeGIF renders={renders} />
      <InteractiveSegConfirmActions
        onAcceptClick={onInteractiveAccept}
        onCancelClick={onInteractiveCancel}
@@ -1514,17 +1564,6 @@ export default function Editor() {
          onClick={() => setShowRefBrush(false)}
        />
        <div className="editor-toolkit-btns">
-          <Button
-            toolTip="Interactive Segmentation"
-            icon={<CursorArrowRaysIcon />}
-            disabled={isInteractiveSeg || isInpainting || !isOriginalLoaded}
-            onClick={() => {
-              setIsInteractiveSeg(true)
-              if (interactiveSegMask !== null) {
-                setShowInteractiveSegModal(true)
-              }
-            }}
-          />
          <Button
            toolTip="Reset Zoom & Pan"
            icon={<ArrowsPointingOutIcon />}
@@ -1591,7 +1630,6 @@ export default function Editor() {
            }}
            disabled={renders.length === 0}
          />
-          <MakeGIF renders={renders} />
          <Button
            toolTip="Save Image"
            icon={<ArrowDownTrayIcon />}
@@ -1617,8 +1655,7 @@ export default function Editor() {
                </svg>
              }
              disabled={
-                isInpainting ||
-                isInteractiveSeg ||
+                isProcessing ||
                (!hadDrawSomething() && interactiveSegMask === null)
              }
              onClick={() => {
--- a/lama_cleaner/app/src/components/Editor/MakeGIF.tsx
+++ b/lama_cleaner/app/src/components/Editor/MakeGIF.tsx
@@ -1,5 +1,4 @@
-import React, { useState } from 'react'
-import { GifIcon } from '@heroicons/react/24/outline'
+import React, { useEffect, useState } from 'react'
 import { useRecoilState, useRecoilValue, useSetRecoilState } from 'recoil'
 import Button from '../shared/Button'
 import { fileState, gifImageState, toastState } from '../../store/Atoms'
@@ -7,6 +6,8 @@ import { makeGif } from '../../adapters/inpainting'
 import Modal from '../shared/Modal'
 import { LoadingIcon } from '../shared/Toast'
 import { downloadImage } from '../../utils'
+import emitter from '../../event'
+import { PluginName } from '../Plugins/Plugins'

 interface Props {
  renders: HTMLImageElement[]
@@ -30,84 +31,94 @@ const MakeGIF = (props: Props) => {
    }
  }

-  return (
-    <div>
-      <Button
-        toolTip="Make Gif"
-        icon={<GifIcon />}
-        disabled={!renders.length}
-        onClick={async () => {
-          setShow(true)
-          setGifImg(null)
-          try {
-            const gif = await makeGif(
-              file,
-              renders[renders.length - 1],
-              file.name,
-              file.type
-            )
-            if (gif) {
-              setGifImg(gif)
-            }
-          } catch (e: any) {
-            setToastState({
-              open: true,
-              desc: e.message ? e.message : e.toString(),
-              state: 'error',
-              duration: 2000,
-            })
-          }
-        }}
-      />
-      <Modal
-        onClose={handleOnClose}
-        title="GIF"
-        className="modal-setting"
-        show={show}
-      >
-        <div
-          style={{
-            display: 'flex',
-            alignItems: 'center',
-            justifyContent: 'center',
-            flexDirection: 'column',
-            gap: 16,
-          }}
-        >
-          {gifImg ? (
-            <img src={gifImg.src} style={{ borderRadius: 8 }} alt="gif" />
-          ) : (
-            <div
-              style={{
-                display: 'flex',
-                alignItems: 'center',
-                justifyContent: 'center',
-                gap: 8,
-              }}
-            >
-              <LoadingIcon />
-              Generating GIF...
-            </div>
-          )}
+  useEffect(() => {
+    emitter.on(PluginName.MakeGIF, async () => {
+      if (renders.length === 0) {
+        setToastState({
+          open: true,
+          desc: 'No render found',
+          state: 'error',
+          duration: 2000,
+        })
+        return
+      }

-          {gifImg && (
-            <div
-              style={{
-                display: 'flex',
-                width: '100%',
-                justifyContent: 'flex-end',
-                alignItems: 'center',
-                gap: '12px',
-              }}
-            >
-              <Button onClick={handleDownload} border>
-                Download
-              </Button>
-            </div>
-          )}
-        </div>
-      </Modal>
-    </div>
+      setShow(true)
+      setGifImg(null)
+      try {
+        const gif = await makeGif(
+          file,
+          renders[renders.length - 1],
+          file.name,
+          file.type
+        )
+        if (gif) {
+          setGifImg(gif)
+        }
+      } catch (e: any) {
+        setToastState({
+          open: true,
+          desc: e.message ? e.message : e.toString(),
+          state: 'error',
+          duration: 2000,
+        })
+        setShow(false)
+      }
+    })
+    return () => {
+      emitter.off(PluginName.MakeGIF)
+    }
+  }, [setGifImg, renders, file, setShow])
+
+  return (
+    <Modal
+      onClose={handleOnClose}
+      title="GIF"
+      className="modal-setting"
+      show={show}
+    >
+      <div
+        style={{
+          display: 'flex',
+          alignItems: 'center',
+          justifyContent: 'center',
+          flexDirection: 'column',
+          gap: 16,
+        }}
+      >
+        {gifImg ? (
+          <img src={gifImg.src} style={{ borderRadius: 8 }} alt="gif" />
+        ) : (
+          <div
+            style={{
+              display: 'flex',
+              alignItems: 'center',
+              justifyContent: 'center',
+              gap: 8,
+            }}
+          >
+            <LoadingIcon />
+            Generating GIF...
+          </div>
+        )}
+
+        {gifImg && (
+          <div
+            style={{
+              display: 'flex',
+              width: '100%',
+              justifyContent: 'flex-end',
+              alignItems: 'center',
+              gap: '12px',
+            }}
+          >
+            <Button onClick={handleDownload} border>
+              Download
+            </Button>
+          </div>
+        )}
+      </div>
+    </Modal>
  )
 }

--- a/lama_cleaner/app/src/components/Plugins/Plugins.scss
+++ b/lama_cleaner/app/src/components/Plugins/Plugins.scss
@@ -0,0 +1,92 @@
+@use '../../styles/Mixins/' as *;
+
+.plugins {
+  position: absolute;
+  top: 68px;
+  left: 1rem;
+  padding: 0.1rem 0.3rem;
+  z-index: 4;
+
+  border-radius: 0.8rem;
+  border-style: solid;
+  border-color: var(--border-color);
+  border-width: 1px;
+}
+
+.plugins-trigger {
+  font-family: 'WorkSans', sans-serif;
+  font-size: 16px;
+  border: 0px;
+}
+
+.plugins-content {
+  outline: none;
+  position: relative;
+  font-family: 'WorkSans', sans-serif;
+  font-size: 14px;
+  top: 8px;
+  left: 1rem;
+  padding: 0.8rem 0.5rem;
+  z-index: 9;
+
+  // backdrop-filter: blur(12px);
+  color: var(--text-color);
+  background-color: var(--page-bg);
+
+  border-radius: 0.8rem;
+  border-style: solid;
+  border-color: var(--border-color);
+  border-width: 1px;
+
+  display: flex;
+  flex-direction: column;
+  gap: 12px;
+
+  .setting-block-content {
+    gap: 1rem;
+  }
+
+  // input {
+  //   height: 24px;
+  //   // border-radius: 4px;
+  // }
+
+  // button {
+  //   height: 28px;
+  //   // border-radius: 4px;
+  // }
+}
+
+.negative-prompt {
+  all: unset;
+  border-width: 0;
+  border-radius: 0.5rem;
+  min-height: 150px;
+  max-width: 200px;
+  width: 100%;
+  padding: 12px 0.8rem;
+  outline: 1px solid var(--border-color);
+
+  &:focus-visible {
+    border-width: 0;
+    outline: 1px solid var(--yellow-accent);
+  }
+
+  &:-webkit-input-placeholder {
+    padding-top: 10px;
+  }
+
+  &:-moz-input-placeholder {
+    padding-top: 10px;
+  }
+
+  &:-ms-input-placeholder {
+    padding-top: 10px;
+  }
+}
+
+.resize-title-tile {
+  width: 86px;
+  font-size: 0.5rem;
+  color: var(--text-color-gray);
+}
--- a/lama_cleaner/app/src/components/Plugins/Plugins.tsx
+++ b/lama_cleaner/app/src/components/Plugins/Plugins.tsx
@@ -0,0 +1,91 @@
+import React, { FormEvent } from 'react'
+import { useRecoilValue } from 'recoil'
+import { CursorArrowRaysIcon, GifIcon } from '@heroicons/react/24/outline'
+import { BoxModelIcon, MarginIcon, HobbyKnifeIcon } from '@radix-ui/react-icons'
+import { useToggle } from 'react-use'
+import * as PopoverPrimitive from '@radix-ui/react-popover'
+import {
+  fileState,
+  isInpaintingState,
+  isPluginRunningState,
+  isProcessingState,
+  serverConfigState,
+} from '../../store/Atoms'
+import emitter from '../../event'
+import Button from '../shared/Button'
+
+export enum PluginName {
+  RemoveBG = 'RemoveBG',
+  RealESRGAN = 'RealESRGAN',
+  InteractiveSeg = 'InteractiveSeg',
+  MakeGIF = 'MakeGIF',
+}
+
+const pluginMap = {
+  [PluginName.RemoveBG]: {
+    IconClass: HobbyKnifeIcon,
+    showName: 'RemoveBG',
+  },
+  [PluginName.RealESRGAN]: {
+    IconClass: BoxModelIcon,
+    showName: 'RealESRGAN 4x',
+  },
+  [PluginName.InteractiveSeg]: {
+    IconClass: CursorArrowRaysIcon,
+    showName: 'Interactive Seg',
+  },
+  [PluginName.MakeGIF]: {
+    IconClass: GifIcon,
+    showName: 'Make GIF',
+  },
+}
+
+const Plugins = () => {
+  const [open, toggleOpen] = useToggle(true)
+  const serverConfig = useRecoilValue(serverConfigState)
+  const file = useRecoilValue(fileState)
+  const isProcessing = useRecoilValue(isProcessingState)
+
+  const onPluginClick = (pluginName: string) => {
+    if (isProcessing) {
+      return
+    }
+    emitter.emit(pluginName)
+  }
+
+  const renderPlugins = () => {
+    return serverConfig.plugins.map((plugin: string) => {
+      const { IconClass } = pluginMap[plugin as PluginName]
+      return (
+        <Button
+          style={{ gap: 6 }}
+          icon={<IconClass style={{ width: 15 }} />}
+          onClick={() => onPluginClick(plugin)}
+          disabled={!file || isProcessing}
+        >
+          {pluginMap[plugin as PluginName].showName}
+        </Button>
+      )
+    })
+  }
+
+  return (
+    <div className="plugins">
+      <PopoverPrimitive.Root open={open}>
+        <PopoverPrimitive.Trigger
+          className="btn-primary plugins-trigger"
+          onClick={() => toggleOpen()}
+        >
+          Plugins
+        </PopoverPrimitive.Trigger>
+        <PopoverPrimitive.Portal>
+          <PopoverPrimitive.Content className="plugins-content">
+            {renderPlugins()}
+          </PopoverPrimitive.Content>
+        </PopoverPrimitive.Portal>
+      </PopoverPrimitive.Root>
+    </div>
+  )
+}
+
+export default Plugins
--- a/lama_cleaner/app/src/components/Settings/ModelSettingBlock.tsx
+++ b/lama_cleaner/app/src/components/Settings/ModelSettingBlock.tsx
@@ -1,11 +1,9 @@
 import React, { ReactNode, useEffect, useState } from 'react'
 import { useRecoilState, useRecoilValue } from 'recoil'
-import { getIsDisableModelSwitch } from '../../adapters/inpainting'
 import {
  AIModel,
  CV2Flag,
  isDisableModelSwitchState,
-  SDSampler,
  settingState,
 } from '../../store/Atoms'
 import Selector from '../shared/Selector'
--- a/lama_cleaner/app/src/components/SidePanel/SidePanel.scss
+++ b/lama_cleaner/app/src/components/SidePanel/SidePanel.scss
@@ -4,7 +4,7 @@
  position: absolute;
  top: 68px;
  right: 1.5rem;
-  padding: 0.3rem 0.3rem;
+  padding: 0.1rem 0.3rem;
  z-index: 4;

  border-radius: 0.8rem;
@@ -20,10 +20,11 @@
 }

 .side-panel-content {
+  outline: none;
  position: relative;
  font-family: 'WorkSans', sans-serif;
  font-size: 14px;
-  top: 1rem;
+  top: 8px;
  right: 1.5rem;
  padding: 1rem 1rem;
  z-index: 9;
--- a/lama_cleaner/app/src/components/Workspace.tsx
+++ b/lama_cleaner/app/src/components/Workspace.tsx
@@ -24,6 +24,7 @@ import SidePanel from './SidePanel/SidePanel'
 import PESidePanel from './SidePanel/PESidePanel'
 import FileManager from './FileManager/FileManager'
 import P2PSidePanel from './SidePanel/P2PSidePanel'
+import Plugins from './Plugins/Plugins'

 const Workspace = () => {
  const setFile = useSetRecoilState(fileState)
@@ -102,6 +103,7 @@ const Workspace = () => {
      {isSD ? <SidePanel /> : <></>}
      {isPaintByExample ? <PESidePanel /> : <></>}
      {isPix2Pix ? <P2PSidePanel /> : <></>}
+      <Plugins />
      <FileManager
        photoWidth={256}
        show={showFileManager}