diff --git a/browser_tests/assets/workflow.webm b/browser_tests/assets/workflow.webm new file mode 100644 index 0000000000..4eea34d618 Binary files /dev/null and b/browser_tests/assets/workflow.webm differ diff --git a/browser_tests/fixtures/ComfyPage.ts b/browser_tests/fixtures/ComfyPage.ts index 0cb047990c..6ee0d21395 100644 --- a/browser_tests/fixtures/ComfyPage.ts +++ b/browser_tests/fixtures/ComfyPage.ts @@ -468,6 +468,7 @@ export class ComfyPage { const getFileType = (fileName: string) => { if (fileName.endsWith('.png')) return 'image/png' if (fileName.endsWith('.webp')) return 'image/webp' + if (fileName.endsWith('.webm')) return 'video/webm' if (fileName.endsWith('.json')) return 'application/json' return 'application/octet-stream' } diff --git a/browser_tests/loadWorkflowInMedia.spec.ts b/browser_tests/loadWorkflowInMedia.spec.ts index e4a89b2708..40e939d2e3 100644 --- a/browser_tests/loadWorkflowInMedia.spec.ts +++ b/browser_tests/loadWorkflowInMedia.spec.ts @@ -7,7 +7,8 @@ test.describe('Load Workflow in Media', () => { 'workflow.webp', 'edited_workflow.webp', 'no_workflow.webp', - 'large_workflow.webp' + 'large_workflow.webp', + 'workflow.webm' ].forEach(async (fileName) => { test(`Load workflow in ${fileName}`, async ({ comfyPage }) => { await comfyPage.dragAndDropFile(fileName) diff --git a/browser_tests/loadWorkflowInMedia.spec.ts-snapshots/workflow-webm-chromium-2x-linux.png b/browser_tests/loadWorkflowInMedia.spec.ts-snapshots/workflow-webm-chromium-2x-linux.png new file mode 100644 index 0000000000..28b892d024 Binary files /dev/null and b/browser_tests/loadWorkflowInMedia.spec.ts-snapshots/workflow-webm-chromium-2x-linux.png differ diff --git a/browser_tests/loadWorkflowInMedia.spec.ts-snapshots/workflow-webm-chromium-linux.png b/browser_tests/loadWorkflowInMedia.spec.ts-snapshots/workflow-webm-chromium-linux.png new file mode 100644 index 0000000000..1ea8f18192 Binary files /dev/null and b/browser_tests/loadWorkflowInMedia.spec.ts-snapshots/workflow-webm-chromium-linux.png differ diff --git a/src/scripts/app.ts b/src/scripts/app.ts index b8fc903b94..fcf1279eff 100644 --- a/src/scripts/app.ts +++ b/src/scripts/app.ts @@ -20,6 +20,7 @@ import { validateComfyWorkflow } from '@/schemas/comfyWorkflowSchema' import type { ComfyNodeDef } from '@/schemas/nodeDefSchema' +import { getFromWebmFile } from '@/scripts/metadata/ebml' import { useDialogService } from '@/services/dialogService' import { useExtensionService } from '@/services/extensionService' import { useLitegraphService } from '@/services/litegraphService' @@ -1408,6 +1409,15 @@ export class ComfyApp { } else { this.showErrorOnFileLoad(file) } + } else if (file.type === 'video/webm') { + const webmInfo = await getFromWebmFile(file) + if (webmInfo.workflow) { + this.loadGraphData(webmInfo.workflow, true, true, fileName) + } else if (webmInfo.prompt) { + this.loadApiJson(webmInfo.prompt, fileName) + } else { + this.showErrorOnFileLoad(file) + } } else if ( file.type === 'application/json' || file.name?.endsWith('.json') diff --git a/src/scripts/metadata/ebml.ts b/src/scripts/metadata/ebml.ts new file mode 100644 index 0000000000..edbee97532 --- /dev/null +++ b/src/scripts/metadata/ebml.ts @@ -0,0 +1,358 @@ +import { + type ComfyApiWorkflow, + type ComfyWorkflowJSON +} from '@/schemas/comfyWorkflowSchema' +import { + ComfyMetadata, + ComfyMetadataTags, + EbmlElementRange, + EbmlTagPosition, + TextRange, + VInt +} from '@/types/metadataTypes' + +const WEBM_SIGNATURE = [0x1a, 0x45, 0xdf, 0xa3] +const MAX_READ_BYTES = 2 * 1024 * 1024 +const EBML_ID = { + SIMPLE_TAG: new Uint8Array([0x67, 0xc8]), + TAG_NAME: new Uint8Array([0x45, 0xa3]), + TAG_VALUE: new Uint8Array([0x44, 0x87]) +} +const ASCII = { + OPEN_BRACE: 0x7b, + NULL: 0, + PRINTABLE_MIN: 32, + PRINTABLE_MAX: 126 +} +const MASK = { + MSB: 0x80, + ALL_BITS_SET: -1 +} + +const hasWebmSignature = (data: Uint8Array): boolean => + WEBM_SIGNATURE.every((byte, index) => data[index] === byte) + +const readVint = (data: Uint8Array, pos: number): VInt | null => { + if (pos >= data.length) return null + + const byte = data[pos] + + // Fast path for common case (1-byte vint) + if ((byte & MASK.MSB) === MASK.MSB) { + return { value: byte & ~MASK.MSB, length: 1 } + } + + const length = findFirstSetBitPosition(byte) + if (length === 0 || pos + length > data.length) { + return null + } + + return { + value: calculateVintValue(data, pos, length), + length + } +} + +const calculateVintValue = ( + data: Uint8Array, + pos: number, + length: number +): number => { + let value = data[pos] & (0xff >> length) + + for (let i = 1; i < length; i++) { + value = (value << 8) | data[pos + i] + } + + const allBitsSet = Math.pow(2, 7 * length) - 1 + return value === allBitsSet ? MASK.ALL_BITS_SET : value +} + +const findFirstSetBitPosition = (byte: number): number => { + for (let mask = MASK.MSB, position = 1; mask !== 0; mask >>= 1, position++) { + if ((byte & mask) !== 0) return position + } + return 0 +} + +const matchesId = (data: Uint8Array, pos: number, id: Uint8Array): boolean => { + if (pos + id.length > data.length) return false + + for (let i = 0; i < id.length; i++) { + if (data[pos + i] !== id[i]) return false + } + + return true +} + +const findNextTag = ( + data: Uint8Array, + pos: number +): { tagEnd: number; contents: EbmlTagPosition | null } | null => { + if (!matchesId(data, pos, EBML_ID.SIMPLE_TAG)) return null + + const tagSize = readVint(data, pos + 2) + if (!tagSize || tagSize.value <= 0) return null + + const tagEnd = pos + 2 + tagSize.length + tagSize.value + if (tagEnd > data.length) return null + + const contents = extractTagContents(data, pos + 2 + tagSize.length, tagEnd) + return { tagEnd, contents } +} + +const extractTagContents = ( + data: Uint8Array, + start: number, + end: number +): EbmlTagPosition | null => { + const nameInfo = findElementsInTag(data, start, end, EBML_ID.TAG_NAME) + const valueInfo = findElementsInTag(data, start, end, EBML_ID.TAG_VALUE) + + if (!nameInfo || !valueInfo) return null + + return { + name: nameInfo, + value: valueInfo + } +} + +const findElementsInTag = ( + data: Uint8Array, + start: number, + end: number, + id: Uint8Array +): EbmlElementRange | null => { + for (let pos = start; pos < end - 1; ) { + if (matchesId(data, pos, id)) { + const size = readVint(data, pos + 2) + if (size && size.value > 0) { + const elementPos = pos + 2 + size.length + return { pos: elementPos, len: size.value } + } + } + pos++ + } + + return null +} + +const processTagContents = ( + data: Uint8Array, + contents: EbmlTagPosition, + meta: ComfyMetadata +) => { + try { + const name = extractEbmlTagName(data, contents.name.pos, contents.name.len) + if (!name) return + + const value = extractEbmlTagValue( + data, + name, + contents.value.pos, + contents.value.len + ) + if (value !== null) { + meta[name.toLowerCase()] = value + } + } catch { + // Silently continue on error + } +} + +const extractEbmlTagValue = ( + data: Uint8Array, + name: string, + pos: number, + len: number +): string | ComfyWorkflowJSON | ComfyApiWorkflow | null => { + if ( + name === ComfyMetadataTags.PROMPT || + name === ComfyMetadataTags.WORKFLOW + ) { + return readJson(data, pos, len) + } + + return ebmlToString(data, pos, len) +} + +const extractEbmlTagName = ( + data: Uint8Array, + start: number, + length: number +): string | null => { + if (length <= 0) return null + + const textRange = findReadableTextRange(data, start, length) + if (!textRange) return null + + return new TextDecoder() + .decode(data.subarray(textRange.start, textRange.end)) + .trim() +} + +const findReadableTextRange = ( + data: Uint8Array, + start: number, + length: number +): TextRange | null => { + const isPrintableAscii = (byte: number) => + byte >= ASCII.PRINTABLE_MIN && byte <= ASCII.PRINTABLE_MAX + + let textStart = start + while (textStart < start + length && !isPrintableAscii(data[textStart])) { + textStart++ + } + + if (textStart >= start + length) return null + + let textEnd = textStart + while (textEnd < start + length && data[textEnd] !== ASCII.NULL) { + textEnd++ + } + + return { start: textStart, end: textEnd } +} + +const readJson = ( + data: Uint8Array, + start: number, + length: number +): ComfyWorkflowJSON | ComfyApiWorkflow | null => { + if (length <= 0) return null + + const nullTerminatorPos = findNullTerminator(data, start, length) + const jsonStartPos = findJsonStart(data, start, nullTerminatorPos - start) + + if (jsonStartPos === null) return null + + const jsonText = decodeJsonText(data, jsonStartPos, nullTerminatorPos) + return parseJsonText(jsonText) +} + +const decodeJsonText = ( + data: Uint8Array, + start: number, + end: number +): string => { + return new TextDecoder().decode(data.subarray(start, end)) +} + +const parseJsonText = ( + jsonText: string +): ComfyWorkflowJSON | ComfyApiWorkflow | null => { + const jsonEndPos = findJsonEnd(jsonText) + if (jsonEndPos === null) return null + + try { + return JSON.parse(jsonText.substring(0, jsonEndPos)) + } catch { + return null + } +} + +const findNullTerminator = ( + data: Uint8Array, + start: number, + maxLength: number +): number => { + const end = Math.min(start + maxLength, data.length) + + for (let pos = start; pos < end; pos++) { + if (data[pos] === ASCII.NULL) return pos + } + + return end +} + +const findJsonStart = ( + data: Uint8Array, + start: number, + length: number +): number | null => { + for (let pos = start; pos < start + length; pos++) { + if (data[pos] === ASCII.OPEN_BRACE) return pos + } + + return null +} + +const findJsonEnd = (text: string): number | null => { + let braceCount = 1 + let pos = 1 // Start after the opening brace + + while (braceCount > 0 && pos < text.length) { + if (text[pos] === '{') braceCount++ + if (text[pos] === '}') braceCount-- + pos++ + } + + return braceCount === 0 ? pos : null +} + +const ebmlToString = ( + data: Uint8Array, + start: number, + length: number +): string => { + if (length <= 0) return '' + + const endPos = findNullTerminator(data, start, length) + return new TextDecoder().decode(data.subarray(start, endPos)).trim() +} + +const parseMetadata = (data: Uint8Array): ComfyMetadata => { + const meta: ComfyMetadata = {} + + for (let pos = 0; pos < data.length - 2; ) { + const tagInfo = findNextTag(data, pos) + if (!tagInfo) { + pos++ + continue + } + + const { tagEnd, contents } = tagInfo + if (contents) { + processTagContents(data, contents, meta) + } + + pos = tagEnd + } + + return meta +} + +const handleFileLoad = ( + event: ProgressEvent, + resolve: (value: ComfyMetadata) => void +) => { + if (!event.target?.result) { + resolve({}) + return + } + + try { + const data = new Uint8Array(event.target.result as ArrayBuffer) + if (data.length < 4 || !hasWebmSignature(data)) { + resolve({}) + return + } + + resolve(parseMetadata(data)) + } catch { + resolve({}) + } +} + +/** + * Extracts ComfyUI Workflow metadata from a WebM file + * @param file - The WebM file to extract metadata from + */ +export function getFromWebmFile(file: File): Promise { + return new Promise((resolve) => { + const reader = new FileReader() + reader.onload = (event) => handleFileLoad(event, resolve) + reader.onerror = () => resolve({}) + reader.readAsArrayBuffer(file.slice(0, MAX_READ_BYTES)) + }) +} diff --git a/src/types/metadataTypes.ts b/src/types/metadataTypes.ts new file mode 100644 index 0000000000..0923fcb7c7 --- /dev/null +++ b/src/types/metadataTypes.ts @@ -0,0 +1,45 @@ +import type { + ComfyApiWorkflow, + ComfyWorkflowJSON +} from '@/schemas/comfyWorkflowSchema' + +/** + * Tag names used in ComfyUI metadata + */ +export enum ComfyMetadataTags { + PROMPT = 'PROMPT', + WORKFLOW = 'WORKFLOW' +} + +/** + * Metadata extracted from ComfyUI output files + */ +export interface ComfyMetadata { + workflow?: ComfyWorkflowJSON + prompt?: ComfyApiWorkflow + [key: string]: string | ComfyWorkflowJSON | ComfyApiWorkflow | undefined +} + +export type EbmlElementRange = { + /** Position in the buffer */ + pos: number + /** Length of the element in bytes */ + len: number +} + +export type EbmlTagPosition = { + name: EbmlElementRange + value: EbmlElementRange +} + +export type VInt = { + /** The value of the variable-length integer */ + value: number + /** The length of the variable-length integer in bytes */ + length: number +} + +export type TextRange = { + start: number + end: number +}