diff --git a/browser_tests/assets/workflow.glb b/browser_tests/assets/workflow.glb new file mode 100644 index 000000000..725727e65 Binary files /dev/null and b/browser_tests/assets/workflow.glb differ diff --git a/browser_tests/fixtures/ComfyPage.ts b/browser_tests/fixtures/ComfyPage.ts index 52f03629f..72971c63d 100644 --- a/browser_tests/fixtures/ComfyPage.ts +++ b/browser_tests/fixtures/ComfyPage.ts @@ -471,6 +471,7 @@ export class ComfyPage { if (fileName.endsWith('.webp')) return 'image/webp' if (fileName.endsWith('.webm')) return 'video/webm' if (fileName.endsWith('.json')) return 'application/json' + if (fileName.endsWith('.glb')) return 'model/gltf-binary' return 'application/octet-stream' } diff --git a/browser_tests/loadWorkflowInMedia.spec.ts b/browser_tests/loadWorkflowInMedia.spec.ts index 40e939d2e..d060befcb 100644 --- a/browser_tests/loadWorkflowInMedia.spec.ts +++ b/browser_tests/loadWorkflowInMedia.spec.ts @@ -8,7 +8,8 @@ test.describe('Load Workflow in Media', () => { 'edited_workflow.webp', 'no_workflow.webp', 'large_workflow.webp', - 'workflow.webm' + 'workflow.webm', + 'workflow.glb' ].forEach(async (fileName) => { test(`Load workflow in ${fileName}`, async ({ comfyPage }) => { await comfyPage.dragAndDropFile(fileName) diff --git a/browser_tests/loadWorkflowInMedia.spec.ts-snapshots/workflow-glb-chromium-linux.png b/browser_tests/loadWorkflowInMedia.spec.ts-snapshots/workflow-glb-chromium-linux.png new file mode 100644 index 000000000..faadbbda9 Binary files /dev/null and b/browser_tests/loadWorkflowInMedia.spec.ts-snapshots/workflow-glb-chromium-linux.png differ diff --git a/src/scripts/app.ts b/src/scripts/app.ts index a06a0d83b..1789db352 100644 --- a/src/scripts/app.ts +++ b/src/scripts/app.ts @@ -23,6 +23,7 @@ import { import { type ComfyNodeDef as ComfyNodeDefV2 } from '@/schemas/nodeDef/nodeDefSchemaV2' import type { ComfyNodeDef as ComfyNodeDefV1 } from '@/schemas/nodeDefSchema' import { getFromWebmFile } from '@/scripts/metadata/ebml' +import { getGltfBinaryMetadata } from '@/scripts/metadata/gltf' import { useDialogService } from '@/services/dialogService' import { useExtensionService } from '@/services/extensionService' import { useLitegraphService } from '@/services/litegraphService' @@ -1391,6 +1392,15 @@ export class ComfyApp { } else { this.showErrorOnFileLoad(file) } + } else if (file.type === 'model/gltf-binary') { + const gltfInfo = await getGltfBinaryMetadata(file) + if (gltfInfo.workflow) { + this.loadGraphData(gltfInfo.workflow, true, true, fileName) + } else if (gltfInfo.prompt) { + this.loadApiJson(gltfInfo.prompt, fileName) + } else { + this.showErrorOnFileLoad(file) + } } else if ( file.type === 'application/json' || file.name?.endsWith('.json') diff --git a/src/scripts/metadata/gltf.ts b/src/scripts/metadata/gltf.ts new file mode 100644 index 000000000..37e65cb00 --- /dev/null +++ b/src/scripts/metadata/gltf.ts @@ -0,0 +1,170 @@ +import { + ComfyApiWorkflow, + ComfyWorkflowJSON +} from '@/schemas/comfyWorkflowSchema' +import { + ASCII, + ComfyMetadata, + ComfyMetadataTags, + GltfChunkHeader, + GltfHeader, + GltfJsonData, + GltfSizeBytes +} from '@/types/metadataTypes' + +const MAX_READ_BYTES = 1 << 20 + +const isJsonChunk = (chunk: GltfChunkHeader | null): boolean => + !!chunk && chunk.chunkTypeIdentifier === ASCII.JSON + +const isValidChunkRange = ( + start: number, + length: number, + bufferSize: number +): boolean => start + length <= bufferSize + +const byteArrayToString = (bytes: Uint8Array): string => + new TextDecoder().decode(bytes) + +const parseGltfBinaryHeader = (dataView: DataView): GltfHeader | null => { + if (dataView.byteLength < GltfSizeBytes.HEADER) return null + + const magicNumber = dataView.getUint32(0, true) + if (magicNumber !== ASCII.GLTF) return null + + return { + magicNumber, + gltfFormatVersion: dataView.getUint32(4, true), + totalLengthBytes: dataView.getUint32(8, true) + } +} + +const parseChunkHeaderAtOffset = ( + dataView: DataView, + offset: number +): GltfChunkHeader | null => { + if (offset + GltfSizeBytes.CHUNK_HEADER > dataView.byteLength) return null + + return { + chunkLengthBytes: dataView.getUint32(offset, true), + chunkTypeIdentifier: dataView.getUint32(offset + 4, true) + } +} + +const extractJsonChunk = ( + buffer: ArrayBuffer +): { start: number; length: number } | null => { + const dataView = new DataView(buffer) + + const header = parseGltfBinaryHeader(dataView) + if (!header) return null + + const chunkOffset = GltfSizeBytes.HEADER + const firstChunk = parseChunkHeaderAtOffset(dataView, chunkOffset) + if (!firstChunk || !isJsonChunk(firstChunk)) return null + + const jsonStart = chunkOffset + GltfSizeBytes.CHUNK_HEADER + const isValid = isValidChunkRange( + jsonStart, + firstChunk.chunkLengthBytes, + dataView.byteLength + ) + if (!isValid) return null + + return { start: jsonStart, length: firstChunk.chunkLengthBytes } +} + +const extractJsonChunkData = (buffer: ArrayBuffer): Uint8Array | null => { + const chunkLocation = extractJsonChunk(buffer) + if (!chunkLocation) return null + + return new Uint8Array(buffer, chunkLocation.start, chunkLocation.length) +} + +const parseJson = (text: string): ReturnType | null => { + try { + return JSON.parse(text) + } catch { + return null + } +} + +const parseJsonBytes = ( + bytes: Uint8Array +): ReturnType | null => { + const jsonString = byteArrayToString(bytes) + return parseJson(jsonString) +} + +const parseMetadataValue = ( + value: string | object +): ComfyWorkflowJSON | ComfyApiWorkflow | undefined => { + if (typeof value !== 'string') + return value as ComfyWorkflowJSON | ComfyApiWorkflow + + const parsed = parseJson(value) + if (!parsed) return undefined + + return parsed as ComfyWorkflowJSON | ComfyApiWorkflow +} + +const extractComfyMetadata = (jsonData: GltfJsonData): ComfyMetadata => { + const metadata: ComfyMetadata = {} + + if (!jsonData?.asset?.extras) return metadata + + const { extras } = jsonData.asset + + if (extras.workflow) { + const parsedValue = parseMetadataValue(extras.workflow) + if (parsedValue) { + metadata[ComfyMetadataTags.WORKFLOW.toLowerCase()] = parsedValue + } + } + + if (extras.prompt) { + const parsedValue = parseMetadataValue(extras.prompt) + if (parsedValue) { + metadata[ComfyMetadataTags.PROMPT.toLowerCase()] = parsedValue + } + } + + return metadata +} + +const processGltfFileBuffer = (buffer: ArrayBuffer): ComfyMetadata => { + const jsonChunk = extractJsonChunkData(buffer) + if (!jsonChunk) return {} + + const parsedJson = parseJsonBytes(jsonChunk) + if (!parsedJson) return {} + + return extractComfyMetadata(parsedJson) +} + +/** + * Extract ComfyUI metadata from a GLTF binary file (GLB) + */ +export function getGltfBinaryMetadata(file: File): Promise { + return new Promise((resolve) => { + if (!file) return Promise.resolve({}) + + const bytesToRead = Math.min(file.size, MAX_READ_BYTES) + + const reader = new FileReader() + reader.onload = (event) => { + try { + if (!event.target?.result) { + resolve({}) + return + } + + resolve(processGltfFileBuffer(event.target.result as ArrayBuffer)) + } catch { + resolve({}) + } + } + reader.onerror = () => resolve({}) + reader.readAsArrayBuffer(file.slice(0, bytesToRead)) + }) +} diff --git a/src/types/metadataTypes.ts b/src/types/metadataTypes.ts index 0923fcb7c..27449a8d5 100644 --- a/src/types/metadataTypes.ts +++ b/src/types/metadataTypes.ts @@ -43,3 +43,38 @@ export type TextRange = { start: number end: number } + +export enum ASCII { + GLTF = 0x46546c67, + JSON = 0x4e4f534a +} + +export enum GltfSizeBytes { + HEADER = 12, + CHUNK_HEADER = 8 +} + +export type GltfHeader = { + magicNumber: number + gltfFormatVersion: number + totalLengthBytes: number +} + +export type GltfChunkHeader = { + chunkLengthBytes: number + chunkTypeIdentifier: number +} + +export type GltfExtras = { + workflow?: string | object + prompt?: string | object + [key: string]: any +} + +export type GltfJsonData = { + asset?: { + extras?: GltfExtras + [key: string]: any + } + [key: string]: any +} diff --git a/tests-ui/tests/scripts/metadata/gltf.test.ts b/tests-ui/tests/scripts/metadata/gltf.test.ts new file mode 100644 index 000000000..1968973c2 --- /dev/null +++ b/tests-ui/tests/scripts/metadata/gltf.test.ts @@ -0,0 +1,163 @@ +import { describe, expect, it } from 'vitest' + +import { ASCII, GltfSizeBytes } from '@/types/metadataTypes' + +import { getGltfBinaryMetadata } from '../../../../src/scripts/metadata/gltf' + +describe('GLTF binary metadata parser', () => { + const createGLTFFileStructure = () => { + const header = new ArrayBuffer(GltfSizeBytes.HEADER) + const headerView = new DataView(header) + return { header, headerView } + } + + const jsonToBinary = (json: object) => { + const jsonString = JSON.stringify(json) + const jsonData = new TextEncoder().encode(jsonString) + return jsonData + } + + const createJSONChunk = (jsonData: ArrayBuffer) => { + const chunkHeader = new ArrayBuffer(GltfSizeBytes.CHUNK_HEADER) + const chunkView = new DataView(chunkHeader) + chunkView.setUint32(0, jsonData.byteLength, true) + chunkView.setUint32(4, ASCII.JSON, true) + return chunkHeader + } + + const setVersionHeader = (headerView: DataView, version: number) => { + headerView.setUint32(4, version, true) + } + + const setTypeHeader = (headerView: DataView, type: number) => { + headerView.setUint32(0, type, true) + } + + const setTotalLengthHeader = (headerView: DataView, length: number) => { + headerView.setUint32(8, length, true) + } + + const setHeaders = (headerView: DataView, jsonData: ArrayBuffer) => { + setTypeHeader(headerView, ASCII.GLTF) + setVersionHeader(headerView, 2) + setTotalLengthHeader( + headerView, + GltfSizeBytes.HEADER + GltfSizeBytes.CHUNK_HEADER + jsonData.byteLength + ) + } + + function createMockGltfFile(jsonContent: object): File { + const jsonData = jsonToBinary(jsonContent) + const { header, headerView } = createGLTFFileStructure() + + setHeaders(headerView, jsonData) + + const chunkHeader = createJSONChunk(jsonData) + + const fileContent = new Uint8Array( + header.byteLength + chunkHeader.byteLength + jsonData.byteLength + ) + fileContent.set(new Uint8Array(header), 0) + fileContent.set(new Uint8Array(chunkHeader), header.byteLength) + fileContent.set(jsonData, header.byteLength + chunkHeader.byteLength) + + return new File([fileContent], 'test.glb', { type: 'model/gltf-binary' }) + } + + it('should extract workflow metadata from GLTF binary file', async () => { + const testWorkflow = { + nodes: [ + { + id: 1, + type: 'TestNode', + pos: [100, 100] + } + ], + links: [] + } + + const mockFile = createMockGltfFile({ + asset: { + version: '2.0', + generator: 'ComfyUI GLTF Test', + extras: { + workflow: testWorkflow + } + }, + scenes: [] + }) + + const metadata = await getGltfBinaryMetadata(mockFile) + + expect(metadata).toBeDefined() + expect(metadata.workflow).toBeDefined() + + const workflow = metadata.workflow as { + nodes: Array<{ id: number; type: string }> + } + expect(workflow.nodes[0].id).toBe(1) + expect(workflow.nodes[0].type).toBe('TestNode') + }) + + it('should extract prompt metadata from GLTF binary file', async () => { + const testPrompt = { + node1: { + class_type: 'TestNode', + inputs: { + seed: 123456 + } + } + } + + const mockFile = createMockGltfFile({ + asset: { + version: '2.0', + generator: 'ComfyUI GLTF Test', + extras: { + prompt: testPrompt + } + }, + scenes: [] + }) + + const metadata = await getGltfBinaryMetadata(mockFile) + expect(metadata).toBeDefined() + expect(metadata.prompt).toBeDefined() + + const prompt = metadata.prompt as Record + expect(prompt.node1.class_type).toBe('TestNode') + expect(prompt.node1.inputs.seed).toBe(123456) + }) + + it('should handle string JSON content', async () => { + const workflowStr = JSON.stringify({ + nodes: [{ id: 1, type: 'StringifiedNode' }], + links: [] + }) + + const mockFile = createMockGltfFile({ + asset: { + version: '2.0', + extras: { + workflow: workflowStr // As string instead of object + } + } + }) + + const metadata = await getGltfBinaryMetadata(mockFile) + + expect(metadata).toBeDefined() + expect(metadata.workflow).toBeDefined() + + const workflow = metadata.workflow as { + nodes: Array<{ id: number; type: string }> + } + expect(workflow.nodes[0].type).toBe('StringifiedNode') + }) + + it('should handle invalid GLTF binary files gracefully', async () => { + const invalidEmptyFile = new File([], 'invalid.glb') + const metadata = await getGltfBinaryMetadata(invalidEmptyFile) + expect(metadata).toEqual({}) + }) +})