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

Add DistributedValue node with per-worker overrides

Browse Source
This commit is contained in:
Robert Wojciechowski
2026-02-28 04:53:20 +00:00
parent 2cf3d717d1
commit c297b34b94
6 changed files with 332 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
api/orchestration/prompt_transform.py
+15
View File
@@ -299,6 +299,20 @@ def _override_upscale_nodes(
inputs["worker_id"] = participant_id
def _override_value_nodes(prompt_copy, prompt_index, is_master, participant_id, worker_index_map):
"""Configure DistributedValue nodes for master or worker role."""
for node_id in prompt_index.nodes_for_class("DistributedValue"):
node = prompt_copy.get(node_id)
if not isinstance(node, dict):
continue
inputs = node.setdefault("inputs", {})
inputs["is_worker"] = not is_master
if is_master:
inputs["worker_id"] = ""
else:
inputs["worker_id"] = f"worker_{worker_index_map.get(participant_id, 0)}"
def apply_participant_overrides(
prompt_copy,
participant_id,
@@ -314,6 +328,7 @@ def apply_participant_overrides(
enabled_json = json.dumps(enabled_worker_ids)
_override_seed_nodes(prompt_copy, prompt_index, is_master, participant_id, worker_index_map)
_override_value_nodes(prompt_copy, prompt_index, is_master, participant_id, worker_index_map)
_override_collector_nodes(
prompt_copy,
prompt_index,
distributed.py
+1
View File
@@ -20,6 +20,7 @@ from .nodes import (
DistributedCollectorNode,
DistributedSeed,
DistributedModelName,
DistributedValue,
AudioBatchDivider,
DistributedEmptyImage,
AnyType,
nodes/__init__.py
+3
View File
@@ -1,6 +1,7 @@
from .utilities import (
DistributedSeed,
DistributedModelName,
DistributedValue,
ImageBatchDivider,
AudioBatchDivider,
DistributedEmptyImage,
@@ -14,6 +15,7 @@ NODE_CLASS_MAPPINGS = {
"DistributedCollector": DistributedCollectorNode,
"DistributedSeed": DistributedSeed,
"DistributedModelName": DistributedModelName,
"DistributedValue": DistributedValue,
"ImageBatchDivider": ImageBatchDivider,
"AudioBatchDivider": AudioBatchDivider,
"DistributedEmptyImage": DistributedEmptyImage,
@@ -22,6 +24,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
"DistributedCollector": "Distributed Collector",
"DistributedSeed": "Distributed Seed",
"DistributedModelName": "Distributed Model Name",
"DistributedValue": "Distributed Value",
"ImageBatchDivider": "Image Batch Divider",
"AudioBatchDivider": "Audio Batch Divider",
"DistributedEmptyImage": "Distributed Empty Image",
nodes/utilities.py
+80
View File
@@ -74,6 +74,7 @@ class DistributedSeed:
# Fallback: return original seed
return (seed,)
# Define ByPassTypeTuple for flexible return types
class AnyType(str):
def __ne__(self, __value: object) -> bool:
@@ -81,6 +82,85 @@ class AnyType(str):
any_type = AnyType("*")
class DistributedValue:
"""
Outputs a different value per worker.
On master: returns default_value.
On workers: looks up the worker-specific value from a JSON map,
falling back to default_value if not set.
"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"default_value": ("STRING", {"default": ""}),
"worker_values": ("STRING", {"default": "{}"}),
},
"hidden": {
"is_worker": ("BOOLEAN", {"default": False}),
"worker_id": ("STRING", {"default": ""}),
},
}
RETURN_TYPES = (any_type,)
RETURN_NAMES = ("value",)
FUNCTION = "distribute"
CATEGORY = "utils"
@staticmethod
def _coerce(value, value_type):
"""Convert a string value to the requested type."""
if value_type == "INT":
return int(float(value))
if value_type == "FLOAT":
return float(value)
return value # STRING and COMBO stay as strings
@staticmethod
def _coerce_safe(value, value_type):
"""Best-effort coercion with graceful fallback to original value."""
try:
return DistributedValue._coerce(value, value_type)
except (TypeError, ValueError):
return value
def distribute(self, default_value, worker_values="{}", is_worker=False, worker_id=""):
values = {}
value_type = "STRING"
try:
values = json.loads(worker_values) if isinstance(worker_values, str) else worker_values
if not isinstance(values, dict):
values = {}
except json.JSONDecodeError as e:
debug_log(f"DistributedValue - Error parsing worker_values: {e}")
values = {}
value_type = values.get("_type", "STRING")
coerced_default = self._coerce_safe(default_value, value_type)
if not is_worker:
debug_log(f"DistributedValue - Master: returning default '{coerced_default}'")
return (coerced_default,)
try:
if worker_id.startswith("worker_"):
idx = int(worker_id.split("_")[1])
else:
idx = int(worker_id)
key = str(idx + 1) # worker_0 → key "1" (1-indexed)
raw = values.get(key, "")
if raw:
coerced = self._coerce(raw, value_type)
debug_log(f"DistributedValue - Worker {idx}: returning '{coerced}'")
return (coerced,)
except (ValueError, IndexError) as e:
debug_log(f"DistributedValue - Error: {e}")
debug_log(f"DistributedValue - Worker fallback: returning default '{coerced_default}'")
return (coerced_default,)
class DistributedModelName:
@classmethod
def INPUT_TYPES(cls):
tests/test_distributed_value.py
+208
View File
@@ -0,0 +1,208 @@
import json
import unittest
class DistributedValueTests(unittest.TestCase):
"""Unit tests for the DistributedValue node's distribute() method."""
def _make_node(self):
# Import inline to avoid plugin-level imports
import importlib.util
import sys
import types
from pathlib import Path
from unittest.mock import MagicMock
module_path = Path(__file__).resolve().parents[1] / "nodes" / "utilities.py"
pkg_name = "dv_test_pkg"
for mod_name in list(sys.modules):
if mod_name == pkg_name or mod_name.startswith(f"{pkg_name}."):
del sys.modules[mod_name]
# Mock torch if not available
if "torch" not in sys.modules:
sys.modules["torch"] = MagicMock()
root_pkg = types.ModuleType(pkg_name)
root_pkg.__path__ = []
sys.modules[pkg_name] = root_pkg
utils_pkg = types.ModuleType(f"{pkg_name}.utils")
utils_pkg.__path__ = []
sys.modules[f"{pkg_name}.utils"] = utils_pkg
logging_mod = types.ModuleType(f"{pkg_name}.utils.logging")
logging_mod.debug_log = lambda *_a, **_k: None
logging_mod.log = lambda *_a, **_k: None
sys.modules[f"{pkg_name}.utils.logging"] = logging_mod
spec = importlib.util.spec_from_file_location(
f"{pkg_name}.nodes.utilities", module_path
)
mod = importlib.util.module_from_spec(spec)
spec.loader.exec_module(mod)
return mod.DistributedValue()
def setUp(self):
self.node = self._make_node()
def test_master_returns_default(self):
result = self.node.distribute(
default_value="model_a",
worker_values="{}",
is_worker=False,
worker_id="",
)
self.assertEqual(result, ("model_a",))
def test_master_coerces_default_int(self):
values = json.dumps({"_type": "INT"})
result = self.node.distribute(
default_value="42",
worker_values=values,
is_worker=False,
worker_id="",
)
self.assertEqual(result, (42,))
self.assertIsInstance(result[0], int)
def test_master_coerces_default_float(self):
values = json.dumps({"_type": "FLOAT"})
result = self.node.distribute(
default_value="2.5",
worker_values=values,
is_worker=False,
worker_id="",
)
self.assertEqual(result, (2.5,))
self.assertIsInstance(result[0], float)
def test_worker_returns_specific_value(self):
values = json.dumps({"1": "model_x", "2": "model_y"})
result = self.node.distribute(
default_value="default",
worker_values=values,
is_worker=True,
worker_id="worker_0",
)
self.assertEqual(result, ("model_x",))
def test_worker_second_index(self):
values = json.dumps({"1": "model_x", "2": "model_y"})
result = self.node.distribute(
default_value="default",
worker_values=values,
is_worker=True,
worker_id="worker_1",
)
self.assertEqual(result, ("model_y",))
def test_worker_falls_back_to_default_when_key_missing(self):
values = json.dumps({"_type": "INT", "1": "3"})
result = self.node.distribute(
default_value="9",
worker_values=values,
is_worker=True,
worker_id="worker_5",
)
self.assertEqual(result, (9,))
self.assertIsInstance(result[0], int)
def test_worker_falls_back_to_default_on_empty_value(self):
values = json.dumps({"1": ""})
result = self.node.distribute(
default_value="fallback",
worker_values=values,
is_worker=True,
worker_id="worker_0",
)
self.assertEqual(result, ("fallback",))
def test_worker_falls_back_on_invalid_json(self):
result = self.node.distribute(
default_value="safe",
worker_values="not-json",
is_worker=True,
worker_id="worker_0",
)
self.assertEqual(result, ("safe",))
def test_worker_falls_back_on_invalid_worker_id(self):
values = json.dumps({"1": "model_x"})
result = self.node.distribute(
default_value="safe",
worker_values=values,
is_worker=True,
worker_id="bad_id",
)
self.assertEqual(result, ("safe",))
def test_worker_id_as_direct_integer(self):
values = json.dumps({"1": "model_x"})
result = self.node.distribute(
default_value="default",
worker_values=values,
is_worker=True,
worker_id="0",
)
self.assertEqual(result, ("model_x",))
def test_type_int_coerces_value(self):
values = json.dumps({"_type": "INT", "1": "42"})
result = self.node.distribute(
default_value="0",
worker_values=values,
is_worker=True,
worker_id="worker_0",
)
self.assertEqual(result, (42,))
self.assertIsInstance(result[0], int)
def test_type_float_coerces_value(self):
values = json.dumps({"_type": "FLOAT", "1": "3.14"})
result = self.node.distribute(
default_value="0",
worker_values=values,
is_worker=True,
worker_id="worker_0",
)
self.assertAlmostEqual(result[0], 3.14)
self.assertIsInstance(result[0], float)
def test_type_combo_stays_string(self):
values = json.dumps({"_type": "COMBO", "1": "model_v2"})
result = self.node.distribute(
default_value="model_v1",
worker_values=values,
is_worker=True,
worker_id="worker_0",
)
self.assertEqual(result, ("model_v2",))
self.assertIsInstance(result[0], str)
def test_type_string_default_stays_string(self):
values = json.dumps({"1": "hello"})
result = self.node.distribute(
default_value="default",
worker_values=values,
is_worker=True,
worker_id="worker_0",
)
self.assertEqual(result, ("hello",))
self.assertIsInstance(result[0], str)
def test_int_coerce_handles_float_string(self):
"""INT coercion of '3.7' should truncate to 3."""
values = json.dumps({"_type": "INT", "1": "3.7"})
result = self.node.distribute(
default_value="0",
worker_values=values,
is_worker=True,
worker_id="worker_0",
)
self.assertEqual(result, (3,))
if __name__ == "__main__":
unittest.main()
tests/test_prompt_transform.py
+25
View File
@@ -464,5 +464,30 @@ class ApplyOverridesUpscaleTests(unittest.TestCase):
self.assertTrue(result["2"]["inputs"].get("pass_through"))
# ---------------------------------------------------------------------------
# apply_participant_overrides – DistributedValue
# ---------------------------------------------------------------------------
class ApplyOverridesValueTests(unittest.TestCase):
def _value_prompt(self):
return {"1": {"class_type": "DistributedValue", "inputs": {}}}
def test_worker_sets_is_worker_true(self):
result = _apply(self._value_prompt(), "worker-a")
self.assertTrue(result["1"]["inputs"]["is_worker"])
def test_worker_id_reflects_index_in_enabled_list(self):
result = _apply(self._value_prompt(), "worker-b", enabled_worker_ids=["worker-a", "worker-b"])
self.assertEqual(result["1"]["inputs"]["worker_id"], "worker_1")
def test_master_sets_is_worker_false(self):
result = _apply(self._value_prompt(), "master")
self.assertFalse(result["1"]["inputs"]["is_worker"])
def test_master_sets_empty_worker_id(self):
result = _apply(self._value_prompt(), "master")
self.assertEqual(result["1"]["inputs"]["worker_id"], "")
if __name__ == "__main__":
unittest.main()