switched to use firefly for image fill

1 files changed, 85 insertions, 27 deletions

diff --git a/src/server/ApiManagers/FireflyManager.ts b/src/server/ApiManagers/FireflyManager.ts
index 6393a1f74..fbf7d7202 100644
--- a/src/server/ApiManagers/FireflyManager.ts
+++ b/src/server/ApiManagers/FireflyManager.ts
@@ -7,8 +7,9 @@ import { DashUserModel } from '../authentication/DashUserModel';
 import { DashUploadUtils } from '../DashUploadUtils';
 import { _error, _invalid, _success, Method } from '../RouteManager';
 import { Upload } from '../SharedMediaTypes';
-import { Directory, filesDirectory } from '../SocketData';
+import { Directory, filesDirectory, pathToDirectory } from '../SocketData';
 import ApiManager, { Registration } from './ApiManager';
+import * as formidable from 'formidable';
 
 export default class FireflyManager extends ApiManager {
     getBearerToken = () =>
@@ -23,6 +24,44 @@ export default class FireflyManager extends ApiManager {
             return undefined;
         });
 
+    generateImageFillWithMask = (prompt: string = 'a realistic illustration of a cat coding', uploadUrl: string, maskUrl: string | undefined, width: number = 2048, height: number = 2048, variations: number = 1) =>
+        this.getBearerToken().then(response =>
+            response?.json().then((data: { access_token: string }) =>
+                //prettier-ignore
+                fetch('https://firefly-api.adobe.io/v3/images/fill', {
+                        method: 'POST',
+                        headers: [
+                            ['Content-Type', 'application/json'],
+                            ['Accept', 'application/json'],
+                            ['x-api-key', process.env._CLIENT_FIREFLY_CLIENT_ID ?? ''],
+                            ['Authorization', `Bearer ${data.access_token}`],
+                        ],
+                        body: JSON.stringify({
+                            image: {
+                                source: { url: uploadUrl },
+                                mask: {
+                                    invert: true,
+                                    url: maskUrl
+                                },
+                            },
+                            size: {
+                                width: Math.round(width),
+                                height: Math.round(height),
+                            },
+                            prompt: prompt ?? '',
+                            numVariations: variations,
+                        }),
+                    })
+                    .then(response2 => response2.json().then(json => 
+                    {
+                        if (json.outputs?.length)
+                            return (json.outputs as {image: {url:string }}[]).map(output => output.image);
+                        throw new Error(JSON.stringify(json));
+                    })
+                )
+            )
+        );
+
     generateImageFromStructure = (
         prompt: string = 'a realistic illustration of a cat coding',
         width: number = 2048,
@@ -79,6 +118,30 @@ export default class FireflyManager extends ApiManager {
             )
         );
 
+    uploadToDropbox = (dropboxClient: Dropbox, user: DashUserModel | undefined, fileUrl: string, contents: NonSharedBuffer, recursed?: boolean): Promise<string | Error> =>
+        dropboxClient
+            .filesUpload({ path: `/Apps/browndash/${path.basename(fileUrl)}`, contents, mode: { '.tag': 'overwrite' } })
+            .then(response =>
+                dropboxClient
+                    .filesGetTemporaryLink({ path: response.result.path_display ?? '' })
+                    .then(link => link.result.link)
+                    .catch(linkErr => new Error('Failed to get temporary link: ' + linkErr.message))
+            )
+            .catch(uploadErr => {
+                if (user?.dropboxRefresh && !recursed) {
+                    console.log('Attempting to refresh Dropbox token for user:', user.email);
+                    return this.refreshDropboxToken(user)
+                        .then(token => {
+                            if (!token) return new Error('Failed to refresh Dropbox token.' + user.email);
+
+                            const dbxNew = new Dropbox({ accessToken: token });
+                            return this.uploadToDropbox(dbxNew, user, fileUrl, contents, true).catch(finalErr => new Error('Failed to refresh Dropbox token:' + finalErr.message));
+                        })
+                        .catch(refreshErr => new Error('Failed to refresh Dropbox token: ' + refreshErr.message));
+                }
+                return new Error('Dropbox error: ' + uploadErr.message);
+            });
+
     uploadImageToDropbox = (fileUrl: string, user: DashUserModel | undefined, dbx = new Dropbox({ accessToken: user?.dropboxToken || '' })) =>
         new Promise<string>((resolve, reject) => {
             fs.readFile(path.join(filesDirectory, `${Directory.images}/${path.basename(fileUrl)}`), undefined, (err, contents) => {
@@ -86,32 +149,7 @@ export default class FireflyManager extends ApiManager {
                     return reject(new Error('Error reading file:' + err.message));
                 }
 
-                const uploadToDropbox = (dropboxClient: Dropbox) =>
-                    dropboxClient
-                        .filesUpload({ path: `/Apps/browndash/${path.basename(fileUrl)}`, contents })
-                        .then(response =>
-                            dropboxClient
-                                .filesGetTemporaryLink({ path: response.result.path_display ?? '' })
-                                .then(link => resolve(link.result.link))
-                                .catch(linkErr => reject(new Error('Failed to get temporary link: ' + linkErr.message)))
-                        )
-                        .catch(uploadErr => {
-                            if (user?.dropboxRefresh) {
-                                console.log('Attempting to refresh Dropbox token for user:', user.email);
-                                this.refreshDropboxToken(user)
-                                    .then(token => {
-                                        if (!token) return reject(new Error('Failed to refresh Dropbox token.' + user.email));
-
-                                        const dbxNew = new Dropbox({ accessToken: token });
-                                        uploadToDropbox(dbxNew).catch(finalErr => reject(new Error('Failed to refresh Dropbox token:' + finalErr.message)));
-                                    })
-                                    .catch(refreshErr => reject(new Error('Failed to refresh Dropbox token: ' + refreshErr.message)));
-                            } else {
-                                reject(new Error('Dropbox error: ' + uploadErr.message));
-                            }
-                        });
-
-                uploadToDropbox(dbx);
+                this.uploadToDropbox(dbx, user, fileUrl, contents).then(value => (value instanceof Error ? reject(value) : resolve(value)));
             });
         });
 
@@ -318,6 +356,26 @@ export default class FireflyManager extends ApiManager {
                     })
             ), // prettier-ignore
         });
+        register({
+            method: Method.POST,
+            subscription: '/queryFireflyImageFillWithMask',
+            secureHandler: ({ req, res }) =>
+                new Promise<string>(resolve => {
+                    const user = req.user as DashUserModel;
+                    const accessToken = user?.dropboxToken || '';
+                    const dbx = new Dropbox({ accessToken });
+                    const form = new formidable.IncomingForm({ keepExtensions: true, uploadDir: pathToDirectory(Directory.parsed_files) });
+                    form.parse(req, async (err, fields, files) => {
+                        if (files.source && files.mask) {
+                            Promise.all([this.uploadToDropbox(dbx, user, 'source.png', fs.readFileSync(files.source[0].filepath)), 
+                                         this.uploadToDropbox(dbx, user, 'mask.png', fs.readFileSync(files.mask[0].filepath))])
+                            .then(stuff => 
+                                stuff.some(s => s instanceof Error) ? resolve("") : this.generateImageFillWithMask(fields["prompt"]?.[0], stuff[0] as string, stuff[1] as string, 2048, 2048, 1).then(url => resolve(url![0].url))
+                            ).catch(() => resolve("") ); // prettier-ignore
+                        }
+                    });
+                }).then(url => (url ? _success(res, { urls: [url] }) : _invalid(res, 'Failed to fill image'))),
+        });
 
         register({
             method: Method.POST,