1 files changed, 19 insertions, 16 deletions
diff --git a/src/client/views/smartdraw/SmartDrawHandler.tsx b/src/client/views/smartdraw/SmartDrawHandler.tsx
index 0c67c7a13..13a93367c 100644
--- a/src/client/views/smartdraw/SmartDrawHandler.tsx
+++ b/src/client/views/smartdraw/SmartDrawHandler.tsx
@@ -267,20 +267,23 @@ export class SmartDrawHandler extends ObservableReactComponent<object> {
     /**
      * Calls Firefly API to create an image based on user input
      */
-    createImageWithFirefly = (input: string, seed?: number) => {
+    createImageWithFirefly = (input: string, seed?: number, changeInPlace?: boolean) => {
         this._lastInput.text = input;
         const dims = FireflyDimensionsMap[this._imgDims];
         return Networking.PostToServer('/queryFireflyImage', { prompt: input, width: dims.width, height: dims.height, seed: seed }).then(img => {
-            const imgDoc: Doc = Docs.Create.ImageDocument(img.accessPaths.agnostic.client, {
-                title: input.match(/^(.*?)~~~.*$/)?.[1] || input,
-                nativeWidth: dims.width,
-                nativeHeight: dims.height,
-                ai: 'firefly',
-                ai_firefly_seed: img.accessPaths.agnostic.client.match(/\/(\d+)upload/)[1],
-                ai_firefly_prompt: input,
-            });
-            DocumentViewInternal.addDocTabFunc(imgDoc, OpenWhere.addRight);
-            this._selectedDocs.push(imgDoc);
+            if (!changeInPlace) {
+                const imgDoc: Doc = Docs.Create.ImageDocument(img.accessPaths.agnostic.client, {
+                    title: input.match(/^(.*?)~~~.*$/)?.[1] || input,
+                    nativeWidth: dims.width,
+                    nativeHeight: dims.height,
+                    ai: 'firefly',
+                    ai_firefly_seed: img.accessPaths.agnostic.client.match(/\/(\d+)upload/)[1],
+                    ai_firefly_prompt: input,
+                });
+                DocumentViewInternal.addDocTabFunc(imgDoc, OpenWhere.addRight);
+                this._selectedDocs.push(imgDoc);
+            }
+            return img.accessPaths.agnostic.client;
         });
     };
 
@@ -289,22 +292,22 @@ export class SmartDrawHandler extends ObservableReactComponent<object> {
      * @param doc the drawing Docs to regenerate
      */
     @action
-    regenerate = async (drawingDocs: Doc[], lastInput?: DrawingOptions, lastResponse?: string, regenInput?: string) => {
+    regenerate = async (drawingDocs: Doc[], lastInput?: DrawingOptions, lastResponse?: string, regenInput?: string, changeInPlace?: boolean) => {
         if (lastInput) this._lastInput = lastInput;
         if (lastResponse) this._lastResponse = lastResponse;
         if (regenInput) this._regenInput = regenInput;
-        await Promise.all(
+        return await Promise.all(
             drawingDocs.map(async doc => {
                 const docData = doc[DocData];
                 if (docData.type == 'image') {
                     const seed: number = docData?.ai_firefly_seed as number;
                     if (this._regenInput !== '') {
                         // if (this._selectedDoc) {
-                        const newPrompt = `${docData.ai_firefly_prompt}, ${this._regenInput}`;
-                        await this.createImageWithFirefly(newPrompt, seed);
+                        const newPrompt = `${docData.ai_firefly_prompt} ~~~ ${this._regenInput}`;
+                        return this.createImageWithFirefly(newPrompt, seed, changeInPlace);
                         // }
                     } else {
-                        await this.createImageWithFirefly(this._lastInput.text || StrCast(docData.ai_firefly_prompt));
+                        return this.createImageWithFirefly(this._lastInput.text || StrCast(docData.ai_firefly_prompt), undefined, changeInPlace);
                     }
                 }
                 if (docData.type == 'collection') {