ComfyUI_frontend/src/extensions/core/editAttention.ts

import { app } from '../../scripts/app'

// Allows you to edit the attention weight by holding ctrl (or cmd) and using the up/down arrow keys

app.registerExtension({
  name: 'Comfy.EditAttention',
  init() {
    const editAttentionDelta = app.ui.settings.addSetting({
      id: 'Comfy.EditAttention.Delta',
      category: ['Comfy', 'EditTokenWeight', 'Delta'],
      name: 'Ctrl+up/down precision',
      type: 'slider',
      attrs: {
        min: 0.01,
        max: 0.5,
        step: 0.01
      },
      defaultValue: 0.05
    })

    function incrementWeight(weight: string, delta: number): string {
      const floatWeight = parseFloat(weight)
      if (isNaN(floatWeight)) return weight
      const newWeight = floatWeight + delta
      return String(Number(newWeight.toFixed(10)))
    }

    type Enclosure = {
      start: number
      end: number
    }

    function findNearestEnclosure(
      text: string,
      cursorPos: number
    ): Enclosure | null {
      let start = cursorPos,
        end = cursorPos
      let openCount = 0,
        closeCount = 0

      // Find opening parenthesis before cursor
      while (start >= 0) {
        start--
        if (text[start] === '(' && openCount === closeCount) break
        if (text[start] === '(') openCount++
        if (text[start] === ')') closeCount++
      }
      if (start < 0) return null

      openCount = 0
      closeCount = 0

      // Find closing parenthesis after cursor
      while (end < text.length) {
        if (text[end] === ')' && openCount === closeCount) break
        if (text[end] === '(') openCount++
        if (text[end] === ')') closeCount++
        end++
      }
      if (end === text.length) return null

      return { start: start + 1, end: end }
    }

    function addWeightToParentheses(text: string): string {
      const parenRegex = /^\((.*)\)$/
      const parenMatch = text.match(parenRegex)

      const floatRegex = /:([+-]?(\d*\.)?\d+([eE][+-]?\d+)?)/
      const floatMatch = text.match(floatRegex)

      if (parenMatch && !floatMatch) {
        return `(${parenMatch[1]}:1.0)`
      } else {
        return text
      }
    }

    function editAttention(event: KeyboardEvent) {
      // @ts-expect-error Runtime narrowing not impl.
      const inputField: HTMLTextAreaElement = event.composedPath()[0]
      const delta = parseFloat(editAttentionDelta.value)

      if (inputField.tagName !== 'TEXTAREA') return
      if (!(event.key === 'ArrowUp' || event.key === 'ArrowDown')) return
      if (!event.ctrlKey && !event.metaKey) return

      event.preventDefault()

      let start = inputField.selectionStart
      let end = inputField.selectionEnd
      let selectedText = inputField.value.substring(start, end)

      // If there is no selection, attempt to find the nearest enclosure, or select the current word
      if (!selectedText) {
        const nearestEnclosure = findNearestEnclosure(inputField.value, start)
        if (nearestEnclosure) {
          start = nearestEnclosure.start
          end = nearestEnclosure.end
          selectedText = inputField.value.substring(start, end)
        } else {
          // Select the current word, find the start and end of the word
          const delimiters = ' .,\\/!?%^*;:{}=-_`~()\r\n\t'

          while (
            !delimiters.includes(inputField.value[start - 1]) &&
            start > 0
          ) {
            start--
          }

          while (
            !delimiters.includes(inputField.value[end]) &&
            end < inputField.value.length
          ) {
            end++
          }

          selectedText = inputField.value.substring(start, end)
          if (!selectedText) return
        }
      }

      // If the selection ends with a space, remove it
      if (selectedText[selectedText.length - 1] === ' ') {
        selectedText = selectedText.substring(0, selectedText.length - 1)
        end -= 1
      }

      // If there are parentheses left and right of the selection, select them
      if (
        inputField.value[start - 1] === '(' &&
        inputField.value[end] === ')'
      ) {
        start -= 1
        end += 1
        selectedText = inputField.value.substring(start, end)
      }

      // If the selection is not enclosed in parentheses, add them
      if (
        selectedText[0] !== '(' ||
        selectedText[selectedText.length - 1] !== ')'
      ) {
        selectedText = `(${selectedText})`
      }

      // If the selection does not have a weight, add a weight of 1.0
      selectedText = addWeightToParentheses(selectedText)

      // Increment the weight
      const weightDelta = event.key === 'ArrowUp' ? delta : -delta
      const updatedText = selectedText.replace(
        /\((.*):([+-]?\d+(?:\.\d+)?)\)/,
        (_, text, weight) => {
          weight = incrementWeight(weight, weightDelta)
          if (weight == 1) {
            return text
          } else {
            return `(${text}:${weight})`
          }
        }
      )

      inputField.setSelectionRange(start, end)
      // Intentional use of deprecated: https://developer.mozilla.org/docs/Web/API/Document/execCommand#using_inserttext
      document.execCommand('insertText', false, updatedText)
      inputField.setSelectionRange(start, start + updatedText.length)
    }
    window.addEventListener('keydown', editAttention)
  }
})