From 52180243955ffcf3ce6e03a346ab2fb55aebe29b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Zolt=C3=A1n=20D=C3=B3cs?= <bezozo97@gmail.com>
Date: Mon, 30 Dec 2024 04:26:03 +0100
Subject: [PATCH] Show workflow templates from custom nodes (#2032)

---
 .../templates/TemplateWorkflowCard.vue        |  75 ++++++++
 .../templates/TemplateWorkflowsContent.vue    | 174 ++++++++++++------
 src/scripts/api.ts                            |  11 ++
 src/stores/workflowTemplatesStore.ts          |  30 +++
 vite.config.mts                               |   6 +-
 5 files changed, 237 insertions(+), 59 deletions(-)
 create mode 100644 src/components/templates/TemplateWorkflowCard.vue
 create mode 100644 src/stores/workflowTemplatesStore.ts

diff --git a/src/components/templates/TemplateWorkflowCard.vue b/src/components/templates/TemplateWorkflowCard.vue
new file mode 100644
index 000000000..518f51b63
--- /dev/null
+++ b/src/components/templates/TemplateWorkflowCard.vue
@@ -0,0 +1,75 @@
+<template>
+  <Card :data-testid="`template-workflow-${props.workflowName}`">
+    <template #header>
+      <div class="flex items-center justify-center">
+        <div
+          class="relative overflow-hidden rounded-lg cursor-pointer w-64 h-64"
+        >
+          <img
+            v-if="!imageError"
+            :src="
+              props.moduleName === 'default'
+                ? `templates/${props.workflowName}.jpg`
+                : `api/workflow_templates/${props.moduleName}/${props.workflowName}.jpg`
+            "
+            @error="imageError = true"
+            class="w-64 h-64 rounded-lg object-cover thumbnail"
+          />
+          <div v-else class="w-64 h-64 content-center text-center">
+            <i class="pi pi-file" style="font-size: 4rem"></i>
+          </div>
+          <a>
+            <div
+              class="absolute top-0 left-0 w-64 h-64 overflow-hidden opacity-0 transition duration-300 ease-in-out hover:opacity-100 bg-opacity-50 bg-black flex items-center justify-center"
+            >
+              <i class="pi pi-play-circle" style="color: white"></i>
+            </div>
+          </a>
+          <ProgressSpinner
+            v-if="loading"
+            class="absolute inset-0 z-1 w-3/12 h-full"
+          />
+        </div>
+      </div>
+    </template>
+    <template #subtitle>
+      <!--Default templates have translations-->
+      <template v-if="props.moduleName === 'default'">
+        {{
+          $t(
+            `templateWorkflows.template.${props.workflowName}`,
+            props.workflowName
+          )
+        }}
+      </template>
+      <template v-else>
+        {{ props.workflowName }}
+      </template>
+    </template>
+  </Card>
+</template>
+
+<script setup lang="ts">
+import Card from 'primevue/card'
+import ProgressSpinner from 'primevue/progressspinner'
+import { ref } from 'vue'
+
+const props = defineProps<{
+  moduleName: string
+  workflowName: string
+  loading: boolean
+}>()
+
+const imageError = ref(false)
+</script>
+
+<style lang="css" scoped>
+.p-card {
+  --p-card-body-padding: 10px 0 0 0;
+  overflow: hidden;
+}
+
+:deep(.p-card-subtitle) {
+  text-align: center;
+}
+</style>
diff --git a/src/components/templates/TemplateWorkflowsContent.vue b/src/components/templates/TemplateWorkflowsContent.vue
index 75a5af529..02c1d0f3a 100644
--- a/src/components/templates/TemplateWorkflowsContent.vue
+++ b/src/components/templates/TemplateWorkflowsContent.vue
@@ -1,81 +1,139 @@
 <template>
-  <div
-    class="flex flex-wrap content-around justify-around gap-4 mt-4"
-    data-testid="template-workflows-content"
-  >
-    <div
-      v-for="template in templates"
-      :key="template"
-      :data-testid="`template-workflow-${template}`"
-    >
-      <Card>
-        <template #header>
-          <div
-            class="relative overflow-hidden rounded-lg cursor-pointer"
-            @click="loadWorkflow(template)"
-          >
-            <img
-              :src="`templates/${template}.jpg`"
-              class="w-64 h-64 rounded-lg object-cover"
-            />
-            <a>
-              <div
-                class="absolute top-0 left-0 w-64 h-64 overflow-hidden opacity-0 transition duration-300 ease-in-out hover:opacity-100 bg-opacity-50 bg-black flex items-center justify-center"
-              >
-                <i class="pi pi-play-circle"></i>
-              </div>
-            </a>
-            <ProgressSpinner
-              v-if="loading === template"
-              class="absolute inset-0 z-1 w-3/12 h-full"
-            />
-          </div>
-        </template>
-        <template #subtitle>{{
-          $t(`templateWorkflows.template.${template}`)
-        }}</template>
-      </Card>
+  <div class="flex h-96" data-testid="template-workflows-content">
+    <div class="relative">
+      <ProgressSpinner
+        v-if="!workflowTemplatesStore.isLoaded"
+        class="absolute w-8 h-full inset-0"
+      />
+      <Listbox
+        :model-value="selectedTab"
+        @update:model-value="handleTabSelection"
+        :options="tabs"
+        optionLabel="title"
+        scroll-height="auto"
+        class="overflow-y-auto w-64 h-full"
+        listStyle="max-height:unset"
+      />
     </div>
