Add batch size support to ResolutionMaster

Azornes · Azornes · commit 6e1d9fa78fef · 2025-11-07T23:55:22.000+01:00
Introduces a batch_size parameter to the ResolutionMaster node, backend, and frontend UI. Updates the Python backend to generate a latent tensor with the specified batch size, and enhances the JavaScript UI to allow users to set batch size via a new value area and dialog, including input validation and tooltips.
diff --git a/aztoolkit.py b/aztoolkit.py
@@ -1,7 +1,8 @@
 # ComfyUI - azToolkit - Azornes 2025
 
-import math
 import time
+import torch
+import comfy.model_management
 
 
 class AnyType(str):
@@ -18,6 +19,9 @@ class ResolutionMaster:
     # Class variable to store image dimensions for auto-detection
     _image_dimensions_cache = {}
     
+    def __init__(self):
+        self.device = comfy.model_management.intermediate_device()
+    
     @classmethod
     def INPUT_TYPES(cls):
         return {
@@ -30,6 +34,7 @@ def INPUT_TYPES(cls):
                 "auto_detect": ("BOOLEAN", {"default": False, "label_on": "Auto-detect from input", "label_off": "Manual"}),
                 "rescale_mode": ("STRING", {"default": "resolution"}),
                 "rescale_value": ("FLOAT", {"default": 1.0, "step": 0.001, "min": 0.0, "max": 100.0}),
+                "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096, "tooltip": "The number of latent images in the batch."}),
             },
             "optional": {
                 "input_image": ("IMAGE",),
@@ -39,15 +44,15 @@ def INPUT_TYPES(cls):
             },
         }
 
-    RETURN_TYPES = ("INT", "INT", "FLOAT")
-    RETURN_NAMES = ("", "", "")
+    RETURN_TYPES = ("INT", "INT", "FLOAT", "LATENT")
+    RETURN_NAMES = ("width", "height", "rescale_factor", "latent")
     FUNCTION = "main"
     CATEGORY = "utils/azToolkit"
 
-    def main(self, mode, width, height, auto_detect, rescale_mode, rescale_value, input_image=None, unique_id=None):
+    def main(self, mode, width, height, auto_detect, rescale_mode, rescale_value, batch_size=1, input_image=None, unique_id=None):
         detected_width = width
         detected_height = height
-        
+
         # If auto_detect is enabled and we have an input image
         if auto_detect and input_image is not None:
             try:
