selection.js

@@ -324,6 +324,24 @@ async function channelToSelectionExe(channel_name = 'mask') {
     }
 }
 
+function keepRatio(selectionInfo, offset) {
+    // Calculate the current width and height
+    let width = selectionInfo.right - selectionInfo.left
+    let height = selectionInfo.bottom - selectionInfo.top
+
+    // Calculate the new coordinates with the offset
+    selectionInfo.right += offset
+    selectionInfo.left -= offset
+    selectionInfo.bottom += offset
+    selectionInfo.top -= offset
+
+    // Update width and height
+    selectionInfo.width = width + 2 * offset
+    selectionInfo.height = height + 2 * offset
+
+    return selectionInfo
+}
+
 function makeSquare(selectionInfo, offset) {
     // Calculate the current width and height
     let width = selectionInfo.right - selectionInfo.left
@@ -349,12 +367,20 @@ function makeSquare(selectionInfo, offset) {
     return selectionInfo
 }
 
-async function inpaintLassoInitImageAndMask(channel_name = 'mask', offset = 0) {
+async function inpaintLassoInitImageAndMask(
+    channel_name = 'mask',
+    offset = 0,
+    make_square = true
+) {
     const selectionInfo = await psapi.getSelectionInfoExe()
     //convert the selection box into square box so that you have best output results
-    const squareSelection = makeSquare(selectionInfo, offset)
+
+    const newSelection = make_square
+        ? makeSquare(selectionInfo, offset)
+        : keepRatio(selectionInfo, offset)
+
     //correct width and height sliders, since this is lasso mode.
-    await calcWidthHeightFromSelection(squareSelection)
+    await calcWidthHeightFromSelection(newSelection)
 
     async function getImageFromCanvas() {
         const width = html_manip.getWidth()
@@ -363,7 +389,7 @@ async function inpaintLassoInitImageAndMask(channel_name = 'mask', offset = 0) {
         const base64 = await io.IO.getSelectionFromCanvasAsBase64Interface_New(
             width,
             height,
-            squareSelection,
+            newSelection,
             true
         )
         return base64
@@ -396,7 +422,7 @@ async function inpaintLassoInitImageAndMask(channel_name = 'mask', offset = 0) {
             synchronousExecution: true,
         })
         // const selection_info = await psapi.getSelectionInfoExe()
-        mask_base64 = await fillSelectionWhiteOutsideBlack(squareSelection)
+        mask_base64 = await fillSelectionWhiteOutsideBlack(newSelection)
     })
 
     //save laso selection to channel

typescripts/comfyui/comfyui.tsx

@@ -1257,12 +1257,12 @@ class ComfyWorkflowComponent extends React.Component<{}, { value?: number }> {
                                 // Start the progress update
 
                                 util.runRandomSeedScript()
-                                let store_output =
+                                let { outputs, separated_outputs } =
                                     await util.postPromptAndGetBase64JsonResult(
                                         toJS(store.data.current_prompt2)
                                     )
                                 store.data.current_prompt2_output =
-                                    store_output ?? {}
+                                    outputs ?? {}
                             } catch (e) {
                                 console.error(e)
                             } finally {

typescripts/comfyui/main_ui.tsx

@@ -16,6 +16,7 @@ import { store } from './util'
 import { base64UrlToBase64, copyJson } from '../util/ts/general'
 import { session_store } from '../stores'
 import ControlNetStore from '../controlnet/store'
+import { setControlDetectMapSrc } from '../controlnet/entry'
 
 // Function to parse metadata from a title string
 function parseMetadata(title: string) {
@@ -422,16 +423,21 @@ async function generateComfyMode(
 
             delete final_prompt[hire_output_node.id]
         }
+
+        const separated_output_node_ids: string[] = []
+        const node_id_to_controlnet_unit_index: Record<string, number> = {}
         for (const [index, unit] of plugin_settings[
             'controlnet_units'
         ].entries()) {
+            const node_name_id = `preprocessor_output_${index + 1}`
+            const node = getNodeByNameId(nodes, node_name_id)
+            const node_id = node.id.toString()
+            node_id_to_controlnet_unit_index[node_id] = index
+
             if (unit['comfy_enabled'] === 'disable') {
-                index
-                mutePromptNode(
-                    nodes,
-                    final_prompt,
-                    `preprocessor_output_${index + 1}`
-                )
+                mutePromptNode(nodes, final_prompt, node_name_id)
+            } else if (unit['comfy_enabled'] === 'enable') {
+                separated_output_node_ids.push(node_id)
             }
         }
         // for (let i = 0; i < 3; ++i) {
@@ -444,9 +450,11 @@ async function generateComfyMode(
         //     }
         // }
         console.log('final_prompt: ', final_prompt)
-        const outputs = await comfyui_util.postPromptAndGetBase64JsonResult(
-            final_prompt
-        )
+        const { outputs, separated_outputs } =
+            await comfyui_util.postPromptAndGetBase64JsonResult(
+                final_prompt,
+                separated_output_node_ids
+            )
 
         if (outputs) {
             image_url_list = Object.values(outputs).flat()
@@ -454,6 +462,16 @@ async function generateComfyMode(
                 return base64UrlToBase64(image_url)
             })
         }
+        if (separated_outputs) {
+            Object.entries(separated_outputs).forEach(([node_id, images]) => {
+                const controlnet_unit_index =
+                    node_id_to_controlnet_unit_index[node_id]
+                setControlDetectMapSrc(
+                    base64UrlToBase64(images[0]),
+                    controlnet_unit_index
+                )
+            })
+        }
     } catch (e) {
         console.error(e)
     }

typescripts/comfyui/util.ts

@@ -7,6 +7,7 @@ import { AStore } from '../main/astore'
 
 import comfyapi from './comfyapi'
 import { base64UrlToBase64 } from '../util/ts/general'
+import { app } from 'photoshop'
 export enum InputTypeEnum {
     NumberField = 'NumberField',
     TextField = 'TextField',
@@ -211,8 +212,12 @@ async function getHistory(prompt_id: string) {
     return history
 }
 export async function postPromptAndGetBase64JsonResult(
-    prompt: Record<string, any>
-) {
+    prompt: Record<string, any>,
+    separated_output_nodes: string[] = []
+): Promise<{
+    outputs: Record<string, any> | undefined
+    separated_outputs: Record<string, any> | undefined
+}> {
     try {
         const res = await comfyapi.comfy_api.prompt(prompt)
         if (res.error) {
@@ -222,9 +227,11 @@ export async function postPromptAndGetBase64JsonResult(
         const prompt_id = res.prompt_id
         const history = await getHistory(prompt_id)
         const promptInfo = history[prompt_id]
-        const store_output = await mapComfyOutputToStoreOutput(
-            promptInfo.outputs
-        )
+        const { outputs, separated_outputs } =
+            await mapComfyOutputToStoreOutput(
+                promptInfo.outputs,
+                separated_output_nodes
+            )
         //         // [4][0] for output id.
         //         const fileName = promptInfo.outputs[promptInfo.prompt[4][0]].images[0].filename
         //         const resultB64 = await ComfyApi.view(this, fileName);
@@ -233,9 +240,11 @@ export async function postPromptAndGetBase64JsonResult(
         //             try { option.imageFinishCallback(resultB64, index) } catch (e) { }
         //         }
         // }
-        return store_output
+        return { outputs, separated_outputs }
     } catch (e) {
         console.error(e)
+        app.showAlert(`${e}`)
+        return { outputs: undefined, separated_outputs: undefined }
     }
 }
 export const getFileFormat = (fileName: string): string =>
@@ -264,9 +273,11 @@ export function updateOutput(output: any, output_store_obj: any) {
 }
 
 export async function mapComfyOutputToStoreOutput(
-    comfy_output: Record<string, any>
+    comfy_output: Record<string, any>,
+    separated_output_nodes: string[] = []
 ) {
-    const store_output: Record<string, any> = {}
+    const outputs: Record<string, any> = {}
+    const separated_outputs: Record<string, any> = {}
 
     for (let key in comfy_output) {
         let base64_url_list = await Promise.all(
@@ -288,10 +299,18 @@ export async function mapComfyOutputToStoreOutput(
             )
         )
         base64_url_list = base64_url_list.filter((item) => item !== '') // Filter out empty strings
-        store_output[key] = [...(store_output[key] || []), ...base64_url_list]
+
+        if (separated_output_nodes.includes(key)) {
+            separated_outputs[key] = [
+                ...(separated_outputs[key] || []),
+                ...base64_url_list,
+            ]
+        } else {
+            outputs[key] = [...(outputs[key] || []), ...base64_url_list]
+        }
     }
 
-    return store_output
+    return { outputs, separated_outputs }
 }
 
 interface LooseObject {
@@ -462,9 +481,10 @@ async function maskExpansion(
         },
     }
     try {
-        const out = await postPromptAndGetBase64JsonResult(prompt)
-        if (out) {
-            const expanded_mask = out['6'][0]
+        const { outputs, separated_outputs } =
+            await postPromptAndGetBase64JsonResult(prompt)
+        if (outputs) {
+            const expanded_mask = outputs['6'][0]
             return base64UrlToBase64(expanded_mask)
         }
         // html_manip.setInitImageMaskSrc(expanded_mask)

typescripts/controlnet/entry.ts

@@ -3,7 +3,7 @@ import { setControlImageSrc } from '../../utility/html_manip'
 // import { session_ts } from '../entry'
 // import * as session_ts from '../session/session'
 import { store as session_store } from '../session/session_store'
-import { Enum, api, python_replacement } from '../util/oldSystem'
+import { Enum, api, io, python_replacement } from '../util/oldSystem'
 import { GenerationModeEnum } from '../util/ts/enum'
 import store, {
     DefaultControlNetUnitData,
@@ -149,11 +149,11 @@ function getEnableControlNet(index: number) {
         )
     else return store.controlNetUnitData[index || 0].enabled
 }
-function mapPluginSettingsToControlNet(plugin_settings: any) {
+async function mapPluginSettingsToControlNet(plugin_settings: any) {
     const ps = plugin_settings // for shortness
     let controlnet_units: any[] = []
     const controlNetUnits = store.controlNetUnitData
-    function getControlNetInputImage(index: number) {
+    async function getControlNetInputImage(index: number) {
         try {
             const b_sync_input_image = controlNetUnits[index].auto_image
             let input_image = controlNetUnits[index].input_image
@@ -163,10 +163,19 @@ function mapPluginSettingsToControlNet(plugin_settings: any) {
             ) {
                 //conditions: 1) txt2img mode 2)auto image on  3)first generation of session
 
-                input_image = session_store.data.controlnet_input_image ?? ''
-                controlNetUnits[index].input_image = input_image
-                controlNetUnits[index].selection_info =
-                    plugin_settings.selection_info
+                if (
+                    session_store.data.generation_number === 1 &&
+                    session_store.data.controlnet_input_image === ''
+                ) {
+                    session_store.data.controlnet_input_image =
+                        await io.getImg2ImgInitImage()
+                }
+                if (session_store.data.controlnet_input_image !== '') {
+                    input_image = session_store.data.controlnet_input_image
+                    controlNetUnits[index].input_image = input_image
+                    controlNetUnits[index].selection_info =
+                        plugin_settings.selection_info
+                }
             }
             if (
                 b_sync_input_image &&
@@ -212,7 +221,7 @@ function mapPluginSettingsToControlNet(plugin_settings: any) {
     for (let index = 0; index < store.maxControlNet; index++) {
         controlnet_units[index] = {
             enabled: getEnableControlNet(index),
-            input_image: getControlNetInputImage(index),
+            input_image: await getControlNetInputImage(index),
             mask: getControlNetMask(index),
             module: controlNetUnits[index].module,
             model: controlNetUnits[index].model,

typescripts/sd_tab/sd_tab.tsx

@@ -38,8 +38,8 @@ import {
     onHeightSliderInput,
     heightSliderOnChangeEventHandler,
     loadPresetSettings,
-    
     isHiResMode,
+    comfy_mask_content_config,
 } from './util'
 import { general } from '../util/oldSystem'
 import { requestSwapModel, setInpaintMaskWeight } from '../util/ts/sdapi'
@@ -105,6 +105,22 @@ const Modes = observer(() => {
                         >
                             Lasso Mode
                         </SpCheckBox>
+                        <SpCheckBox
+                            style={{
+                                marginLeft: '10px',
+                                display: store.data.is_lasso_mode
+                                    ? void 0
+                                    : 'none',
+                            }}
+                            onChange={() => {
+                                helper_store.data.make_square =
+                                    !helper_store.data.make_square
+                            }}
+                            checked={helper_store.data.make_square}
+                            // id={`chEnableControlNet_${this.props.index}`}
+                        >
+                            Make Square
+                        </SpCheckBox>
 
                         <SpSlider
                             show-value="false"
@@ -1043,30 +1059,31 @@ class SDTab extends React.Component<{}> {
                                 <sp-label class="title" slot="label">
                                     Mask Content:
                                 </sp-label>
-                                {mask_content_config.map(
-                                    (mask_content, index: number) => {
-                                        return (
-                                            <sp-radio
-                                                key={index}
-                                                class="rbMaskContent"
-                                                checked={
-                                                    store.data
-                                                        .inpainting_fill ===
+                                {(settings_tab_ts.store.data
+                                    .selected_backend === 'Automatic1111'
+                                    ? mask_content_config
+                                    : comfy_mask_content_config
+                                ).map((mask_content, index: number) => {
+                                    return (
+                                        <sp-radio
+                                            key={index}
+                                            class="rbMaskContent"
+                                            checked={
+                                                store.data.inpainting_fill ===
+                                                mask_content.value
+                                                    ? true
+                                                    : void 0
+                                            }
+                                            value={mask_content.value}
+                                            onClick={(evt: any) => {
+                                                store.data.inpainting_fill =
                                                     mask_content.value
-                                                        ? true
-                                                        : void 0
-                                                }
-                                                value={mask_content.value}
-                                                onClick={(evt: any) => {
-                                                    store.data.inpainting_fill =
-                                                        mask_content.value
-                                                }}
-                                            >
-                                                {mask_content.name}
-                                            </sp-radio>
-                                        )
-                                    }
-                                )}
+                                            }}
+                                        >
+                                            {mask_content.name}
+                                        </sp-radio>
+                                    )
+                                })}
                             </sp-radio-group>
                         </div>
 

typescripts/sd_tab/util.ts

@@ -77,6 +77,24 @@ export const mask_content_config = [
         value: 3,
     },
 ]
+export const comfy_mask_content_config = [
+    // {
+    //     name: 'fill',
+    //     value: 0,
+    // },
+    {
+        name: 'original',
+        value: 1,
+    },
+    // {
+    //     name: 'latent noise',
+    //     value: 2,
+    // },
+    {
+        name: 'latent nothing',
+        value: 3,
+    },
+]
 export enum SelectionModeEnum {
     Ratio = 'ratio',
     Precise = 'precise',
@@ -131,7 +149,9 @@ export const store = new AStore({
     hr_resize_y: 512,
     hr_second_pass_steps: 20,
     restore_faces: false,
-    inpainting_fill: mask_content_config[0].value,
+    // inpainting_fill: mask_content_config[0].value,
+    inpainting_fill: 1, //set original as default value
+
     hr_upscaler: '',
 
     selection_mode: selection_mode_config[0].value,
@@ -167,7 +187,7 @@ export const default_preset = {
         hr_resize_y: 512,
         hr_second_pass_steps: 0,
         // restore_faces: false,
-        inpainting_fill: 0,
+        inpainting_fill: 1,
         hr_upscaler: '',
         selection_mode: 'ratio',
     },
@@ -186,6 +206,7 @@ export const helper_store = new AStore({
     native_presets: {},
     base_size: 512 as number,
     lasso_offset: 10 as number,
+    make_square: true as boolean,
 })
 export async function refreshModels() {
     let b_result = false

typescripts/session/modes.ts

@@ -25,7 +25,8 @@ import {
 
 import { store as extra_page_store } from '../extra_page/extra_page'
 import { requestPost } from '../util/ts/api'
-import { comfyapi } from '../entry'
+import { comfyapi, settings_tab_ts } from '../entry'
+import { newOutputImageName } from '../util/ts/general'
 const executeAsModal = core.executeAsModal
 
 declare let g_inpaint_mask_layer: any
@@ -38,9 +39,17 @@ interface SessionData {
     mask?: string
     selectionInfo?: any
 }
+interface ImageInfo {
+    path: string
+    base64: string
+    auto_metadata: Record<string, any>
+}
 
-async function saveOutputImagesToDrive(images_info: any, settings: any) {
-    const base64OutputImages = [] //delete all previouse images, Note move this to session end ()
+async function saveOutputImagesToDrive(
+    images_info: ImageInfo[],
+    settings: Record<string, any>
+) {
+    const base64OutputImages = []
     let index = 0
     for (const image_info of images_info) {
         const path = image_info['path']
@@ -58,10 +67,32 @@ async function saveOutputImagesToDrive(images_info: any, settings: any) {
         ) //save the settings
         index += 1
     }
-    session_store.data.last_seed =
-        images_info?.length > 0 ? images_info[0]?.auto_metadata?.Seed : '-1'
+    if (settings_tab_ts.store.data.selected_backend === 'Automatic1111') {
+        session_store.data.last_seed =
+            images_info?.length > 0 ? images_info[0]?.auto_metadata?.Seed : '-1'
+    }
     return base64OutputImages
 }
+async function saveOutputImagesToDriveComfy(
+    base64_images: string[],
+    settings: Record<string, any>
+) {
+    let index = 0
+    const document_name = settings['uniqueDocumentId']
+    delete settings['alwayson_scripts']
+    for (const base64 of base64_images) {
+        const image_name = newOutputImageName()
+        await io.saveFileInSubFolder(base64, document_name, image_name) //save the output image
+        const json_file_name = `${image_name.split('.')[0]}.json`
+
+        await io.saveJsonFileInSubFolder(
+            settings,
+            document_name,
+            json_file_name
+        ) //save the settings
+        index += 1
+    }
+}
 class Mode {
     constructor() {}
 
@@ -159,8 +190,9 @@ export class Txt2ImgMode extends Mode {
 
         const full_url = `${g_sd_url}/sdapi/v1/txt2img`
 
-        const control_net_settings =
-            mapPluginSettingsToControlNet(plugin_settings)
+        const control_net_settings = await mapPluginSettingsToControlNet(
+            plugin_settings
+        )
         let control_networks = []
         // let active_control_networks = 0
         for (let index = 0; index < g_controlnet_max_models; index++) {
@@ -281,21 +313,20 @@ export class Txt2ImgMode extends Mode {
                     )
                 }
                 response_json = await this.requestTxt2Img(settings) //this is automatic1111 txt2img
-            } else if (settings_tab.store.data.selected_backend === 'ComfyUI') {
-                //request Txt2Img from comfyui
-                settings = mapPluginSettingsToControlNet(settings)
-                const { image_base64_list, image_url_list } =
-                    await comfyui_main_ui.generateComfyTxt2Img(settings)
-                output_images = image_base64_list
-            }
 
-            if (settings_tab.store.data.selected_backend === 'Automatic1111') {
                 output_images = await this.processOutput(
                     response_json.images_info,
                     settings
                 )
             } else if (settings_tab.store.data.selected_backend === 'ComfyUI') {
-                // output_images = image_base64_list
+                //request Txt2Img from comfyui
+                settings = await mapPluginSettingsToControlNet(settings)
+                const { image_base64_list, image_url_list } =
+                    await comfyui_main_ui.generateComfyTxt2Img(settings)
+                output_images = image_base64_list
+                if (output_images) {
+                    saveOutputImagesToDriveComfy(output_images, settings)
+                }
             }
         } catch (e) {
             console.warn(e)
@@ -322,8 +353,9 @@ export class Img2ImgMode extends Mode {
     //REFACTOR: reuse the same code for (requestControlNetTxt2Img,requestControlNetImg2Img)
     static async requestControlNetImg2Img(plugin_settings: any) {
         const full_url = `${g_sd_url}/sdapi/v1/img2img`
-        const control_net_settings =
-            mapPluginSettingsToControlNet(plugin_settings)
+        const control_net_settings = await mapPluginSettingsToControlNet(
+            plugin_settings
+        )
 
         // let control_networks = 0
         let control_networks = []
@@ -474,7 +506,7 @@ export class Img2ImgMode extends Mode {
                     )
                 }
             } else if (settings_tab.store.data.selected_backend === 'ComfyUI') {
-                settings = mapPluginSettingsToControlNet(settings)
+                settings = await mapPluginSettingsToControlNet(settings)
                 if (settings?.mode === 'img2img') {
                     const { image_base64_list, image_url_list } =
                         await comfyui_main_ui.generateComfyImg2Img(settings)
@@ -488,6 +520,10 @@ export class Img2ImgMode extends Mode {
                         await comfyui_main_ui.generateComfyInpaint(settings)
                     output_images = image_base64_list
                 }
+
+                if (output_images) {
+                    saveOutputImagesToDriveComfy(output_images, settings)
+                }
             }
         } catch (e) {
             console.warn(e)
@@ -552,7 +588,8 @@ export class LassoInpaintMode extends Img2ImgMode {
         }
         const [init_image, mask] = await selection.inpaintLassoInitImageAndMask(
             'mask',
-            sd_tab_util.helper_store.data.lasso_offset
+            sd_tab_util.helper_store.data.lasso_offset,
+            sd_tab_util.helper_store.data.make_square
         )
 
         const selectionInfo = await psapi.getSelectionInfoExe()

typescripts/util/ts/general.ts

@@ -60,3 +60,10 @@ export function base64UrlToBase64(base64_url: string) {
     const base64_image = base64_url.replace('data:image/png;base64,', '')
     return base64_image
 }
+
+export function newOutputImageName(format = 'png') {
+    const random_id = Math.floor(Math.random() * 100000000000 + 1) // Date.now() doesn't have enough resolution to avoid duplicate
+    const image_name = `output- ${Date.now()}-${random_id}.${format}`
+    console.log('generated image name:', image_name)
+    return image_name
+}