[3d] refactor load3d (#2521)

This commit is contained in:
Terry Jia
2025-02-11 17:43:36 -05:00
committed by GitHub
parent c901d5f659
commit 8cfe814daa
2 changed files with 107 additions and 55 deletions

View File

@@ -4,31 +4,28 @@ import type { IStringWidget } from '@comfyorg/litegraph/dist/types/widgets'
import { nextTick } from 'vue'
import Load3DConfiguration from '@/extensions/core/load3d/Load3DConfiguration'
import Load3d from '@/extensions/core/load3d/Load3d'
import Load3dAnimation from '@/extensions/core/load3d/Load3dAnimation'
import Load3dUtils from '@/extensions/core/load3d/Load3dUtils'
import { app } from '@/scripts/app'
import { useLoad3dService } from '@/services/load3dService'
import { useToastStore } from '@/stores/toastStore'
const containerToLoad3D = new Map()
app.registerExtension({
name: 'Comfy.Load3D',
getCustomWidgets(app) {
return {
LOAD_3D(node, inputName) {
let load3dNode = app.graph._nodes.filter((wi) => wi.type == 'Load3D')
node.addProperty('Camera Info', '')
const container = document.createElement('div')
container.id = `comfy-load-3d-${load3dNode.length}`
container.classList.add('comfy-load-3d')
const load3d = new Load3d(container, { createPreview: true })
containerToLoad3D.set(container.id, load3d)
const load3d = useLoad3dService().registerLoad3d(
node,
container,
'Load3D'
)
node.onMouseEnter = function () {
if (load3d) {
@@ -49,7 +46,7 @@ app.registerExtension({
load3d.remove()
}
containerToLoad3D.delete(container.id)
useLoad3dService().removeLoad3d(node)
origOnRemoved?.apply(this, [])
}
@@ -118,9 +115,7 @@ app.registerExtension({
const sceneWidget = node.widgets.find((w: IWidget) => w.name === 'image')
const container = sceneWidget.element
const load3d = containerToLoad3D.get(container.id)
const load3d = useLoad3dService().getLoad3d(node)
load3d.setNode(node)
@@ -155,8 +150,8 @@ app.registerExtension({
node.properties['Camera Info'] = load3d.getCameraState()
const { scene: imageData, mask: maskData } = await load3d.captureScene(
width.value,
height.value
width.value as number,
height.value as number
)
const [data, dataMask] = await Promise.all([
@@ -178,19 +173,17 @@ app.registerExtension({
getCustomWidgets(app) {
return {
LOAD_3D_ANIMATION(node, inputName) {
let load3dNode = app.graph._nodes.filter(
(wi) => wi.type == 'Load3DAnimation'
)
node.addProperty('Camera Info', '')
const container = document.createElement('div')
container.id = `comfy-load-3d-animation-${load3dNode.length}`
container.classList.add('comfy-load-3d-animation')
const load3d = new Load3dAnimation(container, { createPreview: true })
containerToLoad3D.set(container.id, load3d)
const load3d = useLoad3dService().registerLoad3d(
node,
container,
'Load3DAnimation'
)
node.onMouseEnter = function () {
if (load3d) {
@@ -211,7 +204,7 @@ app.registerExtension({
load3d.remove()
}
containerToLoad3D.delete(container.id)
useLoad3dService().removeLoad3d(node)
origOnRemoved?.apply(this, [])
}
@@ -280,9 +273,7 @@ app.registerExtension({
const sceneWidget = node.widgets.find((w: IWidget) => w.name === 'image')
const container = sceneWidget.element
const load3d = containerToLoad3D.get(container.id)
const load3d = useLoad3dService().getLoad3d(node) as Load3dAnimation
load3d.setNode(node)
@@ -319,8 +310,8 @@ app.registerExtension({
load3d.toggleAnimation(false)
const { scene: imageData, mask: maskData } = await load3d.captureScene(
width.value,
height.value
width.value as number,
height.value as number
)
const [data, dataMask] = await Promise.all([
@@ -351,15 +342,15 @@ app.registerExtension({
getCustomWidgets(app) {
return {
PREVIEW_3D(node, inputName) {
let load3dNode = app.graph._nodes.filter((wi) => wi.type == 'Preview3D')
const container = document.createElement('div')
container.id = `comfy-preview-3d-${load3dNode.length}`
container.classList.add('comfy-preview-3d')
const load3d = new Load3d(container, { createPreview: false })
containerToLoad3D.set(container.id, load3d)
const load3d = useLoad3dService().registerLoad3d(
node,
container,
'Preview3D'
)
node.onMouseEnter = function () {
if (load3d) {
@@ -380,7 +371,7 @@ app.registerExtension({
load3d.remove()
}
containerToLoad3D.delete(container.id)
useLoad3dService().removeLoad3d(node)
origOnRemoved?.apply(this, [])
}
@@ -405,11 +396,7 @@ app.registerExtension({
await nextTick()
const sceneWidget = node.widgets.find((w: IWidget) => w.name === 'image')
const container = sceneWidget.element
const load3d = containerToLoad3D.get(container.id)
const load3d = useLoad3dService().getLoad3d(node)
load3d.setNode(node)
@@ -462,17 +449,15 @@ app.registerExtension({
getCustomWidgets(app) {
return {
PREVIEW_3D_ANIMATION(node, inputName) {
let load3dNode = app.graph._nodes.filter(
(wi) => wi.type == 'Preview3DAnimation'
)
const container = document.createElement('div')
container.id = `comfy-preview-3d-animation-${load3dNode.length}`
container.classList.add('comfy-preview-3d-animation')
const load3d = new Load3dAnimation(container, { createPreview: false })
containerToLoad3D.set(container.id, load3d)
const load3d = useLoad3dService().registerLoad3d(
node,
container,
'Preview3DAnimation'
)
node.onMouseEnter = function () {
if (load3d) {
@@ -493,7 +478,7 @@ app.registerExtension({
load3d.remove()
}
containerToLoad3D.delete(container.id)
useLoad3dService().removeLoad3d(node)
origOnRemoved?.apply(this, [])
}
@@ -522,11 +507,7 @@ app.registerExtension({
await nextTick()
const sceneWidget = node.widgets.find((w: IWidget) => w.name === 'image')
const container = sceneWidget.element
const load3d = containerToLoad3D.get(container.id)
const load3d = useLoad3dService().getLoad3d(node)
load3d.setNode(node)

View File

@@ -0,0 +1,71 @@
import type { LGraphNode } from '@comfyorg/litegraph'
import Load3d from '@/extensions/core/load3d/Load3d'
import Load3dAnimation from '@/extensions/core/load3d/Load3dAnimation'
export class Load3dService {
private static instance: Load3dService
private nodeToLoad3dMap = new Map<LGraphNode, Load3d | Load3dAnimation>()
private constructor() {}
static getInstance(): Load3dService {
if (!Load3dService.instance) {
Load3dService.instance = new Load3dService()
}
return Load3dService.instance
}
registerLoad3d(
node: LGraphNode,
container: HTMLElement,
type: 'Load3D' | 'Load3DAnimation' | 'Preview3D' | 'Preview3DAnimation'
) {
if (this.nodeToLoad3dMap.has(node)) {
this.removeLoad3d(node)
}
const isAnimation = type.includes('Animation')
const Load3dClass = isAnimation ? Load3dAnimation : Load3d
const isPreview = type.includes('Preview')
const instance = new Load3dClass(container, { createPreview: !isPreview })
this.nodeToLoad3dMap.set(node, instance)
return instance
}
getLoad3d(node: LGraphNode): Load3d | Load3dAnimation | null {
return this.nodeToLoad3dMap.get(node) || null
}
getNodeByLoad3d(load3d: Load3d | Load3dAnimation): LGraphNode | null {
for (const [node, instance] of this.nodeToLoad3dMap) {
if (instance === load3d) {
return node
}
}
return null
}
removeLoad3d(node: LGraphNode) {
const instance = this.nodeToLoad3dMap.get(node)
if (instance) {
instance.remove()
this.nodeToLoad3dMap.delete(node)
}
}
clear() {
for (const [node] of this.nodeToLoad3dMap) {
this.removeLoad3d(node)
}
}
}
export const useLoad3dService = () => {
return Load3dService.getInstance()
}