diff --git a/core/contract.py b/core/contract.py index f4bab000..7a6a10c3 100644 --- a/core/contract.py +++ b/core/contract.py @@ -57,6 +57,10 @@ class ModuleSpec: # its own. This is how a model's own module teaches the system to recognise # it, so supporting a new model is a new folder and never an edit to core. traits: Optional[Callable] = None + # Named capabilities this module offers to OTHER modules. Core never reads a + # name out of here and never defines one: it is a lookup, so a node can ask + # "who can build one of these" without core learning what the thing is. + provides: Dict[str, Callable] = field(default_factory=dict) ui: Optional[str] = None # served path to its ui.js status: str = "experimental" # or "proven" source: str = "" # dotted import path diff --git a/core/nodes.py b/core/nodes.py index 4d23e96b..4ffef78b 100644 --- a/core/nodes.py +++ b/core/nodes.py @@ -30,7 +30,7 @@ def _schema_of(node) -> Optional[object]: def collect(registry=None) -> Tuple[List[type], List[Tuple[str, str]]]: """(nodes, rejected) -- rejected carries the reason, for the modules dump.""" - reg = registry if registry is not None else registry_mod.scan() + reg = registry if registry is not None else registry_mod.current() nodes: List[type] = [] rejected: List[Tuple[str, str]] = [] claimed = {} # node_id -> module id diff --git a/core/registry.py b/core/registry.py index 98179e74..7cfbe984 100644 --- a/core/registry.py +++ b/core/registry.py @@ -20,7 +20,7 @@ from .schema import SchemaError, validate FIELDS = { "id": "ID", "title": "TITLE", "mount": "MOUNT", "settings": "SETTINGS", "requires": "REQUIRES", "after": "AFTER", "before": "BEFORE", - "stage": "STAGE", "status": "STATUS", "nodes": "NODES", "traits": "TRAITS", + "stage": "STAGE", "status": "STATUS", "nodes": "NODES", "traits": "TRAITS", "provides": "PROVIDES", } @@ -33,6 +33,19 @@ class Registry: def ids(self) -> List[str]: return sorted(self.specs) + def providers(self, capability: str) -> List[Tuple[ModuleSpec, object]]: + """Every module offering this capability, in a stable order. + + Sorted by id so which provider answers first is a property of the tree + rather than of filesystem order. + """ + found = [] + for spec in sorted(self.specs.values(), key=lambda s: s.id): + fn = spec.provides.get(capability) + if fn is not None: + found.append((spec, fn)) + return found + def add(self, spec: ModuleSpec) -> None: if spec.id in self.specs: # Two modules answering to one id means one wins by import order, @@ -113,3 +126,15 @@ def scan(root: Optional[Path] = None, package: Optional[str] = None) -> Registry registry.failed.append((dotted, str(exc))) return registry + + +# Scanned once. A module list that changed under a running graph would mean two +# nodes in one run disagreeing about what exists. +_current: Optional[Registry] = None + + +def current(rescan: bool = False) -> Registry: + global _current + if _current is None or rescan: + _current = scan() + return _current diff --git a/core/routes.py b/core/routes.py index a1b3fd81..496cf8af 100644 --- a/core/routes.py +++ b/core/routes.py @@ -43,16 +43,10 @@ def _serve_under(req, root, allowed): return _respond(static.serve(root, req.match_info["tail"], allowed)) -# Scanned once at startup. A module list that changed under the UI mid-session -# would mean a panel could refer to something that is no longer there. -_registry = None - - def modules(rescan=False): - global _registry - if _registry is None or rescan: - _registry = registry_mod.scan() - return _registry + """The one shared scan. Held in `registry` so the nodes and the app cannot + end up describing different sets of modules in the same session.""" + return registry_mod.current(rescan) def manifest(traits=None): diff --git a/core/schema.py b/core/schema.py index d126a98e..2dac4fb8 100644 --- a/core/schema.py +++ b/core/schema.py @@ -169,6 +169,23 @@ def validate_traits(value: Any, module_id: str): return value +def validate_provides(value: Any, module_id: str) -> dict: + """Named capabilities: {name: callable}.""" + if value is None: + return {} + if not isinstance(value, dict): + raise SchemaError( + f"module {module_id!r}'s provides must be a dict of name -> callable; " + f"got {type(value).__name__}.") + for name, fn in value.items(): + if not isinstance(name, str) or not name.strip(): + raise SchemaError(f"module {module_id!r} provides a capability with no name.") + if not callable(fn): + raise SchemaError( + f"module {module_id!r}'s {name!r} capability is not callable.") + return dict(value) + + def validate(announcement: Dict[str, Any], source: str = "") -> ModuleSpec: """A validated ModuleSpec, or SchemaError naming what is wrong.""" for required in ("id", "title"): @@ -217,6 +234,7 @@ def validate(announcement: Dict[str, Any], source: str = "") -> ModuleSpec: return ModuleSpec( nodes=validate_nodes(announcement.get("nodes"), announcement["id"]), traits=validate_traits(announcement.get("traits"), announcement["id"]), + provides=validate_provides(announcement.get("provides"), announcement["id"]), id=announcement["id"], title=announcement["title"], mount=announcement.get("mount") or "", diff --git a/core/tests/test_loads_as_comfyui_does.py b/core/tests/test_loads_as_comfyui_does.py new file mode 100644 index 00000000..955f4f3f --- /dev/null +++ b/core/tests/test_loads_as_comfyui_does.py @@ -0,0 +1,85 @@ +"""Loading the pack exactly the way ComfyUI loads it. + +Every other test puts the repo root on sys.path, which makes `import core` work +and hides the one import rule that actually matters: ComfyUI loads a pack by FILE +LOCATION and never adds its directory to sys.path, so an absolute `from core +import ...` inside a module finds nothing. + +That gap has now cost the same bug twice -- a module imported cleanly under +pytest, failed silently under ComfyUI, and its nodes simply did not appear. The +suite could not see it, because the suite is the friendlier of the two +environments. This test is the unfriendly one. +""" + +import asyncio +import subprocess +import sys +from pathlib import Path + +import pytest + +ROOT = Path(__file__).resolve().parents[2] + +# Run in a subprocess with ONLY ComfyUI on the path -- not the pack root. In this +# process the root is already imported, which is exactly the leniency being +# tested for. +PROBE = """ +import asyncio, json, sys +sys.path.insert(0, {comfy!r}) +import nodes +ok = asyncio.run(nodes.load_custom_node({pack!r}, module_parent="custom_nodes")) +print("FUNPACK_RESULT " + json.dumps({{ + "ok": bool(ok), + "nodes": sorted(k for k in nodes.NODE_CLASS_MAPPINGS if k.startswith("FunPack")), +}})) +""" + + +@pytest.fixture(scope="module") +def loaded(comfyui): + # cwd MUST NOT be the pack root. Python puts the working directory on + # sys.path for `-c`, so running from the repo makes `import core` succeed and + # the whole test passes while the bug it exists for is present -- verified by + # reintroducing the bug and watching this go green. + result = subprocess.run( + [sys.executable, "-I", "-c", PROBE.format(comfy=comfyui, pack=str(ROOT))], + capture_output=True, text=True, timeout=600, cwd=comfyui, + ) + line = next((l for l in result.stdout.splitlines() if l.startswith("FUNPACK_RESULT ")), None) + if line is None: + pytest.fail(f"the probe did not report.\nstdout:\n{result.stdout}\nstderr:\n{result.stderr}") + import json + return json.loads(line[len("FUNPACK_RESULT "):]), result.stderr + + +def test_the_pack_loads(loaded): + report, _ = loaded + assert report["ok"], "ComfyUI refused the pack entirely" + + +def test_every_module_that_announces_a_node_actually_registers_it(loaded, comfyui): + """The real assertion: what ComfyUI ends up with must equal what the registry + says exists. A module that fails to import under ComfyUI is absent there and + present here, and only comparing the two can tell.""" + report, stderr = loaded + + from core import nodes as nodes_mod + from core import registry as registry_mod + expected = {n.GET_SCHEMA().node_id + for n in nodes_mod.collect(registry_mod.scan())[0]} + + missing = expected - set(report["nodes"]) + assert not missing, ( + f"these nodes exist under pytest but not under ComfyUI: {sorted(missing)}.\n" + f"Almost always an absolute `import core` in a module.\n" + f"loader stderr:\n{stderr}" + ) + + +def test_no_module_failed_to_import_under_comfyui(loaded): + _, stderr = loaded + # core.routes is expected to fail here: PromptServer has no instance outside + # a running server, and that failure is guarded on purpose. + bad = [line for line in stderr.splitlines() + if "failed to load" in line and "core.routes" not in line] + assert not bad, "a module failed to import under ComfyUI:\n" + "\n".join(bad) diff --git a/core/tests/test_shipped_nodes.py b/core/tests/test_shipped_nodes.py index 49af25c9..bb29733f 100644 --- a/core/tests/test_shipped_nodes.py +++ b/core/tests/test_shipped_nodes.py @@ -46,6 +46,7 @@ FROZEN = { "FunPackCLIPLoader": [ "clip_name1", "clip_name2", "clip_name3", "clip_name4", "type", "device", ], + "FunPackEmptyLatent": ["model", "width", "height", "length", "batch_size"], "FunPackDiffusionModelLoader": [ "model_name", "weight_dtype", "compute_dtype", "attention", "fp16_accumulation", ], diff --git a/modules/_core.py b/modules/_core.py new file mode 100644 index 00000000..76b26847 --- /dev/null +++ b/modules/_core.py @@ -0,0 +1,18 @@ +"""How a module reaches core, in both contexts, decided once. + +A module cannot simply `import core`. Inside ComfyUI the pack is imported by file +location, so `core` is only reachable beneath the pack's own package; under the +tests and the dev server the repo root IS on sys.path, so it is top-level. A +static relative import gets one of those right and breaks the other. + +Every module therefore imports core THROUGH here -- `from ..._core import traits` +-- and this file holds the only place that has to know about the difference. The +leading underscore keeps it out of the module scan. +""" + +try: # root on sys.path: tests, dev server + from core import log, registry, schema, traits +except ImportError: # inside ComfyUI: a subpackage of the pack + from ..core import log, registry, schema, traits + +__all__ = ["log", "registry", "schema", "traits"] diff --git a/modules/latent/__init__.py b/modules/latent/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/modules/latent/empty/__init__.py b/modules/latent/empty/__init__.py new file mode 100644 index 00000000..94bda339 --- /dev/null +++ b/modules/latent/empty/__init__.py @@ -0,0 +1,10 @@ +"""Empty latent. A node-only module.""" + +from .nodes import FunPackEmptyLatent + +ID = "latent_empty" +TITLE = "Empty latent" +STAGE = "latent" +STATUS = "proven" + +NODES = [FunPackEmptyLatent] diff --git a/modules/latent/empty/nodes.py b/modules/latent/empty/nodes.py new file mode 100644 index 00000000..f17a41d2 --- /dev/null +++ b/modules/latent/empty/nodes.py @@ -0,0 +1,106 @@ +"""An empty latent, shaped by whatever model is wired in. + +Two paths, and the order matters: + +1. **A model's own module may claim it.** Some latents cannot be derived from + anything ComfyUI publishes -- MiniMax H3's is a NestedTensor of a video and an + audio branch, its video channel count is not the `latent_channels` its format + reports, and its length snaps to a model-specific frame grid. So a model module + may provide `empty_latent` and build it itself. +2. **Otherwise it is derived** from `latent_format`, which is correct for every + single-tensor model: SDXL (4ch, /8), Wan21 (16ch, /8, /4) and LTXV (128ch, + /32, /8) all come out right, including models that ship after this is written. + +The consequence worth stating: nothing here knows the name of a single model, and +nothing here assumes video. A 2-D latent format yields an image latent because +that is what the model says it wants. +""" + +import comfy.model_management +import torch +from comfy_api.latest import io + +from ..._core import log, registry as registry_mod + +CAPABILITY = "empty_latent" + + +def latent_format_of(model): + inner = getattr(model, "model", None) + config = getattr(inner, "model_config", None) + return getattr(config, "latent_format", None) + + +def derive(model, width, height, length, batch_size): + """The generic shape, from what the model publishes about its own latent.""" + fmt = latent_format_of(model) + if fmt is None: + raise RuntimeError( + "This model does not publish a latent format, so its latent shape " + "cannot be worked out. A model module can supply one.") + + channels = int(getattr(fmt, "latent_channels", 4)) + spatial = int(getattr(fmt, "spacial_downscale_ratio", 8) or 8) + temporal = int(getattr(fmt, "temporal_downscale_ratio", 1) or 1) + rank = getattr(fmt, "latent_dimensions", 2) + + # Never below one: a small size divided by a large ratio floors to zero, and + # a zero-sized latent fails somewhere much further along. + w = max(1, width // spatial) + h = max(1, height // spatial) + device = comfy.model_management.intermediate_device() + + if rank == 3: + frames = ((max(1, length) - 1) // temporal) + 1 + samples = torch.zeros([batch_size, channels, frames, h, w], device=device) + elif rank == 1: + samples = torch.zeros([batch_size, channels, max(1, length)], device=device) + else: + samples = torch.zeros([batch_size, channels, h, w], device=device) + + # These say what ratios the latent was BUILT at, so the sampler can rescale + # if the model wants different ones (comfy.sample.fix_empty_latent_channels). + # Reporting what we actually used is right by construction. + return {"samples": samples, + "downscale_ratio_spacial": spatial, + "downscale_ratio_temporal": temporal} + + +class FunPackEmptyLatent(io.ComfyNode): + @classmethod + def define_schema(cls) -> io.Schema: + return io.Schema( + node_id="FunPackEmptyLatent", + display_name="FunPack Empty Latent", + category="FunPack/Latent", + description="An empty latent shaped for the wired model, image or video.", + inputs=[ + io.Model.Input("model", tooltip="The latent's shape comes from this model."), + io.Int.Input("width", default=768, min=16, max=16384, step=16), + io.Int.Input("height", default=512, min=16, max=16384, step=16), + io.Int.Input("length", default=1, min=1, max=16384, + tooltip="Frames. Ignored by models whose latent has no time axis."), + io.Int.Input("batch_size", default=1, min=1, max=4096, optional=True), + ], + outputs=[ + io.Latent.Output(display_name="latent"), + io.String.Output(display_name="status"), + ], + ) + + @classmethod + def execute(cls, model, width: int, height: int, length: int, + batch_size: int = 1) -> io.NodeOutput: + for spec, build in registry_mod.current().providers(CAPABILITY): + try: + claimed = build(model, width=width, height=height, + length=length, batch_size=batch_size) + except Exception as exc: # noqa: BLE001 + # One module's mistake must not decide the shape for every model. + log.failed(f"{spec.id}.{CAPABILITY}", exc) + continue + if claimed is not None: + return io.NodeOutput(claimed, f"{spec.id} built this latent") + + return io.NodeOutput(derive(model, width, height, length, batch_size), + "shaped from the model's own latent format") diff --git a/modules/latent/tests/test_empty_latent.py b/modules/latent/tests/test_empty_latent.py new file mode 100644 index 00000000..ad504ebb --- /dev/null +++ b/modules/latent/tests/test_empty_latent.py @@ -0,0 +1,145 @@ +"""The empty latent: derived where it can be, claimed where it cannot. + +The derivation is checked against ComfyUI's REAL model configs and against the +shapes core's own nodes produce, because a latent of the wrong shape does not +fail here -- it fails at the first sampling step, or worse, samples something +subtly wrong. +""" + +import pytest + + +@pytest.fixture(autouse=True) +def _needs_comfy(comfyui): + """These import comfy and the node modules.""" + + +class _Model: + """The shape the node reads: patcher -> model -> model_config.""" + + def __init__(self, config): + class Inner: + model_config = config + self.model = Inner() + + +def _config_named(name): + import comfy.supported_models as sm + for cfg in sm.models: + if cfg.__name__ == name: + return cfg + pytest.skip(f"upstream no longer ships {name}") + + +@pytest.mark.parametrize("name,width,height,length,expected", [ + # channels, then the spatial dims the format's own ratio implies. + ("SDXL", 1024, 1024, 1, (1, 4, 128, 128)), + ("WAN21_T2V", 512, 512, 81, (1, 16, 21, 64, 64)), + ("LTXV", 768, 512, 97, (1, 128, 13, 16, 24)), +]) +def test_the_derived_shape_matches_the_models_own_format(name, width, height, length, expected): + from modules.latent.empty.nodes import derive + out = derive(_Model(_config_named(name)), width, height, length, 1) + assert tuple(out["samples"].shape) == expected + + +def test_ltx_matches_what_comfys_own_node_produces(): + """The strongest available check short of sampling: the same numbers core's + EmptyLTXVLatentVideo computes, from a completely different route.""" + import torch + from modules.latent.empty.nodes import derive + + width, height, length = 768, 512, 97 + theirs = torch.zeros([1, 128, ((length - 1) // 8) + 1, height // 32, width // 32]) + ours = derive(_Model(_config_named("LTXV")), width, height, length, 1) + + assert tuple(ours["samples"].shape) == tuple(theirs.shape) + assert ours["downscale_ratio_spacial"] == 32 + + +def test_an_image_model_gets_an_image_latent_with_no_family_branch(): + """The not-video-only property, stated as a test: nothing in the node asks + what kind of model this is.""" + from modules.latent.empty.nodes import derive + out = derive(_Model(_config_named("SDXL")), 512, 512, 99, 1) + assert out["samples"].ndim == 4, "a length was supplied and grew a time axis anyway" + + +def test_a_tiny_canvas_never_yields_a_zero_sized_latent(): + """height // 32 floors to 0 on a small canvas, and a zero-sized latent fails + a long way from here.""" + from modules.latent.empty.nodes import derive + out = derive(_Model(_config_named("LTXV")), 16, 16, 1, 1) + assert all(dim >= 1 for dim in out["samples"].shape) + + +def test_a_model_that_publishes_no_latent_format_says_so(): + from modules.latent.empty.nodes import derive + + class Bare: + class model: + model_config = None + + with pytest.raises(RuntimeError, match="latent format"): + derive(Bare(), 64, 64, 1, 1) + + +def test_the_ratios_reported_are_the_ratios_used(): + """The sampler rescales when these disagree with the model's own, so + reporting anything other than what we built at causes a silent resize.""" + from modules.latent.empty.nodes import derive + cfg = _config_named("WAN21_T2V") + out = derive(_Model(cfg), 512, 512, 81, 1) + assert out["downscale_ratio_spacial"] == cfg.latent_format.spacial_downscale_ratio + assert out["downscale_ratio_temporal"] == cfg.latent_format.temporal_downscale_ratio + + +# --- the claimed path ----------------------------------------------------- + +def test_a_model_module_can_claim_a_latent_core_could_not_derive(): + """H3's latent is a NestedTensor whose video channel count is not the + latent_channels its format reports. Deriving it is not possible; claiming it + is the whole point of a model-support module.""" + from modules.models import minimax_h3 + + built = minimax_h3.empty_latent(_H3Model(), width=1344, height=768, length=124) + samples = built["samples"] + assert getattr(samples, "is_nested", False), "not a nested latent" + + video, audio = samples.unbind() + assert video.shape[1] == 24, "used latent_channels instead of the video channel count" + assert audio.shape[1] == 32 + assert video.shape[-2:] == (768 // 16, 1344 // 16) + + +def test_the_h3_module_declines_models_that_are_not_h3(): + """A provider that claimed everything would take the generic path away from + every other model.""" + from modules.models import minimax_h3 + assert minimax_h3.empty_latent(_Model(_config_named("SDXL")), 512, 512, 1) is None + assert tuple(minimax_h3.traits(_Model(_config_named("SDXL")))) == () + + +class _H3Model: + """A model tree carrying H3's own classes, so the structural probe fires for + the reason it would fire on a real one.""" + + def __init__(self): + import comfy.ldm.minimax.model as h3 + + class Root: + def modules(self_inner): + return [_FakeOf(h3.MiniMaxH3Model), _FakeOf(h3.AdalnProj)] + + class Inner: + diffusion_model = Root() + model_config = None + + self.model = Inner() + + +class _FakeOf: + """An object reporting another class's identity, so no weights are needed.""" + + def __init__(self, cls): + self.__class__ = type(cls.__name__, (object,), {"__module__": cls.__module__}) diff --git a/modules/models/__init__.py b/modules/models/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/modules/models/minimax_h3/__init__.py b/modules/models/minimax_h3/__init__.py new file mode 100644 index 00000000..f4446f94 --- /dev/null +++ b/modules/models/minimax_h3/__init__.py @@ -0,0 +1,80 @@ +"""MiniMax H3 support. + +The first model-support module, and the shape every later one should copy: it +teaches the system to recognise H3 and to build H3's latent, and NOTHING outside +this folder mentions H3 at all. Supporting the next model is another folder like +this one. + +Two things are contributed: + +* **Traits** -- what is true of this model that core cannot read for itself. Core + reads shape and prediction type off any model; that H3 carries an audio stream, + and that it exposes per-modality AdaLN projections, are facts about this + architecture. +* **An empty latent** -- because H3's genuinely cannot be derived. Its latent is a + NestedTensor of a video branch and an audio branch; the video channel count (24) + is not the `latent_channels` its format reports (32); and its length snaps to a + frame grid at a fixed frame rate. None of that is in `latent_format`. + +The frame grid is imported from ComfyUI rather than copied. Copying it would be a +second definition of the same rule, and the two would drift the first time +upstream touched theirs. +""" + +import comfy.model_management +import torch +from comfy.nested_tensor import NestedTensor +from comfy_extras.nodes_minimax_h3 import temporal_shape + +from ..._core import traits as _traits + +has_block = _traits.has_block + +ID = "model_minimax_h3" +TITLE = "MiniMax H3" +STAGE = "load" +STATUS = "experimental" + +# Qualified on purpose: a bare class name matches anything built with a class of +# that name, and this decides which model gets H3's latent. +MODEL_CLASS = "comfy.ldm.minimax.model.MiniMaxH3Model" +ADALN_CLASS = "comfy.ldm.minimax.model.AdalnProj" + +VIDEO_CHANNELS = 24 +AUDIO_CHANNELS = 32 +VIDEO_SPATIAL_RATIO = 16 + + +def is_h3(model) -> bool: + return has_block(model, MODEL_CLASS) + + +def traits(model): + """What core cannot read off the model for itself.""" + if not is_h3(model): + return () + found = ["audio_stream"] + if has_block(model, ADALN_CLASS): + found.append("adaln_modalities") + return found + + +def empty_latent(model, width, height, length, batch_size=1): + """H3's joint video+audio latent, or None if this is not an H3 model.""" + if not is_h3(model): + return None + + frame_count, latent_t, audio_t = temporal_shape(length) + device = comfy.model_management.intermediate_device() + video = torch.zeros( + [batch_size, VIDEO_CHANNELS, latent_t, + max(1, height // VIDEO_SPATIAL_RATIO), max(1, width // VIDEO_SPATIAL_RATIO)], + device=device) + audio = torch.zeros([batch_size, AUDIO_CHANNELS, 2, audio_t], device=device) + # Nested, not concatenated: the two branches have different ranks, and the + # sampler reads them as one latent with two parts. + return {"samples": NestedTensor((video, audio))} + + +TRAITS = traits +PROVIDES = {"empty_latent": empty_latent}