diff --git a/src/extensions/core/load3d.ts b/src/extensions/core/load3d.ts index 001181dc7f..1d0aa7fa97 100644 --- a/src/extensions/core/load3d.ts +++ b/src/extensions/core/load3d.ts @@ -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) diff --git a/src/services/load3dService.ts b/src/services/load3dService.ts new file mode 100644 index 0000000000..9008d8014d --- /dev/null +++ b/src/services/load3dService.ts @@ -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() + + 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() +}