282 lines
12 KiB
JavaScript
282 lines
12 KiB
JavaScript
import {app} from "../../../scripts/app.js";
|
|
|
|
const url = new URL('./', import.meta.url).href;
|
|
let sdxl_unet = null;
|
|
fetch(url + "sdxl_unet.svg")
|
|
.then(response => response.text())
|
|
.then(svg => {
|
|
sdxl_unet = svg;
|
|
});
|
|
|
|
let sd15_unet = null;
|
|
fetch(url + "sd15_unet.svg")
|
|
.then(response => response.text())
|
|
.then(svg => {
|
|
sd15_unet = svg;
|
|
});
|
|
|
|
let dit = null;
|
|
fetch(url + "dit.svg")
|
|
.then(response => response.text())
|
|
.then(svg => {
|
|
dit = svg;
|
|
});
|
|
|
|
|
|
|
|
|
|
app.registerExtension({
|
|
name: "LoRABlockEditor",
|
|
async beforeRegisterNodeDef(nodeType, nodeData) {
|
|
const style = document.createElement('style');
|
|
style.textContent = `
|
|
.noselect {
|
|
user-select: none;
|
|
-webkit-user-select: none;
|
|
-moz-user-select: none;
|
|
-ms-user-select: none;
|
|
}`;
|
|
document.head.appendChild(style);
|
|
|
|
function set_svg_color(block, scale) {
|
|
const c = Math.round(128 + (255 - 128) * scale);
|
|
const fillColor = `rgb(${c}, ${c}, ${c})`;
|
|
const rect = block.querySelector("rect");
|
|
if (rect) rect.style.fill = fillColor;
|
|
|
|
// Find a text element with ".sf" in its name to update scale
|
|
const sfLabel = block.querySelector("[id$='.sf']");
|
|
if (sfLabel) {
|
|
const tspan = sfLabel.querySelector("tspan");
|
|
if (tspan) {
|
|
tspan.textContent = `SF: ${scale.toFixed(2)}`;
|
|
}
|
|
}
|
|
}
|
|
|
|
function load_network(mode, div_unet) {
|
|
if (mode === "sdxl") {
|
|
div_unet.innerHTML = sdxl_unet;
|
|
} else if (mode === "sd15") {
|
|
div_unet.innerHTML = sd15_unet;
|
|
} else if (mode === "dit") {
|
|
div_unet.innerHTML = dit;
|
|
}
|
|
}
|
|
|
|
function wire_block(block, settings, widget) {
|
|
block.style.cursor = "ns-resize"; // Visual cue for up/down drag
|
|
|
|
// Create a floating label element for scale display
|
|
const scaleLabel = document.createElement("div");
|
|
scaleLabel.style.position = "fixed";
|
|
scaleLabel.style.pointerEvents = "none";
|
|
scaleLabel.style.background = "black";
|
|
scaleLabel.style.color = "white";
|
|
scaleLabel.style.padding = "2px 6px";
|
|
scaleLabel.style.borderRadius = "4px";
|
|
scaleLabel.style.fontSize = "12px";
|
|
scaleLabel.style.zIndex = "1000";
|
|
scaleLabel.style.display = "none";
|
|
document.body.appendChild(scaleLabel);
|
|
|
|
let isDragging = false;
|
|
let startY = 0;
|
|
block.addEventListener("mousedown", (e) => {
|
|
isDragging = true;
|
|
startY = e.clientY;
|
|
|
|
const onMouseMove = (e) => {
|
|
if (!isDragging) return;
|
|
|
|
e.preventDefault();
|
|
document.body.classList.add("noselect");
|
|
|
|
const dy = startY - e.clientY; // positive = move up
|
|
let scale = settings.blockScales[block.id] + dy / 100; // adjust sensitivity
|
|
scale = Math.min(1, Math.max(0, scale)); // clamp between 0 and 1
|
|
settings.blockScales[block.id] = scale;
|
|
|
|
widget.options.setValue(JSON.stringify(settings));
|
|
|
|
// Interpolate between grey (128) and white (255)
|
|
set_svg_color(block, scale);
|
|
|
|
startY = e.clientY; // update for next movement
|
|
|
|
// Update label
|
|
scaleLabel.style.display = "block";
|
|
scaleLabel.style.left = `${e.clientX + 10}px`;
|
|
scaleLabel.style.top = `${e.clientY + 10}px`;
|
|
scaleLabel.textContent = "Block weight " + scale.toFixed(2);
|
|
};
|
|
|
|
const onMouseUpOrLeave = () => {
|
|
isDragging = false;
|
|
scaleLabel.style.display = "none";
|
|
document.body.classList.remove("noselect");
|
|
window.removeEventListener("mousemove", onMouseMove);
|
|
window.removeEventListener("mouseup", onMouseUpOrLeave);
|
|
window.removeEventListener("mouseleave", onMouseUpOrLeave);
|
|
};
|
|
|
|
window.addEventListener("mousemove", onMouseMove);
|
|
window.addEventListener("mouseup", onMouseUpOrLeave);
|
|
window.addEventListener("mouseleave", onMouseUpOrLeave);
|
|
});
|
|
}
|
|
|
|
if (nodeData?.name === "PM LoRA Modifier") {
|
|
let settings = {"mode": "sdxl", "blockScales": {}};
|
|
|
|
nodeType.prototype.onNodeCreated = function () {
|
|
const node = this;
|
|
|
|
// Container for both widgets
|
|
const container = document.createElement("div");
|
|
container.style.display = "flex";
|
|
container.style.flexDirection = "column";
|
|
container.style.height = "100%"; // fill node height
|
|
|
|
// Add a dropdown where the use can select the model type
|
|
const modelTypeSelect = document.createElement("select");
|
|
modelTypeSelect.style.margin = "4px";
|
|
modelTypeSelect.style.width = "100%";
|
|
modelTypeSelect.innerHTML = `
|
|
<option value="sdxl">SDXL</option>
|
|
<option value="sd15">SD 1.5</option>
|
|
<option value="dit">DiT (40 layers)</option>
|
|
`;
|
|
|
|
// SVG container - flex grow
|
|
const div_unet = document.createElement("div");
|
|
div_unet.id = "sdxl_unet";
|
|
div_unet.style.overflow = "hidden";
|
|
div_unet.style.flexGrow = "1"; // fill available vertical space
|
|
|
|
// Add model type select and SVG container to the main container
|
|
container.appendChild(modelTypeSelect);
|
|
container.appendChild(div_unet);
|
|
|
|
// Add an event listener to the model type select and update the SVG
|
|
modelTypeSelect.addEventListener("change", (e) => {
|
|
const selectedModel = e.target.value;
|
|
load_network(selectedModel, div_unet);
|
|
// Reset block scales for the new SVG
|
|
// Therefore we need to find out all blocks in the new SVG
|
|
let allBlocks;
|
|
if (selectedModel === "dit") {
|
|
// For DiT, select layer groups
|
|
allBlocks = div_unet.querySelectorAll("[id*='layers_group.']");
|
|
} else {
|
|
// For SD/SDXL, select blocks
|
|
allBlocks = div_unet.querySelectorAll("[id*='blocks.']");
|
|
}
|
|
const blocks = Array.from(allBlocks)
|
|
.filter(el => !el.id.includes("rect") && !el.id.includes("grid") && !el.id.includes(".sf"));
|
|
settings = {"mode": selectedModel, blockScales: {}}
|
|
blocks.forEach(block => {
|
|
const blockId = block.id;
|
|
settings.blockScales[blockId] = 1;
|
|
// Set initial color
|
|
set_svg_color(block, settings.blockScales[blockId]);
|
|
// Create event listeners for the new blocks
|
|
wire_block(block, settings, newWidget);
|
|
});
|
|
// Update the widget with the new block scales
|
|
newWidget.options.setValue(JSON.stringify(settings));
|
|
});
|
|
|
|
|
|
// Add entire container as one DOM widget
|
|
const newWidget = node.addDOMWidget("text", "blocks_store", container, {
|
|
blocksStoreWidget: node.widgets.find(w => w.name === "blocks_store"),
|
|
getValue() {
|
|
return JSON.stringify(settings);
|
|
},
|
|
setValue(v) {
|
|
let parsed = {};
|
|
try {
|
|
parsed = JSON.parse(v || "{}");
|
|
} catch (e) {
|
|
console.warn("Failed to parse blocks_store JSON:", v);
|
|
}
|
|
settings = parsed;
|
|
|
|
// Apply the parsed block scales to the SVG blocks
|
|
if (!settings.blockScales) {
|
|
settings.blockScales = {};
|
|
}
|
|
for (const [blockId, scale] of Object.entries(settings.blockScales)) {
|
|
const el = div_unet.querySelector(`[id="${blockId}"]`);
|
|
if (el)
|
|
set_svg_color(el, scale);
|
|
else
|
|
console.warn(`Block with ID ${blockId} not found in SVG.`);
|
|
}
|
|
|
|
// Update the original text widget if needed (optional)
|
|
if (this.blocksStoreWidget) {
|
|
this.blocksStoreWidget.value = JSON.stringify(parsed, null, 2);
|
|
}
|
|
}
|
|
});
|
|
|
|
|
|
setTimeout(() => {
|
|
if (settings.mode) {
|
|
modelTypeSelect.value = settings.mode;
|
|
load_network(settings.mode, div_unet);
|
|
}
|
|
// If blockScales is empty, initialize it with the new SVG blocks
|
|
if (!settings.blockScales || Object.keys(settings.blockScales).length === 0) {
|
|
modelTypeSelect.dispatchEvent(new Event("change"));
|
|
}else {
|
|
// If blockScales already has values, apply them to the existing blocks
|
|
let allBlocks;
|
|
if (settings.mode === "dit") {
|
|
allBlocks = div_unet.querySelectorAll("[id*='layers_group.']");
|
|
} else {
|
|
allBlocks = div_unet.querySelectorAll("[id*='blocks.']");
|
|
}
|
|
const blocks = Array.from(allBlocks)
|
|
.filter(el => !el.id.includes("rect") && !el.id.includes("grid") && !el.id.includes(".sf"));
|
|
|
|
blocks.forEach(block => {
|
|
let blockId = block.id;
|
|
set_svg_color(block, settings.blockScales[blockId]);
|
|
// Create event listeners for the new blocks
|
|
wire_block(block, settings, newWidget);
|
|
});
|
|
}
|
|
}, 100);
|
|
};
|
|
}
|
|
if (nodeData?.name === "PM LoRA Block Display") {
|
|
nodeType.prototype.onNodeCreated = function () {
|
|
const node = this;
|
|
|
|
// Container for both widgets
|
|
const container = document.createElement("div");
|
|
container.style.display = "flex";
|
|
container.style.flexDirection = "column";
|
|
container.style.height = "100%"; // fill node height
|
|
container.innerHTML = sdxl_unet;
|
|
container.id = "sdxl_unet";
|
|
|
|
const svgDisplayWidget = node.addDOMWidget("text", "svg_display", container, {
|
|
getValue: () => "",
|
|
setValue: (val) => {
|
|
console.log("Setting SVG value:", val);
|
|
},
|
|
});
|
|
|
|
// add onExecuted hook to update the widget when the node is executed
|
|
node.onExecuted = function(output_data, context, graph_canvas) {
|
|
console.log("Node executed, output_data is:", output_data);
|
|
};
|
|
};
|
|
}
|
|
}
|
|
});
|