WildAi
ec9640f34b
The audit artifacts and the plan moved to the local .dev/docs tree, so the references in the code, the tools and the tests follow them there. The license file names the tool that produces the record instead of a path.
344 lines
13 KiB
Python
344 lines
13 KiB
Python
"""P2.2 - one-off equivalence run: vendored vs. pristine diffusers 0.11.0.
|
|
|
|
The plan (P2.2) asks for a **recorded** equivalence run: the pristine
|
|
``diffusers==0.11.0`` classes and the vendored copy must have the same
|
|
``state_dict`` keys and produce identical outputs on a seeded input
|
|
(tolerance 0). The result is written to
|
|
``.dev/docs/audit/2026-10-22-diffusers-equivalence.md``.
|
|
|
|
Why this is a *tool* and not a plain test: the pristine diffusers sources are
|
|
not part of the repository. The hard environment rules of this phase forbid
|
|
creating a venv and forbid installing into the environment, so the pristine
|
|
package is used **as source**:
|
|
|
|
1. ``pip download diffusers==0.11.0 --no-deps --no-binary :all:`` into a
|
|
temp dir (the only pip operation the plan allows), then unpack it;
|
|
2. the unpacked ``src`` directory is put in front of ``sys.path``;
|
|
3. the ``diffusers`` *package* is registered by hand with its real
|
|
``__path__``, so the pristine ``diffusers/models/attention.py`` and its
|
|
real sibling modules are loaded without executing the 0.11.0 pipeline
|
|
imports (the pipelines need a ``transformers`` older than the one
|
|
installed here);
|
|
4. the ``huggingface_hub`` names that recent releases dropped
|
|
(``cached_download``, ``HfFolder``, ...) are injected into the already
|
|
imported ``huggingface_hub`` module - nothing is installed and nothing is
|
|
written outside the temp dir.
|
|
|
|
Usage (embedded python):
|
|
|
|
python tools/diffusers_equivalence.py --diffusers-src <unpacked>/src
|
|
|
|
Without ``--diffusers-src`` the tool exits with code 2 ("pristine source not
|
|
available") - the unit test then skips with a message that names the
|
|
argument.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import importlib.util
|
|
import json
|
|
import re
|
|
import sys
|
|
import types
|
|
from pathlib import Path
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parent.parent
|
|
VENDORED_FILE = REPO_ROOT / "modules" / "deps" / "basicsr" / "archs" / "_diffusers_compat.py"
|
|
DEFAULT_MARKDOWN = REPO_ROOT / ".dev" / "docs" / "audit" / "2026-10-22-diffusers-equivalence.md"
|
|
|
|
#: The classes that are vendored (the helper classes are pulled in by
|
|
#: ``FeedForward`` and are part of the copied block).
|
|
VENDORED_CLASSES = ("FeedForward", "GELU", "GEGLU", "ApproximateGELU", "AdaLayerNorm")
|
|
|
|
#: The ``huggingface_hub`` names diffusers 0.11.0 imports at module level.
|
|
HUB_NAMES = (
|
|
"HfFolder",
|
|
"cached_download",
|
|
"hf_hub_download",
|
|
"model_info",
|
|
"whoami",
|
|
"snapshot_download",
|
|
)
|
|
|
|
|
|
def _install_hub_shims(huggingface_hub) -> list[str]:
|
|
"""Add the missing ``huggingface_hub`` names; return the injected ones."""
|
|
injected: list[str] = []
|
|
for name in HUB_NAMES:
|
|
if hasattr(huggingface_hub, name):
|
|
continue
|
|
|
|
def placeholder(*args, _name=name, **kwargs): # pragma: no cover
|
|
raise RuntimeError(
|
|
f"huggingface_hub.{_name} is not available in this "
|
|
"huggingface_hub version - the diffusers 0.11.0 import only "
|
|
"needs the name to exist"
|
|
)
|
|
|
|
placeholder.__name__ = name
|
|
setattr(huggingface_hub, name, placeholder)
|
|
injected.append(name)
|
|
return injected
|
|
|
|
|
|
def _load_vendored():
|
|
"""Load the vendored compat module directly from its file."""
|
|
spec = importlib.util.spec_from_file_location("keep_diffusers_compat", VENDORED_FILE)
|
|
module = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
def class_blocks(text: str) -> dict[str, str]:
|
|
"""The source of every top-level ``class`` of ``text``."""
|
|
blocks: dict[str, str] = {}
|
|
matches = list(re.finditer(r"^class\s+(\w+)\(", text, re.MULTILINE))
|
|
for index, match in enumerate(matches):
|
|
end = matches[index + 1].start() if index + 1 < len(matches) else len(text)
|
|
blocks[match.group(1)] = text[match.start() : end].rstrip("\n")
|
|
return blocks
|
|
|
|
|
|
def run(diffusers_src: Path) -> dict:
|
|
"""Compare the pristine diffusers 0.11.0 classes with the vendored copy."""
|
|
import torch
|
|
|
|
src = Path(diffusers_src).resolve()
|
|
pristine_file = src / "diffusers" / "models" / "attention.py"
|
|
if not pristine_file.is_file():
|
|
raise FileNotFoundError(
|
|
f"{src} does not contain diffusers/models/attention.py - point "
|
|
"--diffusers-src at the unpacked 'src' directory of diffusers 0.11.0"
|
|
)
|
|
|
|
sys.path.insert(0, str(src))
|
|
|
|
import huggingface_hub # noqa: PLC0415 - optional dependency
|
|
|
|
injected = _install_hub_shims(huggingface_hub)
|
|
|
|
package = types.ModuleType("diffusers")
|
|
package.__path__ = [str(src / "diffusers")] # type: ignore[attr-defined]
|
|
package.__version__ = "0.11.0"
|
|
package.__file__ = str(src / "diffusers" / "__init__.py")
|
|
sys.modules["diffusers"] = package
|
|
|
|
import diffusers.models.attention as pristine_attention # noqa: PLC0415
|
|
from diffusers.models.attention import AdaLayerNorm as RealAdaLayerNorm
|
|
from diffusers.models.attention import FeedForward as RealFeedForward
|
|
|
|
version = getattr(package, "__version__", "unknown")
|
|
if version != "0.11.0":
|
|
raise RuntimeError(
|
|
f"the pristine diffusers 0.11.0 source was expected, got {version}"
|
|
)
|
|
|
|
vendored = _load_vendored()
|
|
|
|
results: dict = {
|
|
"diffusers_version": version,
|
|
"diffusers_file": str(Path(pristine_attention.__file__).resolve()),
|
|
"vendored_file": str(VENDORED_FILE.resolve()),
|
|
"huggingface_hub_shims": injected,
|
|
"source_blocks": [],
|
|
"cases": [],
|
|
}
|
|
|
|
# --- 1. the copied source must be byte-identical ----------------------
|
|
pristine_blocks = class_blocks(pristine_file.read_text(encoding="utf-8"))
|
|
vendored_blocks = class_blocks(VENDORED_FILE.read_text(encoding="utf-8"))
|
|
for name in VENDORED_CLASSES:
|
|
results["source_blocks"].append(
|
|
{
|
|
"class": name,
|
|
"byte_identical": pristine_blocks.get(name) == vendored_blocks.get(name),
|
|
"lines": len((vendored_blocks.get(name) or "").splitlines()),
|
|
}
|
|
)
|
|
|
|
# --- 2. FeedForward: state_dict + forward (tolerance 0) ---------------
|
|
torch.manual_seed(1234)
|
|
real_ff = RealFeedForward(8)
|
|
torch.manual_seed(1234)
|
|
vend_ff = vendored.FeedForward(8)
|
|
|
|
real_keys = list(real_ff.state_dict().keys())
|
|
vend_keys = list(vend_ff.state_dict().keys())
|
|
vend_ff.load_state_dict(real_ff.state_dict())
|
|
|
|
generator = torch.Generator().manual_seed(7)
|
|
sample = torch.randn((2, 8), generator=generator)
|
|
with torch.no_grad():
|
|
real_out = real_ff(sample)
|
|
vend_out = vend_ff(sample)
|
|
|
|
results["cases"].append(
|
|
{
|
|
"class": "FeedForward",
|
|
"state_dict_keys_equal": real_keys == vend_keys,
|
|
"state_dict_keys": real_keys,
|
|
"input_shape": list(sample.shape),
|
|
"output_shape": list(real_out.shape),
|
|
"shape_preserved": list(real_out.shape) == list(sample.shape),
|
|
"max_abs_diff": float((real_out - vend_out).abs().max()),
|
|
"bit_identical": bool(torch.equal(real_out, vend_out)),
|
|
}
|
|
)
|
|
|
|
# --- 3. AdaLayerNorm: state_dict + the inner computation --------------
|
|
torch.manual_seed(1234)
|
|
real_aln = RealAdaLayerNorm(8, 4)
|
|
torch.manual_seed(1234)
|
|
vend_aln = vendored.AdaLayerNorm(8, 4)
|
|
|
|
real_keys = list(real_aln.state_dict().keys())
|
|
vend_keys = list(vend_aln.state_dict().keys())
|
|
vend_aln.load_state_dict(real_aln.state_dict())
|
|
|
|
timestep = torch.tensor([0, 2, 3], dtype=torch.long)
|
|
with torch.no_grad():
|
|
real_emb = real_aln.linear(real_aln.silu(real_aln.emb(timestep)))
|
|
vend_emb = vend_aln.linear(vend_aln.silu(vend_aln.emb(timestep)))
|
|
|
|
# The *forward* of the pristine 0.11.0 class chunks the embedding on
|
|
# dim 0 (``torch.chunk(emb, 2)``), which only broadcasts for very
|
|
# specific shapes - the upstream behaviour is reproduced verbatim, so
|
|
# both classes must fail (or succeed) in exactly the same way.
|
|
probe_input = torch.randn((3, 8), generator=torch.Generator().manual_seed(11))
|
|
real_error = vend_error = None
|
|
try:
|
|
with torch.no_grad():
|
|
real_aln(probe_input, timestep)
|
|
except BaseException as exc: # noqa: BLE001 - the behaviour is the evidence
|
|
real_error = f"{type(exc).__name__}: {exc}"
|
|
try:
|
|
with torch.no_grad():
|
|
vend_aln(probe_input, timestep)
|
|
except BaseException as exc: # noqa: BLE001
|
|
vend_error = f"{type(exc).__name__}: {exc}"
|
|
|
|
results["cases"].append(
|
|
{
|
|
"class": "AdaLayerNorm",
|
|
"state_dict_keys_equal": real_keys == vend_keys,
|
|
"state_dict_keys": real_keys,
|
|
"embedding_output_shape": list(real_emb.shape),
|
|
"max_abs_diff_embedding": float((real_emb - vend_emb).abs().max()),
|
|
"bit_identical_embedding": bool(torch.equal(real_emb, vend_emb)),
|
|
"forward_error": real_error,
|
|
"forward_error_identical": real_error == vend_error,
|
|
"note": (
|
|
"diffusers 0.11.0's AdaLayerNorm.forward chunks the timestep "
|
|
"embedding on dim 0; the vendored copy reproduces that "
|
|
"behaviour byte-for-byte (the KEEP architecture never "
|
|
"enables it: use_ada_layer_norm is False without "
|
|
"num_embeds_ada_norm)"
|
|
),
|
|
}
|
|
)
|
|
|
|
results["equivalent"] = (
|
|
all(block["byte_identical"] for block in results["source_blocks"])
|
|
and results["cases"][0]["state_dict_keys_equal"]
|
|
and results["cases"][0]["bit_identical"]
|
|
and results["cases"][1]["state_dict_keys_equal"]
|
|
and results["cases"][1]["bit_identical_embedding"]
|
|
and results["cases"][1]["forward_error_identical"]
|
|
)
|
|
return results
|
|
|
|
|
|
def render_markdown(results: dict) -> str:
|
|
lines = [
|
|
"# P2.2 - diffusers 0.11.0 equivalence run",
|
|
"",
|
|
"Recorded evidence for the vendored copy of `FeedForward` and",
|
|
"`AdaLayerNorm` (`modules/deps/basicsr/archs/_diffusers_compat.py`).",
|
|
"",
|
|
"The pristine `diffusers==0.11.0` **source** was the reference: the",
|
|
"sdist was downloaded with `pip download diffusers==0.11.0 --no-deps",
|
|
"--no-binary :all:` into a temp dir, unpacked, and put in front of",
|
|
"`sys.path`. No virtual environment was created and nothing was",
|
|
"installed into the environment (the phase rules forbid both).",
|
|
"",
|
|
"Two kinds of evidence are recorded:",
|
|
"",
|
|
"1. **source level** - the copied class blocks are byte-identical to",
|
|
" the upstream 0.11.0 ones;",
|
|
"2. **runtime level** - same `state_dict` keys and bit-identical",
|
|
" outputs (tolerance 0) on a seeded input.",
|
|
"",
|
|
f"* pristine diffusers version: **{results['diffusers_version']}**",
|
|
f"* pristine file: `{results['diffusers_file']}`",
|
|
f"* vendored file: `{results['vendored_file']}`",
|
|
f"* `huggingface_hub` names injected (missing in this env): "
|
|
f"`{', '.join(results['huggingface_hub_shims']) or 'none'}`",
|
|
f"* overall equivalent: **{results['equivalent']}**",
|
|
"",
|
|
"## Source level",
|
|
"",
|
|
"| class | byte-identical to diffusers 0.11.0 | lines |",
|
|
"|---|---|---|",
|
|
]
|
|
for block in results["source_blocks"]:
|
|
lines.append(
|
|
f"| {block['class']} | {block['byte_identical']} | {block['lines']} |"
|
|
)
|
|
lines += [
|
|
"",
|
|
"## Runtime level",
|
|
"",
|
|
"| class | state_dict keys equal | max abs diff | bit-identical |",
|
|
"|---|---|---|---|",
|
|
]
|
|
for case in results["cases"]:
|
|
diff = case.get("max_abs_diff", case.get("max_abs_diff_embedding"))
|
|
identical = case.get("bit_identical", case.get("bit_identical_embedding"))
|
|
lines.append(
|
|
f"| {case['class']} | {case['state_dict_keys_equal']} | {diff} | {identical} |"
|
|
)
|
|
lines.append("")
|
|
for case in results["cases"]:
|
|
lines.append(f"* `{case['class']}` state_dict keys: `{case['state_dict_keys']}`")
|
|
lines.append("")
|
|
for case in results["cases"]:
|
|
if case.get("note"):
|
|
lines.append(f"* `{case['class']}`: {case['note']}")
|
|
lines.append("")
|
|
return "\n".join(lines)
|
|
|
|
|
|
def main(argv: list[str] | None = None) -> int:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument(
|
|
"--diffusers-src",
|
|
dest="diffusers_src",
|
|
default=None,
|
|
help="the unpacked 'src' directory of diffusers 0.11.0",
|
|
)
|
|
parser.add_argument("--markdown", dest="markdown", default=str(DEFAULT_MARKDOWN))
|
|
parser.add_argument("--json", dest="json_path", default=None)
|
|
args = parser.parse_args(argv)
|
|
|
|
if not args.diffusers_src:
|
|
print(
|
|
"no pristine diffusers source given: pass --diffusers-src <unpacked>/src "
|
|
"(the 'src' directory of the unpacked diffusers-0.11.0 sdist)",
|
|
file=sys.stderr,
|
|
)
|
|
return 2
|
|
|
|
results = run(Path(args.diffusers_src))
|
|
markdown = render_markdown(results)
|
|
Path(args.markdown).parent.mkdir(parents=True, exist_ok=True)
|
|
Path(args.markdown).write_text(markdown, encoding="utf-8", newline="\n")
|
|
if args.json_path:
|
|
Path(args.json_path).write_text(json.dumps(results, indent=2) + "\n", encoding="utf-8")
|
|
print(json.dumps(results, indent=2))
|
|
return 0 if results["equivalent"] else 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|