Add support for nested dynamic prompts (#2117)

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)

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]

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')
+}

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('<option2> (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')
+  })
+})