From 5da5ee5031cedd2572d32b28daf94769fdb0a84e Mon Sep 17 00:00:00 2001 From: Terry Jia Date: Fri, 10 Jul 2026 11:39:51 -0400 Subject: [PATCH 01/13] feat: wire up Save 3D (Advanced) node family (CORE-329) (#13330) ## Summary Register the save-side advanced nodes in the Load3D viewer infrastructure: Save3DAdvanced reuses the mesh advanced extension, while SaveGaussianSplat and SavePointCloud reuse the splat/point cloud preview extensions. Parameterize both extension factories with a loadFolder so save nodes load the persisted file from the output folder instead of temp, and add the node types to the lazy-load and viewport-state sets. BE change https://github.com/Comfy-Org/ComfyUI/pull/14701 ## Screenshots (if applicable) Save 3D (Advanced) image Save Splat image --- src/composables/useLoad3d.test.ts | 57 +++- src/composables/useLoad3d.ts | 10 +- src/composables/useLoad3dViewer.ts | 4 +- src/extensions/core/load3d.test.ts | 99 +++++- src/extensions/core/load3d.ts | 322 +++++++++++------- src/extensions/core/load3d/interfaces.ts | 1 + src/extensions/core/load3d/nodeTypes.ts | 13 +- src/extensions/core/load3dLazy.test.ts | 5 +- .../core/load3dPreviewExtensions.test.ts | 109 +++++- .../core/load3dPreviewExtensions.ts | 60 +++- .../core/saveImageExtraOutput.test.ts | 3 + src/extensions/core/saveImageExtraOutput.ts | 3 + 12 files changed, 524 insertions(+), 162 deletions(-) diff --git a/src/composables/useLoad3d.test.ts b/src/composables/useLoad3d.test.ts index 30c550b500..2477a441fe 100644 --- a/src/composables/useLoad3d.test.ts +++ b/src/composables/useLoad3d.test.ts @@ -77,6 +77,14 @@ vi.mock('pinia', async (importOriginal) => { } }) +const { settingGetMock } = vi.hoisted(() => ({ + settingGetMock: vi.fn() +})) + +vi.mock('@/platform/settings/settingStore', () => ({ + useSettingStore: () => ({ get: settingGetMock }) +})) + vi.mock('@/renderer/core/canvas/canvasStore', () => ({ useCanvasStore: vi.fn() })) @@ -95,6 +103,9 @@ describe('useLoad3d', () => { vi.clearAllMocks() nodeToLoad3dMap.clear() vi.mocked(getActivePinia).mockReturnValue(null as unknown as Pinia) + settingGetMock.mockImplementation((key: string) => + key === 'Comfy.Load3D.BackgroundColor' ? '282828' : undefined + ) mockNode = createMockLGraphNode({ properties: { @@ -356,6 +367,20 @@ describe('useLoad3d', () => { expect(composable.isPreview.value).toBe(true) }) + it('should set preview mode for save-viewer nodes despite width/height widgets', async () => { + Object.defineProperty(mockNode, 'constructor', { + value: { comfyClass: 'Save3DAdvanced' }, + configurable: true + }) + + const composable = useLoad3d(mockNode) + const containerRef = document.createElement('div') + + await composable.initializeLoad3d(containerRef) + + expect(composable.isPreview.value).toBe(true) + }) + it('should handle initialization errors', async () => { vi.mocked(createLoad3d).mockImplementationOnce(() => { throw new Error('Load3d creation failed') @@ -383,7 +408,37 @@ describe('useLoad3d', () => { const nodeRef = shallowRef(mockNode) const composable = useLoad3d(nodeRef) - expect(composable.sceneConfig.value.backgroundColor).toBe('#000000') + expect(composable.sceneConfig.value.backgroundColor).toBe('#282828') + }) + + it('defaults background color from the Comfy.Load3D.BackgroundColor setting', () => { + vi.mocked(getActivePinia).mockReturnValue({} as unknown as Pinia) + vi.mocked(useCanvasStore).mockReturnValue( + reactive({ appScalePercentage: 100 }) as unknown as ReturnType< + typeof useCanvasStore + > + ) + settingGetMock.mockImplementation((key: string) => + key === 'Comfy.Load3D.BackgroundColor' ? '123456' : undefined + ) + + const composable = useLoad3d(mockNode) + + expect(composable.sceneConfig.value.backgroundColor).toBe('#123456') + }) + + it('attaches event listeners before running queued ready callbacks', async () => { + const composable = useLoad3d(mockNode) + let listenersAttachedWhenCallbackRan = false + + composable.waitForLoad3d(() => { + listenersAttachedWhenCallbackRan = + vi.mocked(mockLoad3d.addEventListener!).mock.calls.length > 0 + }) + + await composable.initializeLoad3d(document.createElement('div')) + + expect(listenersAttachedWhenCallbackRan).toBe(true) }) it('passes getZoomScale callback to createLoad3d', async () => { diff --git a/src/composables/useLoad3d.ts b/src/composables/useLoad3d.ts index 7dff603576..05d02c739d 100644 --- a/src/composables/useLoad3d.ts +++ b/src/composables/useLoad3d.ts @@ -8,6 +8,7 @@ import { useChainCallback } from '@/composables/functional/useChainCallback' import type Load3d from '@/extensions/core/load3d/Load3d' import Load3dUtils from '@/extensions/core/load3d/Load3dUtils' import { createLoad3d } from '@/extensions/core/load3d/createLoad3d' +import { isLoad3dResultViewerNode } from '@/extensions/core/load3d/nodeTypes' import { isAssetPreviewSupported, persistThumbnail @@ -118,7 +119,9 @@ export const useLoad3d = (nodeOrRef: MaybeRef) => { const sceneConfig = ref({ showGrid: true, - backgroundColor: '#000000', + backgroundColor: getActivePinia() + ? '#' + useSettingStore().get('Comfy.Load3D.BackgroundColor') + : '#282828', backgroundImage: '', backgroundRenderMode: 'tiled' }) @@ -192,6 +195,7 @@ export const useLoad3d = (nodeOrRef: MaybeRef) => { const heightWidget = node.widgets?.find((w) => w.name === 'height') if ( + isLoad3dResultViewerNode(node.constructor.comfyClass ?? '') || node.constructor.comfyClass?.startsWith('Preview') || !(widthWidget && heightWidget) ) { @@ -248,6 +252,8 @@ export const useLoad3d = (nodeOrRef: MaybeRef) => { nodeToLoad3dMap.set(node, load3d) + handleEvents('add') + const callbacks = pendingCallbacks.get(node) if (callbacks && load3d) { @@ -263,8 +269,6 @@ export const useLoad3d = (nodeOrRef: MaybeRef) => { if (load3d) invokeReadyCallback(callback, load3d) }) } - - handleEvents('add') } catch (error) { console.error('Error initializing Load3d:', error) useToastStore().addAlert( diff --git a/src/composables/useLoad3dViewer.ts b/src/composables/useLoad3dViewer.ts index c9c7772624..1920edb37d 100644 --- a/src/composables/useLoad3dViewer.ts +++ b/src/composables/useLoad3dViewer.ts @@ -4,7 +4,7 @@ import QuickLRU from '@alloc/quick-lru' import type Load3d from '@/extensions/core/load3d/Load3d' import Load3dUtils from '@/extensions/core/load3d/Load3dUtils' import { createLoad3d } from '@/extensions/core/load3d/createLoad3d' -import { isLoad3dPreviewNode } from '@/extensions/core/load3d/nodeTypes' +import { isLoad3dResultViewerNode } from '@/extensions/core/load3d/nodeTypes' import type { AnimationItem, BackgroundRenderModeType, @@ -371,7 +371,7 @@ export const useLoad3dViewer = (node?: LGraphNode) => { | LightConfig | undefined - isPreview.value = isLoad3dPreviewNode(node.type ?? '') + isPreview.value = isLoad3dResultViewerNode(node.type ?? '') if (sceneConfig) { backgroundColor.value = diff --git a/src/extensions/core/load3d.test.ts b/src/extensions/core/load3d.test.ts index d1bcbe5a6b..2c327923ec 100644 --- a/src/extensions/core/load3d.test.ts +++ b/src/extensions/core/load3d.test.ts @@ -143,14 +143,23 @@ async function loadExtensionsFresh(): Promise<{ load3DExt: ExtCreated preview3DExt: ExtCreated preview3DAdvancedExt: ExtCreated + save3DAdvancedExt: ExtCreated }> { vi.resetModules() registerExtensionMock.mockClear() await import('@/extensions/core/load3d') + const extByName = (name: string): ExtCreated => { + const call = registerExtensionMock.mock.calls.find( + (c) => (c[0] as ExtCreated).name === name + ) + if (!call) throw new Error(`Extension ${name} was not registered`) + return call[0] as ExtCreated + } return { - load3DExt: registerExtensionMock.mock.calls[0][0] as ExtCreated, - preview3DExt: registerExtensionMock.mock.calls[1][0] as ExtCreated, - preview3DAdvancedExt: registerExtensionMock.mock.calls[2][0] as ExtCreated + load3DExt: extByName('Comfy.Load3D'), + preview3DExt: extByName('Comfy.Preview3D'), + preview3DAdvancedExt: extByName('Comfy.Preview3DAdvanced'), + save3DAdvancedExt: extByName('Comfy.Save3DAdvanced') } } @@ -264,14 +273,15 @@ function setupBaseMocks() { describe('load3d module registration', () => { beforeEach(setupBaseMocks) - it('registers Comfy.Load3D, Comfy.Preview3D, and Comfy.Preview3DAdvanced extensions on import', async () => { - const { load3DExt, preview3DExt, preview3DAdvancedExt } = + it('registers Comfy.Load3D, Comfy.Preview3D, Comfy.Preview3DAdvanced, and Comfy.Save3DAdvanced extensions on import', async () => { + const { load3DExt, preview3DExt, preview3DAdvancedExt, save3DAdvancedExt } = await loadExtensionsFresh() - expect(registerExtensionMock).toHaveBeenCalledTimes(3) + expect(registerExtensionMock).toHaveBeenCalledTimes(4) expect(load3DExt.name).toBe('Comfy.Load3D') expect(preview3DExt.name).toBe('Comfy.Preview3D') expect(preview3DAdvancedExt.name).toBe('Comfy.Preview3DAdvanced') + expect(save3DAdvancedExt.name).toBe('Comfy.Save3DAdvanced') }) }) @@ -711,6 +721,39 @@ describe('Comfy.Preview3D.onNodeOutputsUpdated', () => { }) }) +describe('Comfy.Save3DAdvanced.onNodeOutputsUpdated', () => { + beforeEach(setupBaseMocks) + + it('restores the saved model from the output folder when opened from history', async () => { + const { save3DAdvancedExt } = await loadExtensionsFresh() + const node = makePreview3DAdvancedNode({ comfyClass: 'Save3DAdvanced' }) + getNodeByLocatorIdMock.mockReturnValue(node) + + save3DAdvancedExt.onNodeOutputsUpdated!({ + '7': { result: ['3d\\ComfyUI_00001.glb'] } + } as never) + + expect(node.properties['Last Time Model File']).toBe('3d/ComfyUI_00001.glb') + expect(configureForSaveMeshMock).toHaveBeenCalledWith( + 'output', + '3d/ComfyUI_00001.glb', + expect.objectContaining({ silentOnNotFound: true }) + ) + }) + + it('skips nodes whose comfyClass is not Save3DAdvanced', async () => { + const { save3DAdvancedExt } = await loadExtensionsFresh() + const node = makePreview3DAdvancedNode({ comfyClass: 'Preview3DAdvanced' }) + getNodeByLocatorIdMock.mockReturnValue(node) + + save3DAdvancedExt.onNodeOutputsUpdated!({ + '7': { result: ['mesh.glb'] } + } as never) + + expect(configureForSaveMeshMock).not.toHaveBeenCalled() + }) +}) + describe('Comfy.Preview3DAdvanced.nodeCreated', () => { beforeEach(setupBaseMocks) @@ -1032,6 +1075,50 @@ describe('Comfy.Preview3DAdvanced.getNodeMenuItems', () => { }) }) +describe('Comfy.Save3DAdvanced.nodeCreated', () => { + beforeEach(setupBaseMocks) + + it('skips nodes whose comfyClass is not Save3DAdvanced', async () => { + const { save3DAdvancedExt } = await loadExtensionsFresh() + const node = makePreview3DAdvancedNode({ comfyClass: 'Preview3DAdvanced' }) + + await save3DAdvancedExt.nodeCreated(node) + + expect(waitForLoad3dMock).not.toHaveBeenCalled() + expect(configureForSaveMeshMock).not.toHaveBeenCalled() + }) + + it('restores persisted models from the output folder, not temp', async () => { + const { save3DAdvancedExt } = await loadExtensionsFresh() + const node = makePreview3DAdvancedNode({ + comfyClass: 'Save3DAdvanced', + properties: { 'Last Time Model File': '3d/ComfyUI_00001_.glb' } + }) + + await save3DAdvancedExt.nodeCreated(node) + + expect(configureForSaveMeshMock).toHaveBeenCalledWith( + 'output', + '3d/ComfyUI_00001_.glb', + { silentOnNotFound: true } + ) + }) + + it('onExecuted loads the saved file from the output folder', async () => { + const { save3DAdvancedExt } = await loadExtensionsFresh() + const node = makePreview3DAdvancedNode({ comfyClass: 'Save3DAdvanced' }) + + await save3DAdvancedExt.nodeCreated(node) + node.onExecuted!({ result: ['3d/ComfyUI_00002_.glb'] }) + + expect(configureForSaveMeshMock).toHaveBeenCalledWith( + 'output', + '3d/ComfyUI_00002_.glb', + { silentOnNotFound: true } + ) + }) +}) + describe('Comfy.Load3D scene widget serializeValue caching', () => { beforeEach(setupBaseMocks) diff --git a/src/extensions/core/load3d.ts b/src/extensions/core/load3d.ts index cee600bbfa..d41c19af53 100644 --- a/src/extensions/core/load3d.ts +++ b/src/extensions/core/load3d.ts @@ -15,8 +15,10 @@ import { createExportMenuItems } from '@/extensions/core/load3d/exportMenuHelper import type { CameraConfig, CameraState, + LoadFolder, Model3DInfo } from '@/extensions/core/load3d/interfaces' +import type Load3d from '@/extensions/core/load3d/Load3d' import Load3DConfiguration from '@/extensions/core/load3d/Load3DConfiguration' import { LOAD3D_NONE_MODEL, @@ -48,6 +50,7 @@ import { ComponentWidgetImpl, addWidget } from '@/scripts/domWidget' import { useExtensionService } from '@/services/extensionService' import { useLoad3dService } from '@/services/load3dService' import { useDialogStore } from '@/stores/dialogStore' +import type { ComfyExtension } from '@/types/comfy' import { isLoad3dNode } from '@/utils/litegraphUtil' const inputSpecLoad3D: CustomInputSpec = { @@ -287,8 +290,11 @@ useExtensionService().registerExtension({ getCustomWidgets() { const VIEWPORT_STATE_NODES = new Set([ 'Preview3DAdvanced', + 'Save3DAdvanced', 'PreviewGaussianSplat', - 'PreviewPointCloud' + 'PreviewPointCloud', + 'SaveGaussianSplat', + 'SavePointCloud' ]) return { LOAD_3D(node) { @@ -679,155 +685,215 @@ useExtensionService().registerExtension({ } }) -useExtensionService().registerExtension({ - name: 'Comfy.Preview3DAdvanced', +function applyPreview3DAdvancedResult( + node: LGraphNode, + load3d: Load3d, + result: NonNullable, + loadFolder: LoadFolder, + comfyClass: string +): void { + const filePath = result[0] + if (!filePath) return - getNodeMenuItems(node: LGraphNode): (IContextMenuValue | null)[] { - if (node.constructor.comfyClass !== 'Preview3DAdvanced') return [] + const normalizedPath = filePath.replaceAll('\\', '/') + node.properties['Last Time Model File'] = normalizedPath - const load3d = useLoad3dService().getLoad3d(node) - if (!load3d) return [] + const config = new Load3DConfiguration(load3d, node.properties) + config.configureForSaveMesh(loadFolder, normalizedPath, { + silentOnNotFound: true + }) - if (load3d.isSplatModel()) return [] + const cameraState = result[1] + const modelTransform = result[2]?.[0] + if (!cameraState && !modelTransform) return - return createExportMenuItems(load3d) - }, + const targetGeneration = load3d.currentLoadGeneration + void load3d + .whenLoadIdle() + .then(() => { + if (load3d.currentLoadGeneration !== targetGeneration) return + if (cameraState) load3d.setCameraState(cameraState) + if (modelTransform) load3d.applyModelTransform(modelTransform) + }) + .catch((error) => { + console.error( + `Failed to apply input camera_info / model_3d_info from ${comfyClass}:`, + error + ) + }) +} - async nodeCreated(node: LGraphNode) { - if (node.constructor.comfyClass !== 'Preview3DAdvanced') return +function createPreview3DAdvancedExtension( + comfyClass: string, + extensionName: string, + loadFolder: LoadFolder +): ComfyExtension { + return { + name: extensionName, - const [oldWidth, oldHeight] = node.size + onNodeOutputsUpdated( + nodeOutputs: Record + ) { + for (const [locatorId, output] of Object.entries(nodeOutputs)) { + const result = (output as Preview3DAdvancedOutput).result + if (!result?.[0]) continue - node.setSize([Math.max(oldWidth, 400), Math.max(oldHeight, 550)]) + const node = getNodeByLocatorId(app.rootGraph, locatorId) + if (!node || node.constructor.comfyClass !== comfyClass) continue - await nextTick() - - const onExecuted = node.onExecuted - - useLoad3d(node).onLoad3dReady((load3d) => { - const lastTimeModelFile = node.properties['Last Time Model File'] - if (!lastTimeModelFile) return - - const config = new Load3DConfiguration(load3d, node.properties) - config.configureForSaveMesh('temp', lastTimeModelFile as string, { - silentOnNotFound: true - }) - - const cameraConfig = node.properties['Camera Config'] as - | CameraConfig - | undefined - const cameraState = cameraConfig?.state - if (!cameraState) return - - const targetGeneration = load3d.currentLoadGeneration - void load3d - .whenLoadIdle() - .then(() => { - if (load3d.currentLoadGeneration !== targetGeneration) return - load3d.setCameraState(cameraState) - load3d.forceRender() - }) - .catch((error) => { - console.error( - 'Failed to restore camera state for Preview3DAdvanced:', - error + useLoad3d(node).waitForLoad3d((load3d) => { + applyPreview3DAdvancedResult( + node, + load3d, + result, + loadFolder, + comfyClass ) }) - }) - - useLoad3d(node).waitForLoad3d((load3d) => { - const sceneWidget = node.widgets?.find((w) => w.name === 'viewport_state') - if (!sceneWidget) return - - const resolveLoad3d = () => nodeToLoad3dMap.get(node) ?? load3d - - const widthWidget = node.widgets?.find((w) => w.name === 'width') - const heightWidget = node.widgets?.find((w) => w.name === 'height') - if (widthWidget && heightWidget) { - load3d.setTargetSize( - widthWidget.value as number, - heightWidget.value as number - ) - widthWidget.callback = (value: number) => { - resolveLoad3d().setTargetSize(value, heightWidget.value as number) - } - heightWidget.callback = (value: number) => { - resolveLoad3d().setTargetSize(widthWidget.value as number, value) - } } + }, - sceneWidget.serializeValue = async () => { - const currentLoad3d = nodeToLoad3dMap.get(node) - if (!currentLoad3d) { - console.error('No load3d instance found for node') - return null - } + getNodeMenuItems(node: LGraphNode): (IContextMenuValue | null)[] { + if (node.constructor.comfyClass !== comfyClass) return [] - const cameraConfig: CameraConfig = (node.properties['Camera Config'] as - | CameraConfig - | undefined) || { - cameraType: currentLoad3d.getCurrentCameraType(), - fov: currentLoad3d.cameraManager.perspectiveCamera.fov - } - cameraConfig.state = currentLoad3d.getCameraState() - node.properties['Camera Config'] = cameraConfig + const load3d = useLoad3dService().getLoad3d(node) + if (!load3d) return [] - const modelInfo = currentLoad3d.getModelInfo() - const model_3d_info: Model3DInfo = modelInfo ? [modelInfo] : [] + if (load3d.isSplatModel()) return [] - return { - image: '', - mask: '', - normal: '', - camera_info: cameraConfig.state || null, - recording: '', - model_3d_info - } - } + return createExportMenuItems(load3d) + }, - node.onExecuted = function (output: Preview3DAdvancedOutput) { - onExecuted?.call(this, output) + async nodeCreated(node: LGraphNode) { + if (node.constructor.comfyClass !== comfyClass) return - const result = output.result - const filePath = result?.[0] + const [oldWidth, oldHeight] = node.size - if (!filePath) { - const msg = t('toastMessages.unableToGetModelFilePath') - console.error(msg) - useToastStore().addAlert(msg) - return - } + node.setSize([Math.max(oldWidth, 400), Math.max(oldHeight, 550)]) - const normalizedPath = filePath.replaceAll('\\', '/') - node.properties['Last Time Model File'] = normalizedPath + await nextTick() - const currentLoad3d = resolveLoad3d() - const config = new Load3DConfiguration(currentLoad3d, node.properties) - config.configureForSaveMesh('temp', normalizedPath, { + const onExecuted = node.onExecuted + const { onLoad3dReady, waitForLoad3d } = useLoad3d(node) + + onLoad3dReady((load3d) => { + const lastTimeModelFile = node.properties['Last Time Model File'] + if (!lastTimeModelFile) return + + const config = new Load3DConfiguration(load3d, node.properties) + config.configureForSaveMesh(loadFolder, lastTimeModelFile as string, { silentOnNotFound: true }) - const cameraState = result?.[1] - const modelTransform = result?.[2]?.[0] - if (cameraState || modelTransform) { - const targetGeneration = currentLoad3d.currentLoadGeneration - void currentLoad3d - .whenLoadIdle() - .then(() => { - if (currentLoad3d.currentLoadGeneration !== targetGeneration) - return - if (cameraState) currentLoad3d.setCameraState(cameraState) - if (modelTransform) - currentLoad3d.applyModelTransform(modelTransform) - }) - .catch((error) => { - console.error( - 'Failed to apply input camera_info / model_3d_info from Preview3DAdvanced:', - error - ) - }) + const cameraConfig = node.properties['Camera Config'] as + | CameraConfig + | undefined + const cameraState = cameraConfig?.state + if (!cameraState) return + + const targetGeneration = load3d.currentLoadGeneration + void load3d + .whenLoadIdle() + .then(() => { + if (load3d.currentLoadGeneration !== targetGeneration) return + load3d.setCameraState(cameraState) + load3d.forceRender() + }) + .catch((error) => { + console.error( + `Failed to restore camera state for ${comfyClass}:`, + error + ) + }) + }) + + waitForLoad3d((load3d) => { + const sceneWidget = node.widgets?.find( + (w) => w.name === 'viewport_state' + ) + if (!sceneWidget) return + + const resolveLoad3d = () => nodeToLoad3dMap.get(node) ?? load3d + + const widthWidget = node.widgets?.find((w) => w.name === 'width') + const heightWidget = node.widgets?.find((w) => w.name === 'height') + if (widthWidget && heightWidget) { + load3d.setTargetSize( + widthWidget.value as number, + heightWidget.value as number + ) + widthWidget.callback = (value: number) => { + resolveLoad3d().setTargetSize(value, heightWidget.value as number) + } + heightWidget.callback = (value: number) => { + resolveLoad3d().setTargetSize(widthWidget.value as number, value) + } } - } - }) + + sceneWidget.serializeValue = async () => { + const currentLoad3d = nodeToLoad3dMap.get(node) + if (!currentLoad3d) { + console.error('No load3d instance found for node') + return null + } + + const cameraConfig: CameraConfig = (node.properties[ + 'Camera Config' + ] as CameraConfig | undefined) || { + cameraType: currentLoad3d.getCurrentCameraType(), + fov: currentLoad3d.cameraManager.perspectiveCamera.fov + } + cameraConfig.state = currentLoad3d.getCameraState() + node.properties['Camera Config'] = cameraConfig + + const modelInfo = currentLoad3d.getModelInfo() + const model_3d_info: Model3DInfo = modelInfo ? [modelInfo] : [] + + return { + image: '', + mask: '', + normal: '', + camera_info: cameraConfig.state || null, + recording: '', + model_3d_info + } + } + + node.onExecuted = function (output: Preview3DAdvancedOutput) { + onExecuted?.call(this, output) + + const result = output.result + if (!result?.[0]) { + const msg = t('toastMessages.unableToGetModelFilePath') + console.error(msg) + useToastStore().addAlert(msg) + return + } + + applyPreview3DAdvancedResult( + node, + resolveLoad3d(), + result, + loadFolder, + comfyClass + ) + } + }) + } } -}) +} + +useExtensionService().registerExtension( + createPreview3DAdvancedExtension( + 'Preview3DAdvanced', + 'Comfy.Preview3DAdvanced', + 'temp' + ) +) +useExtensionService().registerExtension( + createPreview3DAdvancedExtension( + 'Save3DAdvanced', + 'Comfy.Save3DAdvanced', + 'output' + ) +) diff --git a/src/extensions/core/load3d/interfaces.ts b/src/extensions/core/load3d/interfaces.ts index 762f8252a5..d38e8e1001 100644 --- a/src/extensions/core/load3d/interfaces.ts +++ b/src/extensions/core/load3d/interfaces.ts @@ -14,6 +14,7 @@ export type MaterialMode = export type UpDirection = 'original' | '-x' | '+x' | '-y' | '+y' | '-z' | '+z' export type CameraType = 'perspective' | 'orthographic' export type BackgroundRenderModeType = 'tiled' | 'panorama' +export type LoadFolder = 'temp' | 'output' interface CameraQuaternion { x: number diff --git a/src/extensions/core/load3d/nodeTypes.ts b/src/extensions/core/load3d/nodeTypes.ts index f5d6b0a8cc..f0b5f3d3ad 100644 --- a/src/extensions/core/load3d/nodeTypes.ts +++ b/src/extensions/core/load3d/nodeTypes.ts @@ -3,21 +3,24 @@ * Adding a new node type that uses the viewer = one line change here. */ -const LOAD3D_PREVIEW_NODES = new Set([ +const LOAD3D_RESULT_VIEWER_NODES = new Set([ 'Preview3D', 'PreviewGaussianSplat', - 'PreviewPointCloud' + 'PreviewPointCloud', + 'Save3DAdvanced', + 'SaveGaussianSplat', + 'SavePointCloud' ]) const LOAD3D_ALL_NODES = new Set([ - ...LOAD3D_PREVIEW_NODES, + ...LOAD3D_RESULT_VIEWER_NODES, 'Load3D', 'Load3DAdvanced', 'SaveGLB' ]) -export const isLoad3dPreviewNode = (nodeType: string): boolean => - LOAD3D_PREVIEW_NODES.has(nodeType) +export const isLoad3dResultViewerNode = (nodeType: string): boolean => + LOAD3D_RESULT_VIEWER_NODES.has(nodeType) export const isLoad3dNode = (nodeType: string): boolean => LOAD3D_ALL_NODES.has(nodeType) diff --git a/src/extensions/core/load3dLazy.test.ts b/src/extensions/core/load3dLazy.test.ts index d7d49aaa7f..249f22a40a 100644 --- a/src/extensions/core/load3dLazy.test.ts +++ b/src/extensions/core/load3dLazy.test.ts @@ -90,7 +90,10 @@ describe('load3dLazy', () => { 'Preview3D', 'PreviewGaussianSplat', 'PreviewPointCloud', - 'SaveGLB' + 'SaveGLB', + 'Save3DAdvanced', + 'SaveGaussianSplat', + 'SavePointCloud' ])( 'recognizes %s as a 3D node type and triggers the lazy-load path', async (nodeType) => { diff --git a/src/extensions/core/load3dPreviewExtensions.test.ts b/src/extensions/core/load3dPreviewExtensions.test.ts index f7e4cbe1be..55b73dec3f 100644 --- a/src/extensions/core/load3dPreviewExtensions.test.ts +++ b/src/extensions/core/load3dPreviewExtensions.test.ts @@ -76,14 +76,24 @@ type ExtCreated = ComfyExtension & { async function loadExtensionsFresh(): Promise<{ splatExt: ExtCreated pointCloudExt: ExtCreated + saveSplatExt: ExtCreated + savePointCloudExt: ExtCreated }> { vi.resetModules() registerExtensionMock.mockClear() await import('@/extensions/core/load3dPreviewExtensions') - const [splatCall, pointCloudCall] = registerExtensionMock.mock.calls + const extByName = (name: string): ExtCreated => { + const call = registerExtensionMock.mock.calls.find( + (c) => (c[0] as ExtCreated).name === name + ) + if (!call) throw new Error(`Extension ${name} was not registered`) + return call[0] as ExtCreated + } return { - splatExt: splatCall[0] as ExtCreated, - pointCloudExt: pointCloudCall[0] as ExtCreated + splatExt: extByName('Comfy.PreviewGaussianSplat'), + pointCloudExt: extByName('Comfy.PreviewPointCloud'), + saveSplatExt: extByName('Comfy.SaveGaussianSplat'), + savePointCloudExt: extByName('Comfy.SavePointCloud') } } @@ -92,6 +102,7 @@ interface FakeLoad3d { isSplatModel: ReturnType forceRender: ReturnType setCameraState: ReturnType + applyModelTransform: ReturnType setTargetSize: ReturnType getCurrentCameraType: ReturnType getCameraState: ReturnType @@ -106,6 +117,7 @@ function makeLoad3dMock(): FakeLoad3d { isSplatModel: vi.fn(() => false), forceRender: vi.fn(), setCameraState: vi.fn(), + applyModelTransform: vi.fn(), setTargetSize: vi.fn(), getCurrentCameraType: vi.fn(() => 'perspective'), getCameraState: vi.fn(() => ({ position: { x: 0, y: 0, z: 0 } })), @@ -151,12 +163,59 @@ function setupBaseMocks() { describe('load3dPreviewExtensions module registration', () => { beforeEach(setupBaseMocks) - it('registers both preview extensions on import', async () => { - const { splatExt, pointCloudExt } = await loadExtensionsFresh() + it('registers preview and save extensions on import', async () => { + const { splatExt, pointCloudExt, saveSplatExt, savePointCloudExt } = + await loadExtensionsFresh() - expect(registerExtensionMock).toHaveBeenCalledTimes(2) + expect(registerExtensionMock).toHaveBeenCalledTimes(4) expect(splatExt.name).toBe('Comfy.PreviewGaussianSplat') expect(pointCloudExt.name).toBe('Comfy.PreviewPointCloud') + expect(saveSplatExt.name).toBe('Comfy.SaveGaussianSplat') + expect(savePointCloudExt.name).toBe('Comfy.SavePointCloud') + }) + + it('save extensions load the saved file from the output folder, not temp', async () => { + const { saveSplatExt, savePointCloudExt } = await loadExtensionsFresh() + const load3d = makeLoad3dMock() + waitForLoad3dMock.mockImplementation((cb: (l: FakeLoad3d) => void) => + cb(load3d) + ) + + const splatNode = makePreviewNode({ comfyClass: 'SaveGaussianSplat' }) + await saveSplatExt.nodeCreated(splatNode) + splatNode.onExecuted!({ result: ['3d/ComfyUI_00001_.ply'] }) + + expect(configureForSaveMeshMock).toHaveBeenLastCalledWith( + 'output', + '3d/ComfyUI_00001_.ply', + expect.objectContaining({ silentOnNotFound: true }) + ) + + const pcNode = makePreviewNode({ comfyClass: 'SavePointCloud' }) + await savePointCloudExt.nodeCreated(pcNode) + pcNode.onExecuted!({ result: ['3d/ComfyUI_00002_.ply'] }) + + expect(configureForSaveMeshMock).toHaveBeenLastCalledWith( + 'output', + '3d/ComfyUI_00002_.ply', + expect.objectContaining({ silentOnNotFound: true }) + ) + }) + + it('restores persisted models from the output folder on nodeCreated, not temp', async () => { + const { saveSplatExt } = await loadExtensionsFresh() + const node = makePreviewNode({ + comfyClass: 'SaveGaussianSplat', + properties: { 'Last Time Model File': '3d/ComfyUI_00001_.ply' } + }) + + await saveSplatExt.nodeCreated(node) + + expect(configureForSaveMeshMock).toHaveBeenCalledWith( + 'output', + '3d/ComfyUI_00001_.ply', + expect.objectContaining({ silentOnNotFound: true }) + ) }) }) @@ -214,6 +273,44 @@ describe('Comfy.PreviewGaussianSplat.nodeCreated', () => { expect(cameraConfig?.state).toEqual(cameraState) }) + it('applies onExecuted results to the remounted instance, not the disposed closure', async () => { + const { splatExt } = await loadExtensionsFresh() + const original = makeLoad3dMock() + waitForLoad3dMock.mockImplementation((cb: (l: FakeLoad3d) => void) => + cb(original) + ) + const node = makePreviewNode() + + await splatExt.nodeCreated(node) + + const remounted = makeLoad3dMock() + nodeToLoad3dMapMock.set(node, remounted) + + node.onExecuted!({ + result: ['scene.ply', { position: { x: 1, y: 2, z: 3 } }] + }) + await new Promise((resolve) => setTimeout(resolve, 0)) + + expect(remounted.forceRender).toHaveBeenCalled() + expect(original.forceRender).not.toHaveBeenCalled() + }) + + it('re-applies the model transform from result[2] on execute', async () => { + const { saveSplatExt } = await loadExtensionsFresh() + const load3d = makeLoad3dMock() + waitForLoad3dMock.mockImplementation((cb: (l: FakeLoad3d) => void) => + cb(load3d) + ) + const node = makePreviewNode({ comfyClass: 'SaveGaussianSplat' }) + const transform = { position: { x: 1, y: 2, z: 3 } } + + await saveSplatExt.nodeCreated(node) + node.onExecuted!({ result: ['scene.ply', undefined, [transform]] }) + await new Promise((resolve) => setTimeout(resolve, 0)) + + expect(load3d.applyModelTransform).toHaveBeenCalledWith(transform) + }) + it('syncs width/height widgets to load3d.setTargetSize and registers callbacks', async () => { const { splatExt } = await loadExtensionsFresh() const load3d = makeLoad3dMock() diff --git a/src/extensions/core/load3dPreviewExtensions.ts b/src/extensions/core/load3dPreviewExtensions.ts index 374ae0933a..dbe8c30ec4 100644 --- a/src/extensions/core/load3dPreviewExtensions.ts +++ b/src/extensions/core/load3dPreviewExtensions.ts @@ -5,6 +5,7 @@ import { createExportMenuItems } from '@/extensions/core/load3d/exportMenuHelper import type { CameraConfig, CameraState, + LoadFolder, Model3DInfo } from '@/extensions/core/load3d/interfaces' import type Load3d from '@/extensions/core/load3d/Load3d' @@ -29,7 +30,9 @@ function applyResultToLoad3d( node: LGraphNode, load3d: Load3d, filePath: string, - cameraState: CameraState | undefined + cameraState: CameraState | undefined, + modelTransform: Model3DInfo[number] | undefined, + loadFolder: LoadFolder ): void { const normalizedPath = filePath.replaceAll('\\', '/') node.properties['Last Time Model File'] = normalizedPath @@ -46,7 +49,7 @@ function applyResultToLoad3d( } const config = new Load3DConfiguration(load3d, node.properties) - config.configureForSaveMesh('temp', normalizedPath, { + config.configureForSaveMesh(loadFolder, normalizedPath, { silentOnNotFound: true }) @@ -54,13 +57,15 @@ function applyResultToLoad3d( void load3d.whenLoadIdle().then(() => { if (load3d.currentLoadGeneration !== targetGeneration) return if (cameraState) load3d.setCameraState(cameraState) + if (modelTransform) load3d.applyModelTransform(modelTransform) load3d.forceRender() }) } function createPreview3DExtension( comfyClass: string, - extensionName: string + extensionName: string, + loadFolder: LoadFolder ): ComfyExtension { const applyPreviewOutput = ( node: LGraphNode, @@ -68,10 +73,18 @@ function createPreview3DExtension( ): void => { const filePath = result[0] const cameraState = result[1] + const modelTransform = result[2]?.[0] if (!filePath) return useLoad3d(node).waitForLoad3d((load3d) => { - applyResultToLoad3d(node, load3d, filePath, cameraState) + applyResultToLoad3d( + node, + load3d, + filePath, + cameraState, + modelTransform, + loadFolder + ) }) } @@ -119,7 +132,7 @@ function createPreview3DExtension( if (!lastTimeModelFile) return const config = new Load3DConfiguration(load3d, node.properties) - config.configureForSaveMesh('temp', lastTimeModelFile as string, { + config.configureForSaveMesh(loadFolder, lastTimeModelFile as string, { silentOnNotFound: true }) @@ -136,6 +149,8 @@ function createPreview3DExtension( }) waitForLoad3d((load3d) => { + const resolveLoad3d = () => nodeToLoad3dMap.get(node) ?? load3d + const sceneWidget = node.widgets?.find( (w) => w.name === 'viewport_state' ) @@ -148,10 +163,10 @@ function createPreview3DExtension( heightWidget.value as number ) widthWidget.callback = (value: number) => { - load3d.setTargetSize(value, heightWidget.value as number) + resolveLoad3d().setTargetSize(value, heightWidget.value as number) } heightWidget.callback = (value: number) => { - load3d.setTargetSize(widthWidget.value as number, value) + resolveLoad3d().setTargetSize(widthWidget.value as number, value) } } @@ -199,7 +214,14 @@ function createPreview3DExtension( return } - applyResultToLoad3d(node, load3d, filePath, result?.[1]) + applyResultToLoad3d( + node, + resolveLoad3d(), + filePath, + result?.[1], + result?.[2]?.[0], + loadFolder + ) } }) } @@ -207,8 +229,26 @@ function createPreview3DExtension( } useExtensionService().registerExtension( - createPreview3DExtension('PreviewGaussianSplat', 'Comfy.PreviewGaussianSplat') + createPreview3DExtension( + 'PreviewGaussianSplat', + 'Comfy.PreviewGaussianSplat', + 'temp' + ) ) useExtensionService().registerExtension( - createPreview3DExtension('PreviewPointCloud', 'Comfy.PreviewPointCloud') + createPreview3DExtension( + 'PreviewPointCloud', + 'Comfy.PreviewPointCloud', + 'temp' + ) +) +useExtensionService().registerExtension( + createPreview3DExtension( + 'SaveGaussianSplat', + 'Comfy.SaveGaussianSplat', + 'output' + ) +) +useExtensionService().registerExtension( + createPreview3DExtension('SavePointCloud', 'Comfy.SavePointCloud', 'output') ) diff --git a/src/extensions/core/saveImageExtraOutput.test.ts b/src/extensions/core/saveImageExtraOutput.test.ts index 86947c10e6..e8c5bf39d8 100644 --- a/src/extensions/core/saveImageExtraOutput.test.ts +++ b/src/extensions/core/saveImageExtraOutput.test.ts @@ -81,6 +81,9 @@ describe('Comfy.SaveImageExtraOutput', () => { 'SaveAudioOpus', 'SaveAudioAdvanced', 'SaveGLB', + 'Save3DAdvanced', + 'SaveGaussianSplat', + 'SavePointCloud', 'SaveAnimatedPNG', 'CLIPSave', 'VAESave', diff --git a/src/extensions/core/saveImageExtraOutput.ts b/src/extensions/core/saveImageExtraOutput.ts index 3c3b044ac3..5ac24b7ae6 100644 --- a/src/extensions/core/saveImageExtraOutput.ts +++ b/src/extensions/core/saveImageExtraOutput.ts @@ -16,6 +16,9 @@ const saveNodeTypes = new Set([ 'SaveAudioOpus', 'SaveAudioAdvanced', 'SaveGLB', + 'Save3DAdvanced', + 'SaveGaussianSplat', + 'SavePointCloud', 'SaveAnimatedPNG', 'CLIPSave', 'VAESave', From 4ed2fe70f3b8920eb3691a4ad5e75f12f747f40b Mon Sep 17 00:00:00 2001 From: Terry Jia Date: Fri, 10 Jul 2026 11:40:17 -0400 Subject: [PATCH 02/13] feat: accept bboxes input and add grid snapping to Create Bounding Boxes (#13376) ## Summary - Seed/override the canvas from an upstream bboxes input - Add dotted grid background and magnetic snap-to-grid controls - Darken the canvas well for contrast BE is https://github.com/Comfy-Org/ComfyUI/pull/14724 ## Screenshots (if applicable) https://github.com/user-attachments/assets/7282f4d5-7cac-46f0-9d73-d0add37f4eb9 --- .../boundingBoxes/WidgetBoundingBoxes.vue | 108 +++++--- .../palette/PaletteSwatchRow.test.ts | 18 +- src/components/palette/PaletteSwatchRow.vue | 44 ++-- .../ui/color-picker/ColorPicker.vue | 98 +++---- .../ui/color-picker/ColorPickerPanel.vue | 5 +- .../boundingBoxes/boundingBoxesUtil.test.ts | 25 +- .../boundingBoxes/boundingBoxesUtil.ts | 20 +- .../boundingBoxes/useBoundingBoxes.test.ts | 240 +++++++++++++++++- .../boundingBoxes/useBoundingBoxes.ts | 172 ++++++++++++- .../palette/usePaletteSwatchRow.test.ts | 41 +-- .../palette/usePaletteSwatchRow.ts | 25 +- .../core/createBoundingBoxes.test.ts | 12 +- src/extensions/core/createBoundingBoxes.ts | 25 +- src/locales/en/main.json | 3 +- .../useBoundingBoxesWidget.test.ts | 2 +- .../composables/useBoundingBoxesWidget.ts | 3 +- 16 files changed, 640 insertions(+), 201 deletions(-) diff --git a/src/components/boundingBoxes/WidgetBoundingBoxes.vue b/src/components/boundingBoxes/WidgetBoundingBoxes.vue index 9873610929..5162556a75 100644 --- a/src/components/boundingBoxes/WidgetBoundingBoxes.vue +++ b/src/components/boundingBoxes/WidgetBoundingBoxes.vue @@ -4,37 +4,67 @@ data-testid="bounding-boxes" @pointerdown.stop > -
- -