1 files changed, 29 insertions, 12 deletions
diff --git a/src/server/ApiManagers/FireflyManager.ts b/src/server/ApiManagers/FireflyManager.ts
index a1f8fab8d..185752404 100644
--- a/src/server/ApiManagers/FireflyManager.ts
+++ b/src/server/ApiManagers/FireflyManager.ts
@@ -20,7 +20,7 @@ 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) =>
         this.getBearerToken().then(response =>
             response?.json().then((data: { access_token: string }) =>
                 //prettier-ignore
@@ -44,7 +44,12 @@ export default class FireflyManager extends ApiManager {
                                   },
                               },
                         //prettier-ignore
-                        style: { presets: styles }
+                        style: { 
+                            imageReference: {
+                                source: { uploadId: styleUrl }
+                            },
+                            presets: styles 
+                        }
                     }),
                 })
                     .then(response2 => response2.json().then(json => JSON.stringify((json.outputs?.[0] as { image: { url: string } })?.image)))
@@ -235,16 +240,28 @@ export default class FireflyManager extends ApiManager {
             method: Method.POST,
             subscription: '/queryFireflyImageFromStructure',
             secureHandler: async ({ req, res }) =>
-                this.uploadImageToDropbox(req.body.structureUrl).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);
-                              })
-                          )
-                ),
+                this.uploadImageToDropbox(req.body.styleRef)
+                    .then(styleUrl => {
+                        if (styleUrl instanceof Error) {
+                            _invalid(res, styleUrl.message);
+                            throw new Error('Error uploading images to dropbox');
+                        }
+                        return this.uploadImageToDropbox(req.body.structure).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, uploads.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,