nodes/robertvoy-ComfyUI-Distributed
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Image Batch Divider

Browse Source
This commit is contained in:
robertvoy
2025-07-15 11:16:34 +10:00
parent 9650280c7f
commit fa09a2da10
4 changed files with 177 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+30
View File
@@ -4,6 +4,36 @@ import sys
# Add the directory to Python path
sys.path.append(os.path.dirname(__file__))
# Patch ComfyUI execution validation for DistributedVideoCollector
try:
import execution
from .distributed import ImageBatchDivider
# Store original validate function if it exists
if hasattr(execution, 'validate_outputs'):
original_validate_outputs = execution.validate_outputs
def patched_validate_outputs(executor, node_id, result, node_class):
if node_class == ImageBatchDivider:
return # Skip validation for our dynamic output node
return original_validate_outputs(executor, node_id, result, node_class)
execution.validate_outputs = patched_validate_outputs
else:
# Fallback: patch the main execute method if validate_outputs doesn't exist
if hasattr(execution.PromptExecutor, 'execute'):
original_execute = execution.PromptExecutor.execute
def patched_execute(self, prompt, prompt_id, extra_data={}, execute_outputs=[]):
# This is a more complex patch - for now just call original
# The ByPassTypeTuple should handle most validation issues
return original_execute(self, prompt, prompt_id, extra_data, execute_outputs)
execution.PromptExecutor.execute = patched_execute
except ImportError:
pass # ComfyUI execution module not available during import
# Import everything needed from the main module
from .distributed import (
NODE_CLASS_MAPPINGS as DISTRIBUTED_CLASS_MAPPINGS,
distributed.py
+78 -2
View File
@@ -1638,11 +1638,87 @@ class DistributedSeed:
# Fallback: return original seed
return (seed,)
# Define ByPassTypeTuple for flexible return types
class AnyType(str):
def __ne__(self, __value: object) -> bool:
return False
any_type = AnyType("*")
class ByPassTypeTuple(tuple):
def __getitem__(self, index):
if index > 0:
index = 0
item = super().__getitem__(index)
if isinstance(item, str):
return any_type
return item
class ImageBatchDivider:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"images": ("IMAGE",),
"divide_by": ("INT", {
"default": 2,
"min": 1,
"max": 10,
"step": 1,
"display": "number",
"tooltip": "Number of parts to divide the batch into"
}),
}
}
RETURN_TYPES = ByPassTypeTuple(("IMAGE", )) # Flexible for variable outputs
RETURN_NAMES = ByPassTypeTuple(tuple([f"batch_{i+1}" for i in range(10)]))
FUNCTION = "divide_batch"
OUTPUT_NODE = True
CATEGORY = "image"
def divide_batch(self, images, divide_by):
import torch
# Use divide_by directly
total_splits = divide_by
if total_splits > 10:
total_splits = 10 # Cap to max
# Get total number of frames
total_frames = images.shape[0]
frames_per_split = total_frames // total_splits
remainder = total_frames % total_splits
outputs = []
start_idx = 0
for i in range(total_splits):
current_frames = frames_per_split + (1 if i < remainder else 0)
end_idx = start_idx + current_frames
split_frames = images[start_idx:end_idx]
outputs.append(split_frames)
start_idx = end_idx
# Pad with empty tensors up to max (10) to match potential RETURN_TYPES len
empty_shape = (1, images.shape[1], images.shape[2], images.shape[3]) if total_frames > 0 else (1, 512, 512, 3)
empty_tensor = torch.zeros(empty_shape, dtype=images.dtype if total_frames > 0 else torch.float32,
device=images.device if total_frames > 0 else 'cpu')
while len(outputs) < 10:
outputs.append(empty_tensor)
return tuple(outputs)
NODE_CLASS_MAPPINGS = {
"DistributedCollector": DistributedCollectorNode,
"DistributedSeed": DistributedSeed
"DistributedSeed": DistributedSeed,
"ImageBatchDivider": ImageBatchDivider
}
NODE_DISPLAY_NAME_MAPPINGS = {
"DistributedCollector": "Distributed Collector",
"DistributedSeed": "Distributed Seed"
"DistributedSeed": "Distributed Seed",
"ImageBatchDivider": "Image Batch Divider"
}
utils/constants.py
+2 -2
View File
@@ -3,8 +3,8 @@ Shared constants for ComfyUI-Distributed.
"""
# Timeouts (in seconds)
WORKER_JOB_TIMEOUT = 10.0
TILE_COLLECTION_TIMEOUT = 60.0
WORKER_JOB_TIMEOUT = 30.0
TILE_COLLECTION_TIMEOUT = 30.0
TILE_WAIT_TIMEOUT = 30.0
PROCESS_TERMINATION_TIMEOUT = 5.0
web/ui.js
+67 -56
View File
@@ -568,50 +568,58 @@ export class DistributedUI {
runpodText.textContent = "Deploy cloud worker with Runpod";
runpodText.style.cssText = "font-size: 12px; color: #4a90e2; text-decoration: none; margin-top: 4px; display: none; cursor: pointer;";
typeSelect.onchange = (e) => {
const workerType = e.target.value;
// Show/hide relevant fields
const hostGroup = document.getElementById(`host-group-${worker.id}`);
const hostInput = document.getElementById(`host-${worker.id}`);
const portGroup = document.getElementById(`port-group-${worker.id}`);
const portInput = document.getElementById(`port-${worker.id}`);
const cudaGroup = document.getElementById(`cuda-group-${worker.id}`);
const argsGroup = document.getElementById(`args-group-${worker.id}`);
const runpodTextElem = document.getElementById(`runpod-text-${worker.id}`);
if (workerType === "local") {
hostGroup.style.display = "none";
portGroup.style.display = "flex";
cudaGroup.style.display = "flex";
argsGroup.style.display = "flex";
runpodTextElem.style.display = "none";
} else if (workerType === "remote") {
hostGroup.style.display = "flex";
portGroup.style.display = "flex";
cudaGroup.style.display = "none";
argsGroup.style.display = "none";
runpodTextElem.style.display = "none";
// Update placeholder for remote workers
hostInput.placeholder = "e.g., 192.168.1.100";
// If switching to remote and host is localhost, clear it
if (hostInput.value === "localhost" || hostInput.value === "127.0.0.1") {
hostInput.value = "";
// Store the onchange function to be assigned later
const createOnChangeHandler = () => {
return (e) => {
const workerType = e.target.value;
// Show/hide relevant fields
const hostGroup = document.getElementById(`host-group-${worker.id}`);
const hostInput = document.getElementById(`host-${worker.id}`);
const portGroup = document.getElementById(`port-group-${worker.id}`);
const portInput = document.getElementById(`port-${worker.id}`);
const cudaGroup = document.getElementById(`cuda-group-${worker.id}`);
const argsGroup = document.getElementById(`args-group-${worker.id}`);
const runpodTextElem = document.getElementById(`runpod-text-${worker.id}`);
// Check if elements exist before accessing them
if (!hostGroup || !portGroup || !cudaGroup || !argsGroup || !runpodTextElem || !hostInput || !portInput) {
return; // Elements not ready yet
}
} else if (workerType === "cloud") {
hostGroup.style.display = "flex";
portGroup.style.display = "flex"; // Keep port visible for cloud workers
cudaGroup.style.display = "none";
argsGroup.style.display = "none";
runpodTextElem.style.display = "block";
// Update placeholder for cloud workers
hostInput.placeholder = "e.g., your-cloud-worker.trycloudflare.com";
// Set port to 443 for cloud workers
portInput.value = "443";
// If switching to cloud and host is localhost, clear it
if (hostInput.value === "localhost" || hostInput.value === "127.0.0.1") {
hostInput.value = "";
if (workerType === "local") {
hostGroup.style.display = "none";
portGroup.style.display = "flex";
cudaGroup.style.display = "flex";
argsGroup.style.display = "flex";
runpodTextElem.style.display = "none";
} else if (workerType === "remote") {
hostGroup.style.display = "flex";
portGroup.style.display = "flex";
cudaGroup.style.display = "none";
argsGroup.style.display = "none";
runpodTextElem.style.display = "none";
// Update placeholder for remote workers
hostInput.placeholder = "e.g., 192.168.1.100";
// If switching to remote and host is localhost, clear it
if (hostInput.value === "localhost" || hostInput.value === "127.0.0.1") {
hostInput.value = "";
}
} else if (workerType === "cloud") {
hostGroup.style.display = "flex";
portGroup.style.display = "flex"; // Keep port visible for cloud workers
cudaGroup.style.display = "none";
argsGroup.style.display = "none";
runpodTextElem.style.display = "block";
// Update placeholder for cloud workers
hostInput.placeholder = "e.g., your-cloud-worker.trycloudflare.com";
// Set port to 443 for cloud workers
portInput.value = "443";
// If switching to cloud and host is localhost, clear it
if (hostInput.value === "localhost" || hostInput.value === "127.0.0.1") {
hostInput.value = "";
}
}
}
};
};
typeGroup.appendChild(typeLabel);
@@ -619,20 +627,6 @@ export class DistributedUI {
typeGroup.appendChild(runpodText);
form.appendChild(typeGroup);
// Set initial value based on worker type
if (worker.type === "cloud") {
typeSelect.value = "cloud";
// Show Runpod text immediately for cloud workers
runpodText.style.display = "block";
} else if (extension.isRemoteWorker(worker)) {
typeSelect.value = "remote";
} else {
typeSelect.value = "local";
}
// Trigger initial state
typeSelect.dispatchEvent(new Event('change'));
// Host field (only for remote workers)
const hostGroup = this.createFormGroup("Host:", worker.host || "", `host-${worker.id}`, "text", "e.g., 192.168.1.100");
hostGroup.group.id = `host-group-${worker.id}`;
@@ -675,6 +669,23 @@ export class DistributedUI {
const buttonGroup = this.createButtonGroup([saveBtn, cancelBtn, deleteBtn], " margin-top: 8px;");
form.appendChild(buttonGroup);
// Assign the onchange handler now that all elements are created
typeSelect.onchange = createOnChangeHandler();
// Set initial value and trigger state after all DOM elements are created
if (worker.type === "cloud") {
typeSelect.value = "cloud";
// Show Runpod text immediately for cloud workers
runpodText.style.display = "block";
} else if (extension.isRemoteWorker(worker)) {
typeSelect.value = "remote";
} else {
typeSelect.value = "local";
}
// Trigger initial state now that all elements exist
typeSelect.dispatchEvent(new Event('change'));
return form;
}