754 lines
30 KiB
JavaScript
754 lines
30 KiB
JavaScript
/**
|
|
* Graph Query Utilities for ComfyUI
|
|
* Allows querying node connections and extracting image data from connected nodes
|
|
*/
|
|
|
|
import { app } from '../../../../scripts/app.js';
|
|
import { applyResizeParams } from './image_resize.js';
|
|
|
|
/**
|
|
* Find the source node connected to a specific input of the current node
|
|
* @param {Object} currentNode - The current node object
|
|
* @param {string} inputName - Name of the input to find connection for (e.g., "ref_image")
|
|
* @returns {Object|null} The source node connected to the input, or null if not connected
|
|
*/
|
|
function findConnectedSourceNode(currentNode, inputName) {
|
|
try {
|
|
// Look for connections where this node is the target
|
|
const graph = app.graph;
|
|
if (!graph || !graph.links) {
|
|
return null;
|
|
}
|
|
|
|
// Find input index for the given input name
|
|
let inputIndex = -1;
|
|
let allInputNames = [];
|
|
|
|
if (currentNode.inputs) {
|
|
for (let i = 0; i < currentNode.inputs.length; i++) {
|
|
const input = currentNode.inputs[i];
|
|
if (input && input.name) {
|
|
allInputNames.push(input.name);
|
|
if (input.name === inputName) {
|
|
inputIndex = i;
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
} else {
|
|
// Try to get inputs from the node configuration
|
|
if (currentNode.constructor?.nodeData?.input) {
|
|
const allInputs = {
|
|
...currentNode.constructor?.nodeData?.input?.required,
|
|
...currentNode.constructor?.nodeData?.input?.optional
|
|
};
|
|
|
|
// Find which position the input is in
|
|
let index = 0;
|
|
for (const [inputNameDef] of Object.entries(allInputs)) {
|
|
allInputNames.push(inputNameDef);
|
|
if (inputNameDef === inputName) {
|
|
inputIndex = index;
|
|
break;
|
|
}
|
|
index++;
|
|
}
|
|
}
|
|
}
|
|
|
|
// DO NOT fallback from ref_image to ref_images - they serve different purposes
|
|
// ref_image is for background canvas, ref_images is for box layer references
|
|
// Fallback removed to prevent ref_images[0] from overwriting bg_image.png
|
|
|
|
if (inputIndex === -1) {
|
|
return null;
|
|
}
|
|
|
|
// Find the link connected to this input (graph.links is a Map, not an array)
|
|
let link = null;
|
|
if (graph.links && graph.links instanceof Map) {
|
|
for (const [linkId, linkObj] of graph.links) {
|
|
if (!linkObj) continue;
|
|
|
|
if (linkObj.target_id === currentNode.id && linkObj.target_slot === inputIndex) {
|
|
link = linkObj;
|
|
break;
|
|
}
|
|
}
|
|
} else if (Array.isArray(graph.links)) {
|
|
// Fallback if links is an array (older format)
|
|
link = graph.links.find(linkObj => {
|
|
if (!linkObj) return false;
|
|
return linkObj.target_id === currentNode.id && linkObj.target_slot === inputIndex;
|
|
});
|
|
} else {
|
|
return null;
|
|
}
|
|
|
|
if (!link) {
|
|
// Check what inputs are connected - iterate through the Map
|
|
const connectedLinks = [];
|
|
if (graph.links instanceof Map) {
|
|
for (const [linkId, linkObj] of graph.links) {
|
|
if (linkObj && linkObj.target_id === currentNode.id) {
|
|
connectedLinks.push(linkObj);
|
|
}
|
|
}
|
|
} else if (Array.isArray(graph.links)) {
|
|
// Fallback for array format
|
|
for (const linkObj of graph.links) {
|
|
if (linkObj && linkObj.target_id === currentNode.id) {
|
|
connectedLinks.push(linkObj);
|
|
}
|
|
}
|
|
}
|
|
|
|
// Alternative approach: try to find any connection that might be the ref_image
|
|
if (connectedLinks.length > 0) {
|
|
// If we find connected links, try to match based on common patterns
|
|
// ref_image is typically an IMAGE type input
|
|
for (const connectedLink of connectedLinks) {
|
|
const sourceNode = graph._nodes?.find(node => node.id === connectedLink.origin_id);
|
|
if (sourceNode) {
|
|
// Check if this is likely an image source by checking the source's output type
|
|
const sourceOutput = sourceNode.outputs?.[connectedLink.origin_slot];
|
|
if (sourceOutput && (sourceOutput.type === 'IMAGE' || (sourceOutput.name && sourceOutput.name.includes('image')))) {
|
|
// Return this connection as a possibility
|
|
return {
|
|
node: sourceNode,
|
|
origin_slot: connectedLink.origin_slot,
|
|
target_slot: connectedLink.target_slot
|
|
};
|
|
}
|
|
}
|
|
}
|
|
|
|
// If we still can't find it by output type, just return the first connection as a fallback
|
|
// as it might be connected but the type matching isn't working
|
|
const firstConnectedLink = connectedLinks[0];
|
|
const firstSourceNode = graph._nodes?.find(node => node.id === firstConnectedLink.origin_id);
|
|
if (firstSourceNode) {
|
|
return {
|
|
node: firstSourceNode,
|
|
origin_slot: firstConnectedLink.origin_slot,
|
|
target_slot: firstConnectedLink.target_slot
|
|
};
|
|
}
|
|
}
|
|
return null;
|
|
} else {
|
|
// The direct match was successful
|
|
// Find the source node based on the link
|
|
const sourceNode = graph._nodes?.find(node => node.id === link.origin_id);
|
|
|
|
if (!sourceNode) {
|
|
return null;
|
|
}
|
|
|
|
// If the directly connected source node is not an image node, try the deep search approach
|
|
if (!isImageNode(sourceNode)) {
|
|
const deepResult = findDeepSourceNode(currentNode, inputName);
|
|
if (deepResult) {
|
|
return deepResult;
|
|
}
|
|
}
|
|
|
|
return {
|
|
node: sourceNode,
|
|
origin_slot: link.origin_slot,
|
|
target_slot: link.target_slot
|
|
};
|
|
}
|
|
|
|
// Find the source node based on the link
|
|
const sourceNode = graph._nodes?.find(node => node.id === link.origin_id);
|
|
|
|
if (!sourceNode) {
|
|
console.error(`Source node with id ${link.origin_id} not found`);
|
|
return null;
|
|
}
|
|
|
|
return {
|
|
node: sourceNode,
|
|
origin_slot: link.origin_slot,
|
|
target_slot: link.target_slot
|
|
};
|
|
} catch (error) {
|
|
console.error(`Error finding connected source node for "${inputName}":`, error);
|
|
console.error(`Stack:`, error.stack);
|
|
return null;
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Normalize a widget value into a URL/data URL we can fetch.
|
|
*/
|
|
function normalizeImageValue(value, fallbackType = 'temp') {
|
|
if (!value) return null;
|
|
if (typeof value === 'string') {
|
|
if (value.startsWith('data:') || value.startsWith('http')) return value;
|
|
return `/view?filename=${encodeURIComponent(value)}&type=${fallbackType}`;
|
|
}
|
|
if (typeof value === 'object') {
|
|
if (value.filename) return `/view?filename=${encodeURIComponent(value.filename)}&type=${fallbackType}`;
|
|
if (value.file) return `/view?filename=${encodeURIComponent(value.file)}&type=${fallbackType}`;
|
|
if (value.url) return value.url;
|
|
}
|
|
return null;
|
|
}
|
|
|
|
/**
|
|
* Extract image data from a source node that may contain image previews (single best-effort).
|
|
* @param {Object} sourceNodeObj - Object containing the source node and connection info
|
|
* @returns {string|null} Base64 image data or URL, or null if no image found
|
|
*/
|
|
async function extractImageFromSourceNode(sourceNodeObj) {
|
|
const images = await extractImagesFromSourceNode(sourceNodeObj, true);
|
|
return images && images.length ? images[0] : null;
|
|
}
|
|
|
|
/**
|
|
* Extract one or many image data URLs/base64 strings from a source node.
|
|
* @param {Object} sourceNodeObj
|
|
* @param {boolean} includeDomFallback whether to also scan DOM for a single preview
|
|
* @returns {Promise<string[]>}
|
|
*/
|
|
async function extractImagesFromSourceNode(sourceNodeObj, includeDomFallback = false) {
|
|
if (!sourceNodeObj || !sourceNodeObj.node) {
|
|
console.error('Invalid source node object');
|
|
return [];
|
|
}
|
|
|
|
const sourceNode = sourceNodeObj.node;
|
|
const collected = [];
|
|
const graph = app.graph;
|
|
|
|
const collectFromWidgets = async (widgets, fallbackType = 'temp') => {
|
|
if (!widgets) return;
|
|
for (const widget of widgets) {
|
|
if ((widget.type === 'image' || widget.name?.toLowerCase?.().includes('image')) && widget.value) {
|
|
// If widget value is an array (batch), gather all
|
|
if (Array.isArray(widget.value)) {
|
|
for (const v of widget.value) {
|
|
const normalized = normalizeImageValue(v, fallbackType);
|
|
if (normalized) {
|
|
const data = await loadImageAsBase64(normalized);
|
|
if (data) collected.push(data);
|
|
}
|
|
}
|
|
} else {
|
|
const normalized = normalizeImageValue(widget.value, fallbackType);
|
|
if (normalized) {
|
|
const data = await loadImageAsBase64(normalized);
|
|
if (data) collected.push(data);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
};
|
|
|
|
const collectFromUpstream = async (node, visited = new Set()) => {
|
|
if (!graph || !graph.links || !node || visited.has(node.id)) return;
|
|
visited.add(node.id);
|
|
const inputs = node.inputs || [];
|
|
for (let i = 0; i < inputs.length; i++) {
|
|
const linkId = node.inputs[i]?.link;
|
|
if (!linkId) continue;
|
|
const link = graph.links.get ? graph.links.get(linkId) : graph.links?.[linkId];
|
|
if (!link) continue;
|
|
const upstream = graph._nodes?.find((n) => n.id === link.origin_id);
|
|
if (upstream) {
|
|
const nested = await extractImagesFromSourceNode({ node: upstream, origin_slot: link.origin_slot, target_slot: link.target_slot }, false);
|
|
nested.forEach((img) => collected.push(img));
|
|
await collectFromUpstream(upstream, visited);
|
|
}
|
|
}
|
|
};
|
|
|
|
try {
|
|
// Different node types store image data differently
|
|
switch (sourceNode.type) {
|
|
case 'LoadImage':
|
|
case 'LoadImageUpload': {
|
|
await collectFromWidgets(sourceNode.widgets, 'input');
|
|
if (collected.length) break;
|
|
if (sourceNode.properties?.value) {
|
|
const normalized = normalizeImageValue(sourceNode.properties.value, 'input');
|
|
if (normalized) {
|
|
const data = await loadImageAsBase64(normalized);
|
|
if (data) collected.push(data);
|
|
}
|
|
}
|
|
break;
|
|
}
|
|
|
|
case 'LoadVideo': {
|
|
console.log('[graph_query DEBUG] Processing LoadVideo node');
|
|
console.log('[graph_query DEBUG] LoadVideo widgets:', sourceNode.widgets?.map(w => ({ name: w.name, value: w.value, type: w.type })));
|
|
|
|
// LoadVideo nodes store frames in widgets
|
|
await collectFromWidgets(sourceNode.widgets, 'input');
|
|
console.log('[graph_query DEBUG] Collected images from widgets:', collected.length);
|
|
if (collected.length) break;
|
|
// Also try upstream connections
|
|
console.log('[graph_query DEBUG] Trying upstream connections');
|
|
await collectFromUpstream(sourceNode);
|
|
console.log('[graph_query DEBUG] Total collected after upstream:', collected.length);
|
|
break;
|
|
}
|
|
|
|
case 'VAEDecode':
|
|
case 'PreviewImage':
|
|
case 'SaveImage': {
|
|
await collectFromWidgets(sourceNode.widgets, 'temp');
|
|
break;
|
|
}
|
|
|
|
case 'ImageResizeKJv2': // Handle image processing nodes
|
|
case 'ImageScale':
|
|
case 'ImageScaleBy':
|
|
case 'ImageUpscale':
|
|
case 'ImageCrop': {
|
|
await collectFromWidgets(sourceNode.widgets, 'temp');
|
|
break;
|
|
}
|
|
|
|
// Batching/list helpers: expect widgets to hold arrays of images
|
|
case 'ImageBatch':
|
|
case 'CreateImageList':
|
|
case 'ImpactMakeImageList':
|
|
case 'ImpactMakeAnyList': {
|
|
await collectFromWidgets(sourceNode.widgets, 'temp');
|
|
// If widgets didn't yield images, try upstream inputs (they may carry batches)
|
|
if (!collected.length) {
|
|
await collectFromUpstream(sourceNode);
|
|
}
|
|
break;
|
|
}
|
|
|
|
default: {
|
|
await collectFromWidgets(sourceNode.widgets, 'temp');
|
|
break;
|
|
}
|
|
}
|
|
|
|
if (collected.length) return collected;
|
|
|
|
if (includeDomFallback) {
|
|
// If no image found in widgets, try to look for image elements in the DOM
|
|
const canvasContainer = document.querySelector('.graphcanvas') || document.querySelector('#graph-canvas') || document.querySelector('.comfyui-body');
|
|
if (canvasContainer) {
|
|
// Look for any image associated with the source node
|
|
const nodeElement = canvasContainer.querySelector(`[data-node-id="${sourceNode.id}"]`);
|
|
if (nodeElement) {
|
|
const imgElements = nodeElement.querySelectorAll('img');
|
|
for (const img of imgElements) {
|
|
if (img.src && (img.src.startsWith('data:') || img.src.startsWith('http'))) {
|
|
collected.push(img.src);
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
return collected;
|
|
} catch (error) {
|
|
console.error(`Error extracting image(s) from source node ${sourceNode.type}:`, error);
|
|
return collected;
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Load image from URL as base64
|
|
* @param {string} imageUrl - URL of the image to load
|
|
* @returns {Promise<string|null>} Promise that resolves to base64 image data or null
|
|
*/
|
|
async function loadImageAsBase64(imageUrl) {
|
|
try {
|
|
if (!imageUrl) {
|
|
return null;
|
|
}
|
|
|
|
// Handle data URLs directly
|
|
if (imageUrl.startsWith('data:')) {
|
|
return imageUrl;
|
|
}
|
|
|
|
// Add cache-busting parameter to ensure fresh image
|
|
let urlToUse = imageUrl;
|
|
if (imageUrl.includes('?')) {
|
|
urlToUse = `${imageUrl}&t=${Date.now()}`;
|
|
} else {
|
|
urlToUse = `${imageUrl}?t=${Date.now()}`;
|
|
}
|
|
|
|
const response = await fetch(urlToUse);
|
|
if (!response.ok) {
|
|
// Suppress 404 errors as they're expected when checking for ref images
|
|
if (response.status !== 404) {
|
|
console.error(`Error loading image (${response.status}):`, urlToUse);
|
|
}
|
|
return null;
|
|
}
|
|
|
|
const blob = await response.blob();
|
|
return new Promise((resolve, reject) => {
|
|
const reader = new FileReader();
|
|
reader.onloadend = () => resolve(reader.result);
|
|
reader.onerror = reject;
|
|
reader.readAsDataURL(blob);
|
|
});
|
|
} catch (error) {
|
|
// Only log unexpected errors (not 404s)
|
|
if (!error.message?.includes('404')) {
|
|
console.error('Error loading image as base64:', error);
|
|
}
|
|
return null;
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Check if ImageResizeKJv2 is in the node chain and extract its parameters
|
|
* @param {Object} startNode - The starting node
|
|
* @param {string} inputName - The input name to trace
|
|
* @returns {Object|null} Resize node parameters or null if not found
|
|
*/
|
|
function findResizeNodeInChain(startNode, inputName) {
|
|
const graph = app.graph;
|
|
if (!graph || !graph.links) return null;
|
|
|
|
const visited = new Set();
|
|
const toVisit = [{ node: startNode, input: inputName }];
|
|
|
|
while (toVisit.length > 0) {
|
|
const { node, input } = toVisit.shift();
|
|
|
|
if (visited.has(node.id)) continue;
|
|
visited.add(node.id);
|
|
|
|
// Check if this node is ImageResizeKJv2
|
|
if (node.type === 'ImageResizeKJv2') {
|
|
// Extract parameters from the node's widgets
|
|
const params = {};
|
|
if (node.widgets) {
|
|
for (const widget of node.widgets) {
|
|
params[widget.name] = widget.value;
|
|
}
|
|
}
|
|
|
|
// Check if width/height are connected from other nodes
|
|
if (node.inputs) {
|
|
for (let i = 0; i < node.inputs.length; i++) {
|
|
const nodeInput = node.inputs[i];
|
|
const inputName = nodeInput.name;
|
|
|
|
// Find if this input has a connection
|
|
let link = null;
|
|
if (graph.links instanceof Map) {
|
|
for (const [linkId, linkObj] of graph.links) {
|
|
if (linkObj && linkObj.target_id === node.id && linkObj.target_slot === i) {
|
|
link = linkObj;
|
|
break;
|
|
}
|
|
}
|
|
} else if (Array.isArray(graph.links)) {
|
|
link = graph.links.find(linkObj =>
|
|
linkObj && linkObj.target_id === node.id && linkObj.target_slot === i
|
|
);
|
|
}
|
|
|
|
if (link) {
|
|
// Found a connection for this input
|
|
const sourceNode = graph._nodes?.find(n => n.id === link.origin_id);
|
|
if (sourceNode) {
|
|
// Try to get the value from the source node's output
|
|
const outputSlot = link.origin_slot;
|
|
|
|
// Check if source node has widgets_values or properties
|
|
if (sourceNode.widgets_values && sourceNode.widgets_values.length > outputSlot) {
|
|
params[inputName] = sourceNode.widgets_values[outputSlot];
|
|
} else if (sourceNode.widgets && sourceNode.widgets[outputSlot]) {
|
|
params[inputName] = sourceNode.widgets[outputSlot].value;
|
|
} else {
|
|
// Try to find matching output name
|
|
const outputInfo = sourceNode.outputs?.[outputSlot];
|
|
if (outputInfo) {
|
|
// If it's a primitive or value node, try to get its value
|
|
if (sourceNode.widgets && sourceNode.widgets.length > 0) {
|
|
const widget = sourceNode.widgets.find(w => w.name === outputInfo.name || w.name === 'value');
|
|
if (widget) {
|
|
params[inputName] = widget.value;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
return params;
|
|
}
|
|
|
|
// Continue traversing upstream
|
|
if (node.inputs) {
|
|
for (const nodeInput of node.inputs) {
|
|
const inputIndex = node.inputs.indexOf(nodeInput);
|
|
|
|
// Find link connected to this input
|
|
let link = null;
|
|
if (graph.links instanceof Map) {
|
|
for (const [linkId, linkObj] of graph.links) {
|
|
if (linkObj && linkObj.target_id === node.id && linkObj.target_slot === inputIndex) {
|
|
link = linkObj;
|
|
break;
|
|
}
|
|
}
|
|
} else if (Array.isArray(graph.links)) {
|
|
link = graph.links.find(linkObj =>
|
|
linkObj && linkObj.target_id === node.id && linkObj.target_slot === inputIndex
|
|
);
|
|
}
|
|
|
|
if (link) {
|
|
const sourceNode = graph._nodes?.find(n => n.id === link.origin_id);
|
|
if (sourceNode) {
|
|
toVisit.push({ node: sourceNode, input: nodeInput.name });
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
return null;
|
|
}
|
|
|
|
/**
|
|
* Main function to get reference image from connected node
|
|
* @param {Object} currentNode - The Power Spline Editor node
|
|
* @returns {Promise<string|null>} Promise that resolves to base64 image data or null
|
|
*/
|
|
async function getReferenceImageFromConnectedNode(currentNode, inputName = 'ref_image') {
|
|
// Step 1: Find the connected source node for ref_image input
|
|
let sourceNodeObj = findConnectedSourceNode(currentNode, inputName);
|
|
if (!sourceNodeObj) {
|
|
// Try deep search as fallback
|
|
sourceNodeObj = findDeepSourceNode(currentNode, inputName);
|
|
}
|
|
|
|
if (!sourceNodeObj) {
|
|
return null;
|
|
}
|
|
|
|
// Check if ImageResizeKJv2 is in the chain and get its parameters
|
|
const resizeParams = findResizeNodeInChain(currentNode, inputName);
|
|
|
|
// Step 2: Extract image data from the source node
|
|
let imageDataUrl = await extractImageFromSourceNode(sourceNodeObj);
|
|
if (!imageDataUrl) {
|
|
// If we couldn't extract image from the found source node, try to find the original image source via deep search
|
|
const deepSourceNodeObj = findDeepSourceNode(currentNode, inputName);
|
|
if (deepSourceNodeObj && deepSourceNodeObj.node.id !== sourceNodeObj.node.id) {
|
|
imageDataUrl = await extractImageFromSourceNode(deepSourceNodeObj);
|
|
if (imageDataUrl) {
|
|
sourceNodeObj = deepSourceNodeObj; // Update the source node for potential logging
|
|
}
|
|
}
|
|
}
|
|
|
|
if (!imageDataUrl) {
|
|
return null;
|
|
}
|
|
|
|
// Step 3: Load the image as base64
|
|
let base64Image = await loadImageAsBase64(imageDataUrl);
|
|
if (!base64Image) {
|
|
console.error('Failed to load image as base64');
|
|
return null;
|
|
}
|
|
|
|
// Step 4: Apply resize parameters if ImageResizeKJv2 was detected
|
|
if (resizeParams) {
|
|
try {
|
|
// Load image to apply resize
|
|
const img = new Image();
|
|
await new Promise((resolve, reject) => {
|
|
img.onload = resolve;
|
|
img.onerror = reject;
|
|
img.src = base64Image;
|
|
});
|
|
|
|
// Apply resize transformation
|
|
const { canvas } = await applyResizeParams(img, resizeParams);
|
|
|
|
// Convert back to base64
|
|
base64Image = canvas.toDataURL('image/jpeg', 0.95);
|
|
} catch (error) {
|
|
console.error('Error applying resize parameters:', error);
|
|
// Continue with original image if resize fails
|
|
}
|
|
}
|
|
|
|
return base64Image;
|
|
}
|
|
|
|
/**
|
|
* Wrapper to fetch one or more reference images (first frame chosen) from the ref_images input.
|
|
* Currently returns at most one image because ComfyUI previews expose only a single frame.
|
|
* @param {Object} currentNode
|
|
* @returns {Promise<string[]>} array of base64 data URLs (may be empty)
|
|
*/
|
|
async function getReferenceImagesFromConnectedNode(currentNode) {
|
|
let sourceNodeObj = findConnectedSourceNode(currentNode, 'ref_images');
|
|
if (!sourceNodeObj) {
|
|
sourceNodeObj = findDeepSourceNode(currentNode, 'ref_images');
|
|
}
|
|
|
|
let images = [];
|
|
if (sourceNodeObj) {
|
|
images = await extractImagesFromSourceNode(sourceNodeObj, true);
|
|
}
|
|
|
|
if ((!images || images.length === 0) && sourceNodeObj) {
|
|
const deep = findDeepSourceNode(currentNode, 'ref_images');
|
|
if (deep && deep.node.id !== sourceNodeObj.node.id) {
|
|
images = await extractImagesFromSourceNode(deep, true);
|
|
}
|
|
}
|
|
|
|
// Fallback to single-image path
|
|
if (!images || images.length === 0) {
|
|
const single = await getReferenceImageFromConnectedNode(currentNode, 'ref_images');
|
|
if (single) return [single];
|
|
return [];
|
|
}
|
|
|
|
return images;
|
|
}
|
|
|
|
/**
|
|
* Recursively find the original source node connected to a specific input
|
|
* @param {Object} currentNode - The node to start from
|
|
* @param {string} inputName - Name of the input to trace (e.g., "ref_image")
|
|
* @param {Set<number>} visited - Internal set to prevent infinite loops
|
|
* @returns {Object|null} The ultimate source node or null
|
|
*/
|
|
function findDeepSourceNode(currentNode, inputName, visited = new Set()) {
|
|
const graph = app.graph;
|
|
if (!graph || !graph.links) return null;
|
|
|
|
if (visited.has(currentNode.id)) return null;
|
|
visited.add(currentNode.id);
|
|
|
|
// Get the link object for the input
|
|
const inputIndex = currentNode.inputs?.findIndex(i => i.name === inputName) ?? -1;
|
|
if (inputIndex < 0) return null;
|
|
|
|
let link = null;
|
|
// Find the link connected to this input (graph.links is a Map, not an array)
|
|
if (graph.links && graph.links instanceof Map) {
|
|
for (const [linkId, linkObj] of graph.links) {
|
|
if (!linkObj) continue;
|
|
if (linkObj.target_id === currentNode.id && linkObj.target_slot === inputIndex) {
|
|
link = linkObj;
|
|
break;
|
|
}
|
|
}
|
|
} else if (Array.isArray(graph.links)) {
|
|
// Fallback if links is an array (older format)
|
|
link = graph.links.find(linkObj => {
|
|
if (!linkObj) return false;
|
|
return linkObj.target_id === currentNode.id && linkObj.target_slot === inputIndex;
|
|
});
|
|
}
|
|
|
|
if (!link) return null;
|
|
|
|
const sourceNode = graph._nodes?.find(node => node.id === link.origin_id);
|
|
if (!sourceNode) return null;
|
|
|
|
// Check if the source node contains image data
|
|
if (isImageNode(sourceNode)) {
|
|
// Prioritize true image loaders over processors
|
|
const imageLoaderTypes = ['LoadImage', 'LoadImageUpload'];
|
|
if (imageLoaderTypes.some(type => sourceNode.type?.includes(type))) {
|
|
return {
|
|
node: sourceNode,
|
|
origin_slot: link.origin_slot,
|
|
target_slot: link.target_slot
|
|
};
|
|
}
|
|
|
|
// If this is an image processor, continue looking upstream for the original source
|
|
if (sourceNode.inputs) {
|
|
for (const input of sourceNode.inputs) {
|
|
const deeper = findDeepSourceNode(sourceNode, input.name, new Set(visited)); // Use a copy of visited set for each branch
|
|
if (deeper) {
|
|
return deeper; // Return the original source found upstream
|
|
}
|
|
}
|
|
}
|
|
|
|
// If no original source found upstream, return this processor as fallback
|
|
return {
|
|
node: sourceNode,
|
|
origin_slot: link.origin_slot,
|
|
target_slot: link.target_slot
|
|
};
|
|
}
|
|
|
|
// If not an image node, recursively check each input of the source node for possible image connections
|
|
if (sourceNode.inputs) {
|
|
for (const input of sourceNode.inputs) {
|
|
const deeper = findDeepSourceNode(sourceNode, input.name, new Set(visited)); // Use a copy of visited set for each branch
|
|
if (deeper) return deeper;
|
|
}
|
|
}
|
|
|
|
return null;
|
|
}
|
|
|
|
/**
|
|
* Check if a node is likely to contain image data
|
|
* @param {Object} node - The node to check
|
|
* @returns {boolean} Whether the node is likely to contain image data
|
|
*/
|
|
function isImageNode(node) {
|
|
// Check node type for common image-related nodes
|
|
const imageNodeTypes = [
|
|
'LoadImage', 'LoadImageUpload', 'VAEDecode', 'PreviewImage', 'SaveImage',
|
|
'ImageScale', 'ImageScaleBy', 'ImageResizeKJv2', 'ImageUpscale', 'ImageCrop',
|
|
'ImagePadForOutpaint', 'ImageBatch', 'ImageBlend', 'ImageBlur', 'ImageColorToBW',
|
|
'ImageFlip', 'ImageOnlyCheckpointLoader', 'ImageApplyProcessing',
|
|
// Additional list/batch helpers that produce image lists
|
|
'CreateImageList', 'ImpactMakeImageList', 'ImpactMakeAnyList'
|
|
];
|
|
if (imageNodeTypes.some(type => node.type && node.type.includes(type))) {
|
|
return true;
|
|
}
|
|
|
|
// Check if node has IMAGE type outputs
|
|
if (node.outputs) {
|
|
for (const output of node.outputs) {
|
|
if (output.type === 'IMAGE' || (output.name && (output.name.includes('IMAGE') || output.name.toLowerCase().includes('image')))) {
|
|
return true;
|
|
}
|
|
}
|
|
}
|
|
|
|
// Check if node has image-related widgets
|
|
if (node.widgets) {
|
|
for (const widget of node.widgets) {
|
|
if (widget.type === 'image' || (widget.name && widget.name.toLowerCase().includes('image'))) {
|
|
return true;
|
|
}
|
|
}
|
|
}
|
|
|
|
return false;
|
|
}
|
|
|
|
// Export functions for use in other modules
|
|
export { findConnectedSourceNode, findDeepSourceNode, extractImageFromSourceNode, extractImagesFromSourceNode, loadImageAsBase64, getReferenceImageFromConnectedNode, getReferenceImagesFromConnectedNode };
|