From ada7dde51c6aba2efffddec5662dff77a8926c3c Mon Sep 17 00:00:00 2001
From: Terry Jia <terryjia88@gmail.com>
Date: Wed, 21 May 2025 21:29:52 -0400
Subject: [PATCH] [3d] performance improve (#3961)

---
 src/components/load3d/Load3DScene.vue        |  88 +++++++++---
 src/extensions/core/load3d.ts                | 140 +++++++++----------
 src/extensions/core/load3d/Load3d.ts         |   8 ++
 src/extensions/core/load3d/PreviewManager.ts |   8 ++
 src/services/load3dService.ts                |  31 ++++
 5 files changed, 186 insertions(+), 89 deletions(-)

diff --git a/src/components/load3d/Load3DScene.vue b/src/components/load3d/Load3DScene.vue
index 86a6b82a9..b9d385143 100644
--- a/src/components/load3d/Load3DScene.vue
+++ b/src/components/load3d/Load3DScene.vue
@@ -6,7 +6,7 @@
 
 <script setup lang="ts">
 import { LGraphNode } from '@comfyorg/litegraph'
-import { onMounted, onUnmounted, ref, toRaw, watch, watchEffect } from 'vue'
+import { onMounted, onUnmounted, ref, toRaw, watch } from 'vue'
 
 import LoadingOverlay from '@/components/load3d/LoadingOverlay.vue'
 import Load3d from '@/extensions/core/load3d/Load3d'
@@ -76,18 +76,71 @@ const eventConfig = {
     emit('recordingStatusChange', value)
 } as const
 
-watchEffect(() => {
-  if (load3d.value) {
-    const rawLoad3d = toRaw(load3d.value) as Load3d
+watch(
+  () => props.showPreview,
+  (newValue) => {
+    if (load3d.value) {
+      const rawLoad3d = toRaw(load3d.value) as Load3d
 
-    rawLoad3d.setBackgroundColor(props.backgroundColor)
-    rawLoad3d.toggleGrid(props.showGrid)
-    rawLoad3d.setLightIntensity(props.lightIntensity)
-    rawLoad3d.setFOV(props.fov)
-    rawLoad3d.toggleCamera(props.cameraType)
-    rawLoad3d.togglePreview(props.showPreview)
+      rawLoad3d.togglePreview(newValue)
+    }
   }
-})
+)
+
+watch(
+  () => props.cameraType,
+  (newValue) => {
+    if (load3d.value) {
+      const rawLoad3d = toRaw(load3d.value) as Load3d
+
+      rawLoad3d.toggleCamera(newValue)
+    }
+  }
+)
+
+watch(
+  () => props.fov,
+  (newValue) => {
+    if (load3d.value) {
+      const rawLoad3d = toRaw(load3d.value) as Load3d
+
+      rawLoad3d.setFOV(newValue)
+    }
+  }
+)
+
+watch(
+  () => props.lightIntensity,
+  (newValue) => {
+    if (load3d.value) {
+      const rawLoad3d = toRaw(load3d.value) as Load3d
+
+      rawLoad3d.setLightIntensity(newValue)
+    }
+  }
+)
+
+watch(
+  () => props.showGrid,
+  (newValue) => {
+    if (load3d.value) {
+      const rawLoad3d = toRaw(load3d.value) as Load3d
+
+      rawLoad3d.toggleGrid(newValue)
+    }
+  }
+)
+
+watch(
+  () => props.backgroundColor,
+  (newValue) => {
+    if (load3d.value) {
+      const rawLoad3d = toRaw(load3d.value) as Load3d
+
+      rawLoad3d.setBackgroundColor(newValue)
+    }
+  }
+)
 
 watch(
   () => props.backgroundImage,
@@ -164,12 +217,13 @@ const handleEvents = (action: 'add' | 'remove') => {
 }
 
 onMounted(() => {
-  load3d.value = useLoad3dService().registerLoad3d(
-    node.value as LGraphNode,
-    // @ts-expect-error fixme ts strict error
-    container.value,
-    props.inputSpec
-  )
+  if (container.value) {
+    load3d.value = useLoad3dService().registerLoad3d(
+      node.value as LGraphNode,
+      container.value,
+      props.inputSpec
+    )
+  }
   handleEvents('add')
 })
 
diff --git a/src/extensions/core/load3d.ts b/src/extensions/core/load3d.ts
index 965c508f3..40c0f4149 100644
--- a/src/extensions/core/load3d.ts
+++ b/src/extensions/core/load3d.ts
@@ -112,7 +112,7 @@ useExtensionService().registerExtension({
       LOAD_3D(node) {
         const fileInput = document.createElement('input')
         fileInput.type = 'file'
-        fileInput.accept = '.gltf,.glb,.obj,.mtl,.fbx,.stl'
+        fileInput.accept = '.gltf,.glb,.obj,.fbx,.stl'
         fileInput.style.display = 'none'
 
         fileInput.onchange = async () => {
@@ -195,9 +195,7 @@ useExtensionService().registerExtension({
 
     await nextTick()
 
-    const load3d = useLoad3dService().getLoad3d(node)
-
-    if (load3d) {
+    useLoad3dService().waitForLoad3d(node, (load3d) => {
       let cameraState = node.properties['Camera Info']
 
       const config = new Load3DConfiguration(load3d)
@@ -256,7 +254,7 @@ useExtensionService().registerExtension({
           return returnVal
         }
       }
-    }
+    })
   }
 })
 
@@ -346,67 +344,65 @@ useExtensionService().registerExtension({
 
     await nextTick()
 
-    const sceneWidget = node.widgets?.find((w) => w.name === 'image')
+    useLoad3dService().waitForLoad3d(node, (load3d) => {
+      const sceneWidget = node.widgets?.find((w) => w.name === 'image')
+      const modelWidget = node.widgets?.find((w) => w.name === 'model_file')
+      let cameraState = node.properties['Camera Info']
+      const width = node.widgets?.find((w) => w.name === 'width')
+      const height = node.widgets?.find((w) => w.name === 'height')
 
-    const load3d = useLoad3dService().getLoad3d(node) as Load3dAnimation
+      if (modelWidget && width && height && sceneWidget && load3d) {
+        const config = new Load3DConfiguration(load3d)
 
-    const modelWidget = node.widgets?.find((w) => w.name === 'model_file')
+        config.configure('input', modelWidget, cameraState, width, height)
 
-    let cameraState = node.properties['Camera Info']
+        sceneWidget.serializeValue = async () => {
+          node.properties['Camera Info'] = load3d.getCameraState()
 
-    const width = node.widgets?.find((w) => w.name === 'width')
-    const height = node.widgets?.find((w) => w.name === 'height')
+          const load3dAnimation = load3d as Load3dAnimation
+          load3dAnimation.toggleAnimation(false)
 
-    if (modelWidget && width && height && sceneWidget && load3d) {
-      const config = new Load3DConfiguration(load3d)
+          if (load3dAnimation.isRecording()) {
+            load3dAnimation.stopRecording()
+          }
 
-      config.configure('input', modelWidget, cameraState, width, height)
+          const {
+            scene: imageData,
+            mask: maskData,
+            normal: normalData
+          } = await load3dAnimation.captureScene(
+            width.value as number,
+            height.value as number
+          )
 
-      sceneWidget.serializeValue = async () => {
-        node.properties['Camera Info'] = load3d.getCameraState()
-
-        load3d.toggleAnimation(false)
-
-        if (load3d.isRecording()) {
-          load3d.stopRecording()
-        }
-
-        const {
-          scene: imageData,
-          mask: maskData,
-          normal: normalData
-        } = await load3d.captureScene(
-          width.value as number,
-          height.value as number
-        )
-
-        const [data, dataMask, dataNormal] = await Promise.all([
-          Load3dUtils.uploadTempImage(imageData, 'scene'),
-          Load3dUtils.uploadTempImage(maskData, 'scene_mask'),
-          Load3dUtils.uploadTempImage(normalData, 'scene_normal')
-        ])
-
-        load3d.handleResize()
-
-        const returnVal = {
-          image: `threed/${data.name} [temp]`,
-          mask: `threed/${dataMask.name} [temp]`,
-          normal: `threed/${dataNormal.name} [temp]`,
-          camera_info: node.properties['Camera Info'],
-          recording: ''
-        }
-
-        const recordingData = load3d.getRecordingData()
-        if (recordingData) {
-          const [recording] = await Promise.all([
-            Load3dUtils.uploadTempImage(recordingData, 'recording', 'mp4')
+          const [data, dataMask, dataNormal] = await Promise.all([
+            Load3dUtils.uploadTempImage(imageData, 'scene'),
+            Load3dUtils.uploadTempImage(maskData, 'scene_mask'),
+            Load3dUtils.uploadTempImage(normalData, 'scene_normal')
           ])
-          returnVal['recording'] = `threed/${recording.name} [temp]`
-        }
 
-        return returnVal
+          load3dAnimation.handleResize()
+
+          const returnVal = {
+            image: `threed/${data.name} [temp]`,
+            mask: `threed/${dataMask.name} [temp]`,
+            normal: `threed/${dataNormal.name} [temp]`,
+            camera_info: node.properties['Camera Info'],
+            recording: ''
+          }
+
+          const recordingData = load3dAnimation.getRecordingData()
+          if (recordingData) {
+            const [recording] = await Promise.all([
+              Load3dUtils.uploadTempImage(recordingData, 'recording', 'mp4')
+            ])
+            returnVal['recording'] = `threed/${recording.name} [temp]`
+          }
+
+          return returnVal
+        }
       }
-    }
+    })
   }
 })
 
@@ -467,19 +463,19 @@ useExtensionService().registerExtension({
         useToastStore().addAlert(msg)
       }
 
-      const load3d = useLoad3dService().getLoad3d(node)
-
       let cameraState = message.result[1]
 
-      const modelWidget = node.widgets?.find((w) => w.name === 'model_file')
+      useLoad3dService().waitForLoad3d(node, (load3d) => {
+        const modelWidget = node.widgets?.find((w) => w.name === 'model_file')
 
-      if (load3d && modelWidget) {
-        modelWidget.value = filePath.replaceAll('\\', '/')
+        if (modelWidget) {
+          modelWidget.value = filePath.replaceAll('\\', '/')
 
-        const config = new Load3DConfiguration(load3d)
+          const config = new Load3DConfiguration(load3d)
 
-        config.configure('output', modelWidget, cameraState)
-      }
+          config.configure('output', modelWidget, cameraState)
+        }
+      })
     }
   }
 })
@@ -543,16 +539,16 @@ useExtensionService().registerExtension({
 
       let cameraState = message.result[1]
 
-      const load3d = useLoad3dService().getLoad3d(node)
+      useLoad3dService().waitForLoad3d(node, (load3d) => {
+        const modelWidget = node.widgets?.find((w) => w.name === 'model_file')
+        if (modelWidget) {
+          modelWidget.value = filePath.replaceAll('\\', '/')
 
-      const modelWidget = node.widgets?.find((w) => w.name === 'model_file')
-      if (load3d && modelWidget) {
-        modelWidget.value = filePath.replaceAll('\\', '/')
+          const config = new Load3DConfiguration(load3d)
 
-        const config = new Load3DConfiguration(load3d)
-
-        config.configure('output', modelWidget, cameraState)
-      }
+          config.configure('output', modelWidget, cameraState)
+        }
+      })
     }
   }
 })
diff --git a/src/extensions/core/load3d/Load3d.ts b/src/extensions/core/load3d/Load3d.ts
index 92cd5bc6c..3a1b76789 100644
--- a/src/extensions/core/load3d/Load3d.ts
+++ b/src/extensions/core/load3d/Load3d.ts
@@ -486,6 +486,14 @@ class Load3d {
   }
 
   public remove(): void {
+    this.renderer.forceContextLoss()
+    const canvas = this.renderer.domElement
+    const event = new Event('webglcontextlost', {
+      bubbles: true,
+      cancelable: true
+    })
+    canvas.dispatchEvent(event)
+
     if (this.animationFrameId !== null) {
       cancelAnimationFrame(this.animationFrameId)
     }
diff --git a/src/extensions/core/load3d/PreviewManager.ts b/src/extensions/core/load3d/PreviewManager.ts
index 52dabb3fd..80bd68c43 100644
--- a/src/extensions/core/load3d/PreviewManager.ts
+++ b/src/extensions/core/load3d/PreviewManager.ts
@@ -62,6 +62,14 @@ export class PreviewManager implements PreviewManagerInterface {
 
   dispose(): void {
     if (this.previewRenderer) {
+      this.previewRenderer.forceContextLoss()
+      const canvas = this.previewRenderer.domElement
+      const event = new Event('webglcontextlost', {
+        bubbles: true,
+        cancelable: true
+      })
+      canvas.dispatchEvent(event)
+
       this.previewRenderer.dispose()
     }
 
diff --git a/src/services/load3dService.ts b/src/services/load3dService.ts
index 95bb92554..515579ef9 100644
--- a/src/services/load3dService.ts
+++ b/src/services/load3dService.ts
@@ -5,9 +5,12 @@ import Load3d from '@/extensions/core/load3d/Load3d'
 import Load3dAnimation from '@/extensions/core/load3d/Load3dAnimation'
 import type { CustomInputSpec } from '@/schemas/nodeDef/nodeDefSchemaV2'
 
+type Load3dReadyCallback = (load3d: Load3d | Load3dAnimation) => void
+
 export class Load3dService {
   private static instance: Load3dService
   private nodeToLoad3dMap = new Map<LGraphNode, Load3d | Load3dAnimation>()
+  private pendingCallbacks = new Map<LGraphNode, Load3dReadyCallback[]>()
 
   private constructor() {}
 
@@ -60,6 +63,13 @@ export class Load3dService {
 
     this.nodeToLoad3dMap.set(rawNode, instance)
 
+    const callbacks = this.pendingCallbacks.get(rawNode)
+
+    if (callbacks) {
+      callbacks.forEach((callback) => callback(instance))
+      this.pendingCallbacks.delete(rawNode)
+    }
+
     return instance
   }
 
@@ -69,6 +79,24 @@ export class Load3dService {
     return this.nodeToLoad3dMap.get(rawNode) || null
   }
 
+  waitForLoad3d(node: LGraphNode, callback: Load3dReadyCallback): void {
+    const rawNode = toRaw(node)
+
+    const existingInstance = this.nodeToLoad3dMap.get(rawNode)
+
+    if (existingInstance) {
+      callback(existingInstance)
+
+      return
+    }
+
+    if (!this.pendingCallbacks.has(rawNode)) {
+      this.pendingCallbacks.set(rawNode, [])
+    }
+
+    this.pendingCallbacks.get(rawNode)!.push(callback)
+  }
+
   getNodeByLoad3d(load3d: Load3d | Load3dAnimation): LGraphNode | null {
     for (const [node, instance] of this.nodeToLoad3dMap) {
       if (instance === load3d) {
@@ -88,12 +116,15 @@ export class Load3dService {
 
       this.nodeToLoad3dMap.delete(rawNode)
     }
+
+    this.pendingCallbacks.delete(rawNode)
   }
 
   clear() {
     for (const [node] of this.nodeToLoad3dMap) {
       this.removeLoad3d(node)
     }
+    this.pendingCallbacks.clear()
   }
 }