1 files changed, 18 insertions, 14 deletions
diff --git a/src/client/views/nodes/generativeFill/GenerativeFill.tsx b/src/client/views/nodes/generativeFill/GenerativeFill.tsx
index 1ec6d6e3f..3093287e9 100644
--- a/src/client/views/nodes/generativeFill/GenerativeFill.tsx
+++ b/src/client/views/nodes/generativeFill/GenerativeFill.tsx
@@ -20,6 +20,8 @@ import { CursorData, ImageDimensions, Point } from './generativeFillUtils/genera
 import { APISuccess, ImageUtility } from './generativeFillUtils/ImageHandler';
 import { PointerHandler } from './generativeFillUtils/PointerHandler';
 import React = require('react');
+import { DocumentManager } from '../../../util/DocumentManager';
+import { CollectionFreeFormView } from '../../collections/collectionFreeForm';
 
 enum BrushStyle {
     ADD,
@@ -303,19 +305,21 @@ const GenerativeFill = ({ imageEditorOpen, imageEditorSource, imageRootDoc, addD
             originalImg.current = currImg.current;
             originalDoc.current = parentDoc.current;
             const { urls } = res as APISuccess;
-            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 as Doc];
-                })
-            );
-            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][1] as Doc;
+            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 as Doc];
+                    })
+                );
+                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][1] as Doc;
+            }
         } catch (err) {
             console.log(err);
         }
@@ -418,7 +422,7 @@ const GenerativeFill = ({ imageEditorOpen, imageEditorSource, imageRootDoc, addD
         ImageBox.setImageEditorOpen(false);
         ImageBox.setImageEditorSource('');
         if (newCollectionRef.current) {
-            newCollectionRef.current.fitContentOnce = true;
+            DocumentManager.Instance.AddViewRenderedCb(newCollectionRef.current, dv => (dv.ComponentView as CollectionFreeFormView)?.fitContentOnce());
         }
         setEdits([]);
     };