created image regeneration with dialogue

1 files changed, 130 insertions, 67 deletions

diff --git a/src/client/views/smartdraw/SmartDrawHandler.tsx b/src/client/views/smartdraw/SmartDrawHandler.tsx
index 036ac5983..fb1a5771e 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 { Slider, Switch } from '@mui/material';
+import { Checkbox, Slider, Switch } from '@mui/material';
 import { Button, IconButton } from 'browndash-components';
 import { action, makeObservable, observable, runInAction } from 'mobx';
 import { observer } from 'mobx-react';
@@ -74,6 +74,8 @@ export class SmartDrawHandler extends ObservableReactComponent<object> {
     @observable private _autoColor: boolean = true;
     @observable private _regenInput: string = '';
     @observable private _canInteract: boolean = true;
+    @observable private _generateDrawing: boolean = true;
+    @observable private _generateImage: boolean = true;
 
     @observable public ShowRegenerate: boolean = false;
 
@@ -195,6 +197,7 @@ export class SmartDrawHandler extends ObservableReactComponent<object> {
      */
     @action
     handleSendClick = async () => {
+        if (!this._generateImage && !this._generateDrawing) return;
         this._isLoading = true;
         this._canInteract = false;
         if (this.ShowRegenerate) {
@@ -212,7 +215,12 @@ export class SmartDrawHandler extends ObservableReactComponent<object> {
                 this._showOptions = false;
             });
             try {
-                await this.drawWithGPT({ X: this._pageX, Y: this._pageY }, this._userInput, this._complexity, this._size, this._autoColor);
+                if (this._generateImage) {
+                    await this.createImageWithFirefly(this._userInput);
+                }
+                if (this._generateDrawing) {
+                    await this.drawWithGPT({ X: this._pageX, Y: this._pageY }, this._userInput, this._complexity, this._size, this._autoColor);
+                }
                 this.hideSmartDrawHandler();
 
                 runInAction(() => {
@@ -240,15 +248,12 @@ export class SmartDrawHandler extends ObservableReactComponent<object> {
     drawWithGPT = async (startPt: { X: number; Y: number }, input: string, complexity: number, size: number, autoColor: boolean) => {
         if (input) {
             this._lastInput = { text: input, complexity: complexity, size: size, autoColor: autoColor, x: startPt.X, y: startPt.Y };
-
-            Networking.PostToServer('/queryFireflyImage', { prompt: input }).then(img => DocumentViewInternal.addDocTabFunc(Docs.Create.ImageDocument(img.accessPaths.agnostic.client, { title: input }), OpenWhere.addRight));
-
             const res = await gptAPICall(`"${input}", "${complexity}", "${size}"`, GPTCallType.DRAW, undefined, true);
             if (res) {
                 const strokeData = await this.parseSvg(res, startPt, false, autoColor);
                 const drawingDoc = strokeData && this.CreateDrawingDoc(strokeData.data, strokeData.lastInput, strokeData.lastRes);
                 drawingDoc && this.AddDrawing(drawingDoc, this._lastInput, res);
-
+                this._selectedDoc = drawingDoc;
                 this._errorOccurredOnce = false;
                 return strokeData;
             } else {
@@ -259,6 +264,23 @@ export class SmartDrawHandler extends ObservableReactComponent<object> {
     };
 
     /**
+     * Calls Firefly API to create an image based on user input
+     */
+    createImageWithFirefly = (input: string, seed?: number) => {
+        this._lastInput.text = input;
+        return Networking.PostToServer('/queryFireflyImage', { prompt: input, seed: seed }).then(img => {
+            const imgDoc: Doc = Docs.Create.ImageDocument(img.accessPaths.agnostic.client, {
+                title: input.match(/^(.*?)~~~.*$/)?.[1] || input,
+                ai_generated: true,
+                firefly_seed: img.accessPaths.agnostic.client.match(/\/(\d+)upload/)[1],
+                firefly_prompt: input,
+            });
+            DocumentViewInternal.addDocTabFunc(imgDoc, OpenWhere.addRight);
+            this._selectedDoc = imgDoc;
+        });
+    };
+
+    /**
      * Regenerates drawings with the option to add a specific regenerate prompt/request.
      */
     @action
@@ -266,27 +288,39 @@ export class SmartDrawHandler extends ObservableReactComponent<object> {
         if (lastInput) this._lastInput = lastInput;
         if (lastResponse) this._lastResponse = lastResponse;
         if (regenInput) this._regenInput = regenInput;
-
-        try {
-            let res;
+        if (this._generateDrawing) {
+            try {
+                let res;
+                if (this._regenInput !== '') {
+                    const prompt: string = `This is your previously generated svg code: ${this._lastResponse} for the user input "${this._lastInput.text}". Please regenerate it with the provided specifications.`;
+                    res = await gptAPICall(`"${this._regenInput}"`, GPTCallType.DRAW, prompt, true);
+                    this._lastInput.text = `${this._lastInput.text} ~~~ ${this._regenInput}`;
+                } else {
+                    res = await gptAPICall(`"${this._lastInput.text}", "${this._lastInput.complexity}", "${this._lastInput.size}"`, GPTCallType.DRAW, undefined, true);
+                }
+                if (!res) {
+                    console.error('GPT call failed');
+                    return;
+                }
+                const strokeData = await this.parseSvg(res, { X: this._lastInput.x, Y: this._lastInput.y }, true, lastInput?.autoColor || this._autoColor);
+                this.RemoveDrawing !== unimplementedFunction && this.RemoveDrawing(true, this._selectedDoc);
+                const drawingDoc = strokeData && this.CreateDrawingDoc(strokeData.data, strokeData.lastInput, strokeData.lastRes);
+                drawingDoc && this.AddDrawing(drawingDoc, this._lastInput, res);
+            } catch (err) {
+                console.error('Error regenerating drawing', err);
+            }
+        }
+        if (this._generateImage) {
             if (this._regenInput !== '') {
-                const prompt: string = `This is your previously generated svg code: ${this._lastResponse} for the user input "${this._lastInput.text}". Please regenerate it with the provided specifications.`;
-                res = await gptAPICall(`"${this._regenInput}"`, GPTCallType.DRAW, prompt, true);
-                this._lastInput.text = `${this._lastInput.text} ~~~ ${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 {
-                res = await gptAPICall(`"${this._lastInput.text}", "${this._lastInput.complexity}", "${this._lastInput.size}"`, GPTCallType.DRAW, undefined, true);
+                await this.createImageWithFirefly(this._lastInput.text);
             }
-            if (!res) {
-                console.error('GPT call failed');
-                return;
-            }
-            const strokeData = await this.parseSvg(res, { X: this._lastInput.x, Y: this._lastInput.y }, true, lastInput?.autoColor || this._autoColor);
-            this.RemoveDrawing !== unimplementedFunction && this.RemoveDrawing(true, this._selectedDoc);
-            const drawingDoc = strokeData && this.CreateDrawingDoc(strokeData.data, strokeData.lastInput, strokeData.lastRes);
-            drawingDoc && this.AddDrawing(drawingDoc, this._lastInput, res);
-            return strokeData;
-        } catch (err) {
-            console.error('Error regenerating drawing', err);
         }
     };
 
@@ -397,58 +431,87 @@ export class SmartDrawHandler extends ObservableReactComponent<object> {
                 </div>
                 {this._showOptions && (
                     <div>
-                        <div className="smartdraw-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
+                        <div className="smartdraw-output-options">
+                            <div className="drawing-checkbox">
+                                Generate Ink
+                                <Checkbox
                                     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` } },
+                                        color: 'white',
+                                        '&.Mui-checked': {
+                                            color: SettingsManager.userVariantColor,
+                                        },
                                     }}
-                                    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"
+                                    checked={this._generateDrawing}
+                                    onChange={() => this._canInteract && (this._generateDrawing = !this._generateDrawing)}
                                 />
                             </div>
-                            <div className="size">
-                                Size (in pixels)
-                                <Slider
-                                    className="size-slider"
+                            <div className="image-checkbox">
+                                Generate Image
+                                <Checkbox
                                     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` } },
+                                        color: 'white',
+                                        '&.Mui-checked': {
+                                            color: SettingsManager.userVariantColor,
+                                        },
                                     }}
-                                    min={50}
-                                    max={700}
-                                    step={10}
-                                    size="small"
-                                    value={this._size}
-                                    onChange={action((e, val) => this._canInteract && (this._size = val as number))}
-                                    valueLabelDisplay="auto"
+                                    checked={this._generateImage}
+                                    onChange={() => this._canInteract && (this._generateImage = !this._generateImage)}
                                 />
                             </div>
                         </div>
-                        <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>
+                                <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>
+                            </div>
+                        )}
                     </div>
                 )}
             </div>