375 lines
16 KiB
JavaScript
375 lines
16 KiB
JavaScript
import { SettingUtils } from './sn0w.js';
|
|
import { widgets } from '../settings/scheduler_settings.js';
|
|
import { app } from "../../../scripts/app.js";
|
|
import { api } from '../../../scripts/api.js';
|
|
import { ComfyWidgets } from "../../../scripts/widgets.js";
|
|
|
|
app.registerExtension({
|
|
name: "sn0w.SimpleSamplerCustom",
|
|
async beforeRegisterNodeDef(nodeType, nodeData, app) {
|
|
const allSettings = new Set();
|
|
|
|
if (nodeData.name === "Simple Sampler Custom") {
|
|
// Set positive and negative connection colors
|
|
const onConnectInput = nodeType.prototype.onConnectInput;
|
|
nodeType.prototype.onConnectInput = function (targetSlot, type, output, originNode, originSlot) {
|
|
const slot = type[0]
|
|
const inputType = type[1]
|
|
|
|
const initialCount = countWidgetsOfType(targetSlot.widgets, "converted-widget");
|
|
const originalSize = [targetSlot.size[0], targetSlot.size[1]];
|
|
|
|
if (slot === 3 || slot === 4) {
|
|
if (inputType == "STRING" || inputType == "CONDITIONING") {
|
|
targetSlot.inputs[slot].color_on = app.canvas.default_connection_color_byType[inputType];
|
|
} else {
|
|
console.error(`The input type has to be STRING or CONDITIONING, it cannot be ${inputType}.`)
|
|
targetSlot.inputs[slot].color_on = app.canvas.default_connection_color_byType["VAE"];
|
|
}
|
|
} else if (slot === 5) {
|
|
hideAllSchedulerWidgets(targetSlot);
|
|
SettingUtils.hideWidget(targetSlot, findWidget(targetSlot, "scheduler_name"));
|
|
|
|
resizeNode(targetSlot, initialCount, originalSize);
|
|
} else if (slot === 6) {
|
|
hideAllLatentWidgets(targetSlot);
|
|
resizeNode(targetSlot, initialCount, originalSize);
|
|
}
|
|
|
|
onConnectInput?.apply(targetSlot, type, output, originNode, originSlot);
|
|
}
|
|
|
|
// Handle when connection is disconnected
|
|
const onConnectionsChange = nodeType.prototype.onConnectionsChange;
|
|
nodeType.prototype.onConnectionsChange = function (side, slot, connect, link_info, output) {
|
|
if (slot === 5) {
|
|
if (this.inputs[slot].link === null) {
|
|
const initialCount = countWidgetsOfType(this.widgets, "converted-widget");
|
|
const originalSize = [this.size[0], this.size[1]];
|
|
|
|
const schedulerWidget = findWidget(this, "scheduler_name");
|
|
showSchedulerInputs(this, schedulerWidget.value);
|
|
SettingUtils.showWidget(schedulerWidget);
|
|
if (schedulerWidget.type === undefined)
|
|
schedulerWidget.type = "combo";
|
|
|
|
resizeNode(this, initialCount, originalSize);
|
|
}
|
|
} else if (slot === 6) {
|
|
const initialCount = countWidgetsOfType(this.widgets, "converted-widget");
|
|
const originalSize = [this.size[0], this.size[1]];
|
|
|
|
if (this.inputs[slot].link === null)
|
|
showAllLatentWidgets(this);
|
|
|
|
resizeNode(this, initialCount, originalSize);
|
|
}
|
|
onConnectionsChange?.apply(side, slot, connect, link_info, output);
|
|
}
|
|
|
|
nodeType.prototype.onNodeCreated = function () {
|
|
const schedulerWidget = findWidget(this, "scheduler_name");
|
|
schedulerWidget.value = schedulerWidget.value;
|
|
|
|
// Set up all inputs
|
|
createEverything(this);
|
|
this.inputs[3].type = ["STRING", "CONDITIONING"];
|
|
this.inputs[4].type = ["STRING", "CONDITIONING"];
|
|
|
|
if (schedulerWidget && schedulerWidget.callback) {
|
|
const originalCallback = schedulerWidget.callback;
|
|
schedulerWidget.callback = (newValue) => {
|
|
showSchedulerInputs(this, newValue);
|
|
if (originalCallback) originalCallback.call(this, newValue);
|
|
};
|
|
}
|
|
|
|
api.addEventListener('get_widget_values', (event) => {
|
|
const data = event.detail
|
|
const output = getWidgetOutputs(this, data.widgets_needed);
|
|
if (this.id == data.id) {
|
|
api.fetchApi(`${SettingUtils.API_PREFIX}/widget_values`, {
|
|
method: "POST",
|
|
headers: {
|
|
"Content-Type": "application/json",
|
|
},
|
|
body: JSON.stringify(
|
|
{
|
|
node_id: data.id,
|
|
outputs: output,
|
|
}
|
|
),
|
|
})
|
|
}
|
|
})
|
|
|
|
api.addEventListener('should_decode_image', (event) => {
|
|
const data = event.detail
|
|
const output = checkImageOutput(this);
|
|
if (this.id == data.id) {
|
|
api.fetchApi(`${SettingUtils.API_PREFIX}/should_decode_image`, {
|
|
method: "POST",
|
|
headers: {
|
|
"Content-Type": "application/json",
|
|
},
|
|
body: JSON.stringify(
|
|
{
|
|
node_id: data.id,
|
|
outputs: output,
|
|
}
|
|
),
|
|
})
|
|
}
|
|
})
|
|
|
|
api.addEventListener('sampler_get_sigmas', (event) => {
|
|
const data = event.detail;
|
|
if (this.id == data.id) {
|
|
const imageBase64 = SettingUtils.drawSigmas(data.sigmas);
|
|
|
|
// Send the generated image data back to the server
|
|
api.fetchApi(`${SettingUtils.API_PREFIX}/get_sigmas`, {
|
|
method: "POST",
|
|
headers: {
|
|
"Content-Type": "application/json",
|
|
},
|
|
body: JSON.stringify(
|
|
{
|
|
node_id: data.id,
|
|
outputs: {
|
|
output: imageBase64
|
|
}
|
|
}
|
|
),
|
|
})
|
|
}
|
|
});
|
|
}
|
|
|
|
nodeType.prototype.onConfigure = function () {
|
|
const inputName = findWidget(this, "scheduler_name");
|
|
const desiredWidgets = widgets[inputName.value] ? Object.keys(widgets[inputName.value]) : [];
|
|
showSchedulerInputs(this, inputName, desiredWidgets);
|
|
|
|
|
|
if (this.inputs[5].link !== null) {
|
|
hideAllSchedulerWidgets(this);
|
|
SettingUtils.hideWidget(this, findWidget(this, "scheduler_name"));
|
|
}
|
|
if (this.inputs[6].link !== null) {
|
|
hideAllLatentWidgets(this);
|
|
}
|
|
};
|
|
|
|
function findWidget(node, name) {
|
|
const widget = node.widgets.find(widget => widget.name === name);
|
|
return widget;
|
|
}
|
|
|
|
function countWidgetsOfType(widgets, type) {
|
|
return widgets.filter(widget => widget.type === type).length;
|
|
}
|
|
|
|
function resizeNode(node, initialCount, originalSize) {
|
|
const finalCount = countWidgetsOfType(node.widgets, "converted-widget");
|
|
node.size[0] = originalSize[0];
|
|
node.size[1] = originalSize[1] + (initialCount - finalCount) * (70 / 3);
|
|
}
|
|
|
|
function getWidgetOutputs(node, WidgetsToGet) {
|
|
const widgets = node.widgets;
|
|
|
|
// Filter out only the widgets that are in the WidgetsToGet array
|
|
const outputWidgets = widgets.filter(widget => WidgetsToGet.includes(widget.name));
|
|
|
|
// Create an object with the widget names as keys and their values wrapped in an object
|
|
return Object.fromEntries(
|
|
outputWidgets.map(widget => [
|
|
widget.name,
|
|
{ value: widget.value }
|
|
])
|
|
);
|
|
}
|
|
|
|
function checkImageOutput(node) {
|
|
try {
|
|
if (node.outputs[0].links.length === 0)
|
|
return false;
|
|
return true;
|
|
} catch {
|
|
return false;
|
|
}
|
|
}
|
|
|
|
function createWidgetsToRemove(widgets) {
|
|
// Use a set to avoid duplicates
|
|
const widgetNames = new Set();
|
|
|
|
// Iterate over each category in the widgets object
|
|
for (const category in widgets) {
|
|
// Get all keys from the category and add them to the set
|
|
Object.keys(widgets[category]).forEach(key => widgetNames.add(key));
|
|
}
|
|
|
|
// Convert the set back to an array
|
|
return Array.from(widgetNames);
|
|
}
|
|
|
|
function rearrangeWidgets(node, moveWidgetsBehind, totalWidgetsToMove) {
|
|
let widgets = node.widgets;
|
|
|
|
// Calculate the starting index for removing the last N widgets
|
|
let startIndex = widgets.length - totalWidgetsToMove;
|
|
|
|
// Check if there are enough widgets to move
|
|
if (startIndex < 0) {
|
|
console.warning("Not enough widgets to move.");
|
|
return widgets;
|
|
}
|
|
|
|
// Extract the last N widgets to be moved
|
|
let widgetsToMove = widgets.splice(startIndex, totalWidgetsToMove);
|
|
|
|
// Find the scheduler widget and ensure it is not duplicated
|
|
const schedulerWidget = findWidget(node, "scheduler_name");
|
|
|
|
// If the schedulerWidget is found, remove existing occurrences from the current widgets array
|
|
if (schedulerWidget) {
|
|
widgets = widgets.filter(widget => widget !== schedulerWidget);
|
|
}
|
|
|
|
// Ensure the insertion index is within the bounds of the current array
|
|
moveWidgetsBehind = Math.max(0, Math.min(moveWidgetsBehind, widgets.length));
|
|
|
|
// Prepare the new widgets array by inserting schedulerWidget followed by widgetsToMove
|
|
let newWidgets = [
|
|
...widgets.slice(0, moveWidgetsBehind),
|
|
schedulerWidget,
|
|
...widgetsToMove,
|
|
...widgets.slice(moveWidgetsBehind)
|
|
];
|
|
|
|
// Update the node's widgets with the newly ordered widgets
|
|
node.widgets = newWidgets;
|
|
|
|
return newWidgets;
|
|
}
|
|
|
|
function showSchedulerInputs(node, schedulerName, desiredWidgets = undefined) {
|
|
const shouldResizeNode = !desiredWidgets;
|
|
const initialCount = countWidgetsOfType(node.widgets, "converted-widget");
|
|
const originalSize = [node.size[0], node.size[1]];
|
|
const originalWidgetTypes = new Map();
|
|
|
|
for (const widget of node.widgets) {
|
|
originalWidgetTypes.set(widget.name, widget.type);
|
|
}
|
|
|
|
// Determine which widgets should be present based on schedulerName
|
|
if (!desiredWidgets) {
|
|
desiredWidgets = widgets[schedulerName] ? new Set(Object.keys(widgets[schedulerName])) : new Set();
|
|
} else {
|
|
desiredWidgets = new Set(desiredWidgets);
|
|
}
|
|
|
|
const widgetsToRemove = createWidgetsToRemove(widgets);
|
|
|
|
for (const widget of widgetsToRemove) {
|
|
if (!desiredWidgets.has(widget)) {
|
|
SettingUtils.hideWidget(node, findWidget(node, widget));
|
|
}
|
|
}
|
|
|
|
for (const widgetName of desiredWidgets) {
|
|
const widget = findWidget(node, widgetName);
|
|
SettingUtils.showWidget(widget);
|
|
if (widget.type === undefined) {
|
|
widget.type = originalWidgetTypes.get(widget.name);
|
|
}
|
|
}
|
|
|
|
let addedWidgets = 0;
|
|
let removedWidgets = 0;
|
|
|
|
for (const widget of node.widgets) {
|
|
if (originalWidgetTypes.get(widget.name) === "converted-widget") {
|
|
if (widget.type !== "converted-widget")
|
|
addedWidgets++;
|
|
} else {
|
|
if (widget.type === "converted-widget")
|
|
removedWidgets++;
|
|
}
|
|
}
|
|
|
|
// Adjust the size based on widgets added/removed
|
|
if (shouldResizeNode) {
|
|
resizeNode(node, initialCount, originalSize);
|
|
}
|
|
}
|
|
|
|
function hideAllSchedulerWidgets(node) {
|
|
let nodesHidden = 0;
|
|
for (const widget of node.widgets) {
|
|
if (allSettings.has(widget.name)) {
|
|
SettingUtils.hideWidget(node, widget);
|
|
nodesHidden++;
|
|
}
|
|
}
|
|
return nodesHidden;
|
|
}
|
|
|
|
function hideAllLatentWidgets(node) {
|
|
let heightWidget = findWidget(node, "height");
|
|
let widthWidget = findWidget(node, "width");
|
|
|
|
SettingUtils.hideWidget(node, heightWidget);
|
|
SettingUtils.hideWidget(node, widthWidget);
|
|
|
|
// Remove inputs if they are there
|
|
if (node.inputs.some(obj => obj.name === "width")) {
|
|
node.inputs = node.inputs.filter(obj => obj.name !== "width");
|
|
}
|
|
if (node.inputs.some(obj => obj.name === "height")) {
|
|
node.inputs = node.inputs.filter(obj => obj.name !== "height");
|
|
}
|
|
}
|
|
|
|
function showAllLatentWidgets(node) {
|
|
let heightWidget = findWidget(node, "height");
|
|
let widthWidget = findWidget(node, "width");
|
|
|
|
// Check if the height or width is an input instead of a widget
|
|
if (!node.inputs.some(obj => obj.name === "width")) {
|
|
SettingUtils.showWidget(widthWidget);
|
|
if (widthWidget.type === undefined)
|
|
widthWidget.type = "number";
|
|
}
|
|
if (!node.inputs.some(obj => obj.name === "height")) {
|
|
SettingUtils.showWidget(heightWidget);
|
|
if (heightWidget.type === undefined)
|
|
heightWidget.type = "number";
|
|
}
|
|
}
|
|
|
|
function createEverything(node) {
|
|
for (const [widgetKey, widgetProps] of Object.entries(widgets)) {
|
|
for (const [prop, [type, defaultValue, min, max, step, round]] of Object.entries(widgetProps)) {
|
|
if (type === "FLOAT") {
|
|
ComfyWidgets[type](node, prop, [type, { default: defaultValue, min: min, max: max, step: step, round: round }], app).widget;
|
|
}
|
|
allSettings.add(prop);
|
|
}
|
|
}
|
|
|
|
const originalSize = [node.size[0], node.size[1]];
|
|
|
|
rearrangeWidgets(node, 7, allSettings.size);
|
|
const nodesHidden = hideAllSchedulerWidgets(node);
|
|
|
|
node.size[0] = originalSize[0];
|
|
node.size[1] = originalSize[1] + (-nodesHidden * (70 / 3));
|
|
}
|
|
}
|
|
},
|
|
});
|