332 lines
10 KiB
JavaScript
332 lines
10 KiB
JavaScript
import { app } from "/scripts/app.js";
|
|
import { api } from "/scripts/api.js";
|
|
|
|
const NODE_CLASS = "ChatImageBridge";
|
|
const FETCH_BUTTON = "Fetch Models";
|
|
const CATAPI_NODE_CLASS = "CatAPIGeminiImage";
|
|
const MAX_CATAPI_IMAGES = 14;
|
|
const CATAPI_CUSTOM_RATIO = "Custom Ratio";
|
|
const CATAPI_RESOLUTIONS = ["1K", "2K", "4K"];
|
|
|
|
function findWidget(node, name) {
|
|
return node.widgets?.find((widget) => widget.name === name);
|
|
}
|
|
|
|
function markCanvasDirty() {
|
|
app.graph?.setDirtyCanvas?.(true, true);
|
|
app.canvas?.setDirty?.(true, true);
|
|
}
|
|
|
|
function setWidgetValue(widget, value) {
|
|
if (!widget) {
|
|
return;
|
|
}
|
|
widget.value = value;
|
|
widget.callback?.(value);
|
|
}
|
|
|
|
function prepareModelWidget(node) {
|
|
const modelWidget = findWidget(node, "model");
|
|
if (!modelWidget) {
|
|
return;
|
|
}
|
|
|
|
const currentModel = `${modelWidget.value || ""}`.trim();
|
|
const existingValues = Array.isArray(modelWidget.options?.values) ? modelWidget.options.values : [];
|
|
const values = ["", currentModel, ...existingValues]
|
|
.map((value) => `${value || ""}`.trim())
|
|
.filter((value, index, array) => index === array.indexOf(value));
|
|
|
|
modelWidget.options = modelWidget.options || {};
|
|
modelWidget.options.values = values;
|
|
modelWidget.type = "combo";
|
|
}
|
|
|
|
function updateModelDropdown(node, models) {
|
|
const modelWidget = findWidget(node, "model");
|
|
if (!modelWidget) {
|
|
throw new Error("model widget was not found");
|
|
}
|
|
|
|
const currentModel = `${modelWidget?.value || ""}`.trim();
|
|
const selected = currentModel && models.includes(currentModel) ? currentModel : models[0];
|
|
|
|
modelWidget.options = modelWidget.options || {};
|
|
modelWidget.options.values = models.includes("") ? models : ["", ...models];
|
|
modelWidget.type = "combo";
|
|
setWidgetValue(modelWidget, selected);
|
|
|
|
markCanvasDirty();
|
|
return selected;
|
|
}
|
|
|
|
async function fetchModels(node) {
|
|
const apiKey = `${findWidget(node, "api_key")?.value || ""}`.trim();
|
|
const baseUrl = `${findWidget(node, "base_url")?.value || findWidget(node, "endpoint_url")?.value || ""}`.trim();
|
|
const timeoutSeconds = Number(findWidget(node, "timeout_seconds")?.value || 30);
|
|
|
|
if (!baseUrl) {
|
|
window.alert("base_url is required");
|
|
return;
|
|
}
|
|
if (!apiKey) {
|
|
window.alert("api_key is required");
|
|
return;
|
|
}
|
|
|
|
try {
|
|
const response = await api.fetchApi("/chat_image_bridge/models", {
|
|
method: "POST",
|
|
headers: { "Content-Type": "application/json" },
|
|
body: JSON.stringify({
|
|
base_url: baseUrl,
|
|
api_key: apiKey,
|
|
timeout_seconds: Math.max(5, Math.min(timeoutSeconds, 120)),
|
|
}),
|
|
});
|
|
const data = await response.json();
|
|
|
|
if (!response.ok || data.error) {
|
|
throw new Error(data.error || `HTTP ${response.status}`);
|
|
}
|
|
|
|
const models = Array.isArray(data.models) ? data.models.filter(Boolean) : [];
|
|
if (!models.length) {
|
|
throw new Error("No models were returned by the provider");
|
|
}
|
|
|
|
updateModelDropdown(node, models);
|
|
} catch (error) {
|
|
console.error("[Chat Image Bridge] Failed to fetch models:", error);
|
|
window.alert(`Failed to fetch models: ${error.message || error}`);
|
|
}
|
|
}
|
|
|
|
app.registerExtension({
|
|
name: "Comfy.ChatImageBridge.Models",
|
|
async beforeRegisterNodeDef(nodeType, nodeData) {
|
|
if (nodeData.name !== NODE_CLASS) {
|
|
return;
|
|
}
|
|
|
|
const onNodeCreated = nodeType.prototype.onNodeCreated;
|
|
nodeType.prototype.onNodeCreated = function () {
|
|
const result = onNodeCreated?.apply(this, arguments);
|
|
prepareModelWidget(this);
|
|
|
|
if (!findWidget(this, FETCH_BUTTON)) {
|
|
const fetchButton = this.addWidget("button", FETCH_BUTTON, null, () => fetchModels(this));
|
|
fetchButton.serialize = false;
|
|
}
|
|
|
|
return result;
|
|
};
|
|
},
|
|
});
|
|
|
|
function setWidgetVisible(node, widget, visible) {
|
|
if (!widget) {
|
|
return;
|
|
}
|
|
|
|
if (visible && widget.__catapiHidden) {
|
|
widget.type = widget.__catapiHidden.type;
|
|
widget.computeSize = widget.__catapiHidden.computeSize;
|
|
widget.serializeValue = widget.__catapiHidden.serializeValue;
|
|
widget.hidden = widget.__catapiHidden.hidden;
|
|
delete widget.__catapiHidden;
|
|
} else if (!visible && !widget.__catapiHidden) {
|
|
widget.__catapiHidden = {
|
|
type: widget.type,
|
|
computeSize: widget.computeSize,
|
|
serializeValue: widget.serializeValue,
|
|
hidden: widget.hidden,
|
|
};
|
|
widget.type = "converted-widget";
|
|
widget.computeSize = () => [0, -4];
|
|
widget.serializeValue = () => widget.value;
|
|
widget.hidden = true;
|
|
}
|
|
|
|
const computed = node.computeSize?.() || node.size;
|
|
node.setSize?.([Math.max(node.size?.[0] || 360, 360), computed?.[1] || node.size?.[1] || 200]);
|
|
markCanvasDirty();
|
|
}
|
|
|
|
function catapiImageInputs(node) {
|
|
return (node.inputs || []).filter((input) => /^image_\d+$/.test(input.name));
|
|
}
|
|
|
|
function labelCatapiImageInputs(node) {
|
|
for (const input of catapiImageInputs(node)) {
|
|
const index = Number(input.name.split("_")[1]);
|
|
input.label = `Image ${index}`;
|
|
}
|
|
}
|
|
|
|
function updateCatapiImageInputs(node) {
|
|
let inputs = catapiImageInputs(node);
|
|
labelCatapiImageInputs(node);
|
|
|
|
const last = inputs[inputs.length - 1];
|
|
if (last?.link != null && inputs.length < MAX_CATAPI_IMAGES) {
|
|
const nextIndex = inputs.length + 1;
|
|
node.addInput(`image_${nextIndex}`, "IMAGE");
|
|
inputs = catapiImageInputs(node);
|
|
labelCatapiImageInputs(node);
|
|
}
|
|
|
|
while (inputs.length > 1) {
|
|
const lastInput = inputs[inputs.length - 1];
|
|
const previousInput = inputs[inputs.length - 2];
|
|
if (lastInput.link != null || previousInput.link != null) {
|
|
break;
|
|
}
|
|
node.removeInput(node.inputs.indexOf(lastInput));
|
|
inputs = catapiImageInputs(node);
|
|
}
|
|
|
|
const computed = node.computeSize?.() || node.size;
|
|
node.setSize?.([Math.max(node.size?.[0] || 360, 360), computed?.[1] || node.size?.[1] || 200]);
|
|
markCanvasDirty();
|
|
}
|
|
|
|
function repairShiftedCatapiWidgets(node) {
|
|
const exactWidget = findWidget(node, "exact_custom_ratio");
|
|
const resolutionWidget = findWidget(node, "resolution");
|
|
const apiKeyWidget = findWidget(node, "api_key");
|
|
const baseUrlWidget = findWidget(node, "base_url");
|
|
if (!exactWidget || !resolutionWidget || !apiKeyWidget || !baseUrlWidget) {
|
|
return;
|
|
}
|
|
|
|
const currentResolution = `${resolutionWidget.value || ""}`.trim();
|
|
const shiftedBaseUrl = `${apiKeyWidget.value || ""}`.trim();
|
|
const misplacedExactValue = `${exactWidget.value || ""}`.trim();
|
|
if (/^https?:\/\//i.test(misplacedExactValue)) {
|
|
if (!CATAPI_RESOLUTIONS.includes(currentResolution)) {
|
|
setWidgetValue(resolutionWidget, "1K");
|
|
}
|
|
setWidgetValue(exactWidget, false);
|
|
console.warn("[Chat Image Bridge 2] Repaired widget values saved by the incompatible experimental build.");
|
|
return;
|
|
}
|
|
|
|
const looksShifted = !CATAPI_RESOLUTIONS.includes(currentResolution)
|
|
&& /^https?:\/\//i.test(shiftedBaseUrl);
|
|
if (!looksShifted) {
|
|
return;
|
|
}
|
|
|
|
const previousResolution = `${exactWidget.value || ""}`.trim().toUpperCase();
|
|
setWidgetValue(apiKeyWidget, currentResolution);
|
|
setWidgetValue(baseUrlWidget, shiftedBaseUrl);
|
|
setWidgetValue(
|
|
resolutionWidget,
|
|
CATAPI_RESOLUTIONS.includes(previousResolution) ? previousResolution : "1K",
|
|
);
|
|
setWidgetValue(exactWidget, false);
|
|
console.warn("[Chat Image Bridge 2] Repaired widget values saved by the incompatible experimental build.");
|
|
}
|
|
|
|
function prepareCatapiNode(node) {
|
|
node.title = "Chat Image Bridge 2";
|
|
node.setSize?.([Math.max(node.size?.[0] || 380, 380), node.size?.[1] || 240]);
|
|
repairShiftedCatapiWidgets(node);
|
|
|
|
const labels = {
|
|
prompt: "Prompt",
|
|
output_ratio: "Output Ratio",
|
|
custom_ratio: "Custom Ratio",
|
|
exact_custom_ratio: "Exact Custom Ratio (Experimental)",
|
|
resolution: "Resolution",
|
|
api_key: "API Key",
|
|
base_url: "Base URL",
|
|
};
|
|
for (const [name, label] of Object.entries(labels)) {
|
|
const widget = findWidget(node, name);
|
|
if (widget) {
|
|
widget.label = label;
|
|
}
|
|
}
|
|
|
|
const ratioWidget = findWidget(node, "output_ratio");
|
|
const customRatioWidget = findWidget(node, "custom_ratio");
|
|
const exactCustomRatioWidget = findWidget(node, "exact_custom_ratio");
|
|
const updateRatioWidgets = () => {
|
|
const customRatioSelected = ratioWidget?.value === CATAPI_CUSTOM_RATIO;
|
|
setWidgetVisible(node, customRatioWidget, customRatioSelected);
|
|
setWidgetVisible(node, exactCustomRatioWidget, customRatioSelected);
|
|
markCanvasDirty();
|
|
};
|
|
if (ratioWidget && !ratioWidget.__catapiCallbackWrapped) {
|
|
const originalCallback = ratioWidget.callback;
|
|
ratioWidget.callback = function (value) {
|
|
const result = originalCallback?.apply(this, arguments);
|
|
updateRatioWidgets();
|
|
return result;
|
|
};
|
|
ratioWidget.__catapiCallbackWrapped = true;
|
|
}
|
|
const advancedWidgets = ["api_key", "base_url"].map((name) => findWidget(node, name));
|
|
node.properties = node.properties || {};
|
|
if (node.properties.catapiConnectionExpanded == null) {
|
|
node.properties.catapiConnectionExpanded = false;
|
|
}
|
|
let connectionButton = findWidget(node, "connection_settings");
|
|
const updateConnectionWidgets = () => {
|
|
const expanded = Boolean(node.properties.catapiConnectionExpanded);
|
|
for (const widget of advancedWidgets) {
|
|
setWidgetVisible(node, widget, expanded);
|
|
}
|
|
if (connectionButton) {
|
|
connectionButton.label = expanded ? "Hide Connection Settings" : "Connection Settings";
|
|
connectionButton.name = "connection_settings";
|
|
}
|
|
};
|
|
if (!connectionButton) {
|
|
connectionButton = node.addWidget("button", "connection_settings", null, () => {
|
|
node.properties.catapiConnectionExpanded = !node.properties.catapiConnectionExpanded;
|
|
updateConnectionWidgets();
|
|
});
|
|
connectionButton.serialize = false;
|
|
}
|
|
|
|
updateRatioWidgets();
|
|
updateConnectionWidgets();
|
|
updateCatapiImageInputs(node);
|
|
if (node.outputs?.[0]) {
|
|
node.outputs[0].label = "Image";
|
|
}
|
|
}
|
|
|
|
app.registerExtension({
|
|
name: "Comfy.ChatImageBridge.CatAPISimpleUI",
|
|
async beforeRegisterNodeDef(nodeType, nodeData) {
|
|
if (nodeData.name !== CATAPI_NODE_CLASS) {
|
|
return;
|
|
}
|
|
|
|
const onNodeCreated = nodeType.prototype.onNodeCreated;
|
|
nodeType.prototype.onNodeCreated = function () {
|
|
const result = onNodeCreated?.apply(this, arguments);
|
|
prepareCatapiNode(this);
|
|
return result;
|
|
};
|
|
|
|
const onConnectionsChange = nodeType.prototype.onConnectionsChange;
|
|
nodeType.prototype.onConnectionsChange = function () {
|
|
const result = onConnectionsChange?.apply(this, arguments);
|
|
updateCatapiImageInputs(this);
|
|
return result;
|
|
};
|
|
|
|
const onConfigure = nodeType.prototype.onConfigure;
|
|
nodeType.prototype.onConfigure = function () {
|
|
const result = onConfigure?.apply(this, arguments);
|
|
setTimeout(() => prepareCatapiNode(this), 0);
|
|
return result;
|
|
};
|
|
},
|
|
});
|