Skip to content

Commit

Permalink
feat(ui): show batch group in node title
Browse files Browse the repository at this point in the history
  • Loading branch information
psychedelicious committed Jan 14, 2025
1 parent 789540d commit 4a08374
Show file tree
Hide file tree
Showing 5 changed files with 103 additions and 4 deletions.
1 change: 1 addition & 0 deletions invokeai/frontend/web/public/locales/en.json
Original file line number Diff line number Diff line change
Expand Up @@ -854,6 +854,7 @@
"defaultVAE": "Default VAE"
},
"nodes": {
"noBatchGroup": "no group",
"addNode": "Add Node",
"addNodeToolTip": "Add Node (Shift+A, Space)",
"addLinearView": "Add to Linear View",
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,24 @@
import type { TextProps } from '@invoke-ai/ui-library';
import { Text } from '@invoke-ai/ui-library';
import { useBatchGroupColorToken } from 'features/nodes/hooks/useBatchGroupColorToken';
import { memo } from 'react';

type Props = TextProps & {
batchGroupId?: string;
};

export const BatchGroupId = memo(({ batchGroupId, ...rest }: Props) => {
const batchGroupColorToken = useBatchGroupColorToken(batchGroupId);

if (!batchGroupColorToken || !batchGroupId) {
return null;
}

return (
<Text fontWeight="semibold" color={batchGroupColorToken} {...rest}>
{batchGroupId}
</Text>
);
});

BatchGroupId.displayName = 'BatchGroupId';
Original file line number Diff line number Diff line change
@@ -1,12 +1,14 @@
import type { SystemStyleObject } from '@invoke-ai/ui-library';
import { Box, Editable, EditableInput, EditablePreview, Flex, useEditableControls } from '@invoke-ai/ui-library';
import type { SystemStyleObject, TextProps } from '@invoke-ai/ui-library';
import { Box, Editable, EditableInput, Flex, Text, useEditableControls } from '@invoke-ai/ui-library';
import { useAppDispatch } from 'app/store/storeHooks';
import { useBatchGroupColorToken } from 'features/nodes/hooks/useBatchGroupColorToken';
import { useBatchGroupId } from 'features/nodes/hooks/useBatchGroupId';
import { useNodeLabel } from 'features/nodes/hooks/useNodeLabel';
import { useNodeTemplateTitle } from 'features/nodes/hooks/useNodeTemplateTitle';
import { nodeLabelChanged } from 'features/nodes/store/nodesSlice';
import { DRAG_HANDLE_CLASSNAME } from 'features/nodes/types/constants';
import type { MouseEvent } from 'react';
import { memo, useCallback, useEffect, useState } from 'react';
import { memo, useCallback, useEffect, useMemo, useState } from 'react';
import { useTranslation } from 'react-i18next';

type Props = {
Expand All @@ -17,6 +19,8 @@ type Props = {
const NodeTitle = ({ nodeId, title }: Props) => {
const dispatch = useAppDispatch();
const label = useNodeLabel(nodeId);
const batchGroupId = useBatchGroupId(nodeId);
const batchGroupColorToken = useBatchGroupColorToken(batchGroupId);
const templateTitle = useNodeTemplateTitle(nodeId);
const { t } = useTranslation();

Expand All @@ -29,6 +33,16 @@ const NodeTitle = ({ nodeId, title }: Props) => {
[dispatch, nodeId, title, templateTitle, label, t]
);

const localTitleWithBatchGroupId = useMemo(() => {
if (!batchGroupId) {
return localTitle;
}
if (batchGroupId === 'None') {
return `${localTitle} (${t('nodes.noBatchGroup')})`;
}
return `${localTitle} (${batchGroupId})`;
}, [batchGroupId, localTitle, t]);

const handleChange = useCallback((newTitle: string) => {
setLocalTitle(newTitle);
}, []);
Expand All @@ -50,7 +64,16 @@ const NodeTitle = ({ nodeId, title }: Props) => {
w="full"
h="full"
>
<EditablePreview fontSize="sm" p={0} w="full" noOfLines={1} />
<Preview
fontSize="sm"
p={0}
w="full"
noOfLines={1}
color={batchGroupColorToken}
fontWeight={batchGroupId ? 'semibold' : undefined}
>
{localTitleWithBatchGroupId}
</Preview>
<EditableInput className="nodrag" fontSize="sm" sx={editableInputStyles} />
<EditableControls />
</Editable>
Expand All @@ -60,6 +83,16 @@ const NodeTitle = ({ nodeId, title }: Props) => {

export default memo(NodeTitle);

const Preview = (props: TextProps) => {
const { isEditing } = useEditableControls();

if (isEditing) {
return null;
}

return <Text {...props} />;
};

function EditableControls() {
const { isEditing, getEditButtonProps } = useEditableControls();
const handleDoubleClick = useCallback(
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,22 @@
import { useMemo } from 'react';

export const useBatchGroupColorToken = (batchGroupId?: string) => {
const batchGroupColorToken = useMemo(() => {
switch (batchGroupId) {
case 'Group 1':
return 'invokeGreen.300';
case 'Group 2':
return 'invokeBlue.300';
case 'Group 3':
return 'invokePurple.200';
case 'Group 4':
return 'invokeRed.300';
case 'Group 5':
return 'invokeYellow.300';
default:
return undefined;
}
}, [batchGroupId]);

return batchGroupColorToken;
};
19 changes: 19 additions & 0 deletions invokeai/frontend/web/src/features/nodes/hooks/useBatchGroupId.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,19 @@
import { useNode } from 'features/nodes/hooks/useNode';
import { isBatchNode, isInvocationNode } from 'features/nodes/types/invocation';
import { useMemo } from 'react';

export const useBatchGroupId = (nodeId: string) => {
const node = useNode(nodeId);

const batchGroupId = useMemo(() => {
if (!isInvocationNode(node)) {
return;
}
if (!isBatchNode(node)) {
return;
}
return node.data.inputs['batch_group_id']?.value as string;
}, [node]);

return batchGroupId;
};

0 comments on commit 4a08374

Please sign in to comment.