From a376503a7b0b2afa54c41c1de9e28a46a31aa98e Mon Sep 17 00:00:00 2001
From: Terry Jia <terryjia88@gmail.com>
Date: Mon, 16 Dec 2024 14:56:20 -0500
Subject: [PATCH] [3d] improve Preview 3d node (#1918)

---
 src/extensions/core/load3d.ts | 41 +++++++++++++++++++++++++----------
 1 file changed, 29 insertions(+), 12 deletions(-)

diff --git a/src/extensions/core/load3d.ts b/src/extensions/core/load3d.ts
index 260c618c8..a77acf7ae 100644
--- a/src/extensions/core/load3d.ts
+++ b/src/extensions/core/load3d.ts
@@ -1638,6 +1638,15 @@ app.registerExtension({
 app.registerExtension({
   name: 'Comfy.Preview3D',
 
+  async beforeRegisterNodeDef(nodeType, nodeData) {
+    if (
+      // @ts-expect-error ComfyNode
+      ['Preview3D'].includes(nodeType.comfyClass)
+    ) {
+      nodeData.input.required.image = ['PREVIEW_3D']
+    }
+  },
+
   getCustomWidgets(app) {
     return {
       PREVIEW_3D(node, inputName) {
@@ -1734,17 +1743,25 @@ app.registerExtension({
       (w: IWidget) => w.name === 'up_direction'
     )
 
-    configureLoad3D(
-      load3d,
-      'output',
-      modelWidget,
-      showGrid,
-      cameraType,
-      view,
-      material,
-      bgColor,
-      lightIntensity,
-      upDirection
-    )
+    const onExecuted = node.onExecuted
+
+    node.onExecuted = function (message: any) {
+      onExecuted?.apply(this, arguments)
+
+      modelWidget.value = message.model_file[0]
+
+      configureLoad3D(
+        load3d,
+        'output',
+        modelWidget,
+        showGrid,
+        cameraType,
+        view,
+        material,
+        bgColor,
+        lightIntensity,
+        upDirection
+      )
+    }
   }
 })