switched to use firefly for image fill

1 files changed, 35 insertions, 34 deletions

diff --git a/src/client/views/nodes/imageEditor/ImageEditor.tsx b/src/client/views/nodes/imageEditor/ImageEditor.tsx
index abe235ad5..b56490bc3 100644
--- a/src/client/views/nodes/imageEditor/ImageEditor.tsx
+++ b/src/client/views/nodes/imageEditor/ImageEditor.tsx
@@ -269,7 +269,7 @@ const ImageEditor = ({ imageEditorOpen, imageEditorSource, imageRootDoc, addDoc
     };
 
     // Get AI Edit for Generative Fill
-    const getEdit = async () => {
+    const getEdit = async (useFirefly = true) => {
         const img = currImg.current;
         if (!img) return;
         const canvas = canvasRef.current;
@@ -285,40 +285,41 @@ const ImageEditor = ({ imageEditorOpen, imageEditorSource, imageRootDoc, addDoc
             if (!canvasMask) return;
             const maskBlob = await ImageUtility.canvasToBlob(canvasMask);
             const imgBlob = await ImageUtility.canvasToBlob(canvasOriginalImg);
-            const res = await ImageUtility.getEdit(imgBlob, maskBlob, input || 'Fill in the image in the same style', 2);
-            if (res.status == 'error') {
-                alert(res.message);
-            }
-
-            // create first image
-            if (!newCollectionRef.current) {
-                createNewCollection();
+            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 {
-                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;
-            const { urls } = res as APISuccess;
-            if (res.status !== 'error') {
-                const imgUrls = await Promise.all(urls.map(url => ImageUtility.convertImgToCanvasUrl(url, canvasDims.width, canvasDims.height)));
-                const imgRes = await Promise.all(
-                    imgUrls.map(async url => {
-                        const saveRes = await onSave(url);
-                        return { url, saveRes };
-                    })
-                );
-                setEdits(imgRes);
-                const image = new Image();
-                image.src = imgUrls[0];
-                ImageUtility.drawImgToCanvas(image, canvasRef, canvasDims.width, canvasDims.height);
-                currImg.current = image;
-                parentDoc.current = imgRes[0].saveRes ?? null;
+                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)));
             }
-        } catch (err) {
-            console.log(err);
+
+            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);
     };
@@ -561,7 +562,7 @@ const ImageEditor = ({ imageEditorOpen, imageEditorSource, imageRootDoc, addDoc
     }
 
     // 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 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>(ImageToolType.GenerativeFill);