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

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 => {

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:

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;