From d49815fcb401cc3bb79035c76f9459cb240a985d Mon Sep 17 00:00:00 2001
From: Terry Jia <terryjia88@gmail.com>
Date: Wed, 19 Mar 2025 22:00:59 -0400
Subject: [PATCH] [3d] add preview 3d for saveGlb (#3156)

---
 src/extensions/core/index.ts                  |  1 +
 .../core/load3d/Load3DConfiguration.ts        | 16 ++++
 src/extensions/core/saveMesh.ts               | 79 +++++++++++++++++++
 3 files changed, 96 insertions(+)
 create mode 100644 src/extensions/core/saveMesh.ts

diff --git a/src/extensions/core/index.ts b/src/extensions/core/index.ts
index 517d6794a..1d6b08367 100644
--- a/src/extensions/core/index.ts
+++ b/src/extensions/core/index.ts
@@ -12,6 +12,7 @@ import './nodeTemplates'
 import './noteNode'
 import './rerouteNode'
 import './saveImageExtraOutput'
+import './saveMesh'
 import './simpleTouchSupport'
 import './slotDefaults'
 import './uploadAudio'
diff --git a/src/extensions/core/load3d/Load3DConfiguration.ts b/src/extensions/core/load3d/Load3DConfiguration.ts
index 96f20c837..c848972b8 100644
--- a/src/extensions/core/load3d/Load3DConfiguration.ts
+++ b/src/extensions/core/load3d/Load3DConfiguration.ts
@@ -7,6 +7,11 @@ import { api } from '@/scripts/api'
 class Load3DConfiguration {
   constructor(private load3d: Load3d) {}
 
+  configureForSaveMesh(loadFolder: 'input' | 'output', filePath: string) {
+    this.setupModelHandlingForSaveMesh(filePath, loadFolder)
+    this.setupDefaultProperties()
+  }
+
   configure(
     loadFolder: 'input' | 'output',
     modelWidget: IWidget,
@@ -33,6 +38,17 @@ class Load3DConfiguration {
     }
   }
 
+  private setupModelHandlingForSaveMesh(
+    filePath: string,
+    loadFolder: 'input' | 'output'
+  ) {
+    const onModelWidgetUpdate = this.createModelUpdateHandler(loadFolder)
+
+    if (filePath) {
+      onModelWidgetUpdate(filePath)
+    }
+  }
+
   private setupModelHandling(
     modelWidget: IWidget,
     loadFolder: 'input' | 'output',
diff --git a/src/extensions/core/saveMesh.ts b/src/extensions/core/saveMesh.ts
new file mode 100644
index 000000000..8224f5a71
--- /dev/null
+++ b/src/extensions/core/saveMesh.ts
@@ -0,0 +1,79 @@
+import { IWidget } from '@comfyorg/litegraph'
+import { nextTick } from 'vue'
+
+import Load3D from '@/components/load3d/Load3D.vue'
+import Load3DConfiguration from '@/extensions/core/load3d/Load3DConfiguration'
+import { CustomInputSpec } from '@/schemas/nodeDef/nodeDefSchemaV2'
+import { ComponentWidgetImpl, addWidget } from '@/scripts/domWidget'
+import { useExtensionService } from '@/services/extensionService'
+import { useLoad3dService } from '@/services/load3dService'
+import { generateUUID } from '@/utils/formatUtil'
+
+useExtensionService().registerExtension({
+  name: 'Comfy.SaveGLB',
+
+  async beforeRegisterNodeDef(_nodeType, nodeData) {
+    if ('SaveGLB' === nodeData.name) {
+      // @ts-expect-error InputSpec is not typed correctly
+      nodeData.input.required.image = ['PREVIEW_3D']
+    }
+  },
+
+  getCustomWidgets() {
+    return {
+      PREVIEW_3D(node) {
+        const inputSpec: CustomInputSpec = {
+          name: 'image',
+          type: 'Preview3D'
+        }
+
+        const widget = new ComponentWidgetImpl({
+          id: generateUUID(),
+          node,
+          name: inputSpec.name,
+          component: Load3D,
+          inputSpec,
+          options: {}
+        })
+
+        addWidget(node, widget)
+
+        return { widget }
+      }
+    }
+  },
+
+  async nodeCreated(node) {
+    if (node.constructor.comfyClass !== 'SaveGLB') return
+
+    const [oldWidth, oldHeight] = node.size
+
+    node.setSize([Math.max(oldWidth, 400), Math.max(oldHeight, 550)])
+
+    await nextTick()
+
+    const onExecuted = node.onExecuted
+
+    node.onExecuted = function (message: any) {
+      onExecuted?.apply(this, arguments as any)
+
+      const fileInfo = message['3d'][0]
+
+      const load3d = useLoad3dService().getLoad3d(node)
+
+      const modelWidget = node.widgets?.find((w: IWidget) => w.name === 'image')
+
+      if (load3d && modelWidget) {
+        const filePath = fileInfo['subfolder'] + '/' + fileInfo['filename']
+
+        console.log(filePath)
+
+        modelWidget.value = filePath
+
+        const config = new Load3DConfiguration(load3d)
+
+        config.configureForSaveMesh(fileInfo['type'], filePath)
+      }
+    }
+  }
+})