+    <Carousel
+      class="carousel justify-center"
+      :value="selectedTab.templates"
+      :responsive-options="responsiveOptions"
+      :numVisible="4"
+      :numScroll="3"
+      :key="selectedTab.moduleName"
+    >
+      <template #item="slotProps">
+        <div @click="loadWorkflow(slotProps.data)">
+          <TemplateWorkflowCard
+            :moduleName="selectedTab.moduleName"
+            :workflowName="slotProps.data"
+            :loading="slotProps.data === workflowLoading"
+          />
+        </div>
+      </template>
+    </Carousel>
   </div>
 </template>
 
 <script setup lang="ts">
-import { useDialogStore } from '@/stores/dialogStore'
-import Card from 'primevue/card'
+import Carousel from 'primevue/carousel'
+import Listbox from 'primevue/listbox'
 import ProgressSpinner from 'primevue/progressspinner'
-import { ref } from 'vue'
+import { computed, onMounted, ref } from 'vue'
+import { useDialogStore } from '@/stores/dialogStore'
 import { app } from '@/scripts/app'
 import { api } from '@/scripts/api'
 import { useI18n } from 'vue-i18n'
+import TemplateWorkflowCard from '@/components/templates/TemplateWorkflowCard.vue'
+import { useWorkflowTemplatesStore } from '@/stores/workflowTemplatesStore'
+
+interface WorkflowTemplatesTab {
+  moduleName: string
+  title: string
+  templates: string[]
+}
+
 const { t } = useI18n()
 