@@ -58,15 +63,15 @@ def main(self, mode, width, height, auto_detect, rescale_mode, rescale_value, in
                 elif input_image.dim() == 3:  # [height, width, channels]
                     detected_height = int(input_image.shape[0])
                     detected_width = int(input_image.shape[1])
-                
+
                 # Store dimensions in cache for frontend access
                 if unique_id:
                     ResolutionMaster._image_dimensions_cache[str(unique_id)] = {
                         'width': detected_width,
                         'height': detected_height,
                         'timestamp': time.time()
                     }
-                
+
                 # Only override with detected dimensions if the widget values match the detected values
                 # This allows manually set values (like from auto-fit) to take precedence
                 if width == detected_width and height == detected_height:
@@ -77,12 +82,15 @@ def main(self, mode, width, height, auto_detect, rescale_mode, rescale_value, in
                 else:
                     # Widget values differ from detected values, use widget values (manually set)
                     print(f"[ResolutionMaster] Using manual dimensions: {width}x{height} (detected: {detected_width}x{detected_height})")
-                
+
             except Exception as e:
                 print(f"[ResolutionMaster] Error detecting dimensions: {str(e)}")
                 # Fall back to manual dimensions
-        
+
         # The rescale_factor is calculated on the frontend and passed here
         rescale_factor = rescale_value
-        
-        return (width, height, rescale_factor)
+
+        # Generate latent tensor
+        latent = torch.zeros([batch_size, 4, height // 8, width // 8], device=self.device)
+
+        return (width, height, rescale_factor, {"samples": latent})
diff --git a/js/DialogManager.js b/js/DialogManager.js
@@ -23,7 +23,7 @@ export class DialogManager {
         log.debug(`Clicked on value area: ${valueAreaKey}`);
         
         // Determine the type and current value based on the control key
-        let valueType, currentValue, propertyName, minValue = 0.01;
+        let valueType, currentValue, propertyName, minValue = 0.01, integerOnly = false;
         
         if (valueAreaKey === 'scaleValueArea') {
             valueType = 'Scale Factor';
@@ -42,23 +42,32 @@ export class DialogManager {
             currentValue = this.rm.node.properties.snapValue;
             propertyName = 'snapValue';
             minValue = 1;
+            integerOnly = true;
         } else if (valueAreaKey === 'widthValueArea') {
             valueType = 'Width';
             currentValue = this.rm.widthWidget ? this.rm.widthWidget.value : this.rm.node.properties.valueX;
             propertyName = 'width';
             minValue = 64;
+            integerOnly = true;
         } else if (valueAreaKey === 'heightValueArea') {
             valueType = 'Height';
             currentValue = this.rm.heightWidget ? this.rm.heightWidget.value : this.rm.node.properties.valueY;
             propertyName = 'height';
             minValue = 64;
+            integerOnly = true;
+        } else if (valueAreaKey === 'batchSizeValueArea') {
+            valueType = 'Batch Size';
+            currentValue = this.rm.batchSizeWidget ? this.rm.batchSizeWidget.value : 1;
+            propertyName = 'batch_size';
+            minValue = 1;
+            integerOnly = true;
         } else {
             log.debug(`Unknown value area key: ${valueAreaKey}`);
             return;
         }
         
         log.debug(`Opening dialog for ${valueType} with current value: ${currentValue}`);
-        this.createCustomInputDialog(valueType, currentValue, propertyName, minValue, e);
+        this.createCustomInputDialog(valueType, currentValue, propertyName, minValue, integerOnly, e);
     }
     
     /**
@@ -67,9 +76,10 @@ export class DialogManager {
      * @param {number} currentValue - The current value
      * @param {string} propertyName - The property name to update
      * @param {number} minValue - Minimum allowed value
+     * @param {boolean} integerOnly - Whether to allow only integer values
      * @param {Event} e - The mouse event for positioning
      */
-    createCustomInputDialog(valueType, currentValue, propertyName, minValue, e) {
+    createCustomInputDialog(valueType, currentValue, propertyName, minValue, integerOnly, e) {
         this.inputDialogActive = true;
         log.debug(`Creating dialog for ${valueType}, current: ${currentValue}`);
         
@@ -103,11 +113,12 @@ export class DialogManager {
         dialog.style.top = `${Math.max(10, Math.min(y, window.innerHeight - 200))}px`;
         
         // Create dialog content
+        const inputStep = integerOnly ? '1' : '0.01';
         dialog.innerHTML = `
             <div style="color: #fff; font-size: 16px; font-weight: bold; margin-bottom: 15px; text-align: center;">Set Custom ${valueType}</div>
             <div style="margin-bottom: 10px;">
                 <label style="color: #ccc; font-size: 12px; display: block; margin-bottom: 5px;">Current: ${this.formatValueForDisplay(currentValue, valueType)}</label>
-            <input type="${valueType === 'Scale Factor' ? 'text' : 'number'}" id="customValueInput" value="${currentValue}" step="${valueType === 'Width' || valueType === 'Height' ? '1' : '0.01'}" min="${minValue}"
+            <input type="${valueType === 'Scale Factor' ? 'text' : 'number'}" id="customValueInput" value="${currentValue}" step="${inputStep}" min="${minValue}"
                        style="width: 100%; padding: 8px; border: 1px solid #555; border-radius: 4px; background: #333; color: #fff; font-size: 14px; box-sizing: border-box;">
             </div>
             <div id="validationMessage" style="color: #f55; font-size: 11px; margin-bottom: 5px; min-height: 15px;"></div>
@@ -131,6 +142,37 @@ export class DialogManager {
             infoMsg.textContent = 'Tip: Use /2 for 0.5x, /4 for 0.25x, etc.';
         }
         
+        // Block decimal characters for integer-only inputs
+        if (integerOnly) {
+            input.addEventListener('keydown', (e) => {
+                // Allow: backspace, delete, tab, escape, enter, arrows, home, end
+                const allowedKeys = ['Backspace', 'Delete', 'Tab', 'Escape', 'Enter', 'ArrowLeft', 'ArrowRight', 'ArrowUp', 'ArrowDown', 'Home', 'End'];
+                
+                // Block: decimal point, comma, e, E, +, - (except for navigation)
+                const blockedChars = ['.', ',', 'e', 'E', '+', '-'];
+                
+                if (allowedKeys.includes(e.key)) {
+                    return; // Allow these keys
+                }
+                
+                // Allow Ctrl/Cmd combinations (copy, paste, select all, etc.)
+                if (e.ctrlKey || e.metaKey) {
+                    return;
+                }
+                
+                // Block decimal and scientific notation characters
+                if (blockedChars.includes(e.key)) {
+                    e.preventDefault();
+                    return;
+                }
+                
+                // Allow only digits 0-9
+                if (!/^\d$/.test(e.key)) {
+                    e.preventDefault();
+                }
+            });
+        }
+        
         // Focus and select input
         setTimeout(() => { input.focus(); input.select(); }, 50);
         
@@ -146,6 +188,10 @@ export class DialogManager {
                 validationMsg.textContent = errorMsg;
                 applyBtn.disabled = true; applyBtn.style.opacity = '0.5';
                 return false;
+            } else if (integerOnly && !Number.isInteger(value)) {
+                validationMsg.textContent = 'Value must be a whole number';
+                applyBtn.disabled = true; applyBtn.style.opacity = '0.5';
+                return false;
             } else {
                 validationMsg.textContent = '';
                 applyBtn.disabled = false; applyBtn.style.opacity = '1';
@@ -243,6 +289,12 @@ export class DialogManager {
             const newHeight = Math.round(value);
             const currentWidth = this.rm.widthWidget ? this.rm.widthWidget.value : props.valueX;
             this.rm.setDimensions(currentWidth, newHeight);
+        } else if (propertyName === 'batch_size') {
+            const newBatchSize = Math.max(1, Math.min(4096, Math.round(value)));
+            props.batch_size = newBatchSize;
+            if (this.rm.batchSizeWidget) {
+                this.rm.batchSizeWidget.value = newBatchSize;
+            }
         }
         
         this.closeCustomInputDialog();
@@ -266,6 +318,8 @@ export class DialogManager {
             return value.toFixed(1) + 'MP';
         } else if (valueType === 'Width' || valueType === 'Height') {
             return value.toString() + 'px';
+        } else if (valueType === 'Batch Size') {
+            return value.toString();
         } else {
             return value.toString();
         }
diff --git a/js/ResolutionMaster.js b/js/ResolutionMaster.js
@@ -132,6 +132,7 @@ class ResolutionMasterCanvas {
             mode: "Manual",
             valueX: 512,
             valueY: 512,
+            batch_size: 1,
             canvas_min_x: 0,
             canvas_max_x: 2048,
             canvas_step_x: 64,
@@ -213,6 +214,7 @@ class ResolutionMasterCanvas {
         const autoDetectWidget = node.widgets?.find(w => w.name === 'auto_detect');
         const rescaleModeWidget = node.widgets?.find(w => w.name === 'rescale_mode');
         const rescaleValueWidget = node.widgets?.find(w => w.name === 'rescale_value');
+        const batchSizeWidget = node.widgets?.find(w => w.name === 'batch_size');
         
         // Initialize rescale widgets with proper values
         if (rescaleModeWidget) {
@@ -221,9 +223,8 @@ class ResolutionMasterCanvas {
         if (rescaleValueWidget) {
             rescaleValueWidget.value = node.properties.rescaleValue;
         }
-    
         
-        // Initialize values from widgets
+        // Initialize values from widgets (widget is source of truth for backend parameters)
         if (widthWidget && heightWidget) {
             node.properties.valueX = widthWidget.value;
             node.properties.valueY = heightWidget.value;
@@ -233,12 +234,18 @@ class ResolutionMasterCanvas {
             node.intpos.y = (heightWidget.value - node.properties.canvas_min_y) / (node.properties.canvas_max_y - node.properties.canvas_min_y);
         }
         
+        // Initialize batch_size from widget to property (same as width/height)
+        if (batchSizeWidget) {
+            node.properties.batch_size = batchSizeWidget.value;
+        }
+        
         
         // Store widget references
         this.widthWidget = widthWidget;
         this.heightWidget = heightWidget;
         this.rescaleModeWidget = rescaleModeWidget;
         this.rescaleValueWidget = rescaleValueWidget;
+        this.batchSizeWidget = batchSizeWidget;
         
         
         // Override onDrawForeground
@@ -319,7 +326,7 @@ class ResolutionMasterCanvas {
         };
 
                 // Hide all backend widgets
-        [widthWidget, heightWidget, modeWidget, autoDetectWidget, rescaleModeWidget, rescaleValueWidget].forEach(widget => {
+        [widthWidget, heightWidget, modeWidget, autoDetectWidget, rescaleModeWidget, rescaleValueWidget, batchSizeWidget].forEach(widget => {
             if (widget) {
                 widget.hidden = true;
                 widget.type = "hidden";
@@ -468,16 +475,19 @@ class ResolutionMasterCanvas {
         ctx.textAlign = "right";
         ctx.textBaseline = "middle";
         
-        if (this.widthWidget && this.heightWidget) {
+        if (this.widthWidget && this.heightWidget && this.batchSizeWidget) {
             // Shift values up slightly to better match visual center of slots
             const y_offset_1 = 5 + (LiteGraph.NODE_SLOT_HEIGHT * 0.5);
             const y_offset_2 = 5 + (LiteGraph.NODE_SLOT_HEIGHT * 1.5);
             const y_offset_3 = 5 + (LiteGraph.NODE_SLOT_HEIGHT * 2.5);
+            const y_offset_4 = 5 + (LiteGraph.NODE_SLOT_HEIGHT * 3.5);
 
             // Calculate clickable area dimensions
             const valueAreaWidth = 60; // Wider area for better clicking
+            const batchSizeAreaWidth = 35; // Smaller area for batch size (small numbers)
             const valueAreaHeight = 20;
             const valueAreaX = node.size[0] - valueAreaWidth - 5;
+            const batchSizeAreaX = node.size[0] - batchSizeAreaWidth - 5;
 
             // Width value area
             this.controls.widthValueArea = {
@@ -523,8 +533,31 @@ class ResolutionMasterCanvas {
             ctx.fillStyle = this.hoverElement === 'heightValueArea' ? "#F89" : "#F89";
             ctx.fillText(this.heightWidget.value.toString(), node.size[0] - 20, y_offset_2);
             
+            // Rescale value (non-clickable)
             ctx.fillStyle = "#9F8";
             ctx.fillText(props.rescaleValue.toFixed(2), node.size[0] - 20, y_offset_3);
+            
+            // Batch size value area
+            this.controls.batchSizeValueArea = {
+                x: batchSizeAreaX,
+                y: y_offset_4 - valueAreaHeight/2,
+                w: batchSizeAreaWidth,
+                h: valueAreaHeight
+            };
+            
+            // Draw background for batch size value area if hovered
+            if (this.hoverElement === 'batchSizeValueArea') {
+                ctx.fillStyle = "rgba(255, 136, 187, 0.2)";
+                ctx.strokeStyle = "rgba(255, 136, 187, 0.5)";
+                ctx.lineWidth = 1;
+                ctx.beginPath();
+                ctx.roundRect(batchSizeAreaX, y_offset_4 - valueAreaHeight/2, batchSizeAreaWidth, valueAreaHeight, 4);
+                ctx.fill();
+                ctx.stroke();
+            }
+            
+            ctx.fillStyle = this.hoverElement === 'batchSizeValueArea' ? "#FAB" : "#F8B";
+            ctx.fillText(this.batchSizeWidget.value.toString(), node.size[0] - 20, y_offset_4);
         }
     }
     
diff --git a/js/utils/ResolutionMasterConfig.js b/js/utils/ResolutionMasterConfig.js
@@ -10,6 +10,7 @@ export const tooltips = {
     // Output value areas (editable)
     widthValueArea: "Click to set custom width value",
     heightValueArea: "Click to set custom height value",
+    batchSizeValueArea: "Click to set custom batch size value",
     
     // Scaling controls (buttons and dropdowns only)
     scaleBtn: "Apply manual scaling factor and reset to 1.0x",
