initial commit
This commit is contained in:
403
extensions-builtin/sd_forge_controlnet/javascript/active_units.js
Executable file
403
extensions-builtin/sd_forge_controlnet/javascript/active_units.js
Executable file
@@ -0,0 +1,403 @@
|
||||
/**
|
||||
* Give a badge on ControlNet Accordion indicating total number of active
|
||||
* units.
|
||||
* Make active unit's tab name green.
|
||||
* Append control type to tab name.
|
||||
* Disable resize mode selection when A1111 img2img input is used.
|
||||
*/
|
||||
(function () {
|
||||
const cnetAllAccordions = new Set();
|
||||
onUiUpdate(() => {
|
||||
const ImgChangeType = {
|
||||
NO_CHANGE: 0,
|
||||
REMOVE: 1,
|
||||
ADD: 2,
|
||||
SRC_CHANGE: 3,
|
||||
};
|
||||
|
||||
function imgChangeObserved(mutationsList) {
|
||||
// Iterate over all mutations that just occured
|
||||
for (let mutation of mutationsList) {
|
||||
// Check if the mutation is an addition or removal of a node
|
||||
if (mutation.type === 'childList') {
|
||||
// Check if nodes were added
|
||||
if (mutation.addedNodes.length > 0) {
|
||||
for (const node of mutation.addedNodes) {
|
||||
if (node.tagName === 'IMG') {
|
||||
return ImgChangeType.ADD;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Check if nodes were removed
|
||||
if (mutation.removedNodes.length > 0) {
|
||||
for (const node of mutation.removedNodes) {
|
||||
if (node.tagName === 'IMG') {
|
||||
return ImgChangeType.REMOVE;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
// Check if the mutation is a change of an attribute
|
||||
else if (mutation.type === 'attributes') {
|
||||
if (mutation.target.tagName === 'IMG' && mutation.attributeName === 'src') {
|
||||
return ImgChangeType.SRC_CHANGE;
|
||||
}
|
||||
}
|
||||
}
|
||||
return ImgChangeType.NO_CHANGE;
|
||||
}
|
||||
|
||||
function childIndex(element) {
|
||||
// Get all child nodes of the parent
|
||||
let children = Array.from(element.parentNode.childNodes);
|
||||
|
||||
// Filter out non-element nodes (like text nodes and comments)
|
||||
children = children.filter(child => child.nodeType === Node.ELEMENT_NODE);
|
||||
|
||||
return children.indexOf(element);
|
||||
}
|
||||
|
||||
function imageInputDisabledAlert() {
|
||||
alert('Inpaint control type must use a1111 input in img2img mode.');
|
||||
}
|
||||
|
||||
class ControlNetUnitTab {
|
||||
constructor(tab, accordion) {
|
||||
this.tab = tab;
|
||||
this.tabOpen = false; // Whether the tab is open.
|
||||
this.accordion = accordion;
|
||||
this.isImg2Img = tab.querySelector('.cnet-mask-upload').id.includes('img2img');
|
||||
|
||||
this.enabledAccordionCheckbox = tab.querySelector('.input-accordion-checkbox');
|
||||
this.enabledCheckbox = tab.querySelector('.cnet-unit-enabled input');
|
||||
this.inputImage = tab.querySelector('.cnet-input-image-group .cnet-image input[type="file"]');
|
||||
this.inputImageContainer = tab.querySelector('.cnet-input-image-group .cnet-image');
|
||||
this.generatedImageGroup = tab.querySelector('.cnet-generated-image-group');
|
||||
this.maskImageGroup = tab.querySelector('.cnet-mask-image-group');
|
||||
this.inputImageGroup = tab.querySelector('.cnet-input-image-group');
|
||||
this.controlTypeRadios = tab.querySelectorAll('.controlnet_control_type_filter_group input[type="radio"]');
|
||||
this.resizeModeRadios = tab.querySelectorAll('.controlnet_resize_mode_radio input[type="radio"]');
|
||||
this.runPreprocessorButton = tab.querySelector('.cnet-run-preprocessor');
|
||||
|
||||
this.tabs = tab.parentNode;
|
||||
this.tabIndex = childIndex(tab);
|
||||
|
||||
// By default the InputAccordion checkbox is linked with the state
|
||||
// of accordion's open/close state. To disable this link, we can
|
||||
// simulate click to check the checkbox and uncheck it.
|
||||
this.enabledAccordionCheckbox.click();
|
||||
this.enabledAccordionCheckbox.click();
|
||||
|
||||
this.sync_enabled_checkbox();
|
||||
this.attachEnabledButtonListener();
|
||||
this.attachControlTypeRadioListener();
|
||||
this.attachImageUploadListener();
|
||||
this.attachImageStateChangeObserver();
|
||||
this.attachA1111SendInfoObserver();
|
||||
this.attachAccordionStateObserver();
|
||||
}
|
||||
|
||||
/**
|
||||
* Sync the states of enabledCheckbox and enabledAccordionCheckbox.
|
||||
*/
|
||||
sync_enabled_checkbox() {
|
||||
this.enabledCheckbox.addEventListener("change", () => {
|
||||
if (this.enabledAccordionCheckbox.checked != this.enabledCheckbox.checked) {
|
||||
this.enabledAccordionCheckbox.click();
|
||||
}
|
||||
});
|
||||
this.enabledAccordionCheckbox.addEventListener("change", () => {
|
||||
if (this.enabledCheckbox.checked != this.enabledAccordionCheckbox.checked) {
|
||||
this.enabledCheckbox.click();
|
||||
}
|
||||
});
|
||||
}
|
||||
/**
|
||||
* Get the span that has text "Unit {X}".
|
||||
*/
|
||||
getUnitHeaderTextElement() {
|
||||
return this.tab.querySelector(
|
||||
`button > span:nth-child(1)`
|
||||
);
|
||||
}
|
||||
|
||||
getActiveControlType() {
|
||||
for (let radio of this.controlTypeRadios) {
|
||||
if (radio.checked) {
|
||||
return radio.value;
|
||||
}
|
||||
}
|
||||
return undefined;
|
||||
}
|
||||
|
||||
updateActiveState() {
|
||||
const unitHeader = this.getUnitHeaderTextElement();
|
||||
if (!unitHeader) return;
|
||||
|
||||
if (this.enabledCheckbox.checked) {
|
||||
unitHeader.classList.add('cnet-unit-active');
|
||||
} else {
|
||||
unitHeader.classList.remove('cnet-unit-active');
|
||||
}
|
||||
}
|
||||
|
||||
updateActiveUnitCount() {
|
||||
function getActiveUnitCount(checkboxes) {
|
||||
let activeUnitCount = 0;
|
||||
for (const checkbox of checkboxes) {
|
||||
if (checkbox.checked)
|
||||
activeUnitCount++;
|
||||
}
|
||||
return activeUnitCount;
|
||||
}
|
||||
|
||||
const checkboxes = this.accordion.querySelectorAll('.cnet-unit-enabled input');
|
||||
const span = this.accordion.querySelector('.label-wrap span');
|
||||
|
||||
// Remove existing badge.
|
||||
if (span.childNodes.length !== 1) {
|
||||
span.removeChild(span.lastChild);
|
||||
}
|
||||
// Add new badge if necessary.
|
||||
const activeUnitCount = getActiveUnitCount(checkboxes);
|
||||
if (activeUnitCount > 0) {
|
||||
const div = document.createElement('div');
|
||||
div.classList.add('cnet-badge');
|
||||
div.classList.add('primary');
|
||||
div.innerHTML = `${activeUnitCount} unit${activeUnitCount > 1 ? 's' : ''}`;
|
||||
span.appendChild(div);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Add the active control type to tab displayed text.
|
||||
*/
|
||||
updateActiveControlType() {
|
||||
const unitHeader = this.getUnitHeaderTextElement();
|
||||
if (!unitHeader) return;
|
||||
|
||||
// Remove the control if exists
|
||||
const controlTypeSuffix = unitHeader.querySelector('.control-type-suffix');
|
||||
if (controlTypeSuffix) controlTypeSuffix.remove();
|
||||
|
||||
// Add new suffix.
|
||||
const controlType = this.getActiveControlType();
|
||||
if (controlType === 'All') return;
|
||||
|
||||
const span = document.createElement('span');
|
||||
span.innerHTML = `[${controlType}]`;
|
||||
span.classList.add('control-type-suffix');
|
||||
unitHeader.appendChild(span);
|
||||
}
|
||||
getInputImageSrc() {
|
||||
const img = this.inputImageGroup.querySelector('.cnet-image .forge-image');
|
||||
return (img && img.src.startsWith('data')) ? img.src : null;
|
||||
}
|
||||
getPreprocessorPreviewImageSrc() {
|
||||
const img = this.generatedImageGroup.querySelector('.cnet-image .forge-image');
|
||||
return (img && img.src.startsWith('data')) ? img.src : null;
|
||||
}
|
||||
getMaskImageSrc() {
|
||||
function isEmptyCanvas(canvas) {
|
||||
if (!canvas) return true;
|
||||
if (canvas.width == 0 || canvas.height ==0) return true;
|
||||
const ctx = canvas.getContext('2d');
|
||||
// Get the image data
|
||||
const imageData = ctx.getImageData(0, 0, canvas.width, canvas.height);
|
||||
const data = imageData.data; // This is a Uint8ClampedArray
|
||||
// Check each pixel
|
||||
let isPureBlack = true;
|
||||
for (let i = 0; i < data.length; i += 4) {
|
||||
if (data[i] !== 0 || data[i + 1] !== 0 || data[i + 2] !== 0) { // Check RGB values
|
||||
isPureBlack = false;
|
||||
break;
|
||||
}
|
||||
}
|
||||
return isPureBlack;
|
||||
}
|
||||
const maskImg = this.maskImageGroup.querySelector('.cnet-mask-image .forge-image');
|
||||
// Hand-drawn mask on mask upload.
|
||||
const handDrawnMaskCanvas = this.maskImageGroup.querySelector('.cnet-mask-image .forge-drawing-canvas');
|
||||
// Hand-drawn mask on input image upload.
|
||||
const inputImageHandDrawnMaskCanvas = this.inputImageGroup.querySelector('.cnet-image .forge-drawing-canvas');
|
||||
if (!isEmptyCanvas(handDrawnMaskCanvas)) {
|
||||
return handDrawnMaskCanvas.toDataURL();
|
||||
} else if (maskImg && maskImg.src.startsWith('data')) {
|
||||
return maskImg.src;
|
||||
} else if (!isEmptyCanvas(inputImageHandDrawnMaskCanvas)) {
|
||||
return inputImageHandDrawnMaskCanvas.toDataURL();
|
||||
} else {
|
||||
return null;
|
||||
}
|
||||
}
|
||||
setThumbnail(imgSrc, maskSrc) {
|
||||
if (!imgSrc) return;
|
||||
const unitHeader = this.getUnitHeaderTextElement();
|
||||
if (!unitHeader) return;
|
||||
const img = document.createElement('img');
|
||||
img.src = imgSrc;
|
||||
img.classList.add('cnet-thumbnail');
|
||||
unitHeader.appendChild(img);
|
||||
|
||||
if (maskSrc) {
|
||||
const mask = document.createElement('img');
|
||||
mask.src = maskSrc;
|
||||
mask.classList.add('cnet-thumbnail');
|
||||
unitHeader.appendChild(mask);
|
||||
}
|
||||
}
|
||||
removeThumbnail() {
|
||||
const unitHeader = this.getUnitHeaderTextElement();
|
||||
if (!unitHeader) return;
|
||||
const imgs = unitHeader.querySelectorAll('.cnet-thumbnail');
|
||||
for (const img of imgs) {
|
||||
img.remove();
|
||||
}
|
||||
}
|
||||
/**
|
||||
* When the accordion is folded, display a thumbnail of input image
|
||||
* and mask on the accordion header.
|
||||
*/
|
||||
updateInputImageThumbnail() {
|
||||
if (!opts.controlnet_input_thumbnail) return;
|
||||
if (this.tabOpen) {
|
||||
this.removeThumbnail();
|
||||
} else {
|
||||
this.setThumbnail(this.getInputImageSrc(), this.getMaskImageSrc());
|
||||
}
|
||||
}
|
||||
|
||||
attachEnabledButtonListener() {
|
||||
this.enabledCheckbox.addEventListener('change', () => {
|
||||
this.updateActiveState();
|
||||
this.updateActiveUnitCount();
|
||||
});
|
||||
}
|
||||
|
||||
attachControlTypeRadioListener() {
|
||||
for (const radio of this.controlTypeRadios) {
|
||||
radio.addEventListener('change', () => {
|
||||
this.updateActiveControlType();
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
attachImageUploadListener() {
|
||||
// Automatically check `enable` checkbox when image is uploaded.
|
||||
this.inputImage.addEventListener('change', (event) => {
|
||||
if (!event.target.files) return;
|
||||
if (!this.enabledCheckbox.checked)
|
||||
this.enabledCheckbox.click();
|
||||
});
|
||||
|
||||
// Automatically check `enable` checkbox when JSON pose file is uploaded.
|
||||
this.tab.querySelector('.cnet-upload-pose input').addEventListener('change', (event) => {
|
||||
if (!event.target.files) return;
|
||||
if (!this.enabledCheckbox.checked)
|
||||
this.enabledCheckbox.click();
|
||||
});
|
||||
}
|
||||
|
||||
attachImageStateChangeObserver() {
|
||||
new MutationObserver((mutationsList) => {
|
||||
const changeObserved = imgChangeObserved(mutationsList);
|
||||
|
||||
if (changeObserved === ImgChangeType.ADD) {
|
||||
// enabling the run preprocessor button
|
||||
this.runPreprocessorButton.removeAttribute("disabled");
|
||||
this.runPreprocessorButton.title = 'Run preprocessor';
|
||||
}
|
||||
|
||||
if (changeObserved === ImgChangeType.REMOVE) {
|
||||
// disabling the run preprocessor button
|
||||
this.runPreprocessorButton.setAttribute("disabled", true);
|
||||
this.runPreprocessorButton.title = "No ControlNet input image available";
|
||||
}
|
||||
}).observe(this.inputImageContainer, {
|
||||
childList: true,
|
||||
subtree: true,
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Observe send PNG info buttons in A1111, as they can also directly
|
||||
* set states of ControlNetUnit.
|
||||
*/
|
||||
attachA1111SendInfoObserver() {
|
||||
const pasteButtons = gradioApp().querySelectorAll('#paste');
|
||||
const pngButtons = gradioApp().querySelectorAll(
|
||||
this.isImg2Img ?
|
||||
'#img2img_tab, #inpaint_tab' :
|
||||
'#txt2img_tab'
|
||||
);
|
||||
|
||||
for (const button of [...pasteButtons, ...pngButtons]) {
|
||||
button.addEventListener('click', () => {
|
||||
// The paste/send img generation info feature goes
|
||||
// though gradio, which is pretty slow. Ideally we should
|
||||
// observe the event when gradio has done the job, but
|
||||
// that is not an easy task.
|
||||
// Here we just do a 2 second delay until the refresh.
|
||||
setTimeout(() => {
|
||||
this.updateActiveState();
|
||||
this.updateActiveUnitCount();
|
||||
}, 2000);
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Observer that triggers when the ControlNetUnit's accordion(tab) closes.
|
||||
*/
|
||||
attachAccordionStateObserver() {
|
||||
new MutationObserver((mutationsList) => {
|
||||
for(const mutation of mutationsList) {
|
||||
if (mutation.type === 'attributes' && mutation.attributeName === 'class') {
|
||||
const newState = mutation.target.classList.contains('open');
|
||||
if (this.tabOpen != newState) {
|
||||
this.tabOpen = newState;
|
||||
if (newState) {
|
||||
this.onAccordionOpen();
|
||||
} else {
|
||||
this.onAccordionClose();
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}).observe(this.tab.querySelector('.label-wrap'), { attributes: true, attributeFilter: ['class'] });
|
||||
}
|
||||
|
||||
onAccordionOpen() {
|
||||
this.updateInputImageThumbnail();
|
||||
}
|
||||
|
||||
onAccordionClose() {
|
||||
this.updateInputImageThumbnail();
|
||||
}
|
||||
}
|
||||
|
||||
gradioApp().querySelectorAll('#controlnet').forEach(accordion => {
|
||||
if (cnetAllAccordions.has(accordion)) return;
|
||||
const tabs = [...accordion.querySelectorAll('.input-accordion')]
|
||||
.map(tab => new ControlNetUnitTab(tab, accordion));
|
||||
|
||||
// On open of main extension accordion, if no unit is enabled,
|
||||
// open unit 0 for edit.
|
||||
const labelWrap = accordion.querySelector('.label-wrap');
|
||||
const observerAccordionOpen = new MutationObserver(function (mutations) {
|
||||
for (const mutation of mutations) {
|
||||
if (mutation.target.classList.contains('open') &&
|
||||
tabs.every(tab => !tab.enabledCheckbox.checked &&
|
||||
!tab.tab.querySelector('.label-wrap').classList.contains('open'))
|
||||
) {
|
||||
tabs[0].tab.querySelector('.label-wrap').click();
|
||||
}
|
||||
}
|
||||
});
|
||||
observerAccordionOpen.observe(labelWrap, { attributes: true, attributeFilter: ['class'] });
|
||||
|
||||
cnetAllAccordions.add(accordion);
|
||||
});
|
||||
});
|
||||
})();
|
||||
Reference in New Issue
Block a user