latent: an empty latent the model decides the shape of
Derived from latent_format where that is possible, which is every single-tensor model: SDXL, Wan21 and LTXV all come out matching the numbers core's own nodes compute, including models that ship after this. H3 cannot be derived and is the reason modules exist. Its latent 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 frame grid. So modules/models/minimax_h3 claims it and builds it, importing the grid from ComfyUI rather than copying a second definition of it. Nothing outside that folder names a model, and nothing here assumes video: a 2-D latent format yields an image latent with no branch anywhere. Modules reach core through modules/_core.py. ComfyUI imports a pack by file location, so an absolute `import core` in a module resolves under pytest and not under ComfyUI -- the node then never registers, silently. That cost the same bug twice, so there is now a test that loads the pack the way ComfyUI does, in an isolated subprocess whose cwd is NOT the repo. The first version of that test passed with the bug reintroduced, because cwd was on sys.path. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 5
parent
20a4c39db3
commit
fa2caa8860
@@ -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
|
||||
|
||||
+1
-1
@@ -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
|
||||
|
||||
+26
-1
@@ -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
|
||||
|
||||
+3
-9
@@ -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):
|
||||
|
||||
@@ -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 "",
|
||||
|
||||
@@ -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)
|
||||
@@ -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",
|
||||
],
|
||||
|
||||
@@ -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"]
|
||||
@@ -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]
|
||||
@@ -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")
|
||||
@@ -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__})
|
||||
@@ -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}
|
||||
Reference in New Issue
Block a user