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