From 7f7a62f832c865a13b9181daee79d3cfc21581e2 Mon Sep 17 00:00:00 2001
From: christian-byrne <abolkonsky.rem@gmail.com>
Date: Fri, 3 Jan 2025 14:04:44 -0700
Subject: [PATCH] =?UTF-8?q?feat:=20=E2=9C=A8=20live=20update=20outputs=20g?=
 =?UTF-8?q?rid?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 types/typedefs.js               | 47 +++++++++++++++++++
 web/mtb_input_output_sidebar.js | 80 ++++++++++++++++++++++++++++++++-
 2 files changed, 125 insertions(+), 2 deletions(-)

diff --git a/types/typedefs.js b/types/typedefs.js
index 4883fd0..ac3348f 100644
--- a/types/typedefs.js
+++ b/types/typedefs.js
@@ -16,3 +16,50 @@
  * @typedef {import("./shared.d.ts").INodeOutputSlot} INodeOutputSlot
  */
 
+/**
+ * @typedef {Object} ResultItem
+ * @property {string} [filename] - The filename of the item.
+ * @property {string} [subfolder] - The subfolder of the item.
+ * @property {string} [type] - The type of the item.
+ */
+
+/**
+ * @typedef {Object} Outputs
+ * @property {ResultItem[]} [audio] - Audio result items.
+ * @property {ResultItem[]} [images] - Image result items.
+ * @property {ResultItem[]} [animated] - Animated result items.
+ */
+
+/**
+ * @typedef {Record<string, Outputs>} TaskOutput
+ * - A record mapping Node IDs to their Outputs.
+ */
+
+/**
+ * @typedef {Array} TaskPrompt
+ * @property {QueueIndex} [0] - The queue index.
+ * @property {PromptId} [1] - The unique prompt ID.
+ * @property {PromptInputs} [2] - The prompt inputs.
+ * @property {ExtraData} [3] - Extra data.
+ * @property {OutputsToExecute} [4] - The outputs to execute.
+ */
+
+/**
+ * @typedef {Object} HistoryTaskItem
+ * @property {'History'} taskType - The type of task.
+ * @property {TaskPrompt} prompt - The task prompt.
+ * @property {Status} [status] - The status of the task.
+ * @property {TaskOutput} outputs - The task outputs.
+ * @property {TaskMeta} [meta] - Optional task metadata.
+ */
+
+/**
+ * @typedef {Object} ExecInfo
+ * @property {number} queue_remaining - The number of items remaining in the queue.
+ */
+
+/**
+ * @typedef {Object} StatusWsMessageStatus
+ * @property {ExecInfo} exec_info - Execution information.
+ */
+
diff --git a/web/mtb_input_output_sidebar.js b/web/mtb_input_output_sidebar.js
index 0b18aec..a980c30 100644
--- a/web/mtb_input_output_sidebar.js
+++ b/web/mtb_input_output_sidebar.js
@@ -1,3 +1,5 @@
+/// <reference path="../types/typedefs.js" />
+
 import { app } from '../../scripts/app.js'
 import { api } from '../../scripts/api.js'
 
@@ -20,6 +22,7 @@ let currentSort = 'None'
 
 const IMAGE_NODES = ['LoadImage', 'VHS_LoadImagePath']
 const VIDEO_NODES = ['VHS_LoadVideo']
+const PROCESSED_PROMPT_IDS = new Set()
 
 const updateImage = (node, image) => {
   if (IMAGE_NODES.includes(node.type)) {
@@ -38,7 +41,74 @@ const updateImage = (node, image) => {
   }
 }
 
-const getImgsFromUrls = (urls, target) => {
+/**
+ * Converts a result item to a request url.
+ * @param {ResultItem} resultItem
+ * @returns {string} - The request URL.
+ */
+const resultItemToQuery = (resultItem) =>
+  [
+    `/mtb/view?filename=${resultItem.filename}`,
+    `width=512`,
+    `type=${resultItem.type}`,
+    `subfolder=${resultItem.subfolder}`,
+    `preview=`,
+  ].join('&')
+
+/**
+ * Retrieves the unique prompt ID from a history task item.
+ * @param {HistoryTaskItem} historyTaskItem
+ * @returns {string} - The prompt ID.
+ */
+const getPromptId = (historyTaskItem) => `${historyTaskItem.prompt[1]}`
+
+/**
+ * Process and return any new/unseen outputs from the most recent history item.
+ * @param {HistoryTaskItem} mostRecentTask - The most recent history task item.
+ * @returns {Object<string, string>} - A map of task outputs URLs.
+ */
+const getNewOutputUrls = (mostRecentTask) => {
+  if (!mostRecentTask) return
+
+  const promptId = getPromptId(mostRecentTask)
+  if (PROCESSED_PROMPT_IDS.has(promptId)) return
+
+  const urls = {}
+  for (const nodeOutputs of Object.values(mostRecentTask.outputs)) {
+    const { images, audio, animated } = nodeOutputs
+    if (images) {
+      const imageOutputs = Object.values(nodeOutputs.images)
+      imageOutputs.forEach(
+        (resultItem) =>
+          (urls[resultItem.filename] = resultItemToQuery(resultItem))
+      )
+    }
+    // Can process `animated` and `audio` outputs here.
+  }
+
+  const foundNewOutputs = Object.keys(urls).length > 0
+  if (!foundNewOutputs) return null
+
+  PROCESSED_PROMPT_IDS.add(promptId)
+  return urls
+}
+
+/** Fetch history and update the grid with any new ouput images. */
+const updateOutputsGrid = async () => {
+  try {
+    const history = await api.getHistory(/** maxSize: */ 1)
+    const mostRcentTask = history.History[0]
+    const newUrls = getNewOutputUrls(mostRcentTask)
+    if (newUrls) {
+      const imgGrid = document.querySelector('.mtb_img_grid')
+      getImgsFromUrls(newUrls, imgGrid, { prepend: true })
+    }
+  } catch (error) {
+    console.error('Error fetching history:', error)
+  }
+}
+
+const getImgsFromUrls = (urls, target, options = { prepend: false }) => {
   const imgs = []
   if (urls === undefined) {
     return imgs
@@ -123,7 +193,8 @@ const getImgsFromUrls = (urls, target) => {
     imgs.push(a)
   }
   if (target !== undefined) {
-    target.append(...imgs)
+    if (options.prepend) target.prepend(...imgs)
+    else target.append(...imgs)
   }
   return imgs
 }
@@ -324,11 +395,16 @@ if (window?.__COMFYUI_FRONTEND_VERSION__) {
             }
           })
           handle = renderSidebar(el, cont, [selector, imgGrid, imgTools])
+          app.api.addEventListener('status', async () => {
+            if (currentMode !== 'output') return
+            updateOutputsGrid()
+          })
         },
         destroy: () => {
           if (handle) {
             handle.unregister()
             handle = undefined
+            app.api.removeEventListener('status')
           }
         },
       })