1 files changed, 37 insertions, 17 deletions
diff --git a/src/server/ApiManagers/FireflyManager.ts b/src/server/ApiManagers/FireflyManager.ts
index 6039f5675..d3218842f 100644
--- a/src/server/ApiManagers/FireflyManager.ts
+++ b/src/server/ApiManagers/FireflyManager.ts
@@ -22,11 +22,11 @@ export default class FireflyManager extends ApiManager {
             return undefined;
         });
 
-    generateImageFromStructure = (prompt: string = 'a realistic illustration of a cat coding', width: number = 2048, height: number = 2048, structureUrl: string, strength: number = 50, styles: string[]) =>
+    generateImageFromStructure = (prompt: string = 'a realistic illustration of a cat coding', width: number = 2048, height: number = 2048, structureUrl: string, strength: number = 50, styles: string[], styleUrl: string | undefined) =>
         this.getBearerToken().then(response =>
-            response?.json().then((data: { access_token: string }) =>
+            response?.json().then((data: { access_token: string }) => {
                 //prettier-ignore
-                fetch('https://firefly-api.adobe.io/v3/images/generate', {
+                return fetch('https://firefly-api.adobe.io/v3/images/generate', {
                     method: 'POST',
                     headers: [
                         ['Content-Type', 'application/json'],
@@ -45,8 +45,15 @@ export default class FireflyManager extends ApiManager {
                                       source: { url: structureUrl },
                                   },
                               },
-                        //prettier-ignore
-                        style: { presets: styles }
+                        // prettier-ignore
+                        style: {
+                            presets: styles,
+                            ...(styleUrl && {
+                                imageReference: {
+                                    source: { uploadId: styleUrl },
+                                },
+                            }),
+                        }
                     }),
                 })
                     .then(response2 => response2.json().then(json => JSON.stringify((json.outputs?.[0] as { image: { url: string } })?.image)))
@@ -54,7 +61,7 @@ export default class FireflyManager extends ApiManager {
                         console.error('Error:', error);
                         return '';
                     })
-            )
+            })
         );
 
     uploadImageToDropbox = (fileUrl: string, user: DashUserModel | undefined, dbx = new Dropbox({ accessToken: user?.dropboxToken || '' })) =>
@@ -283,17 +290,30 @@ export default class FireflyManager extends ApiManager {
         register({
             method: Method.POST,
             subscription: '/queryFireflyImageFromStructure',
-            secureHandler: async ({ req, res }) =>
-                this.uploadImageToDropbox(req.body.structureUrl, req.user as DashUserModel).then(uploadUrl =>
-                    uploadUrl instanceof Error
-                        ? _invalid(res, uploadUrl.message)
-                        : this.generateImageFromStructure(req.body.prompt, req.body.width, req.body.height, uploadUrl, req.body.strength, req.body.styles).then(fire =>
-                              DashUploadUtils.UploadImage(JSON.parse(fire ?? '').url).then(info => {
-                                  if (info instanceof Error) _invalid(res, info.message);
-                                  else _success(res, info);
-                              })
-                          )
-                ),
+            secureHandler: async ({ req, res }) => {
+                (req.body.styleRef ? this.uploadImageToDropbox(req.body.styleRef, req.user as DashUserModel) : Promise.resolve(undefined))
+                    .then(styleUrl => {
+                        if (styleUrl instanceof Error) {
+                            _invalid(res, styleUrl.message);
+                            throw new Error('Error uploading images to dropbox');
+                        }
+                        return this.uploadImageToDropbox(req.body.structur, req.user as DashUserModel).then(structureUrl => {
+                            if (structureUrl instanceof Error) {
+                                _invalid(res, structureUrl.message);
+                                throw new Error('Error uploading images to dropbox');
+                            }
+                            return { styleUrl, structureUrl };
+                        });
+                    })
+                    .then(uploads =>
+                        this.generateImageFromStructure(req.body.prompt, req.body.width, req.body.height, uploads.structureUrl, req.body.strength, req.body.presets, req.body.styleUrl).then(fire => {
+                            DashUploadUtils.UploadImage(JSON.parse(fire ?? '').url).then(info => {
+                                if (info instanceof Error) _invalid(res, info.message);
+                                else _success(res, info);
+                            });
+                        })
+                    );
+            },
         });
         register({
             method: Method.POST,