diff --git a/__init__.py b/__init__.py index 9bc2dde..18d9765 100644 --- a/__init__.py +++ b/__init__.py @@ -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, diff --git a/distributed.py b/distributed.py index 1a1a14e..bf2b2c1 100644 --- a/distributed.py +++ b/distributed.py @@ -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" } diff --git a/utils/constants.py b/utils/constants.py index a199f24..ff1869d 100644 --- a/utils/constants.py +++ b/utils/constants.py @@ -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 diff --git a/web/ui.js b/web/ui.js index 2421cdf..bdbc566 100644 --- a/web/ui.js +++ b/web/ui.js @@ -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; }