update

web_app/package-lock.json

@@ -50,6 +50,7 @@
         "tailwind-merge": "^2.0.0",
         "tailwindcss-animate": "^1.0.7",
         "zod": "^3.22.4",
+        "zundo": "^2.0.0",
         "zustand": "^4.4.6"
       },
       "devDependencies": {
@@ -6250,6 +6251,18 @@
         "url": "https://github.com/sponsors/colinhacks"
       }
     },
+    "node_modules/zundo": {
+      "version": "2.0.0",
+      "resolved": "https://registry.npmjs.org/zundo/-/zundo-2.0.0.tgz",
+      "integrity": "sha512-XzKDyunmyxvQHKDjgTmOClOQscJAm5NAa1iEazR0DilvV/uwCjnDwlHJuJ+GmG/oj5RMjzsD0ptghZzjEj1w4g==",
+      "funding": {
+        "type": "individual",
+        "url": "https://github.com/sponsors/charkour"
+      },
+      "peerDependencies": {
+        "zustand": "^4.3.0"
+      }
+    },
     "node_modules/zustand": {
       "version": "4.4.6",
       "resolved": "https://registry.npmjs.org/zustand/-/zustand-4.4.6.tgz",

web_app/package.json

@@ -52,6 +52,7 @@
     "tailwind-merge": "^2.0.0",
     "tailwindcss-animate": "^1.0.7",
     "zod": "^3.22.4",
+    "zundo": "^2.0.0",
     "zustand": "^4.4.6"
   },
   "devDependencies": {

web_app/src/App.tsx

@@ -9,6 +9,7 @@ import Workspace from "@/components/Workspace"
 import FileSelect from "@/components/FileSelect"
 import { Toaster } from "./components/ui/toaster"
 import { useStore } from "./lib/states"
+import { useWindowSize } from "react-use"
 
 const SUPPORTED_FILE_TYPE = [
   "image/jpeg",
@@ -18,20 +19,27 @@ const SUPPORTED_FILE_TYPE = [
   "image/tiff",
 ]
 function Home() {
-  const [file, setServerConfig, setFile] = useStore((state) => [
+  const [file, updateAppState, setServerConfig, setFile] = useStore((state) => [
     state.file,
+    state.updateAppState,
     state.setServerConfig,
     state.setFile,
   ])
 
   const userInputImage = useInputImage()
 
+  const windowSize = useWindowSize()
+
   useEffect(() => {
     if (userInputImage) {
       setFile(userInputImage)
     }
   }, [userInputImage, setFile])
 
+  useEffect(() => {
+    updateAppState({ windowSize })
+  }, [windowSize])
+
   // Keeping GUI Window Open
   useEffect(() => {
     const fetchData = async () => {

web_app/src/components/Editor.tsx

@@ -13,9 +13,11 @@ import {
   askWritePermission,
   copyCanvasImage,
   downloadImage,
+  drawLines,
   isMidClick,
   isRightClick,
   loadImage,
+  mouseXY,
   srcToFile,
 } from "@/lib/utils"
 import { Eraser, Eye, Redo, Undo, Expand, Download } from "lucide-react"
@@ -29,7 +31,7 @@ import emitter, {
 } from "@/lib/event"
 import { useImage } from "@/hooks/useImage"
 import { Slider } from "./ui/slider"
-import { PluginName } from "@/lib/types"
+import { Line, LineGroup, PluginName } from "@/lib/types"
 import { useHotkeys } from "react-hotkeys-hook"
 import { useStore } from "@/lib/states"
 import Cropper from "./Cropper"
@@ -38,40 +40,6 @@ const TOOLBAR_HEIGHT = 200
 const MIN_BRUSH_SIZE = 10
 const MAX_BRUSH_SIZE = 200
 const COMPARE_SLIDER_DURATION_MS = 300
-const BRUSH_COLOR = "#ffcc00bb"
-
-interface Line {
-  size?: number
-  pts: { x: number; y: number }[]
-}
-
-type LineGroup = Array<Line>
-
-function drawLines(
-  ctx: CanvasRenderingContext2D,
-  lines: LineGroup,
-  color = BRUSH_COLOR
-) {
-  ctx.strokeStyle = color
-  ctx.lineCap = "round"
-  ctx.lineJoin = "round"
-
-  lines.forEach((line) => {
-    if (!line?.pts.length || !line.size) {
-      return
-    }
-    ctx.lineWidth = line.size
-    ctx.beginPath()
-    ctx.moveTo(line.pts[0].x, line.pts[0].y)
-    line.pts.forEach((pt) => ctx.lineTo(pt.x, pt.y))
-    ctx.stroke()
-  })
-}
-
-function mouseXY(ev: SyntheticEvent) {
-  const mouseEvent = ev.nativeEvent as MouseEvent
-  return { x: mouseEvent.offsetX, y: mouseEvent.offsetY }
-}
 
 interface EditorProps {
   file: File
@@ -85,14 +53,12 @@ export default function Editor(props: EditorProps) {
     isInpainting,
     imageWidth,
     imageHeight,
-    baseBrushSize,
-    brushSizeScale,
     settings,
     enableAutoSaving,
     cropperRect,
     enableManualInpainting,
     setImageSize,
-    setBrushSize,
+    setBaseBrushSize,
     setIsInpainting,
     setSeed,
     interactiveSegState,
@@ -100,18 +66,25 @@ export default function Editor(props: EditorProps) {
     resetInteractiveSegState,
     isPluginRunning,
     setIsPluginRunning,
+    handleCanvasMouseDown,
+    handleCanvasMouseMove,
+    cleanCurLineGroup,
+    updateEditorState,
+    resetRedoState,
+    undo,
+    redo,
+    undoDisabled,
+    redoDisabled,
   ] = useStore((state) => [
     state.isInpainting,
     state.imageWidth,
     state.imageHeight,
-    state.brushSize,
-    state.brushSizeScale,
     state.settings,
     state.serverConfig.enableAutoSaving,
     state.cropperState,
     state.settings.enableManualInpainting,
     state.setImageSize,
-    state.setBrushSize,
+    state.setBaseBrushSize,
     state.setIsInpainting,
     state.setSeed,
     state.interactiveSegState,
@@ -119,8 +92,24 @@ export default function Editor(props: EditorProps) {
     state.resetInteractiveSegState,
     state.isPluginRunning,
     state.setIsPluginRunning,
+    state.handleCanvasMouseDown,
+    state.handleCanvasMouseMove,
+    state.cleanCurLineGroup,
+    state.updateEditorState,
+    state.resetRedoState,
+    state.undo,
+    state.redo,
+    state.undoDisabled(),
+    state.redoDisabled(),
   ])
-  const brushSize = baseBrushSize * brushSizeScale
+  const baseBrushSize = useStore((state) => state.editorState.baseBrushSize)
+  const brushSize = useStore((state) => state.getBrushSize())
+  const renders = useStore((state) => state.editorState.renders)
+  const lineGroups = useStore((state) => state.editorState.lineGroups)
+
+  const lastLineGroup = useStore((state) => state.editorState.lastLineGroup)
+  const curLineGroup = useStore((state) => state.editorState.curLineGroup)
+  const redoLineGroups = useStore((state) => state.editorState.redoLineGroups)
 
   // 纯 local state
   const [showOriginal, setShowOriginal] = useState(false)
@@ -151,14 +140,10 @@ export default function Editor(props: EditorProps) {
     useState<LineGroup>([])
 
   const [original, isOriginalLoaded] = useImage(file)
-  const [renders, setRenders] = useState<HTMLImageElement[]>([])
   const [context, setContext] = useState<CanvasRenderingContext2D>()
   const [maskCanvas] = useState<HTMLCanvasElement>(() => {
     return document.createElement("canvas")
   })
-  const [lineGroups, setLineGroups] = useState<LineGroup[]>([])
-  const [lastLineGroup, setLastLineGroup] = useState<LineGroup>([])
-  const [curLineGroup, setCurLineGroup] = useState<LineGroup>([])
   const [{ x, y }, setCoords] = useState({ x: -1, y: -1 })
   const [showBrush, setShowBrush] = useState(false)
   const [showRefBrush, setShowRefBrush] = useState(false)
@@ -185,11 +170,6 @@ export default function Editor(props: EditorProps) {
 
   const [sliderPos, setSliderPos] = useState<number>(0)
 
-  // redo 相关
-  const [redoRenders, setRedoRenders] = useState<HTMLImageElement[]>([])
-  const [redoCurLines, setRedoCurLines] = useState<Line[]>([])
-  const [redoLineGroups, setRedoLineGroups] = useState<LineGroup[]>([])
-
   const draw = useCallback(
     (render: HTMLImageElement, lineGroup: LineGroup) => {
       if (!context) {
@@ -276,28 +256,51 @@ export default function Editor(props: EditorProps) {
         )
       }
     },
-    [context, maskCanvas, isPix2Pix, imageWidth, imageHeight]
+    [context, maskCanvas, imageWidth, imageHeight]
   )
 
   const hadDrawSomething = useCallback(() => {
-    if (isPix2Pix) {
-      return true
-    }
     return curLineGroup.length !== 0
-  }, [curLineGroup, isPix2Pix])
+  }, [curLineGroup])
 
-  const drawOnCurrentRender = useCallback(
+  // const drawOnCurrentRender = useCallback(
-    (lineGroup: LineGroup) => {
+  //   (lineGroup: LineGroup) => {
-      console.log("[drawOnCurrentRender] draw on current render")
+  //     console.log("[drawOnCurrentRender] draw on current render")
-      if (renders.length === 0) {
+  //     if (renders.length === 0) {
-        draw(original, lineGroup)
+  //       draw(original, lineGroup)
-      } else {
+  //     } else {
-        draw(renders[renders.length - 1], lineGroup)
+  //       draw(renders[renders.length - 1], lineGroup)
+  //     }
+  //   },
+  //   [original, renders, draw]
+  // )
+
+  useEffect(() => {
+    if (!context) {
+      return
     }
-    },
+
-    [original, renders, draw]
+    const render = renders.length === 0 ? original : renders[renders.length - 1]
+
+    console.log(
+      `[draw] render size: ${render.width}x${render.height} image size: ${imageWidth}x${imageHeight} canvas size: ${context.canvas.width}x${context.canvas.height}`
     )
 
+    context.clearRect(0, 0, context.canvas.width, context.canvas.height)
+    context.drawImage(render, 0, 0, imageWidth, imageHeight)
+    // if (interactiveSegState.isInteractiveSeg && tmpInteractiveSegMask) {
+    //   context.drawImage(tmpInteractiveSegMask, 0, 0, imageWidth, imageHeight)
+    // }
+    // if (!interactiveSegState.isInteractiveSeg && interactiveSegMask) {
+    //   context.drawImage(interactiveSegMask, 0, 0, imageWidth, imageHeight)
+    // }
+    // if (dreamButtonHoverSegMask) {
+    //   context.drawImage(dreamButtonHoverSegMask, 0, 0, imageWidth, imageHeight)
+    // }
+    drawLines(context, curLineGroup)
+    // drawLines(context, dreamButtonHoverLineGroup)
+  }, [renders, file, original, context, curLineGroup, imageHeight, imageWidth])
+
   const runInpainting = useCallback(
     async (
       useLastLineGroup?: boolean,
@@ -332,13 +335,13 @@ export default function Editor(props: EditorProps) {
           return
         }
 
-        setLastLineGroup(curLineGroup)
+        // setLastLineGroup(curLineGroup)
         maskLineGroup = curLineGroup
       }
 
       const newLineGroups = [...lineGroups, maskLineGroup]
 
-      setCurLineGroup([])
+      cleanCurLineGroup()
       setIsDraging(false)
       setIsInpainting(true)
       drawLinesOnMask([maskLineGroup], maskImage)
@@ -391,15 +394,17 @@ export default function Editor(props: EditorProps) {
         if (useLastLineGroup === true) {
           const prevRenders = renders.slice(0, -1)
           const newRenders = [...prevRenders, newRender]
-          setRenders(newRenders)
+          // setRenders(newRenders)
+          updateEditorState({ renders: newRenders })
         } else {
           const newRenders = [...renders, newRender]
-          setRenders(newRenders)
+          updateEditorState({ renders: newRenders })
         }
 
         draw(newRender, [])
         // Only append new LineGroup after inpainting success
-        setLineGroups(newLineGroups)
+        // setLineGroups(newLineGroups)
+        updateEditorState({ lineGroups: newLineGroups })
 
         // clear redo stack
         resetRedoState()
@@ -409,7 +414,7 @@ export default function Editor(props: EditorProps) {
           title: "Uh oh! Something went wrong.",
           description: e.message ? e.message : e.toString(),
         })
-        drawOnCurrentRender([])
+        // drawOnCurrentRender([])
       }
       setIsInpainting(false)
       setPrevInteractiveSegMask(maskImage)
@@ -423,7 +428,7 @@ export default function Editor(props: EditorProps) {
       maskCanvas,
       settings,
       cropperRect,
-      drawOnCurrentRender,
+      // drawOnCurrentRender,
       hadDrawSomething,
       drawLinesOnMask,
     ]
@@ -488,7 +493,7 @@ export default function Editor(props: EditorProps) {
     hadDrawSomething,
     interactiveSegMask,
     prevInteractiveSegMask,
-    drawOnCurrentRender,
+    // drawOnCurrentRender,
     lineGroups,
     redoLineGroups,
   ])
@@ -499,13 +504,13 @@ export default function Editor(props: EditorProps) {
       if (!hadDrawSomething() && !interactiveSegMask) {
         setDreamButtonHoverSegMask(null)
         setDreamButtonHoverLineGroup([])
-        drawOnCurrentRender([])
+        // drawOnCurrentRender([])
       }
     })
     return () => {
       emitter.off(DREAM_BUTTON_MOUSE_LEAVE)
     }
-  }, [hadDrawSomething, interactiveSegMask, drawOnCurrentRender])
+  }, [hadDrawSomething, interactiveSegMask])
 
   useEffect(() => {
     emitter.on(EVENT_CUSTOM_MASK, (data: any) => {
@@ -601,9 +606,10 @@ export default function Editor(props: EditorProps) {
         await loadImage(newRender, blob)
         setImageSize(newRender.height, newRender.width)
         const newRenders = [...renders, newRender]
-        setRenders(newRenders)
+        // setRenders(newRenders)
+        updateEditorState({ renders: newRenders })
         const newLineGroups = [...lineGroups, []]
-        setLineGroups(newLineGroups)
+        updateEditorState({ lineGroups: newLineGroups })
 
         const end = new Date()
         const time = end.getTime() - start.getTime()
@@ -632,7 +638,7 @@ export default function Editor(props: EditorProps) {
     },
     [
       renders,
-      setRenders,
+      // setRenders,
       getCurrentRender,
       setIsPluginRunning,
       isProcessing,
@@ -640,7 +646,7 @@ export default function Editor(props: EditorProps) {
       lineGroups,
       viewportRef,
       windowSize,
-      setLineGroups,
+      // setLineGroups,
     ]
   )
 
@@ -737,7 +743,7 @@ export default function Editor(props: EditorProps) {
       context.canvas.width = width
       context.canvas.height = height
       console.log("[on file load] set canvas size && drawOnCurrentRender")
-      drawOnCurrentRender([])
+      // drawOnCurrentRender([])
     }
 
     if (!initialCentered) {
@@ -747,13 +753,13 @@ export default function Editor(props: EditorProps) {
       setInitialCentered(true)
     }
   }, [
-    // context?.canvas,
+    context?.canvas,
     viewportRef,
     original,
     isOriginalLoaded,
     windowSize,
     initialCentered,
-    drawOnCurrentRender,
+    // drawOnCurrentRender,
     getCurrentWidthHeight,
   ])
 
@@ -790,12 +796,6 @@ export default function Editor(props: EditorProps) {
     minScale,
   ])
 
-  const resetRedoState = () => {
-    setRedoCurLines([])
-    setRedoLineGroups([])
-    setRedoRenders([])
-  }
-
   useEffect(() => {
     window.addEventListener("resize", () => {
       resetZoom()
@@ -825,8 +825,8 @@ export default function Editor(props: EditorProps) {
 
     if (isDraging) {
       setIsDraging(false)
-      setCurLineGroup([])
+      // setCurLineGroup([])
-      drawOnCurrentRender([])
+      // drawOnCurrentRender([])
     } else {
       resetZoom()
     }
@@ -836,7 +836,7 @@ export default function Editor(props: EditorProps) {
     isDraging,
     isInpainting,
     resetZoom,
-    drawOnCurrentRender,
+    // drawOnCurrentRender,
   ])
 
   const onMouseMove = (ev: SyntheticEvent) => {
@@ -845,15 +845,15 @@ export default function Editor(props: EditorProps) {
   }
 
   const onMouseDrag = (ev: SyntheticEvent) => {
-    if (isChangingBrushSizeByMouse) {
+    // if (isChangingBrushSizeByMouse) {
-      const initX = changeBrushSizeByMouseInit.x
+    //   const initX = changeBrushSizeByMouseInit.x
-      // move right: increase brush size
+    //   // move right: increase brush size
-      const newSize = changeBrushSizeByMouseInit.brushSize + (x - initX)
+    //   const newSize = changeBrushSizeByMouseInit.brushSize + (x - initX)
-      if (newSize <= MAX_BRUSH_SIZE && newSize >= MIN_BRUSH_SIZE) {
+    //   if (newSize <= MAX_BRUSH_SIZE && newSize >= MIN_BRUSH_SIZE) {
-        setBrushSize(newSize)
+    //     setBaseBrushSize(newSize)
-      }
+    //   }
-      return
+    //   return
-    }
+    // }
     if (interactiveSegState.isInteractiveSeg) {
       return
     }
@@ -866,10 +866,12 @@ export default function Editor(props: EditorProps) {
     if (curLineGroup.length === 0) {
       return
     }
-    const lineGroup = [...curLineGroup]
+
-    lineGroup[lineGroup.length - 1].pts.push(mouseXY(ev))
+    handleCanvasMouseMove(mouseXY(ev))
-    setCurLineGroup(lineGroup)
+    // const lineGroup = [...curLineGroup]
-    drawOnCurrentRender(lineGroup)
+    // lineGroup[lineGroup.length - 1].pts.push(mouseXY(ev))
+    // setCurLineGroup(lineGroup)
+    // drawOnCurrentRender(lineGroup)
   }
 
   const runInteractiveSeg = async (newClicks: number[][]) => {
@@ -1010,166 +1012,29 @@ export default function Editor(props: EditorProps) {
 
     setIsDraging(true)
 
-    let lineGroup: LineGroup = []
+    // let lineGroup: LineGroup = []
-    if (enableManualInpainting) {
+    // if (enableManualInpainting) {
-      lineGroup = [...curLineGroup]
+    //   lineGroup = [...curLineGroup]
-    }
-    lineGroup.push({ size: brushSize, pts: [mouseXY(ev)] })
-    setCurLineGroup(lineGroup)
-    drawOnCurrentRender(lineGroup)
-  }
-
-  const undoStroke = useCallback(() => {
-    if (curLineGroup.length === 0) {
-      return
-    }
-    setLastLineGroup([])
-
-    const lastLine = curLineGroup.pop()!
-    const newRedoCurLines = [...redoCurLines, lastLine]
-    setRedoCurLines(newRedoCurLines)
-
-    const newLineGroup = [...curLineGroup]
-    setCurLineGroup(newLineGroup)
-    drawOnCurrentRender(newLineGroup)
-  }, [curLineGroup, redoCurLines, drawOnCurrentRender])
-
-  const undoRender = useCallback(() => {
-    if (!renders.length) {
-      return
-    }
-
-    // save line Group
-    const latestLineGroup = lineGroups.pop()!
-    setRedoLineGroups([...redoLineGroups, latestLineGroup])
-    // If render is undo, clear strokes
-    setRedoCurLines([])
-
-    setLineGroups([...lineGroups])
-    setCurLineGroup([])
-    setIsDraging(false)
-
-    // save render
-    const lastRender = renders.pop()!
-    setRedoRenders([...redoRenders, lastRender])
-
-    const newRenders = [...renders]
-    setRenders(newRenders)
-    // if (newRenders.length === 0) {
-    //   draw(original, [])
-    // } else {
-    //   draw(newRenders[newRenders.length - 1], [])
     // }
-  }, [
+    // lineGroup.push({ size: brushSize, pts: [mouseXY(ev)] })
-    draw,
+    // setCurLineGroup(lineGroup)
-    renders,
-    redoRenders,
-    redoLineGroups,
-    lineGroups,
-    original,
-    context,
-  ])
 
-  const undo = (keyboardEvent: KeyboardEvent | SyntheticEvent) => {
+    handleCanvasMouseDown(mouseXY(ev))
+
+    // drawOnCurrentRender(lineGroup)
+  }
+
+  const handleUndo = (keyboardEvent: KeyboardEvent | SyntheticEvent) => {
     keyboardEvent.preventDefault()
-    if (enableManualInpainting && curLineGroup.length !== 0) {
+    undo()
-      undoStroke()
-    } else {
-      undoRender()
-    }
   }
+  useHotkeys("meta+z,ctrl+z", handleUndo)
 
-  useHotkeys("meta+z,ctrl+z", undo, undefined, [
+  const handleRedo = (keyboardEvent: KeyboardEvent | SyntheticEvent) => {
-    undoStroke,
-    undoRender,
-    enableManualInpainting,
-    curLineGroup,
-    context?.canvas,
-    renders,
-  ])
-
-  const disableUndo = () => {
-    if (isProcessing) {
-      return true
-    }
-    if (renders.length > 0) {
-      return false
-    }
-
-    if (enableManualInpainting) {
-      if (curLineGroup.length === 0) {
-        return true
-      }
-    } else if (renders.length === 0) {
-      return true
-    }
-
-    return false
-  }
-
-  const redoStroke = useCallback(() => {
-    if (redoCurLines.length === 0) {
-      return
-    }
-    const line = redoCurLines.pop()!
-    setRedoCurLines([...redoCurLines])
-
-    const newLineGroup = [...curLineGroup, line]
-    setCurLineGroup(newLineGroup)
-    drawOnCurrentRender(newLineGroup)
-  }, [curLineGroup, redoCurLines, drawOnCurrentRender])
-
-  const redoRender = useCallback(() => {
-    if (redoRenders.length === 0) {
-      return
-    }
-    const lineGroup = redoLineGroups.pop()!
-    setRedoLineGroups([...redoLineGroups])
-
-    setLineGroups([...lineGroups, lineGroup])
-    setCurLineGroup([])
-    setIsDraging(false)
-
-    const render = redoRenders.pop()!
-    const newRenders = [...renders, render]
-    setRenders(newRenders)
-    // draw(newRenders[newRenders.length - 1], [])
-  }, [draw, renders, redoRenders, redoLineGroups, lineGroups, original])
-
-  const redo = (keyboardEvent: KeyboardEvent | SyntheticEvent) => {
     keyboardEvent.preventDefault()
-    if (enableManualInpainting && redoCurLines.length !== 0) {
+    redo()
-      redoStroke()
-    } else {
-      redoRender()
-    }
-  }
-
-  useHotkeys("shift+ctrl+z,shift+meta+z", redo, undefined, [
-    redoStroke,
-    redoRender,
-    enableManualInpainting,
-    redoCurLines,
-  ])
-
-  const disableRedo = () => {
-    if (isProcessing) {
-      return true
-    }
-    if (redoRenders.length > 0) {
-      return false
-    }
-
-    if (enableManualInpainting) {
-      if (redoCurLines.length === 0) {
-        return true
-      }
-    } else if (redoRenders.length === 0) {
-      return true
-    }
-
-    return false
   }
+  useHotkeys("shift+ctrl+z,shift+meta+z", handleRedo)
 
   useKeyPressEvent(
     "Tab",
@@ -1265,7 +1130,7 @@ export default function Editor(props: EditorProps) {
       if (baseBrushSize <= 10 && baseBrushSize > 0) {
         newBrushSize = baseBrushSize - 5
       }
-      setBrushSize(newBrushSize)
+      setBaseBrushSize(newBrushSize)
     },
     [baseBrushSize]
   )
@@ -1273,7 +1138,7 @@ export default function Editor(props: EditorProps) {
   useHotkeys(
     "]",
     () => {
-      setBrushSize(baseBrushSize + 10)
+      setBaseBrushSize(baseBrushSize + 10)
     },
     [baseBrushSize]
   )
@@ -1366,7 +1231,7 @@ export default function Editor(props: EditorProps) {
   }
 
   const handleSliderChange = (value: number) => {
-    setBrushSize(value)
+    setBaseBrushSize(value)
 
     if (!showRefBrush) {
       setShowRefBrush(true)
@@ -1552,10 +1417,18 @@ export default function Editor(props: EditorProps) {
           >
             <Expand />
           </IconButton>
-          <IconButton tooltip="Undo" onClick={undo} disabled={disableUndo()}>
+          <IconButton
+            tooltip="Undo"
+            onClick={handleUndo}
+            disabled={undoDisabled}
+          >
             <Undo />
           </IconButton>
-          <IconButton tooltip="Redo" onClick={redo} disabled={disableRedo()}>
+          <IconButton
+            tooltip="Redo"
+            onClick={handleRedo}
+            disabled={redoDisabled}
+          >
             <Redo />
           </IconButton>
           <IconButton

web_app/src/components/Plugins.tsx

@@ -8,7 +8,13 @@ import {
   DropdownMenuTrigger,
 } from "./ui/dropdown-menu"
 import { Button } from "./ui/button"
-import { Fullscreen, MousePointerClick, Slice, Smile } from "lucide-react"
+import {
+  Blocks,
+  Fullscreen,
+  MousePointerClick,
+  Slice,
+  Smile,
+} from "lucide-react"
 import { MixIcon } from "@radix-ui/react-icons"
 import { useStore } from "@/lib/states"
 import { InteractiveSeg } from "./InteractiveSeg"
@@ -118,8 +124,8 @@ const Plugins = () => {
         className="border rounded-lg z-10 bg-background"
         tabIndex={-1}
       >
-        <Button variant="ghost" size="icon" asChild>
+        <Button variant="ghost" size="icon" asChild className="p-1.5">
-          <MixIcon className="p-2" />
+          <Blocks strokeWidth={1} />
         </Button>
       </DropdownMenuTrigger>
       <DropdownMenuContent side="bottom" align="start">

web_app/src/components/PromptInput.tsx

@@ -45,7 +45,7 @@ const PromptInput = () => {
   return (
     <div className="flex gap-4 items-center">
       <Input
-        className="min-w-[600px]"
+        className="min-w-[500px]"
         value={prompt}
         onInput={handleOnInput}
         onKeyUp={onKeyUp}

web_app/src/components/SidePanel.tsx

@@ -32,6 +32,8 @@ import {
   SheetTitle,
   SheetTrigger,
 } from "./ui/sheet"
+import { ChevronLeft } from "lucide-react"
+import { Button } from "./ui/button"
 
 const SidePanel = () => {
   const [settings, updateSettings, showSidePanel] = useStore((state) => [
@@ -214,9 +216,11 @@ const SidePanel = () => {
     <Sheet open={open} onOpenChange={toggleOpen} modal={false}>
       <SheetTrigger
         tabIndex={-1}
-        className="z-10 outline-none absolute top-[68px] right-6 px-3 py-2 rounded-lg border-solid border hover:bg-primary hover:text-primary-foreground"
+        className="z-10 outline-none absolute top-[68px] right-6 rounded-lg border bg-background"
       >
-        Config
+        <Button variant="ghost" size="icon" asChild className="p-1.5">
+          <ChevronLeft strokeWidth={1} />
+        </Button>
       </SheetTrigger>
       <SheetContent
         side="right"

web_app/src/components/theme-provider.tsx

@@ -1,73 +0,0 @@
-import { createContext, useContext, useEffect, useState } from "react"
-
-type Theme = "dark" | "light" | "system"
-
-type ThemeProviderProps = {
-  children: React.ReactNode
-  defaultTheme?: Theme
-  storageKey?: string
-}
-
-type ThemeProviderState = {
-  theme: Theme
-  setTheme: (theme: Theme) => void
-}
-
-const initialState: ThemeProviderState = {
-  theme: "system",
-  setTheme: () => null,
-}
-
-const ThemeProviderContext = createContext<ThemeProviderState>(initialState)
-
-export function ThemeProvider({
-  children,
-  defaultTheme = "system",
-  storageKey = "vite-ui-theme",
-  ...props
-}: ThemeProviderProps) {
-  const [theme, setTheme] = useState<Theme>(
-    () => (localStorage.getItem(storageKey) as Theme) || defaultTheme
-  )
-
-  useEffect(() => {
-    const root = window.document.documentElement
-
-    root.classList.remove("light", "dark")
-
-    if (theme === "system") {
-      const systemTheme = window.matchMedia("(prefers-color-scheme: dark)")
-        .matches
-        ? "dark"
-        : "light"
-
-      root.classList.add(systemTheme)
-      return
-    }
-
-    root.classList.add(theme)
-  }, [theme])
-
-  const value = {
-    theme,
-    setTheme: (theme: Theme) => {
-      localStorage.setItem(storageKey, theme)
-      setTheme(theme)
-    },
-  }
-
-  return (
-    <ThemeProviderContext.Provider {...props} value={value}>
-      {children}
-    </ThemeProviderContext.Provider>
-  )
-}
-
-export const useTheme = () => {
-  const context = useContext(ThemeProviderContext)
-
-  if (context === undefined)
-    throw new Error("useTheme must be used within a ThemeProvider")
-
-  return context
-}

web_app/src/globals.css

@@ -70,7 +70,7 @@ html { font-family: "Inter", "system-ui"; }
     --radius: 0.5rem;
   }
  
-  .dark {
+  [data-theme='dark'] {
     --background: 240 10% 3.9%;
     --foreground: 0 0% 98%;
  
@@ -83,8 +83,8 @@ html { font-family: "Inter", "system-ui"; }
     --primary: 48 100.0% 50.0%;
     --primary-foreground: 220.9 39.3% 11%;
  
-    --secondary: 215 27.9% 16.9%;
+    --secondary: 240 3.7% 15.9%;
-    --secondary-foreground: 210 20% 98%;
+    --secondary-foreground: 0 0% 98%;
  
     --muted: 240 3.7% 15.9%;
     --muted-foreground: 240 5% 64.9%;

web_app/src/lib/const.ts

@@ -6,3 +6,4 @@ export const MODEL_TYPE_DIFFUSERS_SDXL = "diffusers_sdxl"
 export const MODEL_TYPE_DIFFUSERS_SD_INPAINT = "diffusers_sd_inpaint"
 export const MODEL_TYPE_DIFFUSERS_SDXL_INPAINT = "diffusers_sdxl_inpaint"
 export const MODEL_TYPE_OTHER = "diffusers_other"
+export const BRUSH_COLOR = "#ffcc00bb"

web_app/src/lib/states.ts

@@ -1,16 +1,22 @@
 import { create } from "zustand"
 import { persist } from "zustand/middleware"
+import { shallow } from "zustand/shallow"
 import { immer } from "zustand/middleware/immer"
+import { createWithEqualityFn } from "zustand/traditional"
 import {
   CV2Flag,
   FreeuConfig,
   LDMSampler,
+  Line,
+  LineGroup,
   ModelInfo,
+  Point,
   SDSampler,
+  Size,
   SortBy,
   SortOrder,
 } from "./types"
-import { DEFAULT_BRUSH_SIZE } from "./const"
+import { DEFAULT_BRUSH_SIZE, MODEL_TYPE_INPAINT } from "./const"
 
 type FileManagerState = {
   sortBy: SortBy
@@ -93,15 +99,28 @@ type InteractiveSegState = {
   clicks: number[][]
 }
 
+type EditorState = {
+  baseBrushSize: number
+  brushSizeScale: number
+  renders: HTMLImageElement[]
+  lineGroups: LineGroup[]
+  lastLineGroup: LineGroup
+  curLineGroup: LineGroup
+  // redo 相关
+  redoRenders: HTMLImageElement[]
+  redoCurLines: Line[]
+  redoLineGroups: LineGroup[]
+}
+
 type AppState = {
   file: File | null
   customMask: File | null
   imageHeight: number
   imageWidth: number
-  brushSize: number
-  brushSizeScale: number
   isInpainting: boolean
   isPluginRunning: boolean
+  windowSize: Size
+  editorState: EditorState
 
   interactiveSegState: InteractiveSegState
   fileManagerState: FileManagerState
@@ -112,11 +131,13 @@ type AppState = {
 }
 
 type AppAction = {
+  updateAppState: (newState: Partial<AppState>) => void
   setFile: (file: File) => void
   setCustomFile: (file: File) => void
   setIsInpainting: (newValue: boolean) => void
   setIsPluginRunning: (newValue: boolean) => void
-  setBrushSize: (newValue: number) => void
+  setBaseBrushSize: (newValue: number) => void
+  getBrushSize: () => number
   setImageSize: (width: number, height: number) => void
 
   setCropperX: (newValue: number) => void
@@ -132,6 +153,18 @@ type AppAction = {
   resetInteractiveSegState: () => void
   showPromptInput: () => boolean
   showSidePanel: () => boolean
+
+  // EditorState
+  updateEditorState: (newState: Partial<EditorState>) => void
+  runMannually: () => boolean
+  handleCanvasMouseDown: (point: Point) => void
+  handleCanvasMouseMove: (point: Point) => void
+  cleanCurLineGroup: () => void
+  resetRedoState: () => void
+  undo: () => void
+  redo: () => void
+  undoDisabled: () => boolean
+  redoDisabled: () => boolean
 }
 
 const defaultValues: AppState = {
@@ -139,10 +172,23 @@ const defaultValues: AppState = {
   customMask: null,
   imageHeight: 0,
   imageWidth: 0,
-  brushSize: DEFAULT_BRUSH_SIZE,
-  brushSizeScale: 1,
   isInpainting: false,
   isPluginRunning: false,
+  windowSize: {
+    height: 600,
+    width: 800,
+  },
+  editorState: {
+    baseBrushSize: DEFAULT_BRUSH_SIZE,
+    brushSizeScale: 1,
+    renders: [],
+    lineGroups: [],
+    lastLineGroup: [],
+    curLineGroup: [],
+    redoRenders: [],
+    redoCurLines: [],
+    redoLineGroups: [],
+  },
 
   interactiveSegState: {
     isInteractiveSeg: false,
@@ -216,12 +262,176 @@ const defaultValues: AppState = {
   },
 }
 
-export const useStore = create<AppState & AppAction>()(
+export const useStore = createWithEqualityFn<AppState & AppAction>()(
-  immer(
   persist(
-      (set, get) => ({
+    immer((set, get) => ({
       ...defaultValues,
 
+      // Edirot State //
+      updateEditorState: (newState: Partial<EditorState>) => {
+        set((state) => {
+          return {
+            ...state,
+            editorState: {
+              ...state.editorState,
+              ...newState,
+            },
+          }
+        })
+      },
+
+      cleanCurLineGroup: () => {
+        get().updateEditorState({ curLineGroup: [] })
+      },
+
+      handleCanvasMouseDown: (point: Point) => {
+        let lineGroup: LineGroup = []
+        const state = get()
+        if (state.runMannually()) {
+          lineGroup = [...state.editorState.curLineGroup]
+        }
+        lineGroup.push({ size: state.getBrushSize(), pts: [point] })
+        set((state) => {
+          state.editorState.curLineGroup = lineGroup
+        })
+      },
+
+      handleCanvasMouseMove: (point: Point) => {
+        set((state) => {
+          const curLineGroup = state.editorState.curLineGroup
+          if (curLineGroup.length) {
+            curLineGroup[curLineGroup.length - 1].pts.push(point)
+          }
+        })
+      },
+
+      runMannually: (): boolean => {
+        const state = get()
+        return (
+          state.settings.enableManualInpainting ||
+          state.settings.model.model_type !== MODEL_TYPE_INPAINT
+        )
+      },
+
+      // undo/redo
+
+      undoDisabled: (): boolean => {
+        const editorState = get().editorState
+        if (editorState.renders.length > 0) {
+          return false
+        }
+        if (get().runMannually()) {
+          if (editorState.curLineGroup.length === 0) {
+            return true
+          }
+        } else if (editorState.renders.length === 0) {
+          return true
+        }
+        return false
+      },
+
+      undo: () => {
+        if (
+          get().runMannually() &&
+          get().editorState.curLineGroup.length !== 0
+        ) {
+          // undoStroke
+          set((state) => {
+            const editorState = state.editorState
+            if (editorState.curLineGroup.length === 0) {
+              return
+            }
+            editorState.lastLineGroup = []
+            const lastLine = editorState.curLineGroup.pop()!
+            editorState.redoCurLines.push(lastLine)
+          })
+        } else {
+          set((state) => {
+            const editorState = state.editorState
+            if (
+              editorState.renders.length === 0 ||
+              editorState.lineGroups.length === 0
+            ) {
+              return
+            }
+            const lastLineGroup = editorState.lineGroups.pop()!
+            editorState.redoLineGroups.push(lastLineGroup)
+            editorState.redoCurLines = []
+            editorState.curLineGroup = []
+
+            const lastRender = editorState.renders.pop()!
+            editorState.redoRenders.push(lastRender)
+          })
+        }
+      },
+
+      redoDisabled: (): boolean => {
+        const editorState = get().editorState
+        if (editorState.redoRenders.length > 0) {
+          return false
+        }
+        if (get().runMannually()) {
+          if (editorState.redoCurLines.length === 0) {
+            return true
+          }
+        } else if (editorState.redoRenders.length === 0) {
+          return true
+        }
+        return false
+      },
+
+      redo: () => {
+        if (
+          get().runMannually() &&
+          get().editorState.redoCurLines.length !== 0
+        ) {
+          set((state) => {
+            const editorState = state.editorState
+            if (editorState.redoCurLines.length === 0) {
+              return
+            }
+            const line = editorState.redoCurLines.pop()!
+            editorState.curLineGroup.push(line)
+          })
+        } else {
+          set((state) => {
+            const editorState = state.editorState
+            if (
+              editorState.redoRenders.length === 0 ||
+              editorState.redoLineGroups.length === 0
+            ) {
+              return
+            }
+            const lastLineGroup = editorState.redoLineGroups.pop()!
+            editorState.lineGroups.push(lastLineGroup)
+            editorState.curLineGroup = []
+
+            const lastRender = editorState.redoRenders.pop()!
+            editorState.renders.push(lastRender)
+          })
+        }
+      },
+
+      resetRedoState: () => {
+        set((state) => {
+          state.editorState.redoCurLines = []
+          state.editorState.redoLineGroups = []
+          state.editorState.redoRenders = []
+        })
+      },
+
+      //****//
+
+      updateAppState: (newState: Partial<AppState>) => {
+        set(() => newState)
+      },
+
+      getBrushSize: (): number => {
+        return (
+          get().editorState.baseBrushSize * get().editorState.brushSizeScale
+        )
+      },
+
       showPromptInput: (): boolean => {
         const model_type = get().settings.model.model_type
         return ["diffusers_sd", "diffusers_sd_inpaint"].includes(model_type)
@@ -233,13 +443,13 @@ export const useStore = create<AppState & AppAction>()(
       },
 
       setServerConfig: (newValue: ServerConfig) => {
-          set((state: AppState) => {
+        set((state) => {
           state.serverConfig = newValue
         })
       },
 
       updateSettings: (newSettings: Partial<Settings>) => {
-          set((state: AppState) => {
+        set((state) => {
           state.settings = {
             ...state.settings,
             ...newSettings,
@@ -248,7 +458,7 @@ export const useStore = create<AppState & AppAction>()(
       },
 
       updateFileManagerState: (newState: Partial<FileManagerState>) => {
-          set((state: AppState) => {
+        set((state) => {
           state.fileManagerState = {
             ...state.fileManagerState,
             ...newState,
@@ -257,7 +467,7 @@ export const useStore = create<AppState & AppAction>()(
       },
 
       updateInteractiveSegState: (newState: Partial<InteractiveSegState>) => {
-          set((state: AppState) => {
+        set((state) => {
           state.interactiveSegState = {
             ...state.interactiveSegState,
             ...newState,
@@ -265,81 +475,88 @@ export const useStore = create<AppState & AppAction>()(
         })
       },
       resetInteractiveSegState: () => {
-          set((state: AppState) => {
+        set((state) => {
           state.interactiveSegState = defaultValues.interactiveSegState
         })
       },
 
       setIsInpainting: (newValue: boolean) =>
-          set((state: AppState) => {
+        set((state) => {
           state.isInpainting = newValue
         }),
 
       setIsPluginRunning: (newValue: boolean) =>
-          set((state: AppState) => {
+        set((state) => {
           state.isPluginRunning = newValue
         }),
 
       setFile: (file: File) =>
-          set((state: AppState) => {
+        set((state) => {
           // TODO: 清空各种状态
           state.file = file
         }),
 
       setCustomFile: (file: File) =>
-          set((state: AppState) => {
+        set((state) => {
           state.customMask = file
         }),
 
-        setBrushSize: (newValue: number) =>
+      setBaseBrushSize: (newValue: number) =>
-          set((state: AppState) => {
+        set((state) => {
-            state.brushSize = newValue
+          state.editorState.baseBrushSize = newValue
         }),
 
       setImageSize: (width: number, height: number) => {
         // 根据图片尺寸调整 brushSize 的 scale
-          set((state: AppState) => {
+        set((state) => {
           state.imageWidth = width
           state.imageHeight = height
-            state.brushSizeScale = Math.max(Math.min(width, height), 512) / 512
+          state.editorState.brushSizeScale =
+            Math.max(Math.min(width, height), 512) / 512
         })
       },
 
       setCropperX: (newValue: number) =>
-          set((state: AppState) => {
+        set((state) => {
           state.cropperState.x = newValue
         }),
 
       setCropperY: (newValue: number) =>
-          set((state: AppState) => {
+        set((state) => {
           state.cropperState.y = newValue
         }),
 
       setCropperWidth: (newValue: number) =>
-          set((state: AppState) => {
+        set((state) => {
           state.cropperState.width = newValue
         }),
 
       setCropperHeight: (newValue: number) =>
-          set((state: AppState) => {
+        set((state) => {
           state.cropperState.height = newValue
         }),
 
       setSeed: (newValue: number) =>
-          set((state: AppState) => {
+        set((state) => {
           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", "settings"].includes(key)
+            ["fileManagerState", "settings"].includes(key)
           )
         ),
     }
+  ),
+  shallow
 )
-  )
+
-)
+// export const useStore = <U>(selector: (state: AppState & AppAction) => U) => {
+//   return createWithEqualityFn(selector, shallow)
+// }
+
+// export const useStore = createWithEqualityFn(useBaseStore, shallow)

web_app/src/lib/types.ts

@@ -69,3 +69,20 @@ export interface FreeuConfig {
   b1: number
   b2: number
 }
+
+export interface Point {
+  x: number
+  y: number
+}
+
+export interface Line {
+  size?: number
+  pts: Point[]
+}
+
+export type LineGroup = Array<Line>
+
+export interface Size {
+  width: number
+  height: number
+}

web_app/src/lib/utils.ts

@@ -1,6 +1,8 @@
 import { type ClassValue, clsx } from "clsx"
 import { SyntheticEvent } from "react"
 import { twMerge } from "tailwind-merge"
+import { LineGroup } from "./types"
+import { BRUSH_COLOR } from "./const"
 
 export function cn(...inputs: ClassValue[]) {
   return twMerge(clsx(inputs))
@@ -131,3 +133,29 @@ export function downloadImage(uri: string, name: string) {
     link.remove()
   }, 100)
 }
+
+export function mouseXY(ev: SyntheticEvent) {
+  const mouseEvent = ev.nativeEvent as MouseEvent
+  return { x: mouseEvent.offsetX, y: mouseEvent.offsetY }
+}
+
+export function drawLines(
+  ctx: CanvasRenderingContext2D,
+  lines: LineGroup,
+  color = BRUSH_COLOR
+) {
+  ctx.strokeStyle = color
+  ctx.lineCap = "round"
+  ctx.lineJoin = "round"
+
+  lines.forEach((line) => {
+    if (!line?.pts.length || !line.size) {
+      return
+    }
+    ctx.lineWidth = line.size
+    ctx.beginPath()
+    ctx.moveTo(line.pts[0].x, line.pts[0].y)
+    line.pts.forEach((pt) => ctx.lineTo(pt.x, pt.y))
+    ctx.stroke()
+  })
+}

web_app/src/main.tsx

@@ -4,14 +4,14 @@ import { QueryClient, QueryClientProvider } from "@tanstack/react-query"
 import "inter-ui/inter.css"
 import App from "./App.tsx"
 import "./globals.css"
-import { ThemeProvider } from "./components/theme-provider.tsx"
+import { ThemeProvider } from "next-themes"
 
 const queryClient = new QueryClient()
 
 ReactDOM.createRoot(document.getElementById("root")!).render(
   <React.StrictMode>
     <QueryClientProvider client={queryClient}>
-      <ThemeProvider defaultTheme="dark" storageKey="vite-ui-theme">
+      <ThemeProvider defaultTheme="dark" disableTransitionOnChange>
         <App />
       </ThemeProvider>
     </QueryClientProvider>