From bafb2bce54231be6716c51aa4203fcee9db56338 Mon Sep 17 00:00:00 2001 From: Chenlei Hu Date: Sun, 4 Feb 2024 03:00:18 +0000 Subject: [PATCH] Display thumbnail of input image and mask on folded unit header (#33) * Add input thumbnail to unit header * handle hand drawn mask * Update mask logic --- .../javascript/active_units.js | 108 ++++++++++++++++++ .../controlnet_ui/controlnet_ui_group.py | 3 - .../sd_forge_controlnet/style.css | 26 ++++- 3 files changed, 131 insertions(+), 6 deletions(-) diff --git a/extensions-builtin/sd_forge_controlnet/javascript/active_units.js b/extensions-builtin/sd_forge_controlnet/javascript/active_units.js index 13343707..3015a028 100644 --- a/extensions-builtin/sd_forge_controlnet/javascript/active_units.js +++ b/extensions-builtin/sd_forge_controlnet/javascript/active_units.js @@ -65,6 +65,7 @@ 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'); @@ -72,6 +73,9 @@ 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'); @@ -92,6 +96,7 @@ this.attachImageStateChangeObserver(); this.attachA1111SendInfoObserver(); this.attachPresetDropdownObserver(); + this.attachAccordionStateObserver(); } /** @@ -186,6 +191,81 @@ span.classList.add('control-type-suffix'); unitHeader.appendChild(span); } + getInputImageSrc() { + const img = this.inputImageGroup.querySelector('.cnet-image img'); + return img ? img.src : null; + } + getPreprocessorPreviewImageSrc() { + const img = this.generatedImageGroup.querySelector('.cnet-image img'); + return img ? img.src : null; + } + getMaskImageSrc() { + function isEmptyCanvas(canvas) { + if (!canvas) 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 img'); + // Hand-drawn mask on mask upload. + const handDrawnMaskCanvas = this.maskImageGroup.querySelector('.cnet-mask-image canvas[key="mask"]'); + // Hand-drawn mask on input image upload. + const inputImageHandDrawnMaskCanvas = this.inputImageGroup.querySelector('.cnet-image canvas[key="mask"]'); + if (!isEmptyCanvas(handDrawnMaskCanvas)) { + return handDrawnMaskCanvas.toDataURL(); + } else if (maskImg) { + 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 (this.tabOpen) { + this.removeThumbnail(); + } else { + this.setThumbnail(this.getInputImageSrc(), this.getMaskImageSrc()); + } + } attachEnabledButtonListener() { this.enabledCheckbox.addEventListener('change', () => { @@ -285,6 +365,34 @@ subtree: true, }); } + /** + * 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 => { diff --git a/extensions-builtin/sd_forge_controlnet/lib_controlnet/controlnet_ui/controlnet_ui_group.py b/extensions-builtin/sd_forge_controlnet/lib_controlnet/controlnet_ui/controlnet_ui_group.py index 63086e81..b917c043 100644 --- a/extensions-builtin/sd_forge_controlnet/lib_controlnet/controlnet_ui/controlnet_ui_group.py +++ b/extensions-builtin/sd_forge_controlnet/lib_controlnet/controlnet_ui/controlnet_ui_group.py @@ -223,9 +223,6 @@ class ControlNetUiGroup(object): self.prevent_next_n_module_update = 0 self.prevent_next_n_slider_value_update = 0 - # API-only fields - self.advanced_weighting = gr.State(None) - ControlNetUiGroup.all_ui_groups.append(self) def render(self, tabname: str, elem_id_tabname: str) -> None: diff --git a/extensions-builtin/sd_forge_controlnet/style.css b/extensions-builtin/sd_forge_controlnet/style.css index 785e9370..527c2e32 100644 --- a/extensions-builtin/sd_forge_controlnet/style.css +++ b/extensions-builtin/sd_forge_controlnet/style.css @@ -11,18 +11,38 @@ .controlnet .input-accordion { flex: 1 1 calc(50% - 10px); /* Adjusts for the gap, default 2 columns */ - /* Additional styling for items */ + display: flex; + align-items: center; } /* Media query for screens smaller than a specific width */ @media (max-width: 600px) { - - /* Adjust the threshold as needed */ .controlnet .input-accordion { flex: 1 1 100%; /* Changes to 1 column when window width is ≤ 600px */ } } +/* Input image thumbnail */ +.cnet-thumbnail { + height: 3rem !important; + border: 1px solid var(--button-secondary-border-color); +} + +.cnet-unit-active { + display: flex; + flex-direction: row; + align-items: center; + gap: 5px; +} + +.controlnet .input-accordion .icon { + height: 1rem; + width: 1rem; +} + +.controlnet .input-accordion .label-wrap { + align-items: center; +} .cnet-modal { display: none;