1 files changed, 24 insertions, 21 deletions
diff --git a/src/client/views/smartdraw/DrawingFillHandler.tsx b/src/client/views/smartdraw/DrawingFillHandler.tsx
index cd15d19d5..6f2d1cbe6 100644
--- a/src/client/views/smartdraw/DrawingFillHandler.tsx
+++ b/src/client/views/smartdraw/DrawingFillHandler.tsx
@@ -1,7 +1,7 @@
 import { imageUrlToBase64 } from '../../../ClientUtils';
 import { Doc } from '../../../fields/Doc';
 import { DocData } from '../../../fields/DocSymbols';
-import { DocCast, ImageCast } from '../../../fields/Types';
+import { DocCast, ImageCast, NumCast } from '../../../fields/Types';
 import { Upload } from '../../../server/SharedMediaTypes';
 import { gptDescribeImage } from '../../apis/gpt/GPT';
 import { Docs } from '../../documents/Documents';
@@ -12,13 +12,15 @@ import { AspectRatioLimits, FireflyDimensionsMap, FireflyImageDimensions, Firefl
 
 const DashDropboxId = '2m86iveqdr9vzsa';
 export class DrawingFillHandler {
-    static drawingToImage = async (drawing: Doc, strength: number, user_prompt: string) => {
+    static drawingToImage = async (drawing: Doc, strength: number, user_prompt: string, styleDoc?: Doc) => {
         const docData = drawing[DocData];
         const tags: string[] = ((docData?.tags as unknown as string[]) ?? []).map(tag => tag.slice(1)) ?? [];
         const styles = tags.filter(tag => FireflyStylePresets.has(tag));
-        const styleDocs = Doc.Links(drawing)
-            .map(link => Doc.getOppositeAnchor(link, drawing))
-            .map(anchor => anchor && DocCast(anchor.embedContainer));
+        const styleDocs = styleDoc
+            ? [styleDoc]
+            : Doc.Links(drawing)
+                  .map(link => Doc.getOppositeAnchor(link, drawing))
+                  .map(anchor => anchor && DocCast(anchor.embedContainer));
         const styleRef = styleDocs.filter(doc => doc !== undefined && doc.type === 'image').lastElement();
         let styleUrl: string | undefined;
         if (styleRef) {
@@ -47,25 +49,26 @@ export class DrawingFillHandler {
                 const structureUrl = `${hrefParts.slice(0, -1).join('.')}_o.${hrefParts.lastElement()}`;
                 return imageUrlToBase64(structureUrl)
                     .then(gptDescribeImage)
-                    .then(prompt =>
-                        Networking.PostToServer('/queryFireflyImageFromStructure', { prompt: `${user_prompt || prompt}`, width: dims.width, height: dims.height, structure: structureUrl, strength, presets: styles, styleUrl })
-                            .then((info: Upload.ImageInformation) =>
-                                DocumentViewInternal.addDocTabFunc(
-                                    Docs.Create.ImageDocument(info.accessPaths.agnostic.client, {
-                                        ai: 'firefly',
-                                        title: user_prompt || prompt,
-                                        ai_firefly_prompt: user_prompt || prompt,
-                                        _width: 500,
-                                        data_nativeWidth: info.nativeWidth,
-                                        data_nativeHeight: info.nativeHeight,
-                                    }),
-                                    OpenWhere.addRight
-                                )
-                            )
+                    .then((prompt, newPrompt = user_prompt || prompt) =>
+                        Networking.PostToServer('/queryFireflyImageFromStructure', { prompt: `${newPrompt}`, width: dims.width, height: dims.height, structure: structureUrl, strength, presets: styles, styleUrl })
+                            .then((info: Upload.ImageInformation) => {
+                                const img = Docs.Create.ImageDocument(info.accessPaths.agnostic.client, {
+                                    ai: 'firefly',
+                                    title: newPrompt,
+                                    ai_firefly_prompt: newPrompt,
+                                    _width: 500,
+                                    data_nativeWidth: info.nativeWidth,
+                                    data_nativeHeight: info.nativeHeight,
+                                });
+                                const genratedDocs = drawing.generatedDocs ? DocCast(drawing.generatedDocs) : Docs.Create.MasonryDocument([], { _width: 400, _height: 400, x: NumCast(drawing.x) + NumCast(drawing.width), y: NumCast(drawing.y) });
+                                drawing[DocData].generatedDocs = genratedDocs;
+                                Doc.AddDocToList(genratedDocs, undefined, img);
+                                if (!DocumentView.getFirstDocumentView(genratedDocs)) DocumentViewInternal.addDocTabFunc(genratedDocs, OpenWhere.addRight);
+                            })
                             .catch(e => {
                                 if (e.toString().includes('Dropbox') && confirm('Create image failed. Try authorizing DropBox?\r\n' + e.toString().replace(/^[^"]*/, ''))) {
                                     window.open(`https://www.dropbox.com/oauth2/authorize?client_id=${DashDropboxId}&response_type=code&token_access_type=offline&redirect_uri=http://localhost:1050/refreshDropbox`, '_blank')?.focus();
-                                }
+                                } else alert(e.toString());
                             })
                     ); // prettier-ignore:q
             }