diff --git a/core/nodes.py b/core/nodes.py index 98fd8df6..4d23e96b 100644 --- a/core/nodes.py +++ b/core/nodes.py @@ -38,9 +38,16 @@ def collect(registry=None) -> Tuple[List[type], List[Tuple[str, str]]]: # Sorted so the set of registered nodes does not depend on filesystem order. for spec in sorted(reg.specs.values(), key=lambda s: s.id): for node in spec.nodes: - name = getattr(node, "__name__", repr(node)) - where = f"{spec.id}.{name}" + # Named inside the guard, and starting from a name that cannot fail. + # `getattr(node, "__name__", repr(node))` looks safe and is not: + # Python evaluates the default eagerly, so a class whose metaclass + # raises in __repr__ threw BEFORE the try and took every other + # module's nodes with it -- the one thing this loop exists to stop. + where = f"{spec.id}." try: + name = getattr(node, "__name__", None) + if isinstance(name, str) and name: + where = f"{spec.id}.{name}" schema = _schema_of(node) except Exception as exc: # noqa: BLE001 -- absence, not a crash log.failed(where, exc) diff --git a/core/schema.py b/core/schema.py index a478aa0c..d126a98e 100644 --- a/core/schema.py +++ b/core/schema.py @@ -136,8 +136,12 @@ def validate_nodes(value: Any, module_id: str) -> list: outside it. The ComfyUI-specific checks (a FunPack-prefixed node_id, and no two modules claiming one id) happen in `core.nodes`, where ComfyUI exists. """ + if value is None: + return [] if isinstance(value, (str, bytes)) or not isinstance(value, (list, tuple)): - raise SchemaError(f"module {module_id!r}'s nodes must be a list.") + raise SchemaError( + f"module {module_id!r}'s nodes must be a list; got {value!r}." + ) for node in value: if not isinstance(node, type): @@ -200,7 +204,9 @@ def validate(announcement: Dict[str, Any], source: str = "") -> ModuleSpec: settings = {key: validate_setting(key, spec, raw) for key, spec in raw.items()} def _ids(name: str) -> list: - value = announcement.get(name) or [] + value = announcement.get(name) + if value is None: + return [] if isinstance(value, str) or not isinstance(value, (list, tuple)): raise SchemaError(f"module {announcement['id']!r}'s {name} must be a list.") for item in value: @@ -209,7 +215,7 @@ def validate(announcement: Dict[str, Any], source: str = "") -> ModuleSpec: return list(value) return ModuleSpec( - nodes=validate_nodes(announcement.get("nodes") or [], announcement["id"]), + nodes=validate_nodes(announcement.get("nodes"), announcement["id"]), traits=validate_traits(announcement.get("traits"), announcement["id"]), id=announcement["id"], title=announcement["title"], diff --git a/core/tests/fixtures/fixture_modules/nodes/clashes/__init__.py b/core/tests/fixtures/fixture_modules/nodes/clashes/__init__.py index 16864438..5b559da9 100644 --- a/core/tests/fixtures/fixture_modules/nodes/clashes/__init__.py +++ b/core/tests/fixtures/fixture_modules/nodes/clashes/__init__.py @@ -1,4 +1,10 @@ -"""Reuses an id another module already declared.""" +"""Reuses an id another module already declared. + +`ID` deliberately sorts AFTER `nodes_good`, because `collect()` resolves a +duplicate in favour of whoever it reaches first. Without that, this fixture won +the id and every test asserting "FunPackGood was collected" was satisfied by +this class instead of the real one -- a test passing for a reason it never named. +""" from ..good import GoodNode @@ -7,7 +13,7 @@ class Clashing(GoodNode): pass # same GET_SCHEMA, so same node_id -ID = "nodes_clashes" +ID = "nodes_zz_clashes" TITLE = "Clashing node" STAGE = "load" NODES = [Clashing] diff --git a/core/tests/test_nodes.py b/core/tests/test_nodes.py index 15453299..0b668874 100644 --- a/core/tests/test_nodes.py +++ b/core/tests/test_nodes.py @@ -28,9 +28,18 @@ def _ids(nodes): return [n.GET_SCHEMA().node_id for n in nodes] +def _good(): + """The real GoodNode class, by identity.""" + from core.tests.fixtures.fixture_modules.nodes.good import GoodNode + return GoodNode + + def test_an_announced_node_is_collected(collected): _, (nodes, _) = collected - assert "FunPackGood" in _ids(nodes) + # By identity. Asserting on the id alone let a DIFFERENT fixture module that + # claims the same id satisfy this, so the test passed without the module + # under test contributing anything. + assert _good() in nodes def test_a_node_outside_the_funpack_namespace_is_refused(collected): @@ -43,14 +52,18 @@ def test_a_node_outside_the_funpack_namespace_is_refused(collected): def test_two_modules_cannot_claim_one_node_id(collected): _, (nodes, rejected) = collected assert _ids(nodes).count("FunPackGood") == 1 - assert any("duplicate node_id" in why for _, why in rejected) + # And the one that survived is the first module to claim it, not whichever + # happened to be reached last. + assert _good() in nodes + assert any("duplicate node_id" in why and "nodes_zz_clashes" in where + for where, why in rejected) def test_something_that_is_not_a_comfy_node_is_refused(collected): _, (nodes, rejected) = collected assert any("not a comfy_api" in why for _, why in rejected) # And it did not stop the good one being collected. - assert "FunPackGood" in _ids(nodes) + assert _good() in nodes def test_a_node_whose_schema_raises_does_not_take_the_others_down(): @@ -74,7 +87,7 @@ def test_a_node_whose_schema_raises_does_not_take_the_others_down(): ) nodes, rejected = nodes_mod.collect(registry) - assert "FunPackGood" in _ids(nodes) + assert _good() in nodes, "a raising sibling took the good node with it" assert any("boom" in why for _, why in rejected) @@ -99,3 +112,61 @@ def test_a_module_with_settings_still_needs_a_mount(): def test_a_malformed_nodes_declaration_refuses_the_module(bad): with pytest.raises(SchemaError): validate({"id": "m", "title": "M", "nodes": bad}) + + +def test_a_node_that_cannot_be_named_does_not_take_the_others_down(): + """`getattr(node, "__name__", repr(node))` evaluates repr() eagerly, so a + class whose metaclass raises there threw outside the guard and killed the + whole collection -- and ComfyUI reports that as the entire pack failing.""" + + class Meta(type): + def __repr__(cls): + raise RuntimeError("unreprable") + + def __getattribute__(cls, item): + if item == "__name__": + raise RuntimeError("unnameable") + return type.__getattribute__(cls, item) + + class Hostile(metaclass=Meta): + @classmethod + def define_schema(cls): + return None + + @classmethod + def execute(cls): + return None + + @classmethod + def GET_SCHEMA(cls): + return None + + registry = registry_mod.scan(FIXTURES, package=PACKAGE) + spec = registry.specs["nodes_good"] + registry.specs["nodes_good"] = type(spec)( + **{**spec.__dict__, "nodes": [Hostile, *spec.nodes]} + ) + + nodes, rejected = nodes_mod.collect(registry) + assert _good() in nodes + assert any("unnameable" in why or "unreprable" in why for _, why in rejected) + + +@pytest.mark.parametrize("falsy", [0, False, ""]) +def test_a_falsy_but_wrong_nodes_declaration_is_refused_not_repaired(falsy): + """`announcement.get("nodes") or []` quietly turned a typo into "this module + ships no nodes". This file refuses rather than repairs, and 0 is not [].""" + with pytest.raises(SchemaError, match="must be a list"): + validate({"id": "m", "title": "M", "nodes": falsy}) + + +@pytest.mark.parametrize("field", ["requires", "after", "before"]) +def test_a_falsy_but_wrong_id_list_is_refused_too(field): + with pytest.raises(SchemaError, match="must be a list"): + validate({"id": "m", "title": "M", field: 0}) + + +def test_an_absent_declaration_is_still_fine(): + # Absent means "none", which is different from a typo. + spec = validate({"id": "m", "title": "M"}) + assert spec.nodes == [] and spec.requires == [] diff --git a/core/tests/test_traits.py b/core/tests/test_traits.py index 492da90e..d13ce32b 100644 --- a/core/tests/test_traits.py +++ b/core/tests/test_traits.py @@ -161,3 +161,39 @@ def test_traits_describe_the_tensor_not_an_assumed_modality(real_configs): assert "temporal_latent" in read(qwen) assert "temporal_latent" in read(ltxv) + + +def test_a_bare_probe_name_collides_across_architectures(): + """`Attention`, `MLP`, `Block` are not distinctive. A bare name matches any + of them, which is why a probe that needs to be sure must qualify.""" + + class Root: + def modules(self): + # A class named Attention, but from an unrelated architecture. + other = type("Attention", (), {}) + other.__module__ = "somewhere.else.entirely" + return [self, other()] + + class Model: + class model: + diffusion_model = Root() + + assert has_block(Model(), "Attention") # the collision, unqualified + assert not has_block(Model(), "comfy.ldm.minimax.model.Attention") + + +def test_a_qualified_probe_matches_its_own_architecture(): + class Root: + def modules(self): + target = type("AdalnProj", (), {}) + target.__module__ = "comfy.ldm.minimax.model" + return [self, target()] + + class Model: + class model: + diffusion_model = Root() + + assert has_block(Model(), "comfy.ldm.minimax.model.AdalnProj") + assert has_block(Model(), "minimax.model.AdalnProj") # tail match + assert has_block(Model(), "AdalnProj") # distinctive enough + assert not has_block(Model(), "other.pack.AdalnProj") diff --git a/core/traits.py b/core/traits.py index e818652f..d0fd16b2 100644 --- a/core/traits.py +++ b/core/traits.py @@ -60,6 +60,14 @@ def has_block(model: Any, class_name: str) -> bool: The structural probe model modules are expected to use. It asks "does this model have the thing I touch", which is the question a modifier actually has -- unlike "is this H3", it stays true for anything built the same way. + + Pass a DOTTED name to require the defining module too + (`"comfy.ldm.minimax.model.AdalnProj"`, matched on the tail of the module + path). A bare name matches on the class name alone, which is fine for + something distinctive like `AdalnProj` and actively wrong for names as + common as `Attention`, `MLP` or `Block` -- those collide across unrelated + architectures, and a false positive here enables a modifier on a model it + was never written for. When in doubt, qualify. """ root = getattr(getattr(model, "model", None), "diffusion_model", None) if root is None: @@ -68,8 +76,20 @@ def has_block(model: Any, class_name: str) -> bool: modules = root.modules() except Exception: # noqa: BLE001 -- not an nn.Module return False + + if "." in class_name: + wanted_module, _, wanted = class_name.rpartition(".") + else: + wanted_module, wanted = "", class_name + for sub in modules: - if type(sub).__name__ == class_name: + kind = type(sub) + if kind.__name__ != wanted: + continue + if not wanted_module: + return True + origin = getattr(kind, "__module__", "") or "" + if origin == wanted_module or origin.endswith("." + wanted_module): return True return False