From c77a5cab5ba8ffc40f61508376879efdc43302e9 Mon Sep 17 00:00:00 2001 From: pythongosssss <125205205+pythongosssss@users.noreply.github.com> Date: Wed, 1 Jan 2025 20:29:17 +0000 Subject: [PATCH] Add support for nested dynamic prompts (#2117) --- src/extensions/core/dynamicPrompts.ts | 33 +------ src/extensions/core/widgetInputs.ts | 3 +- src/utils/formatUtil.ts | 69 +++++++++++++ tests-ui/tests/dynamicPrompts.test.ts | 137 ++++++++++++++++++++++++++ 4 files changed, 213 insertions(+), 29 deletions(-) create mode 100644 tests-ui/tests/dynamicPrompts.test.ts diff --git a/src/extensions/core/dynamicPrompts.ts b/src/extensions/core/dynamicPrompts.ts index 1c1ff6f8f..b06b57599 100644 --- a/src/extensions/core/dynamicPrompts.ts +++ b/src/extensions/core/dynamicPrompts.ts @@ -1,17 +1,10 @@ -// @ts-strict-ignore -import { app } from '../../scripts/app' +import { useExtensionService } from '@/services/extensionService' +import { processDynamicPrompt } from '@/utils/formatUtil' // Allows for simple dynamic prompt replacement // Inputs in the format {a|b} will have a random value of a or b chosen when the prompt is queued. -/* - * Strips C-style line and block comments from a string - */ -function stripComments(str) { - return str.replace(/\/\*[\s\S]*?\*\/|\/\/.*/g, '') -} - -app.registerExtension({ +useExtensionService().registerExtension({ name: 'Comfy.DynamicPrompts', nodeCreated(node) { if (node.widgets) { @@ -23,25 +16,9 @@ app.registerExtension({ // Override the serialization of the value to resolve dynamic prompts for all widgets supporting it in this node // @ts-expect-error hacky override widget.serializeValue = (workflowNode, widgetIndex) => { - let prompt = stripComments(widget.value) - while ( - prompt.replace('\\{', '').includes('{') && - prompt.replace('\\}', '').includes('}') - ) { - const startIndex = prompt.replace('\\{', '00').indexOf('{') - const endIndex = prompt.replace('\\}', '00').indexOf('}') + if (typeof widget.value !== 'string') return widget.value - const optionsString = prompt.substring(startIndex + 1, endIndex) - const options = optionsString.split('|') - - const randomIndex = Math.floor(Math.random() * options.length) - const randomOption = options[randomIndex] - - prompt = - prompt.substring(0, startIndex) + - randomOption + - prompt.substring(endIndex + 1) - } + const prompt = processDynamicPrompt(widget.value) // Overwrite the value in the serialized workflow pnginfo if (workflowNode?.widgets_values) diff --git a/src/extensions/core/widgetInputs.ts b/src/extensions/core/widgetInputs.ts index 644219e1f..f83bc2cf9 100644 --- a/src/extensions/core/widgetInputs.ts +++ b/src/extensions/core/widgetInputs.ts @@ -633,7 +633,8 @@ export function mergeIfValid( k !== 'defaultInput' && k !== 'control_after_generate' && k !== 'multiline' && - k !== 'tooltip' + k !== 'tooltip' && + k !== 'dynamicPrompts' ) { let v1 = config1[1][k] let v2 = config2[1]?.[k] diff --git a/src/utils/formatUtil.ts b/src/utils/formatUtil.ts index fcc13e311..7de9f7164 100644 --- a/src/utils/formatUtil.ts +++ b/src/utils/formatUtil.ts @@ -135,3 +135,72 @@ export function getPathDetails(path: string) { export function normalizeI18nKey(key: string) { return key.replace(/\./g, '_') } + +/** + * Takes a dynamic prompt in the format {opt1|opt2|{optA|optB}|} and randomly replaces groups. Supports C style comments. + * @param input The dynamic prompt to process + * @returns + */ +export function processDynamicPrompt(input: string): string { + /* + * Strips C-style line and block comments from a string + */ + function stripComments(str: string) { + return str.replace(/\/\*[\s\S]*?\*\/|\/\/.*/g, '') + } + + let i = 0 + let result = '' + input = stripComments(input) + + const handleEscape = () => { + const nextChar = input[i++] + return '\\' + nextChar + } + + function parseChoiceBlock() { + // Parse the content inside {} + const options: string[] = [] + let choice = '' + let depth = 0 + + while (i < input.length) { + const char = input[i++] + + if (char === '\\') { + choice += handleEscape() + continue + } else if (char === '{') { + depth++ + } else if (char === '}') { + if (!depth) break + depth-- + } else if (char === '|') { + if (!depth) { + options.push(choice) + choice = '' + continue + } + } + choice += char + } + + options.push(choice) + + const chosenOption = options[Math.floor(Math.random() * options.length)] + return processDynamicPrompt(chosenOption) + } + + while (i < input.length) { + const char = input[i++] + if (char === '\\') { + result += handleEscape() + } else if (char === '{') { + result += parseChoiceBlock() + } else { + result += char + } + } + + return result.replace(/\\([{}|])/g, '$1') +} diff --git a/tests-ui/tests/dynamicPrompts.test.ts b/tests-ui/tests/dynamicPrompts.test.ts new file mode 100644 index 000000000..e6eb49bb9 --- /dev/null +++ b/tests-ui/tests/dynamicPrompts.test.ts @@ -0,0 +1,137 @@ +import { processDynamicPrompt } from '@/utils/formatUtil' + +describe('dynamic prompts', () => { + afterEach(() => { + jest.restoreAllMocks() + }) + + it('handles single and multiline comments', () => { + const input = + '/*\nStart\n*/Hello /* this is a comment */ world!\n// it\nEnd' + expect(processDynamicPrompt(input)).toBe('Hello world!\n\nEnd') + }) + + it('handles simple option groups', () => { + const input = '{option1|option2}' + + jest.spyOn(Math, 'random').mockReturnValue(0) + expect(processDynamicPrompt(input)).toBe('option1') + + jest.spyOn(Math, 'random').mockReturnValue(0.99) + expect(processDynamicPrompt(input)).toBe('option2') + }) + + test('handles trailing empty options', () => { + const input = '{a|}' + jest.spyOn(Math, 'random').mockReturnValue(0.99) + expect(processDynamicPrompt(input)).toBe('') + }) + + test('handles leading empty options', () => { + const input = '{|a}' + jest.spyOn(Math, 'random').mockReturnValue(0) + expect(processDynamicPrompt(input)).toBe('') + }) + + test('handles multiple empty alternatives', () => { + const input = '{||}' + expect(processDynamicPrompt(input)).toBe('') + }) + + test('handles multiple nested empty alternatives', () => { + const input = '{a|{b||c}|}' + jest.spyOn(Math, 'random').mockReturnValue(0.5) + expect(processDynamicPrompt(input)).toBe('') + }) + + test('handles unescaped special characters gracefully', () => { + const input = '{a|\\}' + jest.spyOn(Math, 'random').mockReturnValue(0.99) + expect(processDynamicPrompt(input)).toBe('}') + }) + + it('handles nested option groups', () => { + jest.spyOn(Math, 'random').mockReturnValue(0) // pick the first option at each level + const input = '{a|{b|{c|d}}}' + expect(processDynamicPrompt(input)).toBe('a') + + jest.spyOn(Math, 'random').mockReturnValue(0.99) // pick the last option at each level + expect(processDynamicPrompt(input)).toBe('d') + }) + + test('handles escaped braces', () => { + const input = '\\{a|b\\}' + expect(processDynamicPrompt(input)).toBe('{a|b}') + }) + + test('handles escaped pipe', () => { + const input = 'a\\|b' + expect(processDynamicPrompt(input)).toBe('a|b') + }) + + it('handles escaped characters', () => { + const input = '{\\{escaped\\}\\|escaped pipe}' + expect(processDynamicPrompt(input)).toBe('{escaped}|escaped pipe') + }) + + test('handles deeply nested escaped characters', () => { + const input = '{a|{b|\\{c\\}}}' + jest.spyOn(Math, 'random').mockReturnValue(0.99) + expect(processDynamicPrompt(input)).toBe('{c}') + }) + + it('handles mixed input', () => { + jest.spyOn(Math, 'random').mockReturnValue(0.99) + const input = + '<{option1|option2}>/*comment*/ ({something|else}:2) \\{escaped\\}!' + expect(processDynamicPrompt(input)).toBe(' (else:2) {escaped}!') + }) + + it('handles non-paired braces gracefully', () => { + jest.spyOn(Math, 'random').mockReturnValue(0) + const input = '{option1|option2|{nested1|nested2' + expect(processDynamicPrompt(input)).toBe('option1') + + jest.spyOn(Math, 'random').mockReturnValue(0.4) + expect(processDynamicPrompt(input)).toBe('option2') + + jest.spyOn(Math, 'random').mockReturnValue(0.99) + expect(processDynamicPrompt(input)).toBe('nested2') + }) + + it('handles deep nesting', () => { + const input = '{a|{b|{c|{d|{e|{f|{g}}1}2}3}4}5}' + jest.spyOn(Math, 'random').mockReturnValue(0.99) + expect(processDynamicPrompt(input)).toBe('g12345') + }) + + test('handles empty alternative inside braces', () => { + const input = '{|a||b|}' + jest.spyOn(Math, 'random').mockReturnValue(0) + expect(processDynamicPrompt(input)).toBe('') + jest.spyOn(Math, 'random').mockReturnValue(0.3) + expect(processDynamicPrompt(input)).toBe('a') + jest.spyOn(Math, 'random').mockReturnValue(0.5) + expect(processDynamicPrompt(input)).toBe('') + jest.spyOn(Math, 'random').mockReturnValue(0.75) + expect(processDynamicPrompt(input)).toBe('b') + jest.spyOn(Math, 'random').mockReturnValue(0.999) + expect(processDynamicPrompt(input)).toBe('') + }) + + test('handles no braces', () => { + const input = 'abcdef' + expect(processDynamicPrompt(input)).toBe('abcdef') + }) + + test('handles empty input', () => { + const input = '' + expect(processDynamicPrompt(input)).toBe('') + }) + + test('handles complex mixed cases', () => { + jest.spyOn(Math, 'random').mockReturnValue(0.5) //pick the second option from each group + const input = '1{a|b|{c|d}}2{e|f}3' + expect(processDynamicPrompt(input)).toBe('1b2f3') + }) +})