-const templates = ['default', 'image2image', 'upscale', 'flux_schnell']
-const loading = ref<string | null>(null)
+//These default templates are provided by the frontend
+const comfyUITemplates: WorkflowTemplatesTab = {
+  moduleName: 'default',
+  title: 'ComfyUI',
+  templates: ['default', 'image2image', 'upscale', 'flux_schnell']
+}
+
+const responsiveOptions = ref([
+  {
+    breakpoint: '1660px',
+    numVisible: 3,
+    numScroll: 2
+  },
+  {
+    breakpoint: '1360px',
+    numVisible: 2,
+    numScroll: 1
+  },
+  {
+    breakpoint: '960px',
+    numVisible: 1,
+    numScroll: 1
+  }
+])
+
+const workflowTemplatesStore = useWorkflowTemplatesStore()
+const selectedTab = ref<WorkflowTemplatesTab>(comfyUITemplates)
+const workflowLoading = ref<string | null>(null)
+
+const tabs = computed<WorkflowTemplatesTab[]>(() => {
+  return [
+    comfyUITemplates,
+    ...Object.entries(workflowTemplatesStore.items).map(([key, value]) => ({
+      moduleName: key,
+      title: key,
+      templates: value
+    }))
+  ]
+})
+
+onMounted(async () => {
+  await workflowTemplatesStore.loadWorkflowTemplates()
+})
+
+const handleTabSelection = (selection: WorkflowTemplatesTab | null) => {
+  //Listbox allows deselecting so this special case is ignored here
+  if (selection !== selectedTab.value && selection !== null)
+    selectedTab.value = selection
+}
 
 const loadWorkflow = async (id: string) => {
-  loading.value = id
-  const json = await fetch(api.fileURL(`templates/${id}.json`)).then((r) =>
-    r.json()
-  )
+  workflowLoading.value = id
+  let json
+  if (selectedTab.value.moduleName === 'default') {
+    // Default templates provided by frontend are served on this separate endpoint
+    json = await fetch(api.fileURL(`templates/${id}.json`)).then((r) =>
+      r.json()
+    )
+  } else {
+    json = await fetch(
+      api.apiURL(
+        `/workflow_templates/${selectedTab.value.moduleName}/${id}.json`
+      )
+    ).then((r) => r.json())
+  }
   useDialogStore().closeDialog()
-  await app.loadGraphData(
-    json,
-    true,
-    true,
-    t(`templateWorkflows.template.${id}`)
-  )
+  const workflowName =
+    selectedTab.value.moduleName === 'default'
+      ? t(`templateWorkflows.template.${id}`, id)
+      : id
+  await app.loadGraphData(json, true, true, workflowName)
 
   return false
 }
 </script>
 
 <style lang="css" scoped>
-.p-card {
-  --p-card-body-padding: 10px 0 0 0;
-  overflow: hidden;
-}
-
-:deep(.p-card-subtitle) {
-  text-align: center;
+.carousel {
+  width: 66vw;
 }
 </style>
diff --git a/src/scripts/api.ts b/src/scripts/api.ts
index 686372481..5b5fe857e 100644
--- a/src/scripts/api.ts
+++ b/src/scripts/api.ts
@@ -397,6 +397,17 @@ export class ComfyApi extends EventTarget {
     return await resp.json()
   }
 
+  /**
+   * Gets the available workflow templates from custom nodes.
+   * @returns A map of custom_node names and associated template workflow names.
+   */
+  async getWorkflowTemplates(): Promise<{
+    [customNodesName: string]: string[]
+  }> {
+    const res = await this.fetchApi('/workflow_templates')
+    return await res.json()
+  }
+
   /**
    * Gets a list of embedding names
    */
diff --git a/src/stores/workflowTemplatesStore.ts b/src/stores/workflowTemplatesStore.ts
new file mode 100644
index 000000000..1c5721642
--- /dev/null
+++ b/src/stores/workflowTemplatesStore.ts
@@ -0,0 +1,30 @@
+import { defineStore } from 'pinia'
+import { ref } from 'vue'
+import { api } from '@/scripts/api'
+
+export const useWorkflowTemplatesStore = defineStore(
+  'workflowTemplates',
+  () => {
+    const items = ref<{
+      [customNodesName: string]: string[]
+    }>({})
+    const isLoaded = ref(false)
+
+    async function loadWorkflowTemplates() {
+      try {
+        if (!isLoaded.value) {
+          items.value = await api.getWorkflowTemplates()
+          isLoaded.value = true
+        }
+      } catch (error) {
+        console.error('Error fetching workflow templates:', error)
+      }
+    }
+
+    return {
+      items,
+      isLoaded,
+      loadWorkflowTemplates
+    }
+  }
+)
diff --git a/vite.config.mts b/vite.config.mts
index 7d454479b..9d94fe575 100644
--- a/vite.config.mts
+++ b/vite.config.mts
@@ -119,6 +119,10 @@ export default defineConfig({
         ws: true
       },
 
+      '/workflow_templates': {
+        target: DEV_SERVER_COMFYUI_URL
+      },
+
       '/testsubrouteindex': {
         target: 'http://localhost:5173',
         rewrite: (path) => path.substring('/testsubrouteindex'.length)
@@ -183,4 +187,4 @@ export default defineConfig({
       '@comfyorg/comfyui-electron-types'
     ]
   }
-}) as UserConfigExport
+}) as UserConfigExport
\ No newline at end of file