Draw canvas image preview in a widget (#2952)

src/composables/widgets/useImagePreviewWidget.ts

@@ -0,0 +1,277 @@
+import { type LGraphNode, LiteGraph } from '@comfyorg/litegraph'
+import type {
+  IBaseWidget,
+  ICustomWidget,
+  IWidgetOptions
+} from '@comfyorg/litegraph/dist/types/widgets'
+
+import type { InputSpec } from '@/schemas/nodeDef/nodeDefSchemaV2'
+import { calculateImageGrid } from '@/scripts/ui/imagePreview'
+import { ComfyWidgetConstructorV2 } from '@/scripts/widgets'
+import { is_all_same_aspect_ratio } from '@/utils/imageUtil'
+
+const renderPreview = (
+  ctx: CanvasRenderingContext2D,
+  node: LGraphNode,
+  shiftY: number
+) => {
+  const canvas = app.canvas
+  const mouse = canvas.graph_mouse
+
+  if (!canvas.pointer_is_down && node.pointerDown) {
+    if (
+      mouse[0] === node.pointerDown.pos[0] &&
+      mouse[1] === node.pointerDown.pos[1]
+    ) {
+      node.imageIndex = node.pointerDown.index
+    }
+    node.pointerDown = null
+  }
+
+  const imgs = node.imgs ?? []
+  let { imageIndex } = node
+  const numImages = imgs.length
+  if (numImages === 1 && !imageIndex) {
+    // This skips the thumbnail render section below
+    node.imageIndex = imageIndex = 0
+  }
+
+  const IMAGE_TEXT_SIZE_TEXT_HEIGHT = 15
+  const dw = node.size[0]
+  const dh = node.size[1] - shiftY - IMAGE_TEXT_SIZE_TEXT_HEIGHT
+
+  if (imageIndex == null) {
+    // No image selected; draw thumbnails of all
+    let cellWidth: number
+    let cellHeight: number
+    let shiftX: number
+    let cell_padding: number
+    let cols: number
+
+    const compact_mode = is_all_same_aspect_ratio(imgs)
+    if (!compact_mode) {
+      // use rectangle cell style and border line
+      cell_padding = 2
+      // Prevent infinite canvas2d scale-up
+      const largestDimension = imgs.reduce(
+        (acc, current) =>
+          Math.max(acc, current.naturalWidth, current.naturalHeight),
+        0
+      )
+      const fakeImgs = []
+      fakeImgs.length = imgs.length
+      fakeImgs[0] = {
+        naturalWidth: largestDimension,
+        naturalHeight: largestDimension
+      }
+      ;({ cellWidth, cellHeight, cols, shiftX } = calculateImageGrid(
+        fakeImgs,
+        dw,
+        dh
+      ))
+    } else {
+      cell_padding = 0
+      ;({ cellWidth, cellHeight, cols, shiftX } = calculateImageGrid(
+        imgs,
+        dw,
+        dh
+      ))
+    }
+
+    let anyHovered = false
+    node.imageRects = []
+    for (let i = 0; i < numImages; i++) {
+      const img = imgs[i]
+      const row = Math.floor(i / cols)
+      const col = i % cols
+      const x = col * cellWidth + shiftX
+      const y = row * cellHeight + shiftY
+      if (!anyHovered) {
+        anyHovered = LiteGraph.isInsideRectangle(
+          mouse[0],
+          mouse[1],
+          x + node.pos[0],
+          y + node.pos[1],
+          cellWidth,
+          cellHeight
+        )
+        if (anyHovered) {
+          node.overIndex = i
+          let value = 110
+          if (canvas.pointer_is_down) {
+            if (!node.pointerDown || node.pointerDown.index !== i) {
+              node.pointerDown = { index: i, pos: [...mouse] }
+            }
+            value = 125
+          }
+          ctx.filter = `contrast(${value}%) brightness(${value}%)`
+          canvas.canvas.style.cursor = 'pointer'
+        }
+      }
+      node.imageRects.push([x, y, cellWidth, cellHeight])
+
+      const wratio = cellWidth / img.width
+      const hratio = cellHeight / img.height
+      const ratio = Math.min(wratio, hratio)
+
+      const imgHeight = ratio * img.height
+      const imgY = row * cellHeight + shiftY + (cellHeight - imgHeight) / 2
+      const imgWidth = ratio * img.width
+      const imgX = col * cellWidth + shiftX + (cellWidth - imgWidth) / 2
+
+      ctx.drawImage(
+        img,
+        imgX + cell_padding,
+        imgY + cell_padding,
+        imgWidth - cell_padding * 2,
+        imgHeight - cell_padding * 2
+      )
+      if (!compact_mode) {
+        // rectangle cell and border line style
+        ctx.strokeStyle = '#8F8F8F'
+        ctx.lineWidth = 1
+        ctx.strokeRect(
+          x + cell_padding,
+          y + cell_padding,
+          cellWidth - cell_padding * 2,
+          cellHeight - cell_padding * 2
+        )
+      }
+
+      ctx.filter = 'none'
+    }
+
+    if (!anyHovered) {
+      node.pointerDown = null
+      node.overIndex = null
+    }
+
+    return
+  }
+  // Draw individual
+  const img = imgs[imageIndex]
+  let w = img.naturalWidth
+  let h = img.naturalHeight
+
+  const scaleX = dw / w
+  const scaleY = dh / h
+  const scale = Math.min(scaleX, scaleY, 1)
+
+  w *= scale
+  h *= scale
+
+  const x = (dw - w) / 2
+  const y = (dh - h) / 2 + shiftY
+  ctx.drawImage(img, x, y, w, h)
+
+  // Draw image size text below the image
+  ctx.fillStyle = LiteGraph.NODE_TEXT_COLOR
+  ctx.textAlign = 'center'
+  ctx.font = '10px sans-serif'
+  const sizeText = `${Math.round(img.naturalWidth)} × ${Math.round(img.naturalHeight)}`
+  const textY = y + h + 10
+  ctx.fillText(sizeText, x + w / 2, textY)
+
+  const drawButton = (
+    x: number,
+    y: number,
+    sz: number,
+    text: string
+  ): boolean => {
+    const hovered = LiteGraph.isInsideRectangle(
+      mouse[0],
+      mouse[1],
+      x + node.pos[0],
+      y + node.pos[1],
+      sz,
+      sz
+    )
+    let fill = '#333'
+    let textFill = '#fff'
+    let isClicking = false
+    if (hovered) {
+      canvas.canvas.style.cursor = 'pointer'
+      if (canvas.pointer_is_down) {
+        fill = '#1e90ff'
+        isClicking = true
+      } else {
+        fill = '#eee'
+        textFill = '#000'
+      }
+    }
+
+    ctx.fillStyle = fill
+    ctx.beginPath()
+    ctx.roundRect(x, y, sz, sz, [4])
+    ctx.fill()
+    ctx.fillStyle = textFill
+    ctx.font = '12px Arial'
+    ctx.textAlign = 'center'
+    ctx.fillText(text, x + 15, y + 20)
+
+    return isClicking
+  }
+
+  if (!(numImages > 1)) return
+
+  const imageNum = (node.imageIndex ?? 0) + 1
+  if (drawButton(dw - 40, dh + shiftY - 40, 30, `${imageNum}/${numImages}`)) {
+    const i = imageNum >= numImages ? 0 : imageNum
+    if (!node.pointerDown || node.pointerDown.index !== i) {
+      node.pointerDown = { index: i, pos: [...mouse] }
+    }
+  }
+
+  if (drawButton(dw - 40, shiftY + 10, 30, `x`)) {
+    if (!node.pointerDown || node.pointerDown.index !== null) {
+      node.pointerDown = { index: null, pos: [...mouse] }
+    }
+  }
+}
+
+class ImagePreviewWidget implements ICustomWidget {
+  readonly type: 'custom'
+  readonly name: string
+  readonly options: IWidgetOptions<unknown>
+  // Dummy value to satisfy type requirements
+  value: string
+
+  constructor(name: string, options: IWidgetOptions<unknown>) {
+    this.type = 'custom'
+    this.name = name
+    this.options = options
+    this.value = ''
+  }
+
+  draw(
+    ctx: CanvasRenderingContext2D,
+    node: LGraphNode,
+    widget_width: number,
+    y: number,
+    H: number
+  ): void {
+    renderPreview(ctx, node, y)
+  }
+
+  computeLayoutSize(this: IBaseWidget, node: LGraphNode) {
+    return {
+      minHeight: 220,
+      minWidth: 1
+    }
+  }
+}
+
+export const useImagePreviewWidget = () => {
+  const widgetConstructor: ComfyWidgetConstructorV2 = (
+    node: LGraphNode,
+    inputSpec: InputSpec
+  ) => {
+    return node.addCustomWidget(
+      new ImagePreviewWidget(inputSpec.name, {
+        serialize: false
+      })
+    )
+  }
+
+  return widgetConstructor
+}