added size option to image generation

3 files changed, 157 insertions, 96 deletions

diff --git a/src/client/views/smartdraw/SmartDrawHandler.scss b/src/client/views/smartdraw/SmartDrawHandler.scss
index 513779512..4b21c92a5 100644
--- a/src/client/views/smartdraw/SmartDrawHandler.scss
+++ b/src/client/views/smartdraw/SmartDrawHandler.scss
@@ -1,6 +1,6 @@
 .smart-draw-handler {
     position: absolute;
-    // width: 265px;
+    width: 265px;
 
     .smart-draw-main {
         display: flex;
@@ -18,36 +18,54 @@
         justify-content: center;
     }
 
-    .smartdraw-svg-options {
-        margin-top: 5px;
+    .smartdraw-options-container {
+        padding: 5px;
+        font-weight: bolder;
+        text-align: center;
         display: flex;
-        flex-direction: row;
-        justify-content: space-around;
+        flex-direction: column;
 
-        .auto-color {
+        .smartdraw-options {
+            font-weight: normal;
+            margin-top: 5px;
             display: flex;
-            flex-direction: column;
-            align-items: center;
-            width: 30%;
-        }
+            flex-direction: row;
+            justify-content: space-around;
 
-        .complexity {
-            display: flex;
-            flex-direction: column;
-            align-items: center;
-            width: 31%;
-        }
+            .smartdraw-auto-color {
+                display: flex;
+                flex-direction: column;
+                align-items: center;
+                width: 30%;
+                margin-top: 3px;
+            }
 
-        .size {
-            display: flex;
-            flex-direction: column;
-            align-items: center;
-            width: 39%;
+            .smartdraw-complexity {
+                display: flex;
+                flex-direction: column;
+                align-items: center;
+                width: 31%;
+                margin-top: 3px;
+            }
 
-            .size-slider {
-                width: 80%;
+            .smartdraw-size {
+                display: flex;
+                flex-direction: column;
+                align-items: center;
+                width: 39%;
+                margin-top: 3px;
             }
         }
+
+        .smartdraw-dimensions {
+            font-weight: normal;
+            margin-top: 7px;
+            padding-left: 25px;
+        }
+    }
+
+    .smartdraw-slider {
+        width: 65px;
     }
 
     .regenerate-box,
diff --git a/src/client/views/smartdraw/SmartDrawHandler.tsx b/src/client/views/smartdraw/SmartDrawHandler.tsx
index fb1a5771e..ed2734c8b 100644
--- a/src/client/views/smartdraw/SmartDrawHandler.tsx
+++ b/src/client/views/smartdraw/SmartDrawHandler.tsx
@@ -1,5 +1,5 @@
 import { FontAwesomeIcon } from '@fortawesome/react-fontawesome';
-import { Checkbox, Slider, Switch } from '@mui/material';
+import { Checkbox, FormControlLabel, Radio, RadioGroup, Slider, Switch } from '@mui/material';
 import { Button, IconButton } from 'browndash-components';
 import { action, makeObservable, observable, runInAction } from 'mobx';
 import { observer } from 'mobx-react';
@@ -35,6 +35,13 @@ export interface DrawingOptions {
     y: number;
 }
 
+enum FireflyImageDimensions {
+    Square = 'square',
+    Landscape = 'landscape',
+    Portrait = 'portrait',
+    Widescreen = 'widescreen',
+}
+
 /**
  * The SmartDrawHandler allows users to generate drawings with GPT from text input. Users are able to enter
  * the item to draw, how complex they want the drawing to be, how large the drawing should be, and whether
@@ -66,13 +73,16 @@ export class SmartDrawHandler extends ObservableReactComponent<object> {
     @observable private _scale: number = 0;
     @observable private _yRelativeToTop: boolean = true;
     @observable private _isLoading: boolean = false;
+
     @observable private _userInput: string = '';
+    @observable private _regenInput: string = '';
     @observable private _showOptions: boolean = false;
     @observable private _showEditBox: boolean = false;
     @observable private _complexity: number = 5;
     @observable private _size: number = 200;
     @observable private _autoColor: boolean = true;
-    @observable private _regenInput: string = '';
+    @observable private _imgDims: FireflyImageDimensions = FireflyImageDimensions.Square;
+
     @observable private _canInteract: boolean = true;
     @observable private _generateDrawing: boolean = true;
     @observable private _generateImage: boolean = true;
@@ -197,7 +207,7 @@ export class SmartDrawHandler extends ObservableReactComponent<object> {
      */
     @action
     handleSendClick = async () => {
-        if (!this._generateImage && !this._generateDrawing) return;
+        if ((!this.ShowRegenerate && this._userInput == '') || (!this._generateImage && !this._generateDrawing)) return;
         this._isLoading = true;
         this._canInteract = false;
         if (this.ShowRegenerate) {
@@ -207,10 +217,6 @@ export class SmartDrawHandler extends ObservableReactComponent<object> {
                 this._showEditBox = false;
             });
         } else {
-            if (this._userInput == '') {
-                this._isLoading = false;
-                return;
-            }
             runInAction(() => {
                 this._showOptions = false;
             });
@@ -232,7 +238,7 @@ export class SmartDrawHandler extends ObservableReactComponent<object> {
                     this._errorOccurredOnce = false;
                 } else {
                     this._errorOccurredOnce = true;
-                    await this.drawWithGPT({ X: this._pageX, Y: this._pageY }, this._userInput, this._complexity, this._size, this._autoColor);
+                    await this.handleSendClick();
                 }
             }
         }
@@ -268,7 +274,26 @@ export class SmartDrawHandler extends ObservableReactComponent<object> {
      */
     createImageWithFirefly = (input: string, seed?: number) => {
         this._lastInput.text = input;
-        return Networking.PostToServer('/queryFireflyImage', { prompt: input, seed: seed }).then(img => {
+        let width = 0;
+        let height = 0;
+        switch (this._imgDims) {
+            case FireflyImageDimensions.Square:
+                width = 2048;
+                height = 2048;
+                break;
+            case FireflyImageDimensions.Landscape:
+                width = 2304;
+                height = 1792;
+            case FireflyImageDimensions.Portrait:
+                width = 1792;
+                height = 2304;
+                break;
+            case FireflyImageDimensions.Widescreen:
+                width = 2688;
+                height = 1536;
+                break;
+        }
+        return Networking.PostToServer('/queryFireflyImage', { prompt: input, width: width, height: height, seed: seed }).then(img => {
             const imgDoc: Doc = Docs.Create.ImageDocument(img.accessPaths.agnostic.client, {
                 title: input.match(/^(.*?)~~~.*$/)?.[1] || input,
                 ai_generated: true,
@@ -288,6 +313,18 @@ export class SmartDrawHandler extends ObservableReactComponent<object> {
         if (lastInput) this._lastInput = lastInput;
         if (lastResponse) this._lastResponse = lastResponse;
         if (regenInput) this._regenInput = regenInput;
+        if (this._generateImage) {
+            if (this._regenInput !== '') {
+                if (this._selectedDoc) {
+                    const docData = this._selectedDoc[DocData];
+                    const newPrompt = `${docData.firefly_prompt}, ${this._regenInput}`;
+                    const seed: number = docData?.firefly_seed as number;
+                    await this.createImageWithFirefly(newPrompt, seed);
+                }
+            } else {
+                await this.createImageWithFirefly(this._lastInput.text);
+            }
+        }
         if (this._generateDrawing) {
             try {
                 let res;
@@ -310,18 +347,6 @@ export class SmartDrawHandler extends ObservableReactComponent<object> {
                 console.error('Error regenerating drawing', err);
             }
         }
-        if (this._generateImage) {
-            if (this._regenInput !== '') {
-                if (this._selectedDoc) {
-                    const docData = this._selectedDoc[DocData];
-                    const newPrompt = `${docData.firefly_prompt}, ${this._regenInput}`;
-                    const seed: number = docData?.firefly_seed as number;
-                    await this.createImageWithFirefly(newPrompt, seed);
-                }
-            } else {
-                await this.createImageWithFirefly(this._lastInput.text);
-            }
-        }
     };
 
     /**
@@ -460,55 +485,71 @@ export class SmartDrawHandler extends ObservableReactComponent<object> {
                             </div>
                         </div>
                         {this._generateDrawing && (
-                            <div className="smartdraw-svg-options">
-                                <div className="auto-color">
-                                    Auto color
-                                    <Switch
-                                        sx={{
-                                            '& .MuiSwitch-switchBase.Mui-checked': { color: SettingsManager.userColor },
-                                            '& .MuiSwitch-switchBase.Mui-checked + .MuiSwitch-track': { backgroundColor: SettingsManager.userVariantColor },
-                                        }}
-                                        defaultChecked={true}
-                                        value={this._autoColor}
-                                        size="small"
-                                        onChange={action(() => this._canInteract && (this._autoColor = !this._autoColor))}
-                                    />
-                                </div>
-                                <div className="complexity">
-                                    Complexity
-                                    <Slider
-                                        sx={{
-                                            '& .MuiSlider-track': { color: SettingsManager.userVariantColor },
-                                            '& .MuiSlider-rail': { color: SettingsManager.userColor },
-                                            '& .MuiSlider-thumb': { color: SettingsManager.userColor, '&.Mui-focusVisible, &:hover, &.Mui-active': { boxShadow: `0px 0px 0px 8px${SettingsManager.userColor.slice(0, 7)}10` } },
-                                        }}
-                                        style={{ width: '80%' }}
-                                        min={1}
-                                        max={10}
-                                        step={1}
-                                        size="small"
-                                        value={this._complexity}
-                                        onChange={action((e, val) => this._canInteract && (this._complexity = val as number))}
-                                        valueLabelDisplay="auto"
-                                    />
+                            <div className="smartdraw-options-container">
+                                Drawing Options
+                                <div className="smartdraw-options">
+                                    <div className="smartdraw-auto-color">
+                                        Auto color
+                                        <Switch
+                                            sx={{
+                                                '& .MuiSwitch-switchBase.Mui-checked': { color: SettingsManager.userColor },
+                                                '& .MuiSwitch-switchBase.Mui-checked + .MuiSwitch-track': { backgroundColor: SettingsManager.userVariantColor },
+                                            }}
+                                            defaultChecked={true}
+                                            value={this._autoColor}
+                                            size="small"
+                                            onChange={action(() => this._canInteract && (this._autoColor = !this._autoColor))}
+                                        />
+                                    </div>
+                                    <div className="smartdraw-complexity">
+                                        Complexity
+                                        <Slider
+                                            className="smartdraw-slider"
+                                            sx={{
+                                                '& .MuiSlider-track': { color: SettingsManager.userVariantColor },
+                                                '& .MuiSlider-rail': { color: SettingsManager.userColor },
+                                                '& .MuiSlider-thumb': { color: SettingsManager.userColor, '&.Mui-focusVisible, &:hover, &.Mui-active': { boxShadow: `0px 0px 0px 8px${SettingsManager.userColor.slice(0, 7)}10` } },
+                                            }}
+                                            min={1}
+                                            max={10}
+                                            step={1}
+                                            size="small"
+                                            value={this._complexity}
+                                            onChange={action((e, val) => this._canInteract && (this._complexity = val as number))}
+                                            valueLabelDisplay="auto"
+                                        />
+                                    </div>
+                                    <div className="smartdraw-size">
+                                        Size (in pixels)
+                                        <Slider
+                                            className="smartdraw-slider"
+                                            sx={{
+                                                '& .MuiSlider-track': { color: SettingsManager.userVariantColor },
+                                                '& .MuiSlider-rail': { color: SettingsManager.userColor },
+                                                '& .MuiSlider-thumb': { color: SettingsManager.userColor, '&.Mui-focusVisible, &:hover, &.Mui-active': { boxShadow: `0px 0px 0px 8px${SettingsManager.userColor.slice(0, 7)}20` } },
+                                            }}
+                                            min={50}
+                                            max={700}
+                                            step={10}
+                                            size="small"
+                                            value={this._size}
+                                            onChange={action((e, val) => this._canInteract && (this._size = val as number))}
+                                            valueLabelDisplay="auto"
+                                        />
+                                    </div>
                                 </div>
-                                <div className="size">
-                                    Size (in pixels)
-                                    <Slider
-                                        className="size-slider"
-                                        sx={{
-                                            '& .MuiSlider-track': { color: SettingsManager.userVariantColor },
-                                            '& .MuiSlider-rail': { color: SettingsManager.userColor },
-                                            '& .MuiSlider-thumb': { color: SettingsManager.userColor, '&.Mui-focusVisible, &:hover, &.Mui-active': { boxShadow: `0px 0px 0px 8px${SettingsManager.userColor.slice(0, 7)}20` } },
-                                        }}
-                                        min={50}
-                                        max={700}
-                                        step={10}
-                                        size="small"
-                                        value={this._size}
-                                        onChange={action((e, val) => this._canInteract && (this._size = val as number))}
-                                        valueLabelDisplay="auto"
-                                    />
+                            </div>
+                        )}
+                        {this._generateImage && (
+                            <div className="smartdraw-options-container">
+                                Image Options
+                                <div className="smartdraw-dimensions">
+                                    <RadioGroup row defaultValue="square" sx={{ alignItems: 'center' }}>
+                                        <FormControlLabel sx={{ width: '40%' }} value="square" control={<Radio />} onChange={() => this._canInteract && (this._imgDims = FireflyImageDimensions.Square)} label="Square" />
+                                        <FormControlLabel sx={{ width: '40%' }} value="landscape" control={<Radio />} onChange={() => this._canInteract && (this._imgDims = FireflyImageDimensions.Landscape)} label="Landscape" />
+                                        <FormControlLabel sx={{ width: '40%' }} value="portrait" control={<Radio />} onChange={() => this._canInteract && (this._imgDims = FireflyImageDimensions.Portrait)} label="Portrait" />
+                                        <FormControlLabel sx={{ width: '40%' }} value="widescreen" control={<Radio />} onChange={() => this._canInteract && (this._imgDims = FireflyImageDimensions.Widescreen)} label="Widescreen" />
+                                    </RadioGroup>
                                 </div>
                             </div>
                         )}
diff --git a/src/server/ApiManagers/FireflyManager.ts b/src/server/ApiManagers/FireflyManager.ts
index c5348c7db..a41492745 100644
--- a/src/server/ApiManagers/FireflyManager.ts
+++ b/src/server/ApiManagers/FireflyManager.ts
@@ -65,10 +65,12 @@ export default class FireflyManager extends ApiManager {
             })
         );
 
-    generateImage = (prompt: string = 'a realistic illustration of a cat coding', seed?: number) => {
-        let body = `{  "prompt": "${prompt}" }`;
+    generateImage = (prompt: string = 'a realistic illustration of a cat coding', width: number = 2048, height: number = 2048, seed?: number) => {
+        console.log('DIMENSIONS', width, height);
+        let body = `{  "prompt": "${prompt}", "size": { "width": ${width}, "height": ${height}} }`;
         if (seed) {
-            body = `{  "prompt": "${prompt}", "seeds": [${seed}]}`;
+            console.log('RECEIVED SEED', seed);
+            body = `{  "prompt": "${prompt}", "size": { "width": ${width}, "height": ${height}}, "seeds": [${seed}]}`;
         }
         const fetched = this.getBearerToken().then(response =>
             response?.json().then((data: { access_token: string }) =>
@@ -236,7 +238,7 @@ export default class FireflyManager extends ApiManager {
             method: Method.POST,
             subscription: '/queryFireflyImage',
             secureHandler: ({ req, res }) =>
-                this.generateImage(req.body.prompt, req.body.seed).then(img =>
+                this.generateImage(req.body.prompt, req.body.width, req.body.height, req.body.seed).then(img =>
                     DashUploadUtils.UploadImage(img?.url ?? '', undefined, img?.seed).then(info => {
                         if (info instanceof Error) _invalid(res, info.message);
                         else _success(res, info);