app, core: a node's outputs are visible, and model family detection
Every node's own OUTPUTS were invisible on its own panel in the pipeline editor -- wiring an input already let you pick a source, but opening the node that PRODUCES something told you nothing about what it makes or whether anything reads it. Found live, by a real user just looking at a node, after two full review rounds that only ever checked the consuming side. Fixed with a read-only "Outputs" section per node, naming each output, its type, and who (if anyone) currently wires from it. Also: model family detection. core/probe.py reads a safetensors header (tensor names only, no weights) and asks every installed model module, through the same capability-lookup mechanism already used for empty_latent/decode, whether the file is its own -- core never names a family itself. A real, adversarially-found bug in the ported signature: one of MiniMax H3's three candidate prefixes doesn't correspond to any prefix ComfyUI's own loader actually resolves to, traced through comfy/model_detection.py and comfy/sd.py -- the probe could have called a file H3 that ComfyUI's real loader would refuse. Removed, and pinned with a dedicated test module. A second finding narrowed the folder search to only the folders FunPack's own loaders actually read from. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
digital-garbage
co-authored by
Claude Sonnet 5
Claude Sonnet 5
parent
75e9e9f686
commit
93bdd27b32
@@ -467,9 +467,43 @@ export function mount({ load, describe, check, search, onApply,
|
||||
}));
|
||||
}
|
||||
|
||||
// Outputs are never wired from HERE -- wiring is always asked for by the
|
||||
// input that wants one, through its own Wire button, so the same link
|
||||
// is never described two different ways depending on which end you
|
||||
// opened. But an output that is never SHOWN anywhere on its own node is
|
||||
// a real blind spot, not just a different workflow: open a node that
|
||||
// exists to produce things and there was nothing here saying what it
|
||||
// produces, or whether anything downstream already reads it.
|
||||
if (described.outputs && described.outputs.length) {
|
||||
rows.push(composer.label.section({ text: "Outputs" }));
|
||||
described.outputs.forEach((type, index) => {
|
||||
const name = (described.output_names || [])[index] || type;
|
||||
const consumers = consumersOf(slot.id, index);
|
||||
rows.push(composer.settingsRow.default({
|
||||
label: name,
|
||||
hint: consumers.length
|
||||
? `${type} — feeds ${consumers.join(", ")}`
|
||||
: `${type} — nothing reads this yet`,
|
||||
}));
|
||||
});
|
||||
}
|
||||
|
||||
return rows;
|
||||
}
|
||||
|
||||
/** Every (consumerId.inputName) currently wired to this slot's output. */
|
||||
function consumersOf(slotId, index) {
|
||||
const found = [];
|
||||
for (const other of slots) {
|
||||
for (const [name, value] of Object.entries(other.inputs || {})) {
|
||||
if (isLink(value) && value[0] === slotId && value[1] === index) {
|
||||
found.push(`${other.id}.${name}`);
|
||||
}
|
||||
}
|
||||
}
|
||||
return found;
|
||||
}
|
||||
|
||||
/** Every OTHER slot's output that could legally feed this socket. */
|
||||
function sourcesFor(slot, socket) {
|
||||
const found = [];
|
||||
|
||||
@@ -523,6 +523,58 @@ test("an unfed socket has no Unwire button to press", async () => {
|
||||
win.close();
|
||||
});
|
||||
|
||||
// --- a node's own outputs, shown on its own panel ---------------------------
|
||||
//
|
||||
// Wiring is always asked for by the CONSUMING input, never from the output's
|
||||
// own side -- but an output that is never shown anywhere on its own node's
|
||||
// panel is a real blind spot: open the node that PRODUCES something and
|
||||
// there was nothing here saying what it produces, or whether anything reads
|
||||
// it yet. Missed by two full review rounds because every one of them
|
||||
// reasoned from "does wiring an input work", never "can you see what a node
|
||||
// makes just by looking at it".
|
||||
|
||||
test("a node's own outputs are listed on its panel, named and typed", async () => {
|
||||
const { win } = await opened();
|
||||
win.enter("Loaders");
|
||||
const rows = [...win.node.querySelectorAll(".cx-settings-row")]
|
||||
.filter((r) => r.querySelector(".cx-settings-label")?.textContent === "model");
|
||||
// "model" is both an INPUT widget on other nodes and this node's OWN
|
||||
// output name -- the output row is the one with no input control, only
|
||||
// the read-only hint.
|
||||
const outputRow = rows.find((r) => !r.querySelector("input, select, textarea, button"));
|
||||
assert.ok(outputRow, "the Loader's own 'model' output is not shown at all");
|
||||
assert.match(outputRow.querySelector(".cx-hint")?.textContent ?? "", /MODEL/);
|
||||
win.close();
|
||||
});
|
||||
|
||||
test("an output already consumed says who reads it", async () => {
|
||||
const { win } = await opened();
|
||||
win.enter("Loaders");
|
||||
const rows = [...win.node.querySelectorAll(".cx-settings-row")]
|
||||
.filter((r) => r.querySelector(".cx-settings-label")?.textContent === "model");
|
||||
const outputRow = rows.find((r) => !r.querySelector("input, select, textarea, button"));
|
||||
assert.match(outputRow.querySelector(".cx-hint")?.textContent ?? "", /feeds sampler\.model/);
|
||||
win.close();
|
||||
});
|
||||
|
||||
test("an output nothing reads yet says so", async () => {
|
||||
const custom = server({ slots: [{ id: "s", group: "Sampling", node: "Sampler", inputs: {} }] });
|
||||
const win = openWindow(custom);
|
||||
await win.ready;
|
||||
win.enter("Sampling");
|
||||
const outputRow = rowFor(win, "LATENT");
|
||||
assert.match(outputRow.querySelector(".cx-hint")?.textContent ?? "", /nothing reads this yet/);
|
||||
win.close();
|
||||
});
|
||||
|
||||
test("a node with no outputs shows no Outputs section", async () => {
|
||||
const { win } = await opened();
|
||||
win.enter("Other"); // "save" has no group of its own
|
||||
assert.equal([...win.node.querySelectorAll(".cx-eyebrow")]
|
||||
.some((el) => el.textContent === "Outputs"), false);
|
||||
win.close();
|
||||
});
|
||||
|
||||
// --- wiring a WIDGET input -- "linked inputs" -------------------------------
|
||||
//
|
||||
// A widget-typed input (a number, a string, a combo) can hold a link too --
|
||||
|
||||
@@ -161,6 +161,24 @@ test("a widget input -- a linked input -- can be wired to another node's output,
|
||||
await expect(textRow.locator("textarea, input")).toBeVisible();
|
||||
});
|
||||
|
||||
test("a node's own outputs are shown on its own panel, named and consumed", async ({ page }) => {
|
||||
// Wiring is always asked for by the consuming input; that does not mean an
|
||||
// output should be invisible on the node that PRODUCES it. Missed by two
|
||||
// review rounds, both reasoning only from the consuming side -- found
|
||||
// live, by someone just looking at a node and not seeing what it makes.
|
||||
await openWindow(page);
|
||||
await page.locator(".cx-card", { hasText: "Loaders" }).click();
|
||||
await page.locator('[aria-label="Nodes"] .cx-filter-row', { hasText: "model" }).click();
|
||||
|
||||
const outputsHeading = page.locator(".cx-eyebrow", { hasText: "Outputs" });
|
||||
await expect(outputsHeading).toBeVisible();
|
||||
const modelOutputRow = page.locator(".cx-settings-row")
|
||||
.filter({ has: page.locator(".cx-settings-label", { hasText: /^model$/ }) })
|
||||
.filter({ hasNot: page.locator("input, select, textarea, button") });
|
||||
await expect(modelOutputRow).toContainText("MODEL");
|
||||
await expect(modelOutputRow).toContainText(/feeds/);
|
||||
});
|
||||
|
||||
test("the Add-node picker has one search box, not two disagreeing ones", async ({ page }) => {
|
||||
// filterList draws its own "Filter these results" search box; this picker
|
||||
// already has an outer one that re-queries the server on every keystroke.
|
||||
|
||||
+114
@@ -0,0 +1,114 @@
|
||||
"""Which model family a checkpoint file is, without loading it.
|
||||
|
||||
Core reads the file -- a safetensors header is a length-prefixed JSON object
|
||||
whose keys are tensor names, readable in kilobytes regardless of whether the
|
||||
file itself is 2GB or 40GB -- and that is ALL it knows how to do. Whether a
|
||||
set of tensor names belongs to MiniMax H3, or anything else, is a model
|
||||
module's own business: asking one to identify itself is the same "detect"
|
||||
capability every model module can already offer, the same shape as the
|
||||
`empty_latent` and `decode` capabilities modules/models/minimax_h3 provides.
|
||||
Core aggregates whichever module claims a file and says plainly when none
|
||||
does -- it is never the one deciding what a family IS called.
|
||||
|
||||
Only .safetensors today: no loader in this pack can read a GGUF file yet, so
|
||||
detecting one accurately would be groundwork for a format nothing here can
|
||||
use.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import struct
|
||||
from pathlib import Path
|
||||
from typing import List, Optional
|
||||
|
||||
from . import registry as registry_mod
|
||||
|
||||
#: A header past this size is not a safetensors file -- reading a wild length
|
||||
#: prefix as a real one would try to allocate however many gigabytes a
|
||||
#: corrupt or unrelated file's first eight bytes happen to spell out.
|
||||
MAX_HEADER_BYTES = 256 * 1024 * 1024
|
||||
|
||||
|
||||
def read_safetensors_keys(path: Path) -> Optional[List[str]]:
|
||||
"""Tensor names out of a safetensors header, or None when this is not one.
|
||||
|
||||
Layout: an 8-byte little-endian length, then that many bytes of JSON whose
|
||||
top-level keys are the tensor names (plus `__metadata__`, dropped here).
|
||||
Nothing past the header is read.
|
||||
"""
|
||||
try:
|
||||
with open(path, "rb") as f:
|
||||
raw = f.read(8)
|
||||
if len(raw) < 8:
|
||||
return None
|
||||
(length,) = struct.unpack("<Q", raw)
|
||||
if length <= 0 or length > MAX_HEADER_BYTES:
|
||||
return None
|
||||
header = json.loads(f.read(length).decode("utf-8", errors="replace"))
|
||||
except (OSError, ValueError, json.JSONDecodeError, struct.error):
|
||||
return None
|
||||
if not isinstance(header, dict):
|
||||
return None
|
||||
return [k for k in header if k != "__metadata__"]
|
||||
|
||||
|
||||
def resolve_diffusion_model(filename: str) -> Optional[Path]:
|
||||
"""Where `filename` lives, across the folders a FunPack loader can
|
||||
actually read a diffusion model from -- `modules/loaders/diffusion_model`
|
||||
offers files from "diffusion_models", `modules/loaders/checkpoint` from
|
||||
"checkpoints"; nothing in this app reads a model from a "unet" folder,
|
||||
so probing one would risk answering for a file no loader here would ever
|
||||
actually load under that name. None when it is nowhere findable -- not
|
||||
this function's job to say why.
|
||||
"""
|
||||
try:
|
||||
import folder_paths
|
||||
except ImportError: # not inside ComfyUI
|
||||
return None
|
||||
for folder in ("diffusion_models", "checkpoints"):
|
||||
try:
|
||||
found = folder_paths.get_full_path(folder, filename)
|
||||
except Exception: # noqa: BLE001
|
||||
found = None
|
||||
if found:
|
||||
return Path(found)
|
||||
return None
|
||||
|
||||
|
||||
def detect(path: Path, registry=None) -> dict:
|
||||
"""Inspect one checkpoint file.
|
||||
|
||||
Returns ``{"module", "title", "detected", "reason"}``. `module` is the id
|
||||
of whichever model module claimed it, or None -- and None is a real
|
||||
answer, not a failure to try hard enough: it means this install has no
|
||||
module that recognises the file, and the caller should say so rather
|
||||
than guess a family. `reason` is plain text for a person, always
|
||||
populated.
|
||||
"""
|
||||
p = Path(path)
|
||||
if not p.is_file():
|
||||
return {"module": None, "title": None, "detected": False,
|
||||
"reason": f"{p.name}: file not found"}
|
||||
if p.suffix.lower() != ".safetensors":
|
||||
return {"module": None, "title": None, "detected": False,
|
||||
"reason": f"{p.name}: only .safetensors can be inspected without loading it"}
|
||||
keys = read_safetensors_keys(p)
|
||||
if keys is None:
|
||||
return {"module": None, "title": None, "detected": False,
|
||||
"reason": f"{p.name}: not a readable safetensors header"}
|
||||
|
||||
keyset = set(keys)
|
||||
reg = registry or registry_mod.current()
|
||||
for spec, claims in reg.providers("detect"):
|
||||
try:
|
||||
matched = claims(keyset)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
from . import log
|
||||
log.failed(f"{spec.id}.detect", exc)
|
||||
continue
|
||||
if matched:
|
||||
return {"module": spec.id, "title": spec.title, "detected": True,
|
||||
"reason": f"{p.name}: {spec.title}"}
|
||||
return {"module": None, "title": None, "detected": False,
|
||||
"reason": f"{p.name}: no installed model module recognises this file"}
|
||||
+24
-1
@@ -8,7 +8,8 @@ thin adapters over pure functions in `serve`.
|
||||
import asyncio
|
||||
import json
|
||||
|
||||
from . import (backend_log, config, graph as graph_mod, log, media, nodes_manager, projects,
|
||||
from . import (backend_log, config, graph as graph_mod, log, media, nodes_manager, probe as probe_mod,
|
||||
projects,
|
||||
sysinfo,
|
||||
temp_files,
|
||||
update as update_mod,
|
||||
@@ -447,6 +448,28 @@ def register(routes, prefix=None):
|
||||
payload = await asyncio.to_thread(sysinfo.collect)
|
||||
return web.json_response(payload)
|
||||
|
||||
@routes.get(P + "/api/probe")
|
||||
async def _probe(req):
|
||||
"""Which model family a checkpoint file is, before anything loads it.
|
||||
|
||||
Reads a safetensors header off disk -- to_thread because that is
|
||||
real (if small) file I/O -- and asks every installed model module
|
||||
whether the tensor names are its own. Nothing here names a family:
|
||||
that answer comes back from whichever module claimed it.
|
||||
"""
|
||||
filename = (req.query.get("file") or "").strip()
|
||||
if not filename:
|
||||
return web.json_response(
|
||||
{"problems": ["which file is named by ?file="]}, status=400)
|
||||
path = probe_mod.resolve_diffusion_model(filename)
|
||||
if path is None:
|
||||
return web.json_response({
|
||||
"module": None, "title": None, "detected": False,
|
||||
"reason": f"{filename}: not found in the models folders",
|
||||
})
|
||||
payload = await asyncio.to_thread(probe_mod.detect, path)
|
||||
return web.json_response(payload)
|
||||
|
||||
@routes.post(P + "/api/git/update")
|
||||
async def _git_pull(req):
|
||||
body = await req.json() if req.can_read_body else {}
|
||||
|
||||
@@ -0,0 +1,186 @@
|
||||
"""Model family detection: a safetensors header read cold, then handed to
|
||||
whichever model module claims it. Core never names a family -- these tests
|
||||
prove that split holds, not any one architecture's signature (that belongs
|
||||
to modules/models/minimax_h3's own tests).
|
||||
"""
|
||||
|
||||
import json
|
||||
import struct
|
||||
import sys
|
||||
import types
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from core import probe
|
||||
from core.contract import ModuleSpec
|
||||
from core.registry import Registry
|
||||
|
||||
|
||||
def _write_safetensors(path, keys):
|
||||
header = {k: {"dtype": "F16", "shape": [1], "data_offsets": [0, 0]} for k in keys}
|
||||
header["__metadata__"] = {"format": "pt"}
|
||||
body = json.dumps(header).encode("utf-8")
|
||||
path.write_bytes(struct.pack("<Q", len(body)) + body)
|
||||
|
||||
|
||||
# --- reading the header -----------------------------------------------------
|
||||
|
||||
def test_keys_are_read_without_touching_anything_past_the_header(tmp_path):
|
||||
p = tmp_path / "m.safetensors"
|
||||
_write_safetensors(p, ["video_patch_proj.weight", "unrelated.bias"])
|
||||
keys = probe.read_safetensors_keys(p)
|
||||
assert set(keys) == {"video_patch_proj.weight", "unrelated.bias"}
|
||||
|
||||
|
||||
def test_metadata_is_not_a_tensor_name(tmp_path):
|
||||
p = tmp_path / "m.safetensors"
|
||||
_write_safetensors(p, ["a.weight"])
|
||||
assert "__metadata__" not in probe.read_safetensors_keys(p)
|
||||
|
||||
|
||||
def test_a_file_too_short_to_hold_a_length_is_not_a_header(tmp_path):
|
||||
p = tmp_path / "m.safetensors"
|
||||
p.write_bytes(b"\x01\x02\x03")
|
||||
assert probe.read_safetensors_keys(p) is None
|
||||
|
||||
|
||||
def test_a_wild_length_is_refused_rather_than_allocated(tmp_path):
|
||||
p = tmp_path / "m.safetensors"
|
||||
# A length bigger than the ceiling, however the bytes got that way --
|
||||
# some other file format's first eight bytes, not a real header.
|
||||
p.write_bytes(struct.pack("<Q", probe.MAX_HEADER_BYTES + 1) + b"{}")
|
||||
assert probe.read_safetensors_keys(p) is None
|
||||
|
||||
|
||||
def test_a_header_that_is_not_json_is_not_a_header(tmp_path):
|
||||
p = tmp_path / "m.safetensors"
|
||||
body = b"not json"
|
||||
p.write_bytes(struct.pack("<Q", len(body)) + body)
|
||||
assert probe.read_safetensors_keys(p) is None
|
||||
|
||||
|
||||
def test_a_header_that_is_valid_json_but_not_an_object_is_not_a_header(tmp_path):
|
||||
p = tmp_path / "m.safetensors"
|
||||
body = json.dumps(["a", "list"]).encode()
|
||||
p.write_bytes(struct.pack("<Q", len(body)) + body)
|
||||
assert probe.read_safetensors_keys(p) is None
|
||||
|
||||
|
||||
# --- detect() -----------------------------------------------------------
|
||||
|
||||
def _registry(*claims):
|
||||
"""claims: (id, title, keys_it_recognises) triples."""
|
||||
reg = Registry()
|
||||
for module_id, title, recognises in claims:
|
||||
reg.add(ModuleSpec(id=module_id, title=title, mount="",
|
||||
provides={"detect": lambda keyset, r=recognises: r <= keyset}))
|
||||
return reg
|
||||
|
||||
|
||||
def test_a_file_no_module_recognises_says_so_rather_than_guessing(tmp_path):
|
||||
p = tmp_path / "m.safetensors"
|
||||
_write_safetensors(p, ["some.unrelated.weight"])
|
||||
out = probe.detect(p, registry=_registry(("m", "M", {"video_patch_proj.weight"})))
|
||||
assert out == {"module": None, "title": None, "detected": False,
|
||||
"reason": f"{p.name}: no installed model module recognises this file"}
|
||||
|
||||
|
||||
def test_a_file_a_module_recognises_names_that_module(tmp_path):
|
||||
p = tmp_path / "m.safetensors"
|
||||
_write_safetensors(p, ["video_patch_proj.weight", "audio_patch_proj.weight"])
|
||||
out = probe.detect(p, registry=_registry(("model_minimax_h3", "MiniMax H3",
|
||||
{"video_patch_proj.weight", "audio_patch_proj.weight"})))
|
||||
assert out["module"] == "model_minimax_h3"
|
||||
assert out["title"] == "MiniMax H3"
|
||||
assert out["detected"] is True
|
||||
assert "MiniMax H3" in out["reason"]
|
||||
|
||||
|
||||
def test_a_missing_file_is_not_found_not_undetected(tmp_path):
|
||||
out = probe.detect(tmp_path / "nope.safetensors", registry=_registry())
|
||||
assert out["detected"] is False
|
||||
assert "not found" in out["reason"]
|
||||
|
||||
|
||||
def test_a_non_safetensors_extension_is_refused_without_reading_it(tmp_path):
|
||||
p = tmp_path / "m.ckpt"
|
||||
p.write_bytes(b"whatever a pickle looks like")
|
||||
out = probe.detect(p, registry=_registry())
|
||||
assert out["detected"] is False
|
||||
assert "only .safetensors" in out["reason"]
|
||||
|
||||
|
||||
def test_the_first_module_that_claims_it_wins_in_a_stable_order(tmp_path):
|
||||
"""providers() is sorted by module id -- deterministic, not import order."""
|
||||
p = tmp_path / "m.safetensors"
|
||||
_write_safetensors(p, ["shared.weight"])
|
||||
out = probe.detect(p, registry=_registry(
|
||||
("z_module", "Z", {"shared.weight"}),
|
||||
("a_module", "A", {"shared.weight"}),
|
||||
))
|
||||
assert out["module"] == "a_module"
|
||||
|
||||
|
||||
def test_a_detector_that_raises_is_skipped_not_fatal(tmp_path):
|
||||
p = tmp_path / "m.safetensors"
|
||||
_write_safetensors(p, ["real.weight"])
|
||||
reg = Registry()
|
||||
def explodes(_keys):
|
||||
raise RuntimeError("boom")
|
||||
reg.add(ModuleSpec(id="bad", title="Bad", mount="", provides={"detect": explodes}))
|
||||
reg.add(ModuleSpec(id="good", title="Good", mount="",
|
||||
provides={"detect": lambda keys: "real.weight" in keys}))
|
||||
out = probe.detect(p, registry=reg)
|
||||
assert out["module"] == "good"
|
||||
|
||||
|
||||
# --- resolve_diffusion_model -------------------------------------------------
|
||||
#
|
||||
# `folder_paths` only exists inside a running ComfyUI, and whether it happens
|
||||
# to be importable here depends on which OTHER tests already ran and pulled
|
||||
# ComfyUI's source onto sys.path (core/tests/conftest.py's `comfyui` fixture
|
||||
# does this session-wide, not just for the test that asked). Asserting on
|
||||
# ambient absence would make this test's result depend on run order rather
|
||||
# than on what resolve_diffusion_model() actually does -- so a fake module is
|
||||
# installed instead of trusting whatever is or isn't already importable.
|
||||
|
||||
def _fake_folder_paths(monkeypatch, get_full_path):
|
||||
fake = types.SimpleNamespace(get_full_path=get_full_path)
|
||||
monkeypatch.setitem(sys.modules, "folder_paths", fake)
|
||||
|
||||
|
||||
def test_resolve_finds_a_file_in_the_diffusion_models_folder(monkeypatch):
|
||||
_fake_folder_paths(monkeypatch, lambda folder, name: (
|
||||
"/models/diffusion_models/x.safetensors" if folder == "diffusion_models" else None))
|
||||
assert probe.resolve_diffusion_model("x.safetensors") == Path("/models/diffusion_models/x.safetensors")
|
||||
|
||||
|
||||
def test_resolve_falls_back_to_checkpoints(monkeypatch):
|
||||
_fake_folder_paths(monkeypatch, lambda folder, name: (
|
||||
"/models/checkpoints/x.safetensors" if folder == "checkpoints" else None))
|
||||
assert probe.resolve_diffusion_model("x.safetensors") == Path("/models/checkpoints/x.safetensors")
|
||||
|
||||
|
||||
def test_resolve_never_looks_in_a_unet_folder(monkeypatch):
|
||||
"""No loader in this app reads a model from a "unet" folder --
|
||||
modules/loaders/diffusion_model offers "diffusion_models" files,
|
||||
modules/loaders/checkpoint offers "checkpoints" ones. Probing "unet" too
|
||||
would risk answering for a file no loader here would ever actually load
|
||||
under that name."""
|
||||
asked = []
|
||||
_fake_folder_paths(monkeypatch, lambda folder, name: (asked.append(folder), None)[1])
|
||||
probe.resolve_diffusion_model("x.safetensors")
|
||||
assert "unet" not in asked
|
||||
|
||||
|
||||
def test_resolve_is_none_when_no_folder_has_it(monkeypatch):
|
||||
_fake_folder_paths(monkeypatch, lambda folder, name: None)
|
||||
assert probe.resolve_diffusion_model("nope.safetensors") is None
|
||||
|
||||
|
||||
def test_resolve_is_none_when_folder_paths_itself_raises(monkeypatch):
|
||||
def boom(folder, name):
|
||||
raise KeyError(folder)
|
||||
_fake_folder_paths(monkeypatch, boom)
|
||||
assert probe.resolve_diffusion_model("x.safetensors") is None
|
||||
@@ -0,0 +1,124 @@
|
||||
"""Model family detection, over HTTP -- against the REAL registry, so this
|
||||
proves modules/models/minimax_h3's own detect() is actually reachable
|
||||
through core/registry.py's provider lookup, not just correct in isolation
|
||||
(core/tests/test_probe.py) or in the module's own tests.
|
||||
"""
|
||||
|
||||
import json
|
||||
import socket
|
||||
import struct
|
||||
import threading
|
||||
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("aiohttp")
|
||||
|
||||
from aiohttp import web # noqa: E402
|
||||
|
||||
from core import config, routes # noqa: E402
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def registered(comfyui):
|
||||
"""FunPack's modules importable the way the real registry scan needs --
|
||||
modules/models/minimax_h3 imports `comfy.*` at module level."""
|
||||
import asyncio
|
||||
import nodes as comfy_nodes
|
||||
|
||||
async def load():
|
||||
await comfy_nodes.init_extra_nodes(init_custom_nodes=False)
|
||||
|
||||
asyncio.run(load())
|
||||
return comfy_nodes
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def server(registered, tmp_path_factory):
|
||||
from aiohttp import web as aioweb
|
||||
|
||||
models_dir = tmp_path_factory.mktemp("models")
|
||||
|
||||
import folder_paths
|
||||
folder_paths.folder_names_and_paths["diffusion_models"] = ([str(models_dir)], {".safetensors"})
|
||||
|
||||
app = aioweb.Application()
|
||||
table = aioweb.RouteTableDef()
|
||||
routes.register(table, prefix="/funpack")
|
||||
app.add_routes(table)
|
||||
|
||||
holder = {}
|
||||
|
||||
def run():
|
||||
import asyncio
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
runner = aioweb.AppRunner(app)
|
||||
loop.run_until_complete(runner.setup())
|
||||
site = aioweb.TCPSite(runner, "127.0.0.1", 0)
|
||||
loop.run_until_complete(site.start())
|
||||
holder["port"] = site._server.sockets[0].getsockname()[1]
|
||||
holder["loop"] = loop
|
||||
holder["ready"].set()
|
||||
loop.run_forever()
|
||||
|
||||
holder["ready"] = threading.Event()
|
||||
holder["models_dir"] = models_dir
|
||||
thread = threading.Thread(target=run, daemon=True)
|
||||
thread.start()
|
||||
assert holder["ready"].wait(20), "the test server did not start"
|
||||
yield holder
|
||||
holder["loop"].call_soon_threadsafe(holder["loop"].stop)
|
||||
|
||||
|
||||
def _request(port, method, path):
|
||||
head = (f"{method} {path} HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n").encode()
|
||||
with socket.create_connection(("127.0.0.1", port), timeout=30) as sock:
|
||||
sock.sendall(head)
|
||||
chunks = []
|
||||
while True:
|
||||
got = sock.recv(65536)
|
||||
if not got:
|
||||
break
|
||||
chunks.append(got)
|
||||
header, _, rest = b"".join(chunks).partition(b"\r\n\r\n")
|
||||
status = int(header.split()[1])
|
||||
try:
|
||||
return status, json.loads(rest.decode() or "{}")
|
||||
except ValueError:
|
||||
return status, {"raw": rest.decode(errors="replace")}
|
||||
|
||||
|
||||
def _write_safetensors(path, keys):
|
||||
header = {k: {"dtype": "F16", "shape": [1], "data_offsets": [0, 0]} for k in keys}
|
||||
body = json.dumps(header).encode("utf-8")
|
||||
path.write_bytes(struct.pack("<Q", len(body)) + body)
|
||||
|
||||
|
||||
def test_a_real_h3_signature_is_recognised_through_the_real_registry(server):
|
||||
_write_safetensors(server["models_dir"] / "h3.safetensors",
|
||||
["video_patch_proj.weight", "audio_patch_proj.weight"])
|
||||
status, body = _request(server["port"], "GET", "/funpack/api/probe?file=h3.safetensors")
|
||||
assert status == 200, body
|
||||
assert body["module"] == "model_minimax_h3"
|
||||
assert body["detected"] is True
|
||||
|
||||
|
||||
def test_a_file_matching_no_known_signature_says_so(server):
|
||||
_write_safetensors(server["models_dir"] / "unknown.safetensors", ["some.other.weight"])
|
||||
status, body = _request(server["port"], "GET", "/funpack/api/probe?file=unknown.safetensors")
|
||||
assert status == 200, body
|
||||
assert body["module"] is None
|
||||
assert body["detected"] is False
|
||||
|
||||
|
||||
def test_a_file_that_does_not_exist_says_not_found(server):
|
||||
status, body = _request(server["port"], "GET", "/funpack/api/probe?file=nope.safetensors")
|
||||
assert status == 200
|
||||
assert body["detected"] is False
|
||||
assert "not found" in body["reason"]
|
||||
|
||||
|
||||
def test_no_file_named_is_a_400_not_a_guess(server):
|
||||
status, body = _request(server["port"], "GET", "/funpack/api/probe")
|
||||
assert status == 400
|
||||
assert body["problems"]
|
||||
@@ -44,11 +44,46 @@ VIDEO_CHANNELS = 24
|
||||
AUDIO_CHANNELS = 32
|
||||
VIDEO_SPATIAL_RATIO = 16
|
||||
|
||||
# Mirrors comfy/model_detection.py's own signature for this architecture: key
|
||||
# NAMES only, no shapes or values, so a checkpoint file can be identified by
|
||||
# its safetensors header alone -- without loading a single weight.
|
||||
#
|
||||
# Only the two prefixes ComfyUI's own loading path can actually end up
|
||||
# reading H3's tensors under: bare (a standalone unet-only file, where
|
||||
# `unet_prefix_from_state_dict`'s guess strips nothing and the original
|
||||
# unprefixed keys survive), and "model.diffusion_model." (the first, and
|
||||
# only H3-relevant, candidate that function tries for a full-checkpoint
|
||||
# layout). A bare "diffusion_model." (no "model." in front) was here too,
|
||||
# copied verbatim from the old v4 probe this was ported from -- traced
|
||||
# through comfy/model_detection.py and comfy/sd.py and confirmed it is not
|
||||
# reachable: ComfyUI's own prefix guess for that layout falls through to
|
||||
# "model.", strips nothing (no key starts with "model."), and the ORIGINAL
|
||||
# "diffusion_model."-prefixed keys are what detect_unet_config then searches
|
||||
# under key_prefix="" -- which never finds them, so ComfyUI's real loader
|
||||
# would refuse a file this probe would have called H3. A false "yes" is
|
||||
# exactly the failure this feature exists to prevent, so it is gone rather
|
||||
# than left in on the chance some file happens to use it.
|
||||
_PREFIXES = ("", "model.diffusion_model.")
|
||||
_SIGNATURE_KEYS = ("video_patch_proj.weight", "audio_patch_proj.weight")
|
||||
|
||||
|
||||
def is_h3(model) -> bool:
|
||||
return has_block(model, MODEL_CLASS)
|
||||
|
||||
|
||||
def detect(keys) -> bool:
|
||||
"""Whether a safetensors file's tensor names are MiniMax H3's own.
|
||||
|
||||
Read before anything is loaded: this is what lets core refuse a
|
||||
mismatched checkpoint (or say what a file actually is) before queueing
|
||||
a run that would otherwise fail deep inside sampling, far from the
|
||||
mistake.
|
||||
"""
|
||||
keyset = set(keys)
|
||||
return any(all(f"{prefix}{key}" in keyset for key in _SIGNATURE_KEYS)
|
||||
for prefix in _PREFIXES)
|
||||
|
||||
|
||||
def traits(model):
|
||||
"""What core cannot read off the model for itself."""
|
||||
if not is_h3(model):
|
||||
@@ -112,4 +147,4 @@ def empty_latent(model, width, height, length, batch_size=1):
|
||||
|
||||
|
||||
TRAITS = traits
|
||||
PROVIDES = {"empty_latent": empty_latent, "decode": decode}
|
||||
PROVIDES = {"empty_latent": empty_latent, "decode": decode, "detect": detect}
|
||||
|
||||
@@ -0,0 +1,81 @@
|
||||
"""detect(): whether a set of tensor names is MiniMax H3's own.
|
||||
|
||||
Pinned against ComfyUI's real prefix-resolution algorithm
|
||||
(comfy/model_detection.py's unet_prefix_from_state_dict, and what
|
||||
comfy/sd.py's load_diffusion_model_state_dict does with its answer) rather
|
||||
than only against this module's own idea of what a "reasonable" prefix
|
||||
looks like -- a probe that says yes to a file ComfyUI's real loader would
|
||||
refuse is worse than one that says "not detected", which is exactly the
|
||||
failure a prior review round found and this file exists to pin shut.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _needs_comfy(comfyui):
|
||||
"""Only needed to import the module under test, which imports comfy.* at
|
||||
module level -- detect() itself touches no comfy API."""
|
||||
|
||||
|
||||
def _module():
|
||||
from modules.models import minimax_h3
|
||||
return minimax_h3
|
||||
|
||||
|
||||
def _keys(prefix, *, both=True):
|
||||
names = ["video_patch_proj.weight", "audio_patch_proj.weight"]
|
||||
if not both:
|
||||
names = names[:1]
|
||||
return {f"{prefix}{n}" for n in names}
|
||||
|
||||
|
||||
def test_bare_keys_are_recognised():
|
||||
"""A standalone unet-only file: ComfyUI's own prefix guess strips
|
||||
nothing for this layout (no key starts with "model."), so the ORIGINAL
|
||||
unprefixed keys are what its real loader ends up reading."""
|
||||
m = _module()
|
||||
assert m.detect(_keys(""))
|
||||
|
||||
|
||||
def test_model_diffusion_model_prefix_is_recognised():
|
||||
"""The first, and only H3-relevant, candidate
|
||||
unet_prefix_from_state_dict actually tries for a full-checkpoint
|
||||
layout."""
|
||||
m = _module()
|
||||
assert m.detect(_keys("model.diffusion_model."))
|
||||
|
||||
|
||||
def test_bare_diffusion_model_prefix_is_not_recognised():
|
||||
"""Found by adversarial review: this prefix is not one of the three real
|
||||
candidates ComfyUI's own unet_prefix_from_state_dict tries, and tracing
|
||||
what its fallback guess ("model.") does with a state dict keyed this way
|
||||
shows ComfyUI's real loader would never find these tensors either -- so
|
||||
this probe must not claim it can."""
|
||||
m = _module()
|
||||
assert not m.detect(_keys("diffusion_model."))
|
||||
|
||||
|
||||
def test_only_one_of_the_two_signature_keys_is_not_enough():
|
||||
m = _module()
|
||||
assert not m.detect(_keys("", both=False))
|
||||
assert not m.detect(_keys("model.diffusion_model.", both=False))
|
||||
|
||||
|
||||
def test_the_two_signature_keys_split_across_different_prefixes_do_not_match():
|
||||
"""One tensor under one prefix, the other under a different one -- not a
|
||||
real file's shape, and must not be read as one by accident."""
|
||||
m = _module()
|
||||
mixed = {"video_patch_proj.weight",
|
||||
"model.diffusion_model.audio_patch_proj.weight"}
|
||||
assert not m.detect(mixed)
|
||||
|
||||
|
||||
def test_unrelated_keys_are_not_recognised():
|
||||
m = _module()
|
||||
assert not m.detect({"model.diffusion_model.unrelated.weight", "vae.decoder.weight"})
|
||||
|
||||
|
||||
def test_an_empty_key_set_is_not_recognised():
|
||||
m = _module()
|
||||
assert not m.detect(set())
|
||||
Reference in New Issue
Block a user