From 5f3afa3776b17174e09b7a1fd00b539f774b2f5c Mon Sep 17 00:00:00 2001
From: Chenlei Hu <huchenlei@proton.me>
Date: Tue, 8 Oct 2024 17:10:44 -0400
Subject: [PATCH] Supports VHS video outputs in queue sidebar tab (#1174)

* Properly identify gifs

* Detect VHS video

* Basic video support in queue

* Video in lightbox

* Preview button

* nit

* Fix vitest
---
 .../sidebar/tabs/QueueSidebarTab.vue          |  2 +-
 .../sidebar/tabs/queue/ResultGallery.vue      |  5 ++
 .../sidebar/tabs/queue/ResultItem.vue         | 41 ++++++++--------
 .../queue/__tests__/ResultGallery.spec.ts     | 11 ++---
 src/i18n.ts                                   |  2 +
 src/stores/queueStore.ts                      | 48 +++++++++++++++----
 6 files changed, 72 insertions(+), 37 deletions(-)

diff --git a/src/components/sidebar/tabs/QueueSidebarTab.vue b/src/components/sidebar/tabs/QueueSidebarTab.vue
index 765ead9f5..64539aa3c 100644
--- a/src/components/sidebar/tabs/QueueSidebarTab.vue
+++ b/src/components/sidebar/tabs/QueueSidebarTab.vue
@@ -249,7 +249,7 @@ const menuItems = computed<MenuItem[]>(() => [
   {
     label: t('loadWorkflow'),
     icon: 'pi pi-file-export',
-    command: () => menuTargetTask.value?.loadWorkflow()
+    command: () => menuTargetTask.value?.loadWorkflow(app)
   },
   {
     label: t('goToNode'),
diff --git a/src/components/sidebar/tabs/queue/ResultGallery.vue b/src/components/sidebar/tabs/queue/ResultGallery.vue
index c842976ef..65815c985 100644
--- a/src/components/sidebar/tabs/queue/ResultGallery.vue
+++ b/src/components/sidebar/tabs/queue/ResultGallery.vue
@@ -25,7 +25,12 @@
         :src="item.url"
         :contain="false"
         class="galleria-image"
+        v-if="item.isImage"
       />
+      <video v-else-if="item.isVideo" controls width="100%" height="100%">
+        <source :src="item.url" :type="item.format" />
+        {{ $t('videoFailedToLoad') }}
+      </video>
     </template>
   </Galleria>
 </template>
diff --git a/src/components/sidebar/tabs/queue/ResultItem.vue b/src/components/sidebar/tabs/queue/ResultItem.vue
index b27593372..1fd59779a 100644
--- a/src/components/sidebar/tabs/queue/ResultItem.vue
+++ b/src/components/sidebar/tabs/queue/ResultItem.vue
@@ -1,27 +1,30 @@
 <template>
   <div class="result-container" ref="resultContainer">
-    <template
-      v-if="result.mediaType === 'images' || result.mediaType === 'gifs'"
-    >
-      <ComfyImage
-        :src="result.url"
-        class="task-output-image"
-        :contain="imageFit === 'contain'"
-      />
-      <div class="image-preview-mask">
-        <Button
-          icon="pi pi-eye"
-          severity="secondary"
-          @click="emit('preview', result)"
-          rounded
-        />
-      </div>
+    <ComfyImage
+      v-if="result.isImage"
+      :src="result.url"
+      class="task-output-image"
+      :contain="imageFit === 'contain'"
+    />
+    <template v-else-if="result.isVideo">
+      <video controls width="100%" height="100%">
+        <source :src="result.url" :type="result.format" />
+        {{ $t('videoFailedToLoad') }}
+      </video>
     </template>
-    <!-- TODO: handle more media types -->
     <div v-else class="task-result-preview">
       <i class="pi pi-file"></i>
       <span>{{ result.mediaType }}</span>
     </div>
+
+    <div v-if="result.supportsPreview" class="preview-mask">
+      <Button
+        icon="pi pi-eye"
+        severity="secondary"
+        @click="emit('preview', result)"
+        rounded
+      />
+    </div>
   </div>
 </template>
 
@@ -67,7 +70,7 @@ onMounted(() => {
   align-items: center;
 }
 
-.image-preview-mask {
+.preview-mask {
   position: absolute;
   left: 50%;
   top: 50%;
@@ -80,7 +83,7 @@ onMounted(() => {
   z-index: 1;
 }
 
-.result-container:hover .image-preview-mask {
+.result-container:hover .preview-mask {
   opacity: 1;
 }
 </style>
diff --git a/src/components/sidebar/tabs/queue/__tests__/ResultGallery.spec.ts b/src/components/sidebar/tabs/queue/__tests__/ResultGallery.spec.ts
index 80e1a4cbe..2bfefed45 100644
--- a/src/components/sidebar/tabs/queue/__tests__/ResultGallery.spec.ts
+++ b/src/components/sidebar/tabs/queue/__tests__/ResultGallery.spec.ts
@@ -12,15 +12,12 @@ describe('ResultGallery', () => {
   let mockResultItem: ResultItemImpl
 
   beforeEach(() => {
-    mockResultItem = {
+    mockResultItem = new ResultItemImpl({
       filename: 'test.jpg',
       type: 'images',
-      nodeId: 'test',
-      mediaType: 'images',
-      url: 'https://picsum.photos/200/300',
-      urlWithTimestamp: 'https://picsum.photos/200/300?t=123456',
-      supportsPreview: true
-    }
+      nodeId: 1,
+      mediaType: 'images'
+    })
   })
 
   const mountResultGallery = (props: ResultGalleryProps, options = {}) => {
diff --git a/src/i18n.ts b/src/i18n.ts
index b73289669..5d744b5c6 100644
--- a/src/i18n.ts
+++ b/src/i18n.ts
@@ -2,6 +2,7 @@ import { createI18n } from 'vue-i18n'
 
 const messages = {
   en: {
+    videoFailedToLoad: 'Video failed to load',
     extensionName: 'Extension Name',
     reloadToApplyChanges: 'Reload to apply changes',
     insert: 'Insert',
@@ -112,6 +113,7 @@ const messages = {
     }
   },
   zh: {
+    videoFailedToLoad: '视频加载失败',
     extensionName: '扩展名称',
     reloadToApplyChanges: '重新加载以应用更改',
     insert: '插入',
diff --git a/src/stores/queueStore.ts b/src/stores/queueStore.ts
index 04c7aee55..73625338b 100644
--- a/src/stores/queueStore.ts
+++ b/src/stores/queueStore.ts
@@ -1,5 +1,5 @@
 import { api } from '@/scripts/api'
-import { app } from '@/scripts/app'
+import type { ComfyApp } from '@/scripts/app'
 import type {
   TaskItem,
   TaskType,
@@ -10,7 +10,6 @@ import type {
   ResultItem
 } from '@/types/apiTypes'
 import type { NodeId } from '@/types/comfyWorkflow'
-import { plainToClass } from 'class-transformer'
 import _ from 'lodash'
 import { defineStore } from 'pinia'
 import { toRaw } from 'vue'
@@ -35,6 +34,22 @@ export class ResultItemImpl {
   // 'audio' | 'images' | ...
   mediaType: string
 
+  // VHS output specific fields
+  format?: string
+  frame_rate?: number
+
+  constructor(obj: Record<string, any>) {
+    this.filename = obj.filename
+    this.subfolder = obj.subfolder
+    this.type = obj.type
+
+    this.nodeId = obj.nodeId
+    this.mediaType = obj.mediaType
+
+    this.format = obj.format
+    this.frame_rate = obj.frame_rate
+  }
+
   get url(): string {
     return api.apiURL(`/view?filename=${encodeURIComponent(this.filename)}&type=${this.type}&
 					subfolder=${encodeURIComponent(this.subfolder || '')}`)
@@ -44,8 +59,20 @@ export class ResultItemImpl {
     return `${this.url}&t=${+new Date()}`
   }
 
+  get isVideo(): boolean {
+    return this.format && this.format.startsWith('video/')
+  }
+
+  get isGif(): boolean {
+    return this.filename.endsWith('.gif')
+  }
+
+  get isImage(): boolean {
+    return this.mediaType === 'images' || this.isGif
+  }
+
   get supportsPreview(): boolean {
-    return ['images', 'gifs'].includes(this.mediaType)
+    return this.isImage || this.isVideo
   }
 }
 
@@ -76,12 +103,13 @@ export class TaskItemImpl {
     }
     return Object.entries(this.outputs).flatMap(([nodeId, nodeOutputs]) =>
       Object.entries(nodeOutputs).flatMap(([mediaType, items]) =>
-        (items as ResultItem[]).map((item: ResultItem) =>
-          plainToClass(ResultItemImpl, {
-            ...item,
-            nodeId,
-            mediaType
-          })
+        (items as ResultItem[]).map(
+          (item: ResultItem) =>
+            new ResultItemImpl({
+              ...item,
+              nodeId,
+              mediaType
+            })
         )
       )
     )
@@ -213,7 +241,7 @@ export class TaskItemImpl {
       : undefined
   }
 
-  public async loadWorkflow() {
+  public async loadWorkflow(app: ComfyApp) {
     await app.loadGraphData(toRaw(this.workflow))
     if (this.outputs) {
       app.nodeOutputs = toRaw(this.outputs)