/* eslint-disable no-use-before-define */ import { Checkbox, FormControlLabel, Slider, TextField } from '@mui/material'; import { Button, IconButton, Type } from '@dash/components'; import * as React from 'react'; import { useEffect, useRef, useState } from 'react'; import { CgClose } from 'react-icons/cg'; import { IoMdRedo, IoMdUndo } from 'react-icons/io'; import { ClientUtils } from '../../../../ClientUtils'; import { Doc, DocListCast } from '../../../../fields/Doc'; import { List } from '../../../../fields/List'; import { NumCast } from '../../../../fields/Types'; import { Networking } from '../../../Network'; import { DocUtils } from '../../../documents/DocUtils'; import { Docs } from '../../../documents/Documents'; import { CollectionDockingView } from '../../collections/CollectionDockingView'; import { CollectionFreeFormView } from '../../collections/collectionFreeForm'; import { ImageEditorData } from '../ImageBox'; import { OpenWhereMod } from '../OpenWhere'; import './ImageEditor.scss'; import { ApplyFuncButtons, ImageToolButton } from './ImageEditorButtons'; import { BrushHandler } from './imageEditorUtils/BrushHandler'; import { APISuccess, ImageUtility } from './imageEditorUtils/ImageHandler'; import { PointerHandler } from './imageEditorUtils/PointerHandler'; import { activeColor, bgColor, brushWidthOffset, canvasSize, eraserColor, freeformRenderSize, newCollectionSize, offsetDistanceY, offsetX } from './imageEditorUtils/imageEditorConstants'; import { CutMode, CursorData, ImageDimensions, ImageEditTool, ImageToolType, Point } from './imageEditorUtils/imageEditorInterfaces'; import { DocumentView } from '../DocumentView'; import { SettingsManager } from '../../../util/SettingsManager'; import { Upload } from '../../../../server/SharedMediaTypes'; interface GenerativeFillProps { imageEditorOpen: boolean; imageEditorSource: string; imageRootDoc: Doc | undefined; addDoc: ((doc: Doc | Doc[], annotationKey?: string) => boolean) | undefined; } // Added field on image doc: gen_fill_children: List of children Docs /** * The image editor interface can be accessed by opening a document's context menu, then going to Options --> Open Image Editor. * The image editor supports various operations on images. Currently, there is a Generative Fill feature that allows users to erase * part of an image, add an optional prompt, and send this to GPT. GPT then returns a newly generated image that replaces the erased * portion based on the optional prompt. There is also an image cutting tool that allows users to cut images in different ways to * reshape the images, take out portions of images, and overall use them more creatively (see the header comment for cutImage() for more information). */ const ImageEditor = ({ imageEditorOpen, imageEditorSource, imageRootDoc, addDoc }: GenerativeFillProps) => { const canvasRef = useRef(null); const canvasBackgroundRef = useRef(null); const drawingAreaRef = useRef(null); const [cursorData, setCursorData] = useState({ x: 0, y: 0, width: 150, }); const [isBrushing, setIsBrushing] = useState(false); const [canvasScale, setCanvasScale] = useState(0.5); // format: array of [image source, corresponding image Doc] const [edits, setEdits] = useState<{ url: string; saveRes: Doc | undefined }[]>([]); const [edited, setEdited] = useState(false); const [isFirstDoc, setIsFirstDoc] = useState(true); const [input, setInput] = useState(''); const [loading, setLoading] = useState(false); const [canvasDims, setCanvasDims] = useState({ width: canvasSize, height: canvasSize, }); const [cutType, setCutType] = useState(CutMode.IN); // whether to create a new collection or not const [isNewCollection, setIsNewCollection] = useState(true); // the current image in the main canvas const currImg = useRef(null); // the unedited version of each generation (parent) const originalImg = useRef(null); const originalDoc = useRef(null); // stores history of data urls const undoStack = useRef([]); // stores redo stack const redoStack = useRef([]); // references to keep track of tree structure const newCollectionRef = useRef(null); const parentDoc = useRef(null); const childrenDocs = useRef([]); // constants for image cutting const cutPts = useRef([]); /** * * @param type The new tool type we are changing to */ const changeTool = (type: ImageToolType) => { setCurrToolType(type); setCursorData(prev => ({ ...prev, width: currTool().sliderDefault as number })); }; // Undo and Redo const handleUndo = () => { const ctx = ImageUtility.getCanvasContext(canvasRef); if (!ctx || !currImg.current || !canvasRef.current) return; const target = undoStack.current[undoStack.current.length - 1]; if (!target) { ImageUtility.drawImgToCanvas(currImg.current, canvasRef, canvasDims.width, canvasDims.height); } else { redoStack.current = [...redoStack.current, canvasRef.current.toDataURL()]; const img = new Image(); img.src = target; ImageUtility.drawImgToCanvas(img, canvasRef, canvasDims.width, canvasDims.height); undoStack.current = undoStack.current.slice(0, -1); } }; const handleRedo = () => { const ctx = ImageUtility.getCanvasContext(canvasRef); if (!ctx || !currImg.current || !canvasRef.current) return; const target = redoStack.current[redoStack.current.length - 1]; if (target) { undoStack.current = [...undoStack.current, canvasRef.current?.toDataURL()]; const img = new Image(); img.src = target; ImageUtility.drawImgToCanvas(img, canvasRef, canvasDims.width, canvasDims.height); redoStack.current = redoStack.current.slice(0, -1); } }; // resets any erase strokes const handleReset = () => { if (!canvasRef.current || !currImg.current) return; const ctx = ImageUtility.getCanvasContext(canvasRef); if (!ctx) return; ctx.clearRect(0, 0, canvasSize, canvasSize); undoStack.current = []; redoStack.current = []; cutPts.current.length = 0; ImageUtility.drawImgToCanvas(currImg.current, canvasRef, canvasDims.width, canvasDims.height); }; // initiate brushing const handlePointerDown = (e: React.PointerEvent) => { const canvas = canvasRef.current; if (!canvas) return; const ctx = ImageUtility.getCanvasContext(canvasRef); if (!ctx) return; undoStack.current = [...undoStack.current, canvasRef.current.toDataURL()]; redoStack.current = []; setIsBrushing(true); const { x, y } = PointerHandler.getPointRelativeToElement(canvas, e, canvasScale); BrushHandler.brushCircleOverlay(x, y, cursorData.width / 2 / canvasScale, ctx, eraserColor /* , brushStyle === BrushStyle.SUBTRACT */); }; // stop brushing, push to undo stack const handlePointerUp = () => { const ctx = ImageUtility.getCanvasContext(canvasBackgroundRef); if (!ctx) return; if (!isBrushing) return; setIsBrushing(false); }; // handles brushing on pointer movement useEffect(() => { const canvas = canvasRef.current; if (!isBrushing || !canvas) return undefined; const ctx = ImageUtility.getCanvasContext(canvasRef); if (!ctx) return undefined; const handlePointerMove = (e: PointerEvent) => { const currPoint = PointerHandler.getPointRelativeToElement(canvas, e, canvasScale); const lastPoint: Point = { x: currPoint.x - e.movementX / canvasScale, y: currPoint.y - e.movementY / canvasScale, }; const pts = BrushHandler.createBrushPathOverlay(lastPoint, currPoint, cursorData.width / 2 / canvasScale, ctx, eraserColor); cutPts.current.push(...pts); }; drawingAreaRef.current?.addEventListener('pointermove', handlePointerMove); return () => drawingAreaRef.current?.removeEventListener('pointermove', handlePointerMove); }, [isBrushing]); // first load useEffect(() => { if (imageEditorSource && imageEditorSource) { ImageUtility.urlToBase64(imageEditorSource).then(res => { if (res) { const img = new Image(); img.src = `data:image/png;base64,${res}`; img.onload = () => { currImg.current = img; originalImg.current = img; const imgWidth = img.naturalWidth; const imgHeight = img.naturalHeight; const scale = Math.min(canvasSize / imgWidth, canvasSize / imgHeight); const width = imgWidth * scale; const height = imgHeight * scale; setCanvasDims({ width, height }); }; } }); } // cleanup return () => { setInput(''); setEdited(false); newCollectionRef.current = null; parentDoc.current = null; childrenDocs.current = []; currImg.current = null; originalImg.current = null; originalDoc.current = null; undoStack.current = []; redoStack.current = []; ImageUtility.clearCanvas(canvasRef); }; }, [canvasRef, imageEditorSource]); // once the appropriate dimensions are set, draw the image to the canvas useEffect(() => { if (!currImg.current) return; ImageUtility.drawImgToCanvas(currImg.current, canvasRef, canvasDims.width, canvasDims.height); }, [canvasDims]); // handles brush sizing useEffect(() => { const handleKeyPress = (e: KeyboardEvent) => { if (e.key === 'ArrowUp') { e.preventDefault(); e.stopPropagation(); setCursorData(data => ({ ...data, width: data.width + 5 })); } else if (e.key === 'ArrowDown') { e.preventDefault(); e.stopPropagation(); setCursorData(data => (data.width >= 20 ? { ...data, width: data.width - 5 } : data)); } }; window.addEventListener('keydown', handleKeyPress); return () => window.removeEventListener('keydown', handleKeyPress); }, []); // handle pinch zoom useEffect(() => { const handlePinch = (e: WheelEvent) => { e.preventDefault(); e.stopPropagation(); const delta = e.deltaY; const scaleFactor = delta > 0 ? 0.98 : 1.02; setCanvasScale(prevScale => prevScale * scaleFactor); }; drawingAreaRef.current?.addEventListener('wheel', handlePinch, { passive: false, }); return () => drawingAreaRef.current?.removeEventListener('wheel', handlePinch); }, [drawingAreaRef]); // updates the current position of the cursor const updateCursorData = (e: React.PointerEvent) => { const drawingArea = drawingAreaRef.current; if (!drawingArea) return; const { x, y } = PointerHandler.getPointRelativeToElement(drawingArea, e, 1); setCursorData(data => ({ ...data, x, y, })); }; // Get AI Edit for Generative Fill const getEdit = async (useFirefly = true) => { const img = currImg.current; if (!img) return; const canvas = canvasRef.current; if (!canvas) return; const ctx = ImageUtility.getCanvasContext(canvasRef); if (!ctx) return; setLoading(true); setEdited(true); try { const canvasOriginalImg = ImageUtility.getCanvasImg(img); if (!canvasOriginalImg) return; const canvasMask = ImageUtility.getCanvasMask(canvas, canvas); if (!canvasMask) return; const maskBlob = await ImageUtility.canvasToBlob(canvasMask); const imgBlob = await ImageUtility.canvasToBlob(canvasOriginalImg); let imgUrls: string[] = []; const setupCollection = async () => { // create first image if (!newCollectionRef.current) { createNewCollection(); } else { childrenDocs.current = []; } if (!(originalImg.current && imageRootDoc)) return; // add the doc to the main freeform await createNewImgDoc(originalImg.current, true); originalImg.current = currImg.current; originalDoc.current = parentDoc.current; }; if (useFirefly) { const res = await Networking.PostFormToServer(input || 'Fill in the image in the same style', imgBlob, maskBlob, img.width, img.height); if (!res.ok) throw new Error(await res.text()); const json = (await res.json()) as APISuccess; imgUrls = json.urls ?? []; } else { const res = await ImageUtility.getEdit(imgBlob, maskBlob, input || 'Fill in the image in the same style', 2); if (res.status == 'error') throw new Error(res.message); const json = res as APISuccess; imgUrls = await Promise.all((json.urls ?? []).map(url => ImageUtility.convertImgToCanvasUrl(url, canvasDims.width, canvasDims.height))); } setupCollection(); const imgRes = await Promise.all(imgUrls.map(async url => ({ url, saveRes: await onSave(url) }))); setEdits(imgRes); currImg.current = new Image(); currImg.current.src = imgUrls[0]; ImageUtility.drawImgToCanvas(currImg.current, canvasRef, canvasDims.width, canvasDims.height); parentDoc.current = imgRes[0].saveRes ?? null; } catch (err: unknown) { alert('message' in (err as object) ? (err as { message: string }).message : err); } setLoading(false); }; /** * This function performs image cutting based on the inputted BrushMode. There are currently four ways to cut images: * 1. By outlining the area that should be kept (BrushMode.IN) * 2. By outlining the area that should be removed (BrushMode.OUT) * 3. By drawing in the area that should be kept (where the image is brushed, the image will remain and everything else will be removed) (BrushMode.DRAW_IN) * 4. By drawing the area that she be removed, so this operates as an eraser (BrushMode.ERASE) * @param currCutType BrushMode enum that determines what kind of cutting operation to perform * @param firstDoc boolean for whether it's the first edited image. This is for positioning of the edited images when they render on the canvas. */ const cutImage = async (currCutType: CutMode, brushWidth: number, prevEdits: { url: string; saveRes: Doc | undefined }[], firstDoc: boolean) => { const img = currImg.current; const canvas = canvasRef.current; if (!canvas || !img) return; const ctx = ImageUtility.getCanvasContext(canvasRef); if (!ctx) return; // get the original image const canvasOriginalImg = ImageUtility.getCanvasImg(img); if (!canvasOriginalImg) return; setLoading(true); const currPts = [...cutPts.current]; if (currCutType !== CutMode.ERASE) handleReset(); // gets rid of the visible brush strokes (mostly needed for line_in) unless it's erasing (which depends on the brush strokes) let minX = img.width; let maxX = 0; let minY = img.height; let maxY = 0; // currPts is populated by the brush strokes' points, so this code is drawing a path along the points if (currPts.length) { ctx.beginPath(); ctx.moveTo(currPts[0].x, currPts[0].y); for (let i = 0; i < currPts.length; i++) { ctx.lineTo(currPts[i].x, currPts[i].y); minX = Math.min(currPts[i].x, minX); minY = Math.min(currPts[i].y, minY); maxX = Math.max(currPts[i].x, maxX); maxY = Math.max(currPts[i].y, maxY); } switch ( currCutType // use different canvas operations depending on the type of cutting we're applying ) { case CutMode.IN: ctx.closePath(); ctx.globalCompositeOperation = 'destination-in'; ctx.fill(); break; case CutMode.OUT: ctx.closePath(); ctx.globalCompositeOperation = 'destination-out'; ctx.fill(); break; case CutMode.DRAW_IN: ctx.globalCompositeOperation = 'destination-in'; ctx.lineWidth = brushWidth + brushWidthOffset; // added offset because width gets cut off a little bit ctx.stroke(); break; } } const url = canvas.toDataURL(); if (!newCollectionRef.current) { createNewCollection(); } const image = new Image(); image.src = url; image.onload = async () => { let finalImg: HTMLImageElement | undefined = image; let finalImgURL: string = url; // crop the image for these brush modes to remove excess blank space around the image contents if (currCutType == CutMode.IN || currCutType == CutMode.DRAW_IN) { const croppedData = cropImage(image, Math.max(minX, 0), Math.min(maxX, image.width), Math.max(minY, 0), Math.min(maxY, image.height)); finalImg = croppedData; finalImgURL = croppedData.src; } currImg.current = finalImg; const newImgDoc = await createNewImgDoc(finalImg, firstDoc); if (newImgDoc) { // set the image to transparent to remove the background / brushstrokes newImgDoc.$backgroundColor = 'transparent'; newImgDoc.$disableMixBlend = true; if (firstDoc) setIsFirstDoc(false); setEdits([...prevEdits, { url: finalImgURL, saveRes: undefined }]); } setLoading(false); cutPts.current.length = 0; }; }; /** * Creates a new collection to put the image edits on. Adds to a new tab on the right if "Create New Collection" is checked. * @returns */ const createNewCollection = () => { if (!isNewCollection && imageRootDoc) { // if the parent hasn't been set yet if (!parentDoc.current) parentDoc.current = imageRootDoc; } else { if (!(originalImg.current && imageRootDoc)) return; // create new collection and add it to the view newCollectionRef.current = Docs.Create.FreeformDocument([], { x: NumCast(imageRootDoc.x) + NumCast(imageRootDoc._width) + offsetX, y: NumCast(imageRootDoc.y), _width: newCollectionSize, _height: newCollectionSize, title: 'Image edit collection', }); DocUtils.MakeLink(imageRootDoc, newCollectionRef.current, { link_relationship: 'Image Edit Version History' }); // opening new tab CollectionDockingView.AddSplit(newCollectionRef.current, OpenWhereMod.right); } }; /** * This function crops an image based on the inputted dimensions. This is used to automatically adjust the images that are * edited to be smaller than the original (i.e. for cutting into a small part of the image) */ const cropImage = (image: HTMLImageElement, minX: number, maxX: number, minY: number, maxY: number) => { const croppedCanvas = document.createElement('canvas'); const croppedCtx = croppedCanvas.getContext('2d'); if (!croppedCtx) return image; const cropWidth = Math.abs(maxX - minX); const cropHeight = Math.abs(maxY - minY); croppedCanvas.width = cropWidth; croppedCanvas.height = cropHeight; croppedCtx.globalCompositeOperation = 'source-over'; croppedCtx.clearRect(0, 0, cropWidth, cropHeight); croppedCtx.drawImage(image, minX, minY, cropWidth, cropHeight, 0, 0, cropWidth, cropHeight); const croppedURL = croppedCanvas.toDataURL(); const croppedImage = new Image(); croppedImage.src = croppedURL; return croppedImage; }; // adjusts all the img positions to be aligned const adjustImgPositions = () => { if (!parentDoc.current) return; const startY = NumCast(parentDoc.current.y); const children = DocListCast(parentDoc.current.gen_fill_children); const len = children.length; const initialYPositions: number[] = []; for (let i = 0; i < len; i++) { initialYPositions.push(startY + i * offsetDistanceY); } children.forEach((doc, i) => { if (len % 2 === 1) { doc.y = initialYPositions[i] - Math.floor(len / 2) * offsetDistanceY; } else { doc.y = initialYPositions[i] - (len / 2 - 1 / 2) * offsetDistanceY; } }); }; // creates a new image document and returns its reference const createNewImgDoc = async (img: HTMLImageElement, firstDoc: boolean /*, parent?: Doc */): Promise => { if (!imageRootDoc) return undefined; const { src } = img; const [result] = (await Networking.PostToServer('/uploadRemoteImage', { sources: [src] })) as Upload.ImageInformation[]; const source = ClientUtils.prepend(result.accessPaths.agnostic.client); if (firstDoc) { const x = 0; const initialY = 0; const newImg = Docs.Create.ImageDocument(source, { x: x, y: initialY, _height: freeformRenderSize, _width: freeformRenderSize, data_nativeWidth: result.nativeWidth, data_nativeHeight: result.nativeHeight, }); if (isNewCollection && newCollectionRef.current) { Doc.AddDocToList(newCollectionRef.current, undefined, newImg); } else { addDoc?.(newImg); } parentDoc.current = newImg; return newImg; } if (!parentDoc.current) return undefined; const x = NumCast(parentDoc.current.x) + freeformRenderSize + offsetX; const initialY = 0; const newImg = Docs.Create.ImageDocument(source, { x: x, y: initialY, _height: freeformRenderSize, _width: freeformRenderSize, data_nativeWidth: result.nativeWidth, data_nativeHeight: result.nativeHeight, }); const parentList = DocListCast(parentDoc.current.gen_fill_children); if (parentList.length > 0) { parentList.push(newImg); parentDoc.current.gen_fill_children = new List(parentList); } else { parentDoc.current.gen_fill_children = new List([newImg]); } DocUtils.MakeLink(parentDoc.current, newImg, { link_relationship: `Image edit; Prompt: ${input}` }); adjustImgPositions(); if (isNewCollection && newCollectionRef.current) { Doc.AddDocToList(newCollectionRef.current, undefined, newImg); } else { addDoc?.(newImg); } return newImg; }; // Saves an image to the collection const onSave = async (src: string) => { const img = new Image(); img.src = src; if (!currImg.current || !originalImg.current || !imageRootDoc) return undefined; try { return await createNewImgDoc(img, false); } catch (err) { console.log(err); } return undefined; }; // Closes the editor view const handleViewClose = () => { ImageEditorData.Open = false; ImageEditorData.Source = ''; if (newCollectionRef.current) { DocumentView.addViewRenderedCb(newCollectionRef.current, dv => (dv.ComponentView as CollectionFreeFormView)?.fitContentOnce()); } setEdits([]); setIsFirstDoc(true); }; function currTool() { return imageEditTools.find(tool => tool.type === currToolType) ?? genFillTool; } // defines the tools and sets current tool const genFillTool: ImageEditTool = { type: ImageToolType.GenerativeFill, btnText: 'GET EDITS', icon: 'fill', applyFunc: () => getEdit(), sliderMin: 25, sliderMax: 500, sliderDefault: 150 }; const cutTool: ImageEditTool = { type: ImageToolType.Cut, btnText: 'CUT IMAGE', icon: 'scissors', applyFunc: cutImage, sliderMin: 1, sliderMax: 50, sliderDefault: 5 }; const imageEditTools: ImageEditTool[] = [genFillTool, cutTool]; const [currToolType, setCurrToolType] = useState(ImageToolType.GenerativeFill); // the top controls for making a new collection, resetting, and applying edits, function renderControls() { return (

