sd-parsers
Version:
A library to read metadata from images created by Stable Diffusion
310 lines • 11.4 kB
JavaScript
;
Object.defineProperty(exports, "__esModule", { value: true });
exports.ComfyUIParser = void 0;
const parser_1 = require("./parser");
const data_1 = require("../data");
const exceptions_1 = require("../exceptions");
const SAMPLER_TYPES = ['WanVideoSampler'];
const SAMPLER_PARAMS = new Set(['sampler_name', 'steps', 'cfg']);
const REPLACEMENT_RULES = [];
const POSITIVE_PROMPT_KEYS = ['text', 'positive'];
const NEGATIVE_PROMPT_KEYS = ['text', 'negative'];
const IGNORE_LINK_TYPES_PROMPT = ['CLIP'];
const IGNORE_CLASS_TYPES = ['ConditioningCombine'];
/**
* Parser for images generated by ComfyUI
*/
class ComfyUIParser extends parser_1.Parser {
constructor() {
super(...arguments);
this.generator = data_1.Generators.COMFYUI;
}
async parse(parameters) {
let prompt;
let workflow;
try {
prompt = parameters.prompt;
if (typeof prompt === 'string') {
prompt = JSON.parse(prompt);
}
workflow = parameters.workflow;
if (typeof workflow === 'string') {
workflow = JSON.parse(workflow);
}
}
catch (error) {
throw new exceptions_1.ParserError(`Error reading parameters: ${error}`);
}
const { samplers, metadata } = ImageContext.extract(this, prompt, workflow);
return (0, data_1.createPromptInfo)(this.generator, samplers, metadata, parameters);
}
}
exports.ComfyUIParser = ComfyUIParser;
/**
* Context class for processing ComfyUI image data
*/
class ImageContext {
constructor(parser, prompt, workflow = {}) {
this.parser = parser;
this.processedNodes = new Set();
// Ensure that prompt keys are strings
try {
this.prompt = {};
for (const [k, v] of Object.entries(prompt)) {
this.prompt[String(k)] = v;
}
}
catch (error) {
throw new exceptions_1.ParserError(`Prompt has unexpected format: ${error}`);
}
// Build links dictionary
this.links = {};
try {
if (workflow && workflow.links && Array.isArray(workflow.links)) {
for (const link of workflow.links) {
if (Array.isArray(link) && link.length >= 6) {
const [, outputId, , inputId, , linkType] = link;
const inputIdStr = String(inputId);
const outputIdStr = String(outputId);
if (!this.links[inputIdStr]) {
this.links[inputIdStr] = {};
}
if (!this.links[inputIdStr][outputIdStr]) {
this.links[inputIdStr][outputIdStr] = new Set();
}
this.links[inputIdStr][outputIdStr].add(linkType);
}
}
}
}
catch (error) {
throw new exceptions_1.ParserError(`Workflow has unexpected format: ${error}`);
}
}
static extract(parser, prompt, workflow = {}) {
const context = new ImageContext(parser, prompt, workflow);
const samplers = [];
const metadata = {};
// Pass 1: get samplers and related data
for (const [nodeId, node] of Object.entries(context.prompt)) {
const sampler = context.tryGetSampler(nodeId, node);
if (sampler) {
samplers.push(sampler);
}
}
// Pass 2: put information from unprocessed nodes into metadata
for (const [nodeId, node] of Object.entries(context.prompt)) {
if (context.processedNodes.has(nodeId)) {
continue;
}
try {
if (node.inputs) {
const inputs = context.getInputValues(node.inputs, nodeId);
if (inputs && Object.keys(inputs).length > 0) {
const classType = node.class_type;
if (!metadata[classType]) {
metadata[classType] = [];
}
metadata[classType].push(inputs);
}
}
}
catch (error) {
// Ignore errors for individual nodes
}
}
return { samplers, metadata };
}
tryStepInto(nodeInputs, nodeNames) {
for (const name of nodeNames) {
try {
const targetId = nodeInputs[name]?.[0];
if (targetId && this.prompt[targetId]) {
const targetNode = this.prompt[targetId];
if (targetNode.inputs) {
return { ...targetNode.inputs };
}
}
}
catch (error) {
continue;
}
}
return nodeInputs;
}
tryGetSampler(nodeId, node) {
try {
const inputs = { ...node.inputs };
const classType = node.class_type;
if (!SAMPLER_TYPES.includes(classType) && !this.hasSamplerParams(inputs)) {
return null;
}
}
catch (error) {
return null;
}
if (this.parser.debug) {
console.log(`Found sampler #${nodeId}`);
}
this.processedNodes.add(nodeId);
const inputs = { ...node.inputs };
// Sampler parameters
const samplerName = inputs.sampler_name || inputs.scheduler || 'unknown';
delete inputs.sampler_name;
delete inputs.scheduler;
const samplerParameters = this.parser.normalizeParameters(this.getInputValues(inputs), REPLACEMENT_RULES);
// Get model
const model = this.getModel(nodeId);
// Get prompts
const { prompts, negativePrompts } = this.getPrompts(nodeId);
return (0, data_1.createSampler)(samplerName, samplerParameters, {
samplerId: nodeId,
model,
prompts,
negativePrompts
});
}
hasSamplerParams(inputs) {
const inputKeys = new Set(Object.keys(inputs));
for (const param of SAMPLER_PARAMS) {
if (inputKeys.has(param)) {
return true;
}
}
return false;
}
getModel(initialNodeId) {
// Simplified model extraction - would need more complex logic for full implementation
try {
const trace = this.getTrace([initialNodeId], ['MODEL']);
const metadata = this.getTraceMetadata(trace);
// Look for checkpoint loader or similar
for (const [classType, data] of Object.entries(metadata)) {
if (classType.includes('Checkpoint') || classType.includes('Model')) {
if (Array.isArray(data) && data.length > 0) {
const modelData = data[0];
if (typeof modelData === 'object' && modelData.ckpt_name) {
return (0, data_1.createModel)({ name: modelData.ckpt_name });
}
}
}
}
}
catch (error) {
// Ignore errors
}
return undefined;
}
getPrompts(initialNodeId) {
const prompts = [];
const negativePrompts = [];
try {
const trace = this.getTrace([initialNodeId], ['CONDITIONING']);
const metadata = this.getTraceMetadata(trace);
for (const [classType, data] of Object.entries(metadata)) {
if (IGNORE_CLASS_TYPES.includes(classType)) {
continue;
}
if (Array.isArray(data)) {
for (const item of data) {
if (typeof item === 'object') {
// Check for positive prompts
for (const key of POSITIVE_PROMPT_KEYS) {
if (item[key] && typeof item[key] === 'string') {
prompts.push((0, data_1.createPrompt)(item[key]));
}
}
// Check for negative prompts
for (const key of NEGATIVE_PROMPT_KEYS) {
if (item[key] && typeof item[key] === 'string') {
negativePrompts.push((0, data_1.createPrompt)(item[key]));
}
}
}
}
}
}
}
catch (error) {
// Ignore errors
}
return { prompts, negativePrompts };
}
getTrace(nodeIds, linkTypes) {
const visited = new Set();
const result = [];
const processNode = (nodeId) => {
if (visited.has(nodeId)) {
return;
}
visited.add(nodeId);
result.push(nodeId);
// Follow links of specified types
const nodeLinks = this.links[nodeId];
if (nodeLinks) {
for (const [outputId, types] of Object.entries(nodeLinks)) {
for (const linkType of linkTypes) {
if (types.has(linkType)) {
processNode(outputId);
}
}
}
}
};
for (const nodeId of nodeIds) {
processNode(nodeId);
}
return result;
}
getTraceMetadata(trace) {
const metadata = {};
for (const nodeId of trace) {
try {
const node = this.prompt[nodeId];
if (!node)
continue;
const classType = node.class_type;
if (IGNORE_CLASS_TYPES.includes(classType)) {
continue;
}
const value = this.getInputValues(node.inputs, nodeId);
if (!value || Object.keys(value).length === 0) {
continue;
}
if (metadata[classType]) {
if (Array.isArray(metadata[classType])) {
metadata[classType].push(value);
}
else {
metadata[classType] = [metadata[classType], value];
}
}
else {
metadata[classType] = value;
}
}
catch (error) {
continue;
}
}
return metadata;
}
getInputValues(inputs, nodeId) {
try {
const values = {};
for (const [key, value] of Object.entries(inputs)) {
if (!Array.isArray(value)) {
values[key] = value;
}
}
if (Object.keys(values).length > 0) {
return nodeId ? { id: nodeId, ...values } : values;
}
}
catch (error) {
// Ignore errors
}
return {};
}
}
//# sourceMappingURL=comfyui.js.map