diff --git a/src/core/graph/widgets/dynamicTypes.ts b/src/core/graph/widgets/dynamicTypes.ts new file mode 100644 index 0000000000..92a7850b96 --- /dev/null +++ b/src/core/graph/widgets/dynamicTypes.ts @@ -0,0 +1,35 @@ +import { transformInputSpecV1ToV2 } from '@/schemas/nodeDef/migration' +import { zAutogrowOptions, zMatchTypeOptions } from '@/schemas/nodeDefSchema' +import type { InputSpec } from '@/schemas/nodeDefSchema' +import type { InputSpec as InputSpecV2 } from '@/schemas/nodeDef/nodeDefSchemaV2' + +const dynamicTypeResolvers: Record< + string, + (inputSpec: InputSpecV2) => string[] +> = { + COMFY_AUTOGROW_V3: resolveAutogrowType, + COMFY_MATCHTYPE_V3: (input) => + zMatchTypeOptions + .safeParse(input) + .data?.template?.allowed_types?.split(',') ?? [] +} + +export function resolveInputType(input: InputSpecV2): string[] { + return input.type in dynamicTypeResolvers + ? dynamicTypeResolvers[input.type](input) + : input.type.split(',') +} + +function resolveAutogrowType(rawSpec: InputSpecV2): string[] { + const { input } = zAutogrowOptions.safeParse(rawSpec).data?.template ?? {} + + const inputTypes: (Record | undefined)[] = [ + input?.required, + input?.optional + ] + return inputTypes.flatMap((inputType) => + Object.entries(inputType ?? {}).flatMap(([name, v]) => + resolveInputType(transformInputSpecV1ToV2(v, { name })) + ) + ) +} diff --git a/src/core/graph/widgets/dynamicWidgets.ts b/src/core/graph/widgets/dynamicWidgets.ts index 21913dedd0..6d90c2f56c 100644 --- a/src/core/graph/widgets/dynamicWidgets.ts +++ b/src/core/graph/widgets/dynamicWidgets.ts @@ -16,7 +16,8 @@ import type { ComboInputSpec, InputSpec } from '@/schemas/nodeDefSchema' import type { InputSpec as InputSpecV2 } from '@/schemas/nodeDef/nodeDefSchemaV2' import { zAutogrowOptions, - zDynamicComboInputSpec + zDynamicComboInputSpec, + zMatchTypeOptions } from '@/schemas/nodeDefSchema' import { useLitegraphService } from '@/services/litegraphService' import { app } from '@/scripts/app' @@ -215,6 +216,7 @@ export function applyDynamicInputs( dynamicInputs[inputSpec.type](node, inputSpec) return true } + function spliceInputs( node: LGraphNode, startIndex: number, @@ -329,11 +331,10 @@ function withComfyMatchType(node: LGraphNode): asserts node is MatchTypeNode { function applyMatchType(node: LGraphNode, inputSpec: InputSpecV2) { const { addNodeInput } = useLitegraphService() const name = inputSpec.name - const { allowed_types, template_id } = ( - inputSpec as InputSpecV2 & { - template: { allowed_types: string; template_id: string } - } - ).template + const matchTypeSpec = zMatchTypeOptions.safeParse(inputSpec).data + if (!matchTypeSpec) return + + const { allowed_types, template_id } = matchTypeSpec.template const typedSpec = { ...inputSpec, type: allowed_types } addNodeInput(node, typedSpec) withComfyMatchType(node) diff --git a/src/schemas/nodeDefSchema.ts b/src/schemas/nodeDefSchema.ts index 2fdfeeada9..6120641af2 100644 --- a/src/schemas/nodeDefSchema.ts +++ b/src/schemas/nodeDefSchema.ts @@ -341,6 +341,15 @@ export const zDynamicComboInputSpec = z.tuple([ }) ]) +export const zMatchTypeOptions = z.object({ + ...zBaseInputOptions.shape, + type: z.literal('COMFY_MATCHTYPE_V3'), + template: z.object({ + allowed_types: z.string(), + template_id: z.string() + }) +}) + // `/object_info` export type ComfyInputsSpec = z.infer export type ComfyOutputTypesSpec = z.infer diff --git a/src/services/nodeSearchService.ts b/src/services/nodeSearchService.ts index 75dccf2687..f9ca61e5bc 100644 --- a/src/services/nodeSearchService.ts +++ b/src/services/nodeSearchService.ts @@ -34,10 +34,7 @@ export class NodeSearchService { id: 'input', name: 'Input Type', invokeSequence: 'i', - getItemOptions: (node) => - Object.values(node.inputs ?? []).flatMap((input) => - input.type.split(',') - ), + getItemOptions: (node) => node.inputTypes, fuseOptions }) diff --git a/src/stores/nodeDefStore.ts b/src/stores/nodeDefStore.ts index e43267f50c..2c0703948a 100644 --- a/src/stores/nodeDefStore.ts +++ b/src/stores/nodeDefStore.ts @@ -6,6 +6,7 @@ import { computed, ref, watchEffect } from 'vue' import { t } from '@/i18n' import { isPromotedWidgetView } from '@/core/graph/subgraph/promotedWidgetTypes' import { resolvePromotedWidgetSource } from '@/core/graph/subgraph/resolvePromotedWidgetSource' +import { resolveInputType } from '@/core/graph/widgets/dynamicTypes' import { LiteGraph } from '@/lib/litegraph/src/litegraph' import type { LGraphNode } from '@/lib/litegraph/src/litegraph' import { transformNodeDefV1ToV2 } from '@/schemas/nodeDef/migration' @@ -98,6 +99,7 @@ export class ComfyNodeDefImpl // ComfyNodeDefImpl fields readonly nodeSource: NodeSource + readonly inputTypes: string[] /** * @internal @@ -179,6 +181,9 @@ export class ComfyNodeDefImpl // Initialize node source this.nodeSource = getNodeSource(obj.python_module, this.essentials_category) + this.inputTypes = _.uniq( + Object.values(this.inputs).flatMap(resolveInputType) + ) } get nodePath(): string {