1 files changed, 15 insertions, 9 deletions
diff --git a/src/client/views/nodes/ImageBox.tsx b/src/client/views/nodes/ImageBox.tsx
index 09daea6bb..a1fa9a283 100644
--- a/src/client/views/nodes/ImageBox.tsx
+++ b/src/client/views/nodes/ImageBox.tsx
@@ -45,6 +45,7 @@ import { Button } from 'browndash-components';
 import { SettingsManager } from '../../util/SettingsManager';
 import { AiOutlineSend } from 'react-icons/ai';
 import { returnEmptyDocViewList } from '../StyleProvider';
+import { FireflyImageData } from '../smartdraw/FireflyConstants';
 
 export class ImageEditorData {
     // eslint-disable-next-line no-use-before-define
@@ -533,7 +534,7 @@ export class ImageBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
     @observable private _regenInput = '';
     @observable private _canInteract = true;
     @observable private _regenerateLoading = false;
-    @observable private _prevImgs: { href: string; pathname: string }[] = [];
+    @observable private _prevImgs: FireflyImageData[] = [];
 
     componentAIViewHistory = () => {
         return (
@@ -544,6 +545,8 @@ export class ImageBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
                         src={img.href}
                         onClick={() => {
                             this.dataDoc[this.fieldKey] = new ImageField(img.pathname);
+                            this.dataDoc.ai_firefly_prompt = img.prompt;
+                            this.dataDoc.ai_firefly_seed = img.seed;
                         }}
                     />
                 ))}
@@ -580,14 +583,17 @@ export class ImageBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
                                 onClick={undoable(
                                     action(async () => {
                                         this._regenerateLoading = true;
-                                        SmartDrawHandler.Instance.regenerate([this.Document], undefined, undefined, this._regenInput, true).then(newDocs => {
-                                            const url = newDocs[0];
-                                            const imgField = new ImageField(url);
-                                            this._prevImgs.length === 0 && this._prevImgs.push({ href: this.paths.lastElement(), pathname: field.url.pathname });
-                                            this.dataDoc[this.fieldKey] = imgField;
-                                            this._prevImgs.unshift({ href: this.paths.lastElement(), pathname: url });
-                                            this._regenerateLoading = false;
-                                            this._regenInput = '';
+                                        SmartDrawHandler.Instance.regenerate([this.Document], undefined, undefined, this._regenInput, true).then(newImgs => {
+                                            if (newImgs[0]) {
+                                                const url = newImgs[0].pathname;
+                                                const imgField = new ImageField(url);
+                                                this._prevImgs.length === 0 &&
+                                                    this._prevImgs.push({ prompt: StrCast(this.dataDoc.ai_firefly_prompt), seed: NumCast(this.dataDoc.ai_firefly_seed), href: this.paths.lastElement(), pathname: field.url.pathname });
+                                                this.dataDoc[this.fieldKey] = imgField;
+                                                this._prevImgs.unshift({ prompt: newImgs[0].prompt, seed: newImgs[0].seed, href: this.paths.lastElement(), pathname: url });
+                                                this._regenerateLoading = false;
+                                                this._regenInput = '';
+                                            }
                                         });
                                     }),
                                     'regenerate image'