diff --git a/nodes.py b/nodes.py index e97adb0..d3ee41a 100644 --- a/nodes.py +++ b/nodes.py @@ -460,52 +460,100 @@ def person_dict(result): return dict(zip(RETURN_NAMES, result)) -# -- V1 node (legacy dict API) ------------------------------------------------- +# -- V1 node factory ----------------------------------------------------------- -class RandomPersonNode: +_ALL_KEYS = [k for k, _, _ in CATEGORY_SPECS] - CATEGORY = "utils" - FUNCTION = "generate" - OUTPUT_NODE = False - DESCRIPTION = "Generate a randomised, structured physical person description for image prompts." +# (node_id, display_name, category keys, include_age, output pin names) +NODE_SPECS = [ + ("RandomPersonNode", "Random Person Description", _ALL_KEYS, True, RETURN_NAMES), + ("RandomPersonIdentity", "Random Person: Identity", + ["nationality", "complexion", "skin_texture"], True, + ("description", "sex", "age", "nationality", "complexion", "seed")), + ("RandomPersonFace", "Random Person: Face & Expression", + ["face_shape", "eyes", "eye_shape", "eyebrows", "nose_shape", + "mouth_shape", "face_feature", "expression"], False, + ("description", "sex", "face", "expression", "seed")), + ("RandomPersonHair", "Random Person: Hair", + ["hair_color", "hair_style", "hair_length", "facial_hair"], False, + ("description", "sex", "hair", "facial_hair", "seed")), + ("RandomPersonBody", "Random Person: Body", + ["body_type", "shoulders", "chest", "bust_size", "bust_shape"], False, + ("description", "sex", "body_type", "shoulders", "chest", "bust", "seed")), + ("RandomPersonStyle", "Random Person: Style", + ["accessories", "makeup"], False, + ("description", "sex", "accessories", "makeup", "seed")), +] - RETURN_TYPES = ("STRING",) * (len(RETURN_NAMES) - 1) + ("INT",) - RETURN_NAMES = RETURN_NAMES - - @classmethod - def INPUT_TYPES(cls): - required = { - "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFF, "tooltip": _SEED_TIP}), - "randomize": ("BOOLEAN", {"default": True, "tooltip": _RAND_TIP}), - "sex": (SEX_MODES, {"default": "random", "tooltip": _SEX_TIP}), - } - for key, fixed_labels, placeholder in CATEGORY_SPECS: - required[f"{key}_mode"] = (MODES, {"default": _mode_default(key), "tooltip": _MODE_TIP}) - required[f"{key}_allow_list"] = ("STRING", {"default": "", "multiline": False, - "placeholder": placeholder, "tooltip": _ALLOW_TIP}) - required[f"{key}_fixed"] = (fixed_labels, {"default": "(none)", "tooltip": _FIXED_TIP}) - required["age_mode"] = (AGE_MODES, {"default": "random"}) - required["age_min"] = ("INT", {"default": 0, "min": 0, "max": 90, "tooltip": _AGEMIN_TIP}) - required["age_max"] = ("INT", {"default": 0, "min": 0, "max": 90, "tooltip": _AGEMAX_TIP}) - required["age_fixed"] = ("INT", {"default": 0, "min": 0, "max": 90, "tooltip": _AGEFIX_TIP}) - required["extra_attributes"] = ("STRING", {"default": "", "multiline": True, - "placeholder": "e.g. wearing glasses, tattoo on left arm", - "tooltip": _EXTRA_TIP}) - return {"required": required} - - def generate(self, seed, randomize, sex, age_mode, age_min, age_max, - age_fixed, extra_attributes, **kwargs): - return generate_person(seed, randomize, sex, _collect_category_args(kwargs), - age_mode, age_min, age_max, age_fixed, extra_attributes) - - @classmethod - def IS_CHANGED(cls, randomize, seed, **kwargs): - # randomize=True -> always re-run; otherwise re-run only when seed changes. - return float("nan") if randomize else seed +_NODE_DESCRIPTION = "Generate a randomised, structured physical person description for image prompts." -NODE_CLASS_MAPPINGS = {"RandomPersonNode": RandomPersonNode} -NODE_DISPLAY_NAME_MAPPINGS = {"RandomPersonNode": "Random Person Description"} +def make_v1_node(node_id, display_name, keys, include_age, output_names): + """Build a ComfyUI V1 node class exposing only `keys` and returning `output_names`.""" + key_set = set(keys) + + class _RandomPersonV1: + CATEGORY = "Random Person" + FUNCTION = "generate" + OUTPUT_NODE = False + DESCRIPTION = _NODE_DESCRIPTION + RETURN_NAMES = output_names + RETURN_TYPES = tuple("INT" if n == "seed" else "STRING" for n in output_names) + + @classmethod + def INPUT_TYPES(cls): + required = { + "seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFF, "tooltip": _SEED_TIP}), + "randomize": ("BOOLEAN", {"default": True, "tooltip": _RAND_TIP}), + "sex": (SEX_MODES, {"default": "random", "tooltip": _SEX_TIP}), + } + for key, fixed_labels, placeholder in CATEGORY_SPECS: + if key not in key_set: + continue + required[f"{key}_mode"] = (MODES, {"default": _mode_default(key), "tooltip": _MODE_TIP}) + required[f"{key}_allow_list"] = ("STRING", {"default": "", "multiline": False, + "placeholder": placeholder, "tooltip": _ALLOW_TIP}) + required[f"{key}_fixed"] = (fixed_labels, {"default": "(none)", "tooltip": _FIXED_TIP}) + if include_age: + required["age_mode"] = (AGE_MODES, {"default": "random"}) + required["age_min"] = ("INT", {"default": 0, "min": 0, "max": 90, "tooltip": _AGEMIN_TIP}) + required["age_max"] = ("INT", {"default": 0, "min": 0, "max": 90, "tooltip": _AGEMAX_TIP}) + required["age_fixed"] = ("INT", {"default": 0, "min": 0, "max": 90, "tooltip": _AGEFIX_TIP}) + required["extra_attributes"] = ("STRING", {"default": "", "multiline": True, + "placeholder": "e.g. wearing glasses, tattoo on left arm", + "tooltip": _EXTRA_TIP}) + return {"required": required} + + def generate(self, seed, randomize, sex, **kwargs): + age_mode = kwargs.get("age_mode", "off") if include_age else "off" + age_min = kwargs.get("age_min", 0) + age_max = kwargs.get("age_max", 0) + age_fixed = kwargs.get("age_fixed", 0) + extra = kwargs.get("extra_attributes", "") + cat_args = _collect_category_args(kwargs, keys) + result = generate_person(seed, randomize, sex, cat_args, + age_mode, age_min, age_max, age_fixed, extra) + dmap = person_dict(result) + return tuple(dmap[name] for name in output_names) + + @classmethod + def IS_CHANGED(cls, randomize, seed, **kwargs): + return float("nan") if randomize else seed + + _RandomPersonV1.__name__ = node_id + _RandomPersonV1.__qualname__ = node_id + return _RandomPersonV1 + + +NODE_CLASS_MAPPINGS = {} +NODE_DISPLAY_NAME_MAPPINGS = {} +for _spec in NODE_SPECS: + _cls = make_v1_node(*_spec) + NODE_CLASS_MAPPINGS[_spec[0]] = _cls + NODE_DISPLAY_NAME_MAPPINGS[_spec[0]] = _spec[1] + +# Backwards-compatible alias: tests and external imports reference this name. +RandomPersonNode = NODE_CLASS_MAPPINGS["RandomPersonNode"] # -- V3 node (import-guarded comfy_api schema) --------------------------------- diff --git a/test_random_person.py b/test_random_person.py index b686e4c..3793f52 100644 --- a/test_random_person.py +++ b/test_random_person.py @@ -265,5 +265,57 @@ class TestBodyData(unittest.TestCase): self.assertEqual(out[chest_idx], "", seed) +class TestSegmentNodes(unittest.TestCase): + + def _node(self, node_id): + return core.NODE_CLASS_MAPPINGS[node_id] + + def test_all_nodes_registered(self): + for node_id in ("RandomPersonNode", "RandomPersonIdentity", + "RandomPersonFace", "RandomPersonHair", + "RandomPersonBody", "RandomPersonStyle"): + self.assertIn(node_id, core.NODE_CLASS_MAPPINGS, node_id) + + def test_all_nodes_in_random_person_category(self): + for cls in core.NODE_CLASS_MAPPINGS.values(): + self.assertEqual(cls.CATEGORY, "Random Person") + + def test_identity_exposes_only_its_widgets(self): + req = self._node("RandomPersonIdentity").INPUT_TYPES()["required"] + self.assertIn("nationality_mode", req) + self.assertIn("age_mode", req) + self.assertNotIn("eyes_mode", req) + self.assertNotIn("hair_color_mode", req) + + def test_face_node_has_no_age(self): + req = self._node("RandomPersonFace").INPUT_TYPES()["required"] + self.assertIn("eyes_mode", req) + self.assertNotIn("age_mode", req) + + def test_face_fragment_excludes_other_groups(self): + cls = self._node("RandomPersonFace") + kwargs = {} + req = cls.INPUT_TYPES()["required"] + for name, (typ, opts) in req.items(): + if name.endswith("_mode"): + kwargs[name] = "random" + elif name.endswith("_allow_list"): + kwargs[name] = "" + elif name.endswith("_fixed"): + kwargs[name] = "(none)" + out = cls().generate(seed=3, randomize=False, sex="female", **kwargs) + names = cls.RETURN_NAMES + desc = out[names.index("description")] + self.assertNotIn("build", desc) # body excluded + self.assertNotIn("hair", desc) # hair excluded + self.assertNotIn("year old", desc) # age excluded + + def test_body_node_outputs_bust_pin(self): + cls = self._node("RandomPersonBody") + self.assertIn("bust", cls.RETURN_NAMES) + self.assertIn("shoulders", cls.RETURN_NAMES) + self.assertNotIn("hair", cls.RETURN_NAMES) + + if __name__ == "__main__": unittest.main()