diff --git a/invokeai/frontend/web/src/features/ui/components/tabs/ModelManager/subpanels/MergeModelsPanel.tsx b/invokeai/frontend/web/src/features/ui/components/tabs/ModelManager/subpanels/MergeModelsPanel.tsx index 8c26357720..0cd90a9492 100644 --- a/invokeai/frontend/web/src/features/ui/components/tabs/ModelManager/subpanels/MergeModelsPanel.tsx +++ b/invokeai/frontend/web/src/features/ui/components/tabs/ModelManager/subpanels/MergeModelsPanel.tsx @@ -1,10 +1,260 @@ -import { Flex } from '@chakra-ui/react'; -import MergeModels from 'features/ui/components/tabs/ModelManager/subpanels/MergeModelsPanel/MergeModels'; +import { Flex, Radio, RadioGroup, Text, Tooltip } from '@chakra-ui/react'; +import { RootState } from 'app/store/store'; +import { useAppDispatch, useAppSelector } from 'app/store/storeHooks'; +import IAIButton from 'common/components/IAIButton'; +import IAIInput from 'common/components/IAIInput'; +import IAISelect from 'common/components/IAISelect'; +import IAISimpleCheckbox from 'common/components/IAISimpleCheckbox'; +import IAISlider from 'common/components/IAISlider'; +import { pickBy } from 'lodash-es'; +import { useState } from 'react'; +import { useTranslation } from 'react-i18next'; +import { useListModelsQuery } from 'services/api/endpoints/models'; export default function MergeModelsPanel() { + const { t } = useTranslation(); + + const dispatch = useAppDispatch(); + + const { data } = useListModelsQuery({ + model_type: 'main', + }); + + const diffusersModels = pickBy( + data?.entities, + (value, _) => value?.model_format === 'diffusers' + ); + + const [modelOne, setModelOne] = useState( + Object.keys(diffusersModels)[0] + ); + const [modelTwo, setModelTwo] = useState( + Object.keys(diffusersModels)[1] + ); + const [modelThree, setModelThree] = useState('none'); + + const [mergedModelName, setMergedModelName] = useState(''); + const [modelMergeAlpha, setModelMergeAlpha] = useState(0.5); + + const [modelMergeInterp, setModelMergeInterp] = useState< + 'weighted_sum' | 'sigmoid' | 'inv_sigmoid' | 'add_difference' + >('weighted_sum'); + + const [modelMergeSaveLocType, setModelMergeSaveLocType] = useState< + 'root' | 'custom' + >('root'); + + const [modelMergeCustomSaveLoc, setModelMergeCustomSaveLoc] = + useState(''); + + const [modelMergeForce, setModelMergeForce] = useState(false); + + const modelOneList = Object.keys(diffusersModels).filter( + (model) => model !== modelTwo && model !== modelThree + ); + + const modelTwoList = Object.keys(diffusersModels).filter( + (model) => model !== modelOne && model !== modelThree + ); + + const modelThreeList = [ + { key: t('modelManager.none'), value: 'none' }, + ...Object.keys(diffusersModels) + .filter((model) => model !== modelOne && model !== modelTwo) + .map((model) => ({ key: model, value: model })), + ]; + + const isProcessing = useAppSelector( + (state: RootState) => state.system.isProcessing + ); + + const mergeModelsHandler = () => { + let modelsToMerge: string[] = [modelOne, modelTwo, modelThree]; + modelsToMerge = modelsToMerge.filter((model) => model !== 'none'); + + const mergeModelsInfo: InvokeAI.InvokeModelMergingProps = { + models_to_merge: modelsToMerge, + merged_model_name: + mergedModelName !== '' ? mergedModelName : modelsToMerge.join('-'), + alpha: modelMergeAlpha, + interp: modelMergeInterp, + model_merge_save_path: + modelMergeSaveLocType === 'root' ? null : modelMergeCustomSaveLoc, + force: modelMergeForce, + }; + + dispatch(mergeDiffusersModels(mergeModelsInfo)); + }; + return ( - - + + + {t('modelManager.modelMergeHeaderHelp1')} + + {t('modelManager.modelMergeHeaderHelp2')} + + + + setModelOne(e.target.value)} + /> + setModelTwo(e.target.value)} + /> + { + if (e.target.value !== 'none') { + setModelThree(e.target.value); + setModelMergeInterp('add_difference'); + } else { + setModelThree('none'); + setModelMergeInterp('weighted_sum'); + } + }} + /> + + + setMergedModelName(e.target.value)} + /> + + + setModelMergeAlpha(v)} + withInput + withReset + handleReset={() => setModelMergeAlpha(0.5)} + withSliderMarks + /> + + {t('modelManager.modelMergeAlphaHelp')} + + + + + + {t('modelManager.interpolationType')} + + setModelMergeInterp(v)} + > + + {modelThree === 'none' ? ( + <> + + {t('modelManager.weightedSum')} + + + {t('modelManager.sigmoid')} + + + {t('modelManager.inverseSigmoid')} + + + ) : ( + + + {t('modelManager.addDifference')} + + + )} + + + + + + + + {t('modelManager.mergedModelSaveLocation')} + + setModelMergeSaveLocType(v)} + > + + + {t('modelManager.invokeAIFolder')} + + + + {t('modelManager.custom')} + + + + + + {modelMergeSaveLocType === 'custom' && ( + setModelMergeCustomSaveLoc(e.target.value)} + /> + )} + + + setModelMergeForce(e.target.checked)} + fontWeight="500" + /> + + + {t('modelManager.merge')} + ); } diff --git a/invokeai/frontend/web/src/features/ui/components/tabs/ModelManager/subpanels/MergeModelsPanel/MergeModels.tsx b/invokeai/frontend/web/src/features/ui/components/tabs/ModelManager/subpanels/MergeModelsPanel/MergeModels.tsx deleted file mode 100644 index 219d49d4ee..0000000000 --- a/invokeai/frontend/web/src/features/ui/components/tabs/ModelManager/subpanels/MergeModelsPanel/MergeModels.tsx +++ /dev/null @@ -1,313 +0,0 @@ -import { - Flex, - Modal, - ModalBody, - ModalCloseButton, - ModalContent, - ModalFooter, - ModalHeader, - ModalOverlay, - Radio, - RadioGroup, - Text, - Tooltip, - useDisclosure, -} from '@chakra-ui/react'; -// import { mergeDiffusersModels } from 'app/socketio/actions'; -import { RootState } from 'app/store/store'; -import { useAppDispatch, useAppSelector } from 'app/store/storeHooks'; -import IAIButton from 'common/components/IAIButton'; -import IAIInput from 'common/components/IAIInput'; -import IAISelect from 'common/components/IAISelect'; -import { diffusersModelsSelector } from 'features/system/store/systemSelectors'; -import { useState } from 'react'; -import { useTranslation } from 'react-i18next'; -import * as InvokeAI from 'app/types/invokeai'; -import IAISlider from 'common/components/IAISlider'; -import IAISimpleCheckbox from 'common/components/IAISimpleCheckbox'; - -export default function MergeModels() { - const dispatch = useAppDispatch(); - - const { isOpen, onOpen, onClose } = useDisclosure(); - - const diffusersModels = useAppSelector(diffusersModelsSelector); - - const { t } = useTranslation(); - - const [modelOne, setModelOne] = useState( - Object.keys(diffusersModels)[0] - ); - const [modelTwo, setModelTwo] = useState( - Object.keys(diffusersModels)[1] - ); - const [modelThree, setModelThree] = useState('none'); - - const [mergedModelName, setMergedModelName] = useState(''); - const [modelMergeAlpha, setModelMergeAlpha] = useState(0.5); - - const [modelMergeInterp, setModelMergeInterp] = useState< - 'weighted_sum' | 'sigmoid' | 'inv_sigmoid' | 'add_difference' - >('weighted_sum'); - - const [modelMergeSaveLocType, setModelMergeSaveLocType] = useState< - 'root' | 'custom' - >('root'); - - const [modelMergeCustomSaveLoc, setModelMergeCustomSaveLoc] = - useState(''); - - const [modelMergeForce, setModelMergeForce] = useState(false); - - const modelOneList = Object.keys(diffusersModels).filter( - (model) => model !== modelTwo && model !== modelThree - ); - - const modelTwoList = Object.keys(diffusersModels).filter( - (model) => model !== modelOne && model !== modelThree - ); - - const modelThreeList = [ - { key: t('modelManager.none'), value: 'none' }, - ...Object.keys(diffusersModels) - .filter((model) => model !== modelOne && model !== modelTwo) - .map((model) => ({ key: model, value: model })), - ]; - - const isProcessing = useAppSelector( - (state: RootState) => state.system.isProcessing - ); - - const mergeModelsHandler = () => { - let modelsToMerge: string[] = [modelOne, modelTwo, modelThree]; - modelsToMerge = modelsToMerge.filter((model) => model !== 'none'); - - const mergeModelsInfo: InvokeAI.InvokeModelMergingProps = { - models_to_merge: modelsToMerge, - merged_model_name: - mergedModelName !== '' ? mergedModelName : modelsToMerge.join('-'), - alpha: modelMergeAlpha, - interp: modelMergeInterp, - model_merge_save_path: - modelMergeSaveLocType === 'root' ? null : modelMergeCustomSaveLoc, - force: modelMergeForce, - }; - - dispatch(mergeDiffusersModels(mergeModelsInfo)); - }; - - return ( - <> - - - {t('modelManager.mergeModels')} - - - - - - - {t('modelManager.mergeModels')} - - - - - {t('modelManager.modelMergeHeaderHelp1')} - - {t('modelManager.modelMergeHeaderHelp2')} - - - - setModelOne(e.target.value)} - /> - setModelTwo(e.target.value)} - /> - { - if (e.target.value !== 'none') { - setModelThree(e.target.value); - setModelMergeInterp('add_difference'); - } else { - setModelThree('none'); - setModelMergeInterp('weighted_sum'); - } - }} - /> - - - setMergedModelName(e.target.value)} - /> - - - setModelMergeAlpha(v)} - withInput - withReset - handleReset={() => setModelMergeAlpha(0.5)} - withSliderMarks - /> - - {t('modelManager.modelMergeAlphaHelp')} - - - - - - {t('modelManager.interpolationType')} - - setModelMergeInterp(v)} - > - - {modelThree === 'none' ? ( - <> - - - {t('modelManager.weightedSum')} - - - - {t('modelManager.sigmoid')} - - - - {t('modelManager.inverseSigmoid')} - - - - ) : ( - - - - {t('modelManager.addDifference')} - - - - )} - - - - - - - - {t('modelManager.mergedModelSaveLocation')} - - - setModelMergeSaveLocType(v) - } - > - - - - {t('modelManager.invokeAIFolder')} - - - - - {t('modelManager.custom')} - - - - - - {modelMergeSaveLocType === 'custom' && ( - setModelMergeCustomSaveLoc(e.target.value)} - /> - )} - - - setModelMergeForce(e.target.checked)} - fontWeight="500" - /> - - - {t('modelManager.merge')} - - - - - - - - ); -}