diff --git a/comfy_extras/nodes_audio.py b/comfy_extras/nodes_audio.py index 8efa7c0a..34bcfa96 100644 --- a/comfy_extras/nodes_audio.py +++ b/comfy_extras/nodes_audio.py @@ -3,6 +3,10 @@ import torch import comfy.model_management import folder_paths import os +import io +import json +import struct +from comfy.cli_args import args class EmptyLatentAudio: def __init__(self): @@ -53,6 +57,61 @@ class VAEDecodeAudio: audio = vae.decode(samples["samples"]).movedim(-1, 1) return ({"waveform": audio, "sample_rate": 44100}, ) + +def create_vorbis_comment_block(comment_dict, last_block): + vendor_string = b'ComfyUI' + vendor_length = len(vendor_string) + + comments = [] + for key, value in comment_dict.items(): + comment = f"{key}={value}".encode('utf-8') + comments.append(struct.pack('I', len(comment_data))[1:] + comment_data + + return comment_block + +def insert_or_replace_vorbis_comment(flac_io, comment_dict): + if len(comment_dict) == 0: + return flac_io + + flac_io.seek(4) + + blocks = [] + last_block = False + + while not last_block: + header = flac_io.read(4) + last_block = (header[0] & 0x80) != 0 + block_type = header[0] & 0x7F + block_length = struct.unpack('>I', b'\x00' + header[1:])[0] + block_data = flac_io.read(block_length) + + if block_type == 4 or block_type == 1: + pass + else: + header = bytes([(header[0] & (~0x80))]) + header[1:] + blocks.append(header + block_data) + + blocks.append(create_vorbis_comment_block(comment_dict, last_block=True)) + + new_flac_io = io.BytesIO() + new_flac_io.write(b'fLaC') + for block in blocks: + new_flac_io.write(block) + + new_flac_io.write(flac_io.read()) + return new_flac_io + + class SaveAudio: def __init__(self): self.output_dir = folder_paths.get_output_directory() @@ -78,11 +137,27 @@ class SaveAudio: filename_prefix += self.prefix_append full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir) results = list() + + metadata = {} + if not args.disable_metadata: + if prompt is not None: + metadata["prompt"] = json.dumps(prompt) + if extra_pnginfo is not None: + for x in extra_pnginfo: + metadata[x] = json.dumps(extra_pnginfo[x]) + for (batch_number, waveform) in enumerate(audio["waveform"]): - #TODO: metadata filename_with_batch_num = filename.replace("%batch_num%", str(batch_number)) file = f"{filename_with_batch_num}_{counter:05}_.flac" - torchaudio.save(os.path.join(full_output_folder, file), waveform, audio["sample_rate"], format="FLAC") + + buff = io.BytesIO() + torchaudio.save(buff, waveform, audio["sample_rate"], format="FLAC") + + buff = insert_or_replace_vorbis_comment(buff, metadata) + + with open(os.path.join(full_output_folder, file), 'wb') as f: + f.write(buff.getbuffer()) + results.append({ "filename": file, "subfolder": subfolder, diff --git a/web/scripts/app.js b/web/scripts/app.js index 5fc9b9c6..43df6106 100644 --- a/web/scripts/app.js +++ b/web/scripts/app.js @@ -3,7 +3,7 @@ import { ComfyWidgets, initWidgets } from "./widgets.js"; import { ComfyUI, $el } from "./ui.js"; import { api } from "./api.js"; import { defaultGraph } from "./defaultGraph.js"; -import { getPngMetadata, getWebpMetadata, importA1111, getLatentMetadata } from "./pnginfo.js"; +import { getPngMetadata, getWebpMetadata, getFlacMetadata, importA1111, getLatentMetadata } from "./pnginfo.js"; import { addDomClippingSetting } from "./domWidget.js"; import { createImageHost, calculateImageGrid } from "./ui/imagePreview.js"; import { ComfyAppMenu } from "./ui/menu/index.js"; @@ -2277,6 +2277,19 @@ export class ComfyApp { const workflow = pngInfo?.workflow || pngInfo?.Workflow; const prompt = pngInfo?.prompt || pngInfo?.Prompt; + if (workflow) { + this.loadGraphData(JSON.parse(workflow), true, true, fileName); + } else if (prompt) { + this.loadApiJson(JSON.parse(prompt), fileName); + } else { + this.showErrorOnFileLoad(file); + } + } else if (file.type === "audio/flac") { + const pngInfo = await getFlacMetadata(file); + // Support loading workflows from that webp custom node. + const workflow = pngInfo?.workflow; + const prompt = pngInfo?.prompt; + if (workflow) { this.loadGraphData(JSON.parse(workflow), true, true, fileName); } else if (prompt) { diff --git a/web/scripts/pnginfo.js b/web/scripts/pnginfo.js index 7132fb60..2c03cf74 100644 --- a/web/scripts/pnginfo.js +++ b/web/scripts/pnginfo.js @@ -163,6 +163,78 @@ export function getLatentMetadata(file) { }); } + +function getString(dataView, offset, length) { + let string = ''; + for (let i = 0; i < length; i++) { + string += String.fromCharCode(dataView.getUint8(offset + i)); + } + return string; +} + +// Function to parse the Vorbis Comment block +function parseVorbisComment(dataView) { + let offset = 0; + const vendorLength = dataView.getUint32(offset, true); + offset += 4; + const vendorString = getString(dataView, offset, vendorLength); + offset += vendorLength; + + const userCommentListLength = dataView.getUint32(offset, true); + offset += 4; + const comments = {}; + for (let i = 0; i < userCommentListLength; i++) { + const commentLength = dataView.getUint32(offset, true); + offset += 4; + const comment = getString(dataView, offset, commentLength); + offset += commentLength; + + const [key, value] = comment.split('='); + + comments[key] = value; + } + + return comments; +} + +// Function to read a FLAC file and parse Vorbis comments +export function getFlacMetadata(file) { + return new Promise((r) => { + const reader = new FileReader(); + reader.onload = function(event) { + const arrayBuffer = event.target.result; + const dataView = new DataView(arrayBuffer); + + // Verify the FLAC signature + const signature = String.fromCharCode(...new Uint8Array(arrayBuffer, 0, 4)); + if (signature !== 'fLaC') { + console.error('Not a valid FLAC file'); + return; + } + + // Parse metadata blocks + let offset = 4; + let vorbisComment = null; + while (offset < dataView.byteLength) { + const isLastBlock = dataView.getUint8(offset) & 0x80; + const blockType = dataView.getUint8(offset) & 0x7F; + const blockSize = dataView.getUint32(offset, false) & 0xFFFFFF; + offset += 4; + + if (blockType === 4) { // Vorbis Comment block type + vorbisComment = parseVorbisComment(new DataView(arrayBuffer, offset, blockSize)); + } + + offset += blockSize; + if (isLastBlock) break; + } + + r(vorbisComment); + }; + reader.readAsArrayBuffer(file); + }); +} + export async function importA1111(graph, parameters) { const p = parameters.lastIndexOf("\nSteps:"); if (p > -1) { diff --git a/web/scripts/ui.js b/web/scripts/ui.js index bb094d8b..2c47412c 100644 --- a/web/scripts/ui.js +++ b/web/scripts/ui.js @@ -373,7 +373,7 @@ export class ComfyUI { const fileInput = $el("input", { id: "comfy-file-input", type: "file", - accept: ".json,image/png,.latent,.safetensors,image/webp", + accept: ".json,image/png,.latent,.safetensors,image/webp,audio/flac", style: {display: "none"}, parent: document.body, onchange: () => {