From c297b34b949fefbaf455f29b112dac3e3afd5e4d Mon Sep 17 00:00:00 2001 From: Robert Wojciechowski Date: Sat, 28 Feb 2026 03:04:28 +0000 Subject: [PATCH] Add DistributedValue node with per-worker overrides --- api/orchestration/prompt_transform.py | 15 ++ distributed.py | 1 + nodes/__init__.py | 3 + nodes/utilities.py | 80 ++++++++++ tests/test_distributed_value.py | 208 ++++++++++++++++++++++++++ tests/test_prompt_transform.py | 25 ++++ 6 files changed, 332 insertions(+) create mode 100644 tests/test_distributed_value.py diff --git a/api/orchestration/prompt_transform.py b/api/orchestration/prompt_transform.py index 119f2d1..244bd6c 100644 --- a/api/orchestration/prompt_transform.py +++ b/api/orchestration/prompt_transform.py @@ -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, diff --git a/distributed.py b/distributed.py index 89314dc..f44bc7c 100644 --- a/distributed.py +++ b/distributed.py @@ -20,6 +20,7 @@ from .nodes import ( DistributedCollectorNode, DistributedSeed, DistributedModelName, + DistributedValue, AudioBatchDivider, DistributedEmptyImage, AnyType, diff --git a/nodes/__init__.py b/nodes/__init__.py index 3dc67b4..d7327b3 100644 --- a/nodes/__init__.py +++ b/nodes/__init__.py @@ -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", diff --git a/nodes/utilities.py b/nodes/utilities.py index b38c9e1..3fda3de 100644 --- a/nodes/utilities.py +++ b/nodes/utilities.py @@ -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): diff --git a/tests/test_distributed_value.py b/tests/test_distributed_value.py new file mode 100644 index 0000000..d2551c6 --- /dev/null +++ b/tests/test_distributed_value.py @@ -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() diff --git a/tests/test_prompt_transform.py b/tests/test_prompt_transform.py index 9720a32..0b5bf6f 100644 --- a/tests/test_prompt_transform.py +++ b/tests/test_prompt_transform.py @@ -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()