1 files changed, 559 insertions, 0 deletions
diff --git a/src/client/views/nodes/generativeFill/GenerativeFill.tsx b/src/client/views/nodes/generativeFill/GenerativeFill.tsx
new file mode 100644
index 000000000..f8f9fe077
--- /dev/null
+++ b/src/client/views/nodes/generativeFill/GenerativeFill.tsx
@@ -0,0 +1,559 @@
+import './GenerativeFill.scss';
+import React = require('react');
+import { useEffect, useRef, useState } from 'react';
+import { APISuccess, ImageUtility } from './generativeFillUtils/ImageHandler';
+import { BrushHandler } from './generativeFillUtils/BrushHandler';
+import { Box, IconButton, Slider, TextField } from '@mui/material';
+import { CursorData, Point } from './generativeFillUtils/generativeFillInterfaces';
+import { activeColor, canvasSize, eraserColor, freeformRenderSize, newCollectionSize, offsetDistanceY, offsetX } from './generativeFillUtils/generativeFillConstants';
+import { PointerHandler } from './generativeFillUtils/PointerHandler';
+import { BsEraser, BsX } from 'react-icons/bs';
+import { CiUndo, CiRedo } from 'react-icons/ci';
+import Buttons from './GenerativeFillButtons';
+import { MainView } from '../../MainView';
+import { Doc } from '../../../../fields/Doc';
+import { Networking } from '../../../Network';
+import { Utils } from '../../../../Utils';
+import { DocUtils, Docs } from '../../../documents/Documents';
+import { Cast, DocCast, NumCast } from '../../../../fields/Types';
+import { CollectionDockingView } from '../../collections/CollectionDockingView';
+import { OpenWhere, OpenWhereMod } from '../DocumentView';
+import { Oval } from 'react-loader-spinner';
+
+/**
+ * For images not 1024x1024 fill in the rest in solid black, or a
+ * reflected version of the image.
+ *
+ * add a branch from image directly checkbox
+ */
+
+/**
+ *
+ *
+ * CollectionDockingView.AddSplit(Doc.MakeCopy(DocCast(Doc.UserDoc().emptyPane)), OpenWhereMod.right);
+ * CollectionDockingView.AddSplit(newCollection,OpenWhere.inParent)
+ * mind mapping
+ * this.props.addDocTab();
+ */
+
+enum BrushStyle {
+    ADD,
+    SUBTRACT,
+    MARQUEE,
+}
+
+interface ImageEdit {
+    imgElement: HTMLImageElement;
+    parent: ImageEdit;
+    children: ImageEdit[];
+}
+
+interface GenerativeFillProps {
+    imageEditorOpen: boolean;
+    imageEditorSource: string;
+    imageRootDoc: Doc | undefined;
+    addDoc: ((doc: Doc | Doc[], annotationKey?: string) => boolean) | undefined;
+}
+
+const GenerativeFill = ({ imageEditorOpen, imageEditorSource, imageRootDoc, addDoc }: GenerativeFillProps) => {
+    const canvasRef = useRef<HTMLCanvasElement>(null);
+    const canvasBackgroundRef = useRef<HTMLCanvasElement>(null);
+    const drawingAreaRef = useRef<HTMLDivElement>(null);
+    const fileRef = useRef<HTMLInputElement>(null);
+    const [cursorData, setCursorData] = useState<CursorData>({
+        x: 0,
+        y: 0,
+        width: 150,
+    });
+    const [isBrushing, setIsBrushing] = useState(false);
+    const [canvasScale, setCanvasScale] = useState(0.5);
+    const [edits, setEdits] = useState<string[]>([]);
+    const [brushStyle, setBrushStyle] = useState<BrushStyle>(BrushStyle.ADD);
+    const [input, setInput] = useState('');
+    const [loading, setLoading] = useState(false);
+    const [saveLoading, setSaveLoading] = useState(false);
+    // the current image in the main canvas
+    const currImg = useRef<HTMLImageElement | null>(null);
+    // the unedited version of each generation (parent)
+    const originalImg = useRef<HTMLImageElement | null>(null);
+    // stores history of data urls
+    const undoStack = useRef<string[]>([]);
+    // stores redo stack
+    const redoStack = useRef<string[]>([]);
+
+    // early stage properly, likely will get rid of
+    const freeformPosition = useRef<number[]>([0, 0]);
+
+    // which urls were already saved to canvas
+    const savedSrcs = useRef<Set<string>>(new Set());
+
+    // references to keep track of tree structure
+    const newCollectionRef = useRef<Doc | null>(null);
+    const parentDoc = useRef<Doc | null>(null);
+    const childrenDocs = useRef<Doc[]>([]);
+    const addToExistingCollection = useRef<boolean>(false);
+
+    // 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);
+        } else {
+            redoStack.current = [...redoStack.current, canvasRef.current.toDataURL()];
+            const img = new Image();
+            img.src = target;
+            ImageUtility.drawImgToCanvas(img, canvasRef);
+            undoStack.current = undoStack.current.slice(0, -1);
+        }
+    };
+
+    const handleRedo = () => {
+        // TODO: handle undo as well
+        const target = redoStack.current[redoStack.current.length - 1];
+        if (!target) {
+        } else {
+            const img = new Image();
+            img.src = target;
+            ImageUtility.drawImgToCanvas(img, canvasRef);
+            redoStack.current = redoStack.current.slice(0, -1);
+        }
+    };
+
+    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 = [];
+        ImageUtility.drawImgToCanvas(currImg.current, canvasRef, true);
+    };
+
+    // 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 = (e: React.PointerEvent) => {
+        const ctx = ImageUtility.getCanvasContext(canvasBackgroundRef);
+        if (!ctx) return;
+        if (!isBrushing) return;
+        setIsBrushing(false);
+    };
+
+    // handles brushing on pointer movement
+    useEffect(() => {
+        if (!isBrushing) return;
+        const canvas = canvasRef.current;
+        if (!canvas) return;
+        const ctx = ImageUtility.getCanvasContext(canvasRef);
+        if (!ctx) return;
+
+        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,
+            };
+            BrushHandler.createBrushPathOverlay(lastPoint, currPoint, cursorData.width / 2 / canvasScale, ctx, eraserColor, brushStyle === BrushStyle.SUBTRACT);
+        };
+
+        drawingAreaRef.current?.addEventListener('pointermove', handlePointerMove);
+        return () => {
+            drawingAreaRef.current?.removeEventListener('pointermove', handlePointerMove);
+        };
+    }, [isBrushing]);
+
+    // first load
+    useEffect(() => {
+        console.log('first load');
+        if (!imageEditorSource || imageEditorSource === '') return;
+        const img = new Image();
+        img.src = imageEditorSource;
+        ImageUtility.drawImgToCanvas(img, canvasRef);
+        currImg.current = img;
+        originalImg.current = img;
+        freeformPosition.current = [0, 0];
+
+        return () => {
+            console.log('cleanup');
+            newCollectionRef.current = null;
+            parentDoc.current = null;
+            childrenDocs.current = [];
+            currImg.current = null;
+            originalImg.current = null;
+            freeformPosition.current = [0, 0];
+            undoStack.current = [];
+            redoStack.current = [];
+        };
+    }, [canvasRef, imageEditorSource]);
+
+    // 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; // Adjust the scale factor as per your requirement
+            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
+    const getEdit = async () => {
+        const img = currImg.current;
+        if (!img) return;
+        const canvas = canvasRef.current;
+        if (!canvas) return;
+        const ctx = ImageUtility.getCanvasContext(canvasRef);
+        if (!ctx) return;
+        setLoading(true);
+        // need to adjust later
+        try {
+            const maskBlob = await ImageUtility.canvasToBlob(canvas);
+            const imgBlob = await ImageUtility.canvasToBlob(ImageUtility.getCanvasImg(img));
+            const res = await ImageUtility.getEdit(imgBlob, maskBlob, input !== '' ? input + ' in the same style' : 'Fill in the image in the same style', 2);
+            // const res = await ImageUtility.mockGetEdit(img.src);
+
+            // create first image
+            if (!newCollectionRef.current) {
+                if (addToExistingCollection.current) {
+                }
+                if (!(originalImg.current && imageRootDoc)) return;
+                console.log('creating first image');
+                // 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', link_displayLine: false });
+                // add the doc to the main freeform
+                // addDoc?.(newCollectionRef.current);
+                await createNewImgDoc(originalImg.current, true);
+            } else {
+                parentDoc.current = childrenDocs.current[childrenDocs.current.length - 1];
+                childrenDocs.current = [];
+            }
+
+            originalImg.current = currImg.current;
+
+            const { urls } = res as APISuccess;
+            const image = new Image();
+            image.src = urls[0];
+            setEdits(urls);
+            ImageUtility.drawImgToCanvas(image, canvasRef);
+            currImg.current = image;
+            onSave();
+            freeformPosition.current[0] += 1;
+            freeformPosition.current[1] = 0;
+        } catch (err) {
+            console.log(err);
+        }
+        setLoading(false);
+    };
+
+    // adjusts all the img positions to be aligned
+    const adjustImgPositions = () => {
+        if (!parentDoc.current) return;
+        const startY = NumCast(parentDoc.current.y);
+        const len = childrenDocs.current.length;
+        let initialYPositions: number[] = [];
+        for (let i = 0; i < len; i++) {
+            initialYPositions.push(startY + i * offsetDistanceY);
+        }
+        childrenDocs.current.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): Promise<Doc | undefined> => {
+        if (!newCollectionRef.current || !imageRootDoc) return;
+        const src = img.src;
+        const [result] = await Networking.PostToServer('/uploadRemoteImage', { sources: [src] });
+        const source = Utils.prepend(result.accessPaths.agnostic.client);
+
+        if (firstDoc) {
+            const x = 0;
+            const initialY = 0;
+            console.log('first doc');
+
+            const newImg = Docs.Create.ImageDocument(source, {
+                x: x,
+                y: initialY,
+                _height: freeformRenderSize,
+                _width: freeformRenderSize,
+                data_nativeWidth: result.nativeWidth,
+                data_nativeHeight: result.nativeHeight,
+            });
+
+            Doc.AddDocToList(newCollectionRef.current, undefined, newImg);
+            parentDoc.current = newImg;
+            return newImg;
+        } else {
+            if (!parentDoc.current) return;
+            const x = NumCast(parentDoc.current.x) + freeformRenderSize + offsetX;
+            // dummy position
+            console.log('creating child elements');
+            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,
+            });
+
+            childrenDocs.current.push(newImg);
+            DocUtils.MakeLink(parentDoc.current, newImg, { link_relationship: 'Image Edit', link_displayLine: true });
+            adjustImgPositions();
+
+            Doc.AddDocToList(newCollectionRef.current, undefined, newImg);
+            return newImg;
+        }
+    };
+
+    // need to maybe call on every img click, not just when the save btn is clicked
+    const onSave = async () => {
+        setSaveLoading(true);
+        if (!currImg.current || !originalImg.current || !imageRootDoc) return;
+        try {
+            console.log('creating another image');
+            await createNewImgDoc(currImg.current, false);
+        } catch (err) {
+            console.log(err);
+        }
+        setSaveLoading(false);
+    };
+
+    const handleViewClose = () => {
+        if (newCollectionRef.current) {
+            newCollectionRef.current.fitContentOnce = true;
+            CollectionDockingView.AddSplit(newCollectionRef.current, OpenWhereMod.right);
+        }
+        MainView.Instance.setImageEditorOpen(false);
+        MainView.Instance.setImageEditorSource('');
+        setEdits([]);
+    };
+
+    return (
+        <div className="generativeFillContainer" style={{ display: imageEditorOpen ? 'flex' : 'none' }}>
+            <div className="generativeFillControls">
+                <h1>AI Image Editor</h1>
+                <div style={{ display: 'flex', alignItems: 'center', gap: '1.5rem' }}>
+                    <Buttons canvasRef={canvasRef} currImg={currImg} getEdit={getEdit} loading={loading} onSave={onSave} onReset={handleReset} />
+                    <IconButton onClick={handleViewClose}>
+                        <BsX color={activeColor} />
+                    </IconButton>
+                    {saveLoading && (
+                        <span style={{ height: '100%', display: 'flex', alignItems: 'center', gap: '8px' }}>
+                            Saving image... <Oval height={20} width={20} color="#000000" visible={true} ariaLabel="oval-loading" secondaryColor="#ffffff89" strokeWidth={3} strokeWidthSecondary={3} />
+                        </span>
+                    )}
+                </div>
+            </div>
+            {/* Main canvas for editing */}
+            <div
+                className="drawingArea" // this only works if pointerevents: none is set on the custom pointer
+                ref={drawingAreaRef}
+                onPointerOver={updateCursorData}
+                onPointerMove={updateCursorData}
+                onPointerDown={handlePointerDown}
+                onPointerUp={handlePointerUp}>
+                <canvas ref={canvasRef} width={canvasSize} height={canvasSize} style={{ transform: `scale(${canvasScale})` }} />
+                <canvas ref={canvasBackgroundRef} width={canvasSize} height={canvasSize} style={{ transform: `scale(${canvasScale})` }} />
+                <div
+                    className="pointer"
+                    style={{
+                        left: cursorData.x,
+                        top: cursorData.y,
+                        width: cursorData.width,
+                        height: cursorData.width,
+                    }}>
+                    <div className="innerPointer"></div>
+                </div>
+                {/* Icons */}
+                <div className="iconContainer">
+                    <IconButton
+                        onClick={() => {
+                            setBrushStyle(BrushStyle.ADD);
+                        }}>
+                        <BsEraser color={brushStyle === BrushStyle.ADD ? activeColor : 'inherit'} />
+                    </IconButton>
+                    {/* Undo and Redo */}
+                    <IconButton
+                        onPointerDown={e => {
+                            e.stopPropagation();
+                            handleUndo();
+                        }}
+                        onPointerUp={e => {
+                            e.stopPropagation();
+                        }}>
+                        <CiUndo />
+                    </IconButton>
+                    <IconButton
+                        onPointerDown={e => {
+                            e.stopPropagation();
+                            handleRedo();
+                        }}
+                        onPointerUp={e => {
+                            e.stopPropagation();
+                        }}>
+                        <CiRedo />
+                    </IconButton>
+                    <Box
+                        sx={{
+                            height: 225,
+                            width: '100%',
+                            display: 'flex',
+                            justifyContent: 'center',
+                        }}>
+                        <Slider
+                            sx={{
+                                '& input[type="range"]': {
+                                    WebkitAppearance: 'slider-vertical',
+                                },
+                            }}
+                            orientation="vertical"
+                            min={10}
+                            max={500}
+                            defaultValue={150}
+                            onChange={(e, val) => {
+                                setCursorData(prev => ({ ...prev, width: val as number }));
+                            }}
+                        />
+                    </Box>
+                </div>
+                {/* Edits thumbnails*/}
+                <div className="editsBox">
+                    {edits.map((edit, i) => (
+                        <img
+                            key={i}
+                            width={100}
+                            height={100}
+                            src={edit}
+                            onClick={async () => {
+                                // if (savedSrcs.current.has(edit)) return;
+                                const img = new Image();
+                                img.src = edit;
+                                ImageUtility.drawImgToCanvas(img, canvasRef);
+                                currImg.current = img;
+                                savedSrcs.current.add(edit);
+                                await onSave();
+                            }}
+                        />
+                    ))}
+                    {/* Original img thumbnail */}
+                    {edits.length > 0 && (
+                        <div style={{ position: 'relative' }}>
+                            <label
+                                style={{
+                                    position: 'absolute',
+                                    bottom: 10,
+                                    left: 10,
+                                    color: '#ffffff',
+                                    fontSize: '0.8rem',
+                                    letterSpacing: '1px',
+                                    textTransform: 'uppercase',
+                                }}>
+                                Original
+                            </label>
+                            <img
+                                width={100}
+                                height={100}
+                                src={originalImg.current?.src}
+                                style={{ cursor: 'pointer' }}
+                                onClick={() => {
+                                    if (!originalImg.current) return;
+                                    const img = new Image();
+                                    img.src = originalImg.current.src;
+                                    ImageUtility.drawImgToCanvas(img, canvasRef);
+                                    currImg.current = img;
+                                }}
+                            />
+                        </div>
+                    )}
+                </div>
+            </div>
+            <div>
+                <TextField
+                    value={input}
+                    onChange={e => 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)',
+                    }}
+                />
+            </div>
+        </div>
+    );
+};
+
+export default GenerativeFill;