Image Editor

{/* } /> */}
setIsNewCollection(prev => !prev)} /> } label="Create New Collection" labelPlacement="end" sx={{ whiteSpace: 'nowrap' }} /> currTool().applyFunc(cutType, cursorData.width, edits, isFirstDoc)} loading={loading} onReset={handleReset} btnText={currTool().btnText} /> } onClick={handleViewClose} />
); } // the side icons including tool type, the slider, and undo/redo function renderSideIcons() { return (
{imageEditTools.map(tool => ImageToolButton(tool, tool.type === currTool().type, changeTool))}
{currTool().type == ImageToolType.Cut && (
)}
e.stopPropagation()}> {currTool().type === ImageToolType.GenerativeFill && ( setCursorData(prev => ({ ...prev, width: val as number }))} /> )} {currTool().type === ImageToolType.Cut && ( setCursorData(prev => ({ ...prev, width: val as number }))} /> )}
{/* Undo and Redo */}
{ e.stopPropagation(); handleUndo(); }} onPointerUp={e => e.stopPropagation()} color={activeColor} tooltip="Undo" icon={} /> { e.stopPropagation(); handleRedo(); }} onPointerUp={e => e.stopPropagation()} color={activeColor} tooltip="Redo" icon={} />
); } // circular pointer for drawing/erasing function renderPointer() { return (
); } // the previews for each edit function renderEditThumbnails() { return (
{edits.map(edit => ( image edits { const img = new Image(); img.src = edit.url; ImageUtility.drawImgToCanvas(img, canvasRef, img.width, img.height); currImg.current = img; parentDoc.current = edit.saveRes ?? null; }} /> ))} {/* Original img thumbnail */} {edits.length > 0 && (
image stuff { if (!originalImg.current) return; const img = new Image(); img.src = originalImg.current.src; ImageUtility.drawImgToCanvas(img, canvasRef, canvasDims.width, canvasDims.height); currImg.current = img; if (!parentDoc.current) parentDoc.current = originalDoc.current; }} />
)}
); } // the prompt box for generative fill function renderPromptBox() { return (
setInput(e.target.value)} disabled={isBrushing} type="text" label="Prompt" placeholder="Prompt..." InputLabelProps={{ style: { fontSize: '16px' } }} inputProps={{ style: { fontSize: '16px' } }} sx={{ backgroundColor: '#ffffff', position: 'absolute', bottom: '16px', transform: 'translateX(calc(50vw - 50%))', width: 'calc(100vw - 64px)', }} />
); } return (
{renderControls()} {/* Main canvas for editing */}
{renderPointer()} {renderSideIcons()} {renderEditThumbnails()}
{currTool().type === ImageToolType.GenerativeFill && renderPromptBox()}
); }; export default ImageEditor;