0.2.1: band-coordinate law fixed (+invariant), dispatch-bank/cond classification fixed, schema v1.1 usage card + SCHEMA.md; code now canonical at github.com/AbstractEyes/amoe-lora
Browse files- .github/workflows/invariants.yml +28 -0
- .gitignore +14 -0
- LICENSE +201 -0
- README.md +15 -2
- SCHEMA.md +143 -0
- pyproject.toml +1 -1
- src/amoe/__init__.py +23 -23
- src/amoe/diffusion/train/trainer.py +12 -1
- src/amoe/io/checkpoint.py +74 -8
- src/amoe/testing/diffusion_invariants.py +63 -2
.github/workflows/invariants.yml
ADDED
|
@@ -0,0 +1,28 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# The house laws, enforced on every push.
|
| 2 |
+
#
|
| 3 |
+
# The invariant suites are CPU-only and download nothing (they run on tiny
|
| 4 |
+
# synthetic trunks), so this is a cheap, fast guard on the contracts that
|
| 5 |
+
# cost real GPU hours to establish: the toggle law, bit-exact detach, the
|
| 6 |
+
# zero-init inertness contract, the dtype law, and the band-coordinate law.
|
| 7 |
+
name: invariants
|
| 8 |
+
|
| 9 |
+
on:
|
| 10 |
+
push:
|
| 11 |
+
pull_request:
|
| 12 |
+
workflow_dispatch:
|
| 13 |
+
|
| 14 |
+
jobs:
|
| 15 |
+
invariants:
|
| 16 |
+
runs-on: ubuntu-latest
|
| 17 |
+
steps:
|
| 18 |
+
- uses: actions/checkout@v4
|
| 19 |
+
- uses: actions/setup-python@v5
|
| 20 |
+
with:
|
| 21 |
+
python-version: "3.11"
|
| 22 |
+
- name: Install torch (CPU) + safetensors
|
| 23 |
+
run: |
|
| 24 |
+
python -m pip install --upgrade pip
|
| 25 |
+
pip install torch --index-url https://download.pytorch.org/whl/cpu
|
| 26 |
+
pip install safetensors
|
| 27 |
+
- name: Run the invariant suites
|
| 28 |
+
run: python tests/test_invariants.py
|
.gitignore
ADDED
|
@@ -0,0 +1,14 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
__pycache__/
|
| 2 |
+
*.py[cod]
|
| 3 |
+
*.egg-info/
|
| 4 |
+
build/
|
| 5 |
+
dist/
|
| 6 |
+
.venv/
|
| 7 |
+
venv/
|
| 8 |
+
.idea/
|
| 9 |
+
.vscode/
|
| 10 |
+
.pytest_cache/
|
| 11 |
+
|
| 12 |
+
# adapter weights never live in the code repo — they ship on the hub
|
| 13 |
+
*.pt
|
| 14 |
+
*.safetensors
|
LICENSE
ADDED
|
@@ -0,0 +1,201 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
Apache License
|
| 2 |
+
Version 2.0, January 2004
|
| 3 |
+
http://www.apache.org/licenses/
|
| 4 |
+
|
| 5 |
+
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
| 6 |
+
|
| 7 |
+
1. Definitions.
|
| 8 |
+
|
| 9 |
+
"License" shall mean the terms and conditions for use, reproduction,
|
| 10 |
+
and distribution as defined by Sections 1 through 9 of this document.
|
| 11 |
+
|
| 12 |
+
"Licensor" shall mean the copyright owner or entity authorized by
|
| 13 |
+
the copyright owner that is granting the License.
|
| 14 |
+
|
| 15 |
+
"Legal Entity" shall mean the union of the acting entity and all
|
| 16 |
+
other entities that control, are controlled by, or are under common
|
| 17 |
+
control with that entity. For the purposes of this definition,
|
| 18 |
+
"control" means (i) the power, direct or indirect, to cause the
|
| 19 |
+
direction or management of such entity, whether by contract or
|
| 20 |
+
otherwise, or (ii) ownership of fifty percent (50%) or more of the
|
| 21 |
+
outstanding shares, or (iii) beneficial ownership of such entity.
|
| 22 |
+
|
| 23 |
+
"You" (or "Your") shall mean an individual or Legal Entity
|
| 24 |
+
exercising permissions granted by this License.
|
| 25 |
+
|
| 26 |
+
"Source" form shall mean the preferred form for making modifications,
|
| 27 |
+
including but not limited to software source code, documentation
|
| 28 |
+
source, and configuration files.
|
| 29 |
+
|
| 30 |
+
"Object" form shall mean any form resulting from mechanical
|
| 31 |
+
transformation or translation of a Source form, including but
|
| 32 |
+
not limited to compiled object code, generated documentation,
|
| 33 |
+
and conversions to other media types.
|
| 34 |
+
|
| 35 |
+
"Work" shall mean the work of authorship, whether in Source or
|
| 36 |
+
Object form, made available under the License, as indicated by a
|
| 37 |
+
copyright notice that is included in or attached to the work
|
| 38 |
+
(an example is provided in the Appendix below).
|
| 39 |
+
|
| 40 |
+
"Derivative Works" shall mean any work, whether in Source or Object
|
| 41 |
+
form, that is based on (or derived from) the Work and for which the
|
| 42 |
+
editorial revisions, annotations, elaborations, or other modifications
|
| 43 |
+
represent, as a whole, an original work of authorship. For the purposes
|
| 44 |
+
of this License, Derivative Works shall not include works that remain
|
| 45 |
+
separable from, or merely link (or bind by name) to the interfaces of,
|
| 46 |
+
the Work and Derivative Works thereof.
|
| 47 |
+
|
| 48 |
+
"Contribution" shall mean any work of authorship, including
|
| 49 |
+
the original version of the Work and any modifications or additions
|
| 50 |
+
to that Work or Derivative Works thereof, that is intentionally
|
| 51 |
+
submitted to Licensor for inclusion in the Work by the copyright owner
|
| 52 |
+
or by an individual or Legal Entity authorized to submit on behalf of
|
| 53 |
+
the copyright owner. For the purposes of this definition, "submitted"
|
| 54 |
+
means any form of electronic, verbal, or written communication sent
|
| 55 |
+
to the Licensor or its representatives, including but not limited to
|
| 56 |
+
communication on electronic mailing lists, source code control systems,
|
| 57 |
+
and issue tracking systems that are managed by, or on behalf of, the
|
| 58 |
+
Licensor for the purpose of discussing and improving the Work, but
|
| 59 |
+
excluding communication that is conspicuously marked or otherwise
|
| 60 |
+
designated in writing by the copyright owner as "Not a Contribution."
|
| 61 |
+
|
| 62 |
+
"Contributor" shall mean Licensor and any individual or Legal Entity
|
| 63 |
+
on behalf of whom a Contribution has been received by Licensor and
|
| 64 |
+
subsequently incorporated within the Work.
|
| 65 |
+
|
| 66 |
+
2. Grant of Copyright License. Subject to the terms and conditions of
|
| 67 |
+
this License, each Contributor hereby grants to You a perpetual,
|
| 68 |
+
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
| 69 |
+
copyright license to reproduce, prepare Derivative Works of,
|
| 70 |
+
publicly display, publicly perform, sublicense, and distribute the
|
| 71 |
+
Work and such Derivative Works in Source or Object form.
|
| 72 |
+
|
| 73 |
+
3. Grant of Patent License. Subject to the terms and conditions of
|
| 74 |
+
this License, each Contributor hereby grants to You a perpetual,
|
| 75 |
+
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
| 76 |
+
(except as stated in this section) patent license to make, have made,
|
| 77 |
+
use, offer to sell, sell, import, and otherwise transfer the Work,
|
| 78 |
+
where such license applies only to those patent claims licensable
|
| 79 |
+
by such Contributor that are necessarily infringed by their
|
| 80 |
+
Contribution(s) alone or by combination of their Contribution(s)
|
| 81 |
+
with the Work to which such Contribution(s) was submitted. If You
|
| 82 |
+
institute patent litigation against any entity (including a
|
| 83 |
+
cross-claim or counterclaim in a lawsuit) alleging that the Work
|
| 84 |
+
or a Contribution incorporated within the Work constitutes direct
|
| 85 |
+
or contributory patent infringement, then any patent licenses
|
| 86 |
+
granted to You under this License for that Work shall terminate
|
| 87 |
+
as of the date such litigation is filed.
|
| 88 |
+
|
| 89 |
+
4. Redistribution. You may reproduce and distribute copies of the
|
| 90 |
+
Work or Derivative Works thereof in any medium, with or without
|
| 91 |
+
modifications, and in Source or Object form, provided that You
|
| 92 |
+
meet the following conditions:
|
| 93 |
+
|
| 94 |
+
(a) You must give any other recipients of the Work or
|
| 95 |
+
Derivative Works a copy of this License; and
|
| 96 |
+
|
| 97 |
+
(b) You must cause any modified files to carry prominent notices
|
| 98 |
+
stating that You changed the files; and
|
| 99 |
+
|
| 100 |
+
(c) You must retain, in the Source form of any Derivative Works
|
| 101 |
+
that You distribute, all copyright, patent, trademark, and
|
| 102 |
+
attribution notices from the Source form of the Work,
|
| 103 |
+
excluding those notices that do not pertain to any part of
|
| 104 |
+
the Derivative Works; and
|
| 105 |
+
|
| 106 |
+
(d) If the Work includes a "NOTICE" text file as part of its
|
| 107 |
+
distribution, then any Derivative Works that You distribute must
|
| 108 |
+
include a readable copy of the attribution notices contained
|
| 109 |
+
within such NOTICE file, excluding those notices that do not
|
| 110 |
+
pertain to any part of the Derivative Works, in at least one
|
| 111 |
+
of the following places: within a NOTICE text file distributed
|
| 112 |
+
as part of the Derivative Works; within the Source form or
|
| 113 |
+
documentation, if provided along with the Derivative Works; or,
|
| 114 |
+
within a display generated by the Derivative Works, if and
|
| 115 |
+
wherever such third-party notices normally appear. The contents
|
| 116 |
+
of the NOTICE file are for informational purposes only and
|
| 117 |
+
do not modify the License. You may add Your own attribution
|
| 118 |
+
notices within Derivative Works that You distribute, alongside
|
| 119 |
+
or as an addendum to the NOTICE text from the Work, provided
|
| 120 |
+
that such additional attribution notices cannot be construed
|
| 121 |
+
as modifying the License.
|
| 122 |
+
|
| 123 |
+
You may add Your own copyright statement to Your modifications and
|
| 124 |
+
may provide additional or different license terms and conditions
|
| 125 |
+
for use, reproduction, or distribution of Your modifications, or
|
| 126 |
+
for any such Derivative Works as a whole, provided Your use,
|
| 127 |
+
reproduction, and distribution of the Work otherwise complies with
|
| 128 |
+
the conditions stated in this License.
|
| 129 |
+
|
| 130 |
+
5. Submission of Contributions. Unless You explicitly state otherwise,
|
| 131 |
+
any Contribution intentionally submitted for inclusion in the Work
|
| 132 |
+
by You to the Licensor shall be under the terms and conditions of
|
| 133 |
+
this License, without any additional terms or conditions.
|
| 134 |
+
Notwithstanding the above, nothing herein shall supersede or modify
|
| 135 |
+
the terms of any separate license agreement you may have executed
|
| 136 |
+
with Licensor regarding such Contributions.
|
| 137 |
+
|
| 138 |
+
6. Trademarks. This License does not grant permission to use the trade
|
| 139 |
+
names, trademarks, service marks, or product names of the Licensor,
|
| 140 |
+
except as required for reasonable and customary use in describing the
|
| 141 |
+
origin of the Work and reproducing the content of the NOTICE file.
|
| 142 |
+
|
| 143 |
+
7. Disclaimer of Warranty. Unless required by applicable law or
|
| 144 |
+
agreed to in writing, Licensor provides the Work (and each
|
| 145 |
+
Contributor provides its Contributions) on an "AS IS" BASIS,
|
| 146 |
+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
|
| 147 |
+
implied, including, without limitation, any warranties or conditions
|
| 148 |
+
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
|
| 149 |
+
PARTICULAR PURPOSE. You are solely responsible for determining the
|
| 150 |
+
appropriateness of using or redistributing the Work and assume any
|
| 151 |
+
risks associated with Your exercise of permissions under this License.
|
| 152 |
+
|
| 153 |
+
8. Limitation of Liability. In no event and under no legal theory,
|
| 154 |
+
whether in tort (including negligence), contract, or otherwise,
|
| 155 |
+
unless required by applicable law (such as deliberate and grossly
|
| 156 |
+
negligent acts) or agreed to in writing, shall any Contributor be
|
| 157 |
+
liable to You for damages, including any direct, indirect, special,
|
| 158 |
+
incidental, or consequential damages of any character arising as a
|
| 159 |
+
result of this License or out of the use or inability to use the
|
| 160 |
+
Work (including but not limited to damages for loss of goodwill,
|
| 161 |
+
work stoppage, computer failure or malfunction, or any and all
|
| 162 |
+
other commercial damages or losses), even if such Contributor
|
| 163 |
+
has been advised of the possibility of such damages.
|
| 164 |
+
|
| 165 |
+
9. Accepting Warranty or Additional Liability. While redistributing
|
| 166 |
+
the Work or Derivative Works thereof, You may choose to offer,
|
| 167 |
+
and charge a fee for, acceptance of support, warranty, indemnity,
|
| 168 |
+
or other liability obligations and/or rights consistent with this
|
| 169 |
+
License. However, in accepting such obligations, You may act only
|
| 170 |
+
on Your own behalf and on Your sole responsibility, not on behalf
|
| 171 |
+
of any other Contributor, and only if You agree to indemnify,
|
| 172 |
+
defend, and hold each Contributor harmless for any liability
|
| 173 |
+
incurred by, or claims asserted against, such Contributor by reason
|
| 174 |
+
of your accepting any such warranty or additional liability.
|
| 175 |
+
|
| 176 |
+
END OF TERMS AND CONDITIONS
|
| 177 |
+
|
| 178 |
+
APPENDIX: How to apply the Apache License to your work.
|
| 179 |
+
|
| 180 |
+
To apply the Apache License to your work, attach the following
|
| 181 |
+
boilerplate notice, with the fields enclosed by brackets "[]"
|
| 182 |
+
replaced with your own identifying information. (Don't include
|
| 183 |
+
the brackets!) The text should be enclosed in the appropriate
|
| 184 |
+
comment syntax for the file format. We also recommend that a
|
| 185 |
+
file or class name and description of purpose be included on the
|
| 186 |
+
same "printed page" as the copyright notice for easier
|
| 187 |
+
identification within third-party archives.
|
| 188 |
+
|
| 189 |
+
Copyright 2026 AbstractPhil (AbstractEyes)
|
| 190 |
+
|
| 191 |
+
Licensed under the Apache License, Version 2.0 (the "License");
|
| 192 |
+
you may not use this file except in compliance with the License.
|
| 193 |
+
You may obtain a copy of the License at
|
| 194 |
+
|
| 195 |
+
http://www.apache.org/licenses/LICENSE-2.0
|
| 196 |
+
|
| 197 |
+
Unless required by applicable law or agreed to in writing, software
|
| 198 |
+
distributed under the License is distributed on an "AS IS" BASIS,
|
| 199 |
+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
| 200 |
+
See the License for the specific language governing permissions and
|
| 201 |
+
limitations under the License.
|
README.md
CHANGED
|
@@ -35,10 +35,23 @@ optimized or deep-tested. See *Maturity* below.
|
|
| 35 |
## Install
|
| 36 |
|
| 37 |
```
|
| 38 |
-
pip install
|
| 39 |
-
pip install -
|
| 40 |
```
|
| 41 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 42 |
## The five verbs
|
| 43 |
|
| 44 |
```python
|
|
|
|
| 35 |
## Install
|
| 36 |
|
| 37 |
```
|
| 38 |
+
pip install git+https://github.com/AbstractEyes/amoe-lora # core
|
| 39 |
+
pip install "amoe-lora[diffusion] @ git+https://github.com/AbstractEyes/amoe-lora"
|
| 40 |
```
|
| 41 |
|
| 42 |
+
or from a checkout: `pip install -e .`, extras `.[hf]` (language trunks)
|
| 43 |
+
and `.[diffusion]` (diffusers + safetensors). torch is a dependency, but
|
| 44 |
+
install the CUDA build that matches your machine first — a blind
|
| 45 |
+
`pip install torch` can replace a working one.
|
| 46 |
+
|
| 47 |
+
**Code lives on [GitHub](https://github.com/AbstractEyes/amoe-lora)**
|
| 48 |
+
(canonical, CI-guarded); the Hugging Face repo mirrors it as the card.
|
| 49 |
+
The adapter file format is documented in [SCHEMA.md](SCHEMA.md); trained
|
| 50 |
+
adapters live in
|
| 51 |
+
[aleph-diffusion-adapters](https://huggingface.co/AbstractPhil/aleph-diffusion-adapters)
|
| 52 |
+
and load into ComfyUI via
|
| 53 |
+
[comfyui-geolip-amoe-lora](https://github.com/AbstractEyes/comfyui-geolip-amoe-lora).
|
| 54 |
+
|
| 55 |
## The five verbs
|
| 56 |
|
| 57 |
```python
|
SCHEMA.md
ADDED
|
@@ -0,0 +1,143 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# `amoe.diffusion.anchor` — the adapter file format
|
| 2 |
+
|
| 3 |
+
One file = one adapter stack for one trunk. The file carries its own
|
| 4 |
+
documentation: a loader can tell you what the adapter is, what it was
|
| 5 |
+
trained on, what it measurably does, how strongly to run it, and what it
|
| 6 |
+
does *not* do — without any external database.
|
| 7 |
+
|
| 8 |
+
Two containers, same logical content:
|
| 9 |
+
|
| 10 |
+
| container | when | notes |
|
| 11 |
+
|---|---|---|
|
| 12 |
+
| `.safetensors` | **preferred**, and what ComfyUI loads | meta rides in the safetensors metadata block (str→str) |
|
| 13 |
+
| `.pt` | the campaign's original saves | meta rides as a plain dict under `"meta"` |
|
| 14 |
+
|
| 15 |
+
## Tensor layout
|
| 16 |
+
|
| 17 |
+
Flat, one namespace, site index first:
|
| 18 |
+
|
| 19 |
+
```
|
| 20 |
+
blocks.{site_index}.{param_path} # in the safetensors key space
|
| 21 |
+
{site_index}.{param_path} # in memory / in the .pt payload
|
| 22 |
+
```
|
| 23 |
+
|
| 24 |
+
`site_index` is `0..n_sites-1` in **training enumeration order**.
|
| 25 |
+
|
| 26 |
+
> ### The ordering law (read this before writing a loader)
|
| 27 |
+
>
|
| 28 |
+
> Site order is diffusers `named_modules()` order, which registers
|
| 29 |
+
> `down_blocks` → `up_blocks` → `mid_block`. **The mid block is LAST**, not
|
| 30 |
+
> in the middle. For SD1.5 the width signature is
|
| 31 |
+
>
|
| 32 |
+
> ```
|
| 33 |
+
> [320,320,640,640,1280,1280, 1280,1280,1280, 640,640,640, 320,320,320, 1280]
|
| 34 |
+
> ```
|
| 35 |
+
>
|
| 36 |
+
> A denoiser's *execution* order is different (`input → middle → output`,
|
| 37 |
+
> giving `…1280,1280,1280,1280…` with mid at index 6). Zipping the two
|
| 38 |
+
> positionally misplaces 7 of 16 sites **and still runs**, producing quietly
|
| 39 |
+
> wrong images. Map by site identity, then verify the width signature —
|
| 40 |
+
> the two orders differ at indices 9 and 15, so the signature catches it.
|
| 41 |
+
> `DiffusionAnchorCheckpoint.widths` reads the signature off the tensors.
|
| 42 |
+
|
| 43 |
+
## Metadata
|
| 44 |
+
|
| 45 |
+
### Required (provenance — every file has these)
|
| 46 |
+
|
| 47 |
+
| field | type | meaning |
|
| 48 |
+
|---|---|---|
|
| 49 |
+
| `format` | str | `"amoe.diffusion.anchor"` |
|
| 50 |
+
| `version` | int/str | `1` |
|
| 51 |
+
| `adapter.kind` | str | `relay` · `multiband3` · `mono` · `bank` · `cond` |
|
| 52 |
+
| `substrate.family` | str | `sd15_unet` · `sdxl_unet` · `cosmos_dit` |
|
| 53 |
+
| `substrate.n_sites` | int | must equal the enumerated site count at attach |
|
| 54 |
+
|
| 55 |
+
Only `relay` and `multiband3` are **attachable**. `mono` and `bank` are
|
| 56 |
+
matched controls and falsified-routing evidence; `cond` is the Law-2
|
| 57 |
+
negative. Loaders may read them; `attach()` refuses them by design and
|
| 58 |
+
says why.
|
| 59 |
+
|
| 60 |
+
### Optional (provenance, written when known)
|
| 61 |
+
|
| 62 |
+
`adapter.*` spec (`n_slots`/`K`/`tau`/`hidden` for relay, `rank` for
|
| 63 |
+
multiband3) · `substrate.base_model_id` · `substrate.site_names` ·
|
| 64 |
+
`substrate.widths` · `objective.kind` (`eps`|`flow`|`v`) ·
|
| 65 |
+
`objective.shift` · `blob.lambda` · `dtype` · `seed` · `recipe.*` ·
|
| 66 |
+
`created` · `content_hash_v2` · `imported_from` · `home_reconstructed`
|
| 67 |
+
|
| 68 |
+
### Optional (the usage card — schema v1.1)
|
| 69 |
+
|
| 70 |
+
What a UI renders. All optional, so raw campaign stacks stay valid.
|
| 71 |
+
|
| 72 |
+
| field | type | meaning |
|
| 73 |
+
|---|---|---|
|
| 74 |
+
| `display_name` | str | human name, e.g. `"SD1.5 Relay — Grounding v1 (seed 0)"` |
|
| 75 |
+
| `usage` | str | what it does and when to reach for it |
|
| 76 |
+
| `evidence` | str | the measured verdict **with numbers and seed status** |
|
| 77 |
+
| `recommended_strength` | float | sane default for a strength slider |
|
| 78 |
+
| `band_roles` | list[str] | multiband only: role per band, LOW→HIGH |
|
| 79 |
+
| `caveats` | list[str] | what it costs, where it fails, what is unverified |
|
| 80 |
+
| `license` | str | SPDX-ish string |
|
| 81 |
+
| `nc` | bool | `true` = non-commercial (derived from NC weights) |
|
| 82 |
+
|
| 83 |
+
`DiffusionAnchorCheckpoint.card()` returns exactly this block with safe
|
| 84 |
+
defaults filled in.
|
| 85 |
+
|
| 86 |
+
### Honesty rule for `evidence`
|
| 87 |
+
|
| 88 |
+
`evidence` states what was measured, at how many seeds, against what
|
| 89 |
+
control — never a marketing claim. If a result is single-seed, it says so.
|
| 90 |
+
If a matched control beat it on the aggregate metric, that goes in
|
| 91 |
+
`evidence` or `caveats`, not omitted. Controls and falsified artifacts
|
| 92 |
+
carry `evidence` describing what they *refuted*; that is their value.
|
| 93 |
+
|
| 94 |
+
## Worked example
|
| 95 |
+
|
| 96 |
+
```json
|
| 97 |
+
{
|
| 98 |
+
"format": "amoe.diffusion.anchor",
|
| 99 |
+
"version": 1,
|
| 100 |
+
"display_name": "SD1.5 Multiband — Coarse-to-Fine v1 (seed 0)",
|
| 101 |
+
"adapter": {"kind": "multiband3", "rank": 16},
|
| 102 |
+
"substrate": {"family": "sd15_unet", "n_sites": 16,
|
| 103 |
+
"base_model_id": "stable-diffusion-v1-5/stable-diffusion-v1-5"},
|
| 104 |
+
"objective": {"kind": "eps"},
|
| 105 |
+
"band_roles": ["fidelity/detail (LOW noise)",
|
| 106 |
+
"continuity/semantics (MID)",
|
| 107 |
+
"diversity/structure (HIGH noise)"],
|
| 108 |
+
"usage": "Three sigma-band experts gated per sampling step. Lesion a band to see what it carries.",
|
| 109 |
+
"evidence": "exp008, 2 seeds: band lesions surgical 3/3 — own-band damage 50-200x cross-band. A matched rank-48 monolith still edges it on aggregate eps-MSE under uniform pressure.",
|
| 110 |
+
"recommended_strength": 1.0,
|
| 111 |
+
"caveats": ["Trained on the stock SD1.5 eps trunk; other trunks are untested.",
|
| 112 |
+
"The aggregate win belongs to the monolith control — the value here is the band structure, not the loss number."],
|
| 113 |
+
"license": "MIT", "nc": false,
|
| 114 |
+
"dtype": "float32", "seed": 0
|
| 115 |
+
}
|
| 116 |
+
```
|
| 117 |
+
|
| 118 |
+
## Band gating (multiband3 only)
|
| 119 |
+
|
| 120 |
+
Band windows are cosine crossfades on `s01`, the **normalized discrete
|
| 121 |
+
timestep**:
|
| 122 |
+
|
| 123 |
+
```
|
| 124 |
+
s01 = t / 1000 # eps trunks — t from the model's own sampling
|
| 125 |
+
s01 = sigma # flow trunks — the SHIFT-warped sigma
|
| 126 |
+
edges = (0.35, 0.75) # LAW constants, not tunables
|
| 127 |
+
xfade = 0.06
|
| 128 |
+
```
|
| 129 |
+
|
| 130 |
+
> `s01` is **not** a noise-level proxy. On the real SD1.5 schedule,
|
| 131 |
+
> `1 - alphas_cumprod[t]` puts 316 of 1000 timesteps in a different band
|
| 132 |
+
> than the one the expert was trained on. In ComfyUI, get it from
|
| 133 |
+
> `model_sampling.timestep(sigma) / 1000`.
|
| 134 |
+
|
| 135 |
+
## Loader checklist
|
| 136 |
+
|
| 137 |
+
1. Read meta; reject if `format` is absent or unknown.
|
| 138 |
+
2. Enumerate the host's sites; sort into **training order**.
|
| 139 |
+
3. Assert `n_sites` matches, then assert the width signature matches
|
| 140 |
+
`checkpoint.widths` — refuse loudly on mismatch.
|
| 141 |
+
4. Cast adapters to the trunk's declared dtype (the dtype law).
|
| 142 |
+
5. `kind` in `{relay, multiband3}` to attach; otherwise explain and stop.
|
| 143 |
+
6. Render `card()` so the user sees provenance, evidence, and caveats.
|
pyproject.toml
CHANGED
|
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
|
| 4 |
|
| 5 |
[project]
|
| 6 |
name = "amoe-lora"
|
| 7 |
-
version = "0.2.
|
| 8 |
description = "Aleph mixture-of-experts adapters: train, attach, align, detach — with the honesty diagnostics built in. 0.2 adds the diffusion subsystem (amoe.diffusion) and safetensors I/O."
|
| 9 |
readme = "README.md"
|
| 10 |
requires-python = ">=3.10"
|
|
|
|
| 4 |
|
| 5 |
[project]
|
| 6 |
name = "amoe-lora"
|
| 7 |
+
version = "0.2.1"
|
| 8 |
description = "Aleph mixture-of-experts adapters: train, attach, align, detach — with the honesty diagnostics built in. 0.2 adds the diffusion subsystem (amoe.diffusion) and safetensors I/O."
|
| 9 |
readme = "README.md"
|
| 10 |
requires-python = ">=3.10"
|
src/amoe/__init__.py
CHANGED
|
@@ -1,23 +1,23 @@
|
|
| 1 |
-
"""amoe — aleph mixture-of-experts adapters (train/attach/align/detach).
|
| 2 |
-
|
| 3 |
-
The productized adapter system from the geolip-aleph research line.
|
| 4 |
-
"LoRA-style" refers to the attach/detach usage pattern, not the math:
|
| 5 |
-
these are aleph-addressed patch heads, not low-rank matrices.
|
| 6 |
-
"""
|
| 7 |
-
from .core.adapter import AdapterSpec, RelayPatchwork, BlockWithAdapter
|
| 8 |
-
from .core.address import AlephAddress
|
| 9 |
-
from .core.dispatch import AnchorDispatch, BlockWithDispatch, set_mask
|
| 10 |
-
from .io.checkpoint import (AnchorCheckpoint, DispatchCheckpoint,
|
| 11 |
-
load_anchor, load_dispatch,
|
| 12 |
-
import_legacy_keys)
|
| 13 |
-
from .runtime.attach import attach, detach, AttachHandle
|
| 14 |
-
from .train.config import TrainConfig, AlignConfig
|
| 15 |
-
from .train.trainer import train
|
| 16 |
-
from .train.aligner import align
|
| 17 |
-
from . import laws
|
| 18 |
-
|
| 19 |
-
__version__ = "0.2.
|
| 20 |
-
|
| 21 |
-
# The diffusion subsystem (amoe.diffusion) is imported lazily — its verbs
|
| 22 |
-
# live under `import amoe.diffusion as ad`. Core is pure torch; diffusers
|
| 23 |
-
# is only touched by bindings/samplers/data that need it.
|
|
|
|
| 1 |
+
"""amoe — aleph mixture-of-experts adapters (train/attach/align/detach).
|
| 2 |
+
|
| 3 |
+
The productized adapter system from the geolip-aleph research line.
|
| 4 |
+
"LoRA-style" refers to the attach/detach usage pattern, not the math:
|
| 5 |
+
these are aleph-addressed patch heads, not low-rank matrices.
|
| 6 |
+
"""
|
| 7 |
+
from .core.adapter import AdapterSpec, RelayPatchwork, BlockWithAdapter
|
| 8 |
+
from .core.address import AlephAddress
|
| 9 |
+
from .core.dispatch import AnchorDispatch, BlockWithDispatch, set_mask
|
| 10 |
+
from .io.checkpoint import (AnchorCheckpoint, DispatchCheckpoint,
|
| 11 |
+
load_anchor, load_dispatch,
|
| 12 |
+
import_legacy_keys)
|
| 13 |
+
from .runtime.attach import attach, detach, AttachHandle
|
| 14 |
+
from .train.config import TrainConfig, AlignConfig
|
| 15 |
+
from .train.trainer import train
|
| 16 |
+
from .train.aligner import align
|
| 17 |
+
from . import laws
|
| 18 |
+
|
| 19 |
+
__version__ = "0.2.1"
|
| 20 |
+
|
| 21 |
+
# The diffusion subsystem (amoe.diffusion) is imported lazily — its verbs
|
| 22 |
+
# live under `import amoe.diffusion as ad`. Core is pure torch; diffusers
|
| 23 |
+
# is only touched by bindings/samplers/data that need it.
|
src/amoe/diffusion/train/trainer.py
CHANGED
|
@@ -132,7 +132,18 @@ def train(model, cache: dict, config: "DiffusionTrainConfig | None" = None,
|
|
| 132 |
|
| 133 |
if cfg.objective == "eps":
|
| 134 |
t = torch.randint(0, 1000, (bsz,), generator=gd, device=device)
|
| 135 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 136 |
w = set_w(s01) if cfg.adapter == "multiband3" else None
|
| 137 |
x_t = add_noise(lat, noise, t, acp)
|
| 138 |
pred = ddp_model(x_t.to(dt), t, ehs.to(dt),
|
|
|
|
| 132 |
|
| 133 |
if cfg.objective == "eps":
|
| 134 |
t = torch.randint(0, 1000, (bsz,), generator=gd, device=device)
|
| 135 |
+
# BAND COORDINATE LAW: s01 = t/1000 — the normalized DISCRETE
|
| 136 |
+
# timestep, exactly what every certified bed trained on
|
| 137 |
+
# (dexp008/011/012) and what the proven controller gates on at
|
| 138 |
+
# inference (dexp010, StepGatedSampler). Do NOT substitute a
|
| 139 |
+
# noise-level proxy such as 1 - alphas_cumprod[t]: measured on
|
| 140 |
+
# the real SD1.5 scaled_linear schedule, 316 of 1000 timesteps
|
| 141 |
+
# land in a DIFFERENT band under that proxy (t=300 trains LOW
|
| 142 |
+
# but the proxy says MID; t=700 trains MID, proxy says HIGH),
|
| 143 |
+
# so roughly a third of training would teach the wrong expert
|
| 144 |
+
# and inference would gate on an axis the stack never learned.
|
| 145 |
+
# Pinned by testing.assert_band_coordinate.
|
| 146 |
+
s01 = t.float() / 1000.0
|
| 147 |
w = set_w(s01) if cfg.adapter == "multiband3" else None
|
| 148 |
x_t = add_noise(lat, noise, t, acp)
|
| 149 |
pred = ddp_model(x_t.to(dt), t, ehs.to(dt),
|
src/amoe/io/checkpoint.py
CHANGED
|
@@ -104,9 +104,16 @@ DIFF_ANCHOR_FORMAT = "amoe.diffusion.anchor"
|
|
| 104 |
@dataclass
|
| 105 |
class DiffusionAnchorCheckpoint:
|
| 106 |
"""One diffusion adapter stack: adapters keyed "{site_index}.{param}".
|
| 107 |
-
meta.adapter.kind is 'relay' | 'multiband3'
|
| 108 |
-
('mono', 'bank'). Relay kind requires per-site
|
| 109 |
-
gauge);
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 110 |
adapters: dict[str, torch.Tensor]
|
| 111 |
meta: dict[str, Any] = field(default_factory=dict)
|
| 112 |
|
|
@@ -118,6 +125,45 @@ class DiffusionAnchorCheckpoint:
|
|
| 118 |
def n_sites(self) -> int:
|
| 119 |
return len({k.split(".", 1)[0] for k in self.adapters})
|
| 120 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 121 |
def per_site(self, i: int) -> dict[str, torch.Tensor]:
|
| 122 |
pre = f"{i}."
|
| 123 |
return {k[len(pre):]: v for k, v in self.adapters.items()
|
|
@@ -137,14 +183,23 @@ class DiffusionAnchorCheckpoint:
|
|
| 137 |
|
| 138 |
|
| 139 |
def _sd_kind(sd: dict) -> str:
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 140 |
if any(k.endswith("addr.codebook") for k in sd):
|
| 141 |
return "relay"
|
|
|
|
|
|
|
|
|
|
|
|
|
| 142 |
if "down.0.weight" in sd:
|
| 143 |
return "multiband3"
|
| 144 |
if "down.weight" in sd:
|
| 145 |
return "mono"
|
| 146 |
-
if "key_proj" in sd or "keys" in sd:
|
| 147 |
-
return "bank"
|
| 148 |
return "unknown"
|
| 149 |
|
| 150 |
|
|
@@ -203,11 +258,22 @@ def load_diffusion_anchor(path: str, *, substrate: "dict | None" = None
|
|
| 203 |
key = "banks" if "banks" in blob else "monos"
|
| 204 |
ck = _flatten_stack(list(blob[key]), kind=_sd_kind(blob[key][0]),
|
| 205 |
imported_from=f"legacy_{key}")
|
| 206 |
-
|
| 207 |
-
|
| 208 |
-
|
|
|
|
|
|
|
| 209 |
else:
|
| 210 |
raise ValueError(f"unrecognized diffusion anchor at {path}")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 211 |
if ck.kind == "relay":
|
| 212 |
blocks = {k.split(".", 1)[0] for k in ck.adapters}
|
| 213 |
for b in blocks:
|
|
|
|
| 104 |
@dataclass
|
| 105 |
class DiffusionAnchorCheckpoint:
|
| 106 |
"""One diffusion adapter stack: adapters keyed "{site_index}.{param}".
|
| 107 |
+
meta.adapter.kind is 'relay' | 'multiband3' (attachable) or a
|
| 108 |
+
load-only kind ('mono', 'bank', 'cond'). Relay kind requires per-site
|
| 109 |
+
addr.home (drift gauge); the others have no address, so the
|
| 110 |
+
requirement is kind-aware.
|
| 111 |
+
|
| 112 |
+
The meta block is the SELF-DOCUMENTING half of the format — see
|
| 113 |
+
SCHEMA.md. Everything beyond `adapter`/`substrate` is optional, so
|
| 114 |
+
raw campaign stacks load unchanged while production artifacts can
|
| 115 |
+
carry their own usage card (display_name / usage / evidence /
|
| 116 |
+
recommended_strength / caveats / license)."""
|
| 117 |
adapters: dict[str, torch.Tensor]
|
| 118 |
meta: dict[str, Any] = field(default_factory=dict)
|
| 119 |
|
|
|
|
| 125 |
def n_sites(self) -> int:
|
| 126 |
return len({k.split(".", 1)[0] for k in self.adapters})
|
| 127 |
|
| 128 |
+
@property
|
| 129 |
+
def widths(self) -> list[int]:
|
| 130 |
+
"""Per-site hidden width, read from the tensors themselves — the
|
| 131 |
+
cross-framework ALIGNMENT SIGNATURE. Site order is the training
|
| 132 |
+
order (diffusers named_modules: down -> up -> mid), which is NOT
|
| 133 |
+
a denoiser's execution order; any host that enumerates blocks
|
| 134 |
+
differently must match on this signature, never positionally."""
|
| 135 |
+
out = []
|
| 136 |
+
for i in range(self.n_sites):
|
| 137 |
+
sd = self.per_site(i)
|
| 138 |
+
for key, dim in (("proj.weight", 1), ("down.0.weight", 1),
|
| 139 |
+
("down.weight", 1), ("addr_proj.weight", 0)):
|
| 140 |
+
if key in sd:
|
| 141 |
+
out.append(int(sd[key].shape[dim]))
|
| 142 |
+
break
|
| 143 |
+
else:
|
| 144 |
+
raise ValueError(f"site {i}: cannot infer width from "
|
| 145 |
+
f"{sorted(sd)[:6]}")
|
| 146 |
+
return out
|
| 147 |
+
|
| 148 |
+
def card(self) -> dict:
|
| 149 |
+
"""The human-facing subset of meta (schema v1.1), with safe
|
| 150 |
+
defaults — what a UI should render about this adapter."""
|
| 151 |
+
m = self.meta
|
| 152 |
+
return {
|
| 153 |
+
"display_name": m.get("display_name", m.get("name", "unnamed")),
|
| 154 |
+
"kind": self.kind,
|
| 155 |
+
"n_sites": self.n_sites,
|
| 156 |
+
"trunk": m.get("substrate", {}).get("base_model_id", "unknown"),
|
| 157 |
+
"objective": m.get("objective", {}).get("kind", "unknown"),
|
| 158 |
+
"usage": m.get("usage", ""),
|
| 159 |
+
"evidence": m.get("evidence", ""),
|
| 160 |
+
"recommended_strength": m.get("recommended_strength", 1.0),
|
| 161 |
+
"band_roles": m.get("band_roles", []),
|
| 162 |
+
"caveats": m.get("caveats", []),
|
| 163 |
+
"license": m.get("license", ""),
|
| 164 |
+
"nc": bool(m.get("nc", False)),
|
| 165 |
+
}
|
| 166 |
+
|
| 167 |
def per_site(self, i: int) -> dict[str, torch.Tensor]:
|
| 168 |
pre = f"{i}."
|
| 169 |
return {k[len(pre):]: v for k, v in self.adapters.items()
|
|
|
|
| 183 |
|
| 184 |
|
| 185 |
def _sd_kind(sd: dict) -> str:
|
| 186 |
+
"""Classify one site's state dict.
|
| 187 |
+
|
| 188 |
+
ORDER IS LOAD-BEARING: a dispatch bank carries BOTH `key_proj.weight`
|
| 189 |
+
and `down.0.weight`, so the bank test MUST come before the multiband
|
| 190 |
+
test — otherwise every exp007/exp014/exp015 bank silently loads as a
|
| 191 |
+
multiband3 stack and fails later at module construction.
|
| 192 |
+
"""
|
| 193 |
if any(k.endswith("addr.codebook") for k in sd):
|
| 194 |
return "relay"
|
| 195 |
+
if any(k.startswith("key_proj") for k in sd) or "keys" in sd:
|
| 196 |
+
return "bank" # dispatch bank — must precede
|
| 197 |
+
if "pos_table" in sd:
|
| 198 |
+
return "cond" # AlephCondAdapter (exp002)
|
| 199 |
if "down.0.weight" in sd:
|
| 200 |
return "multiband3"
|
| 201 |
if "down.weight" in sd:
|
| 202 |
return "mono"
|
|
|
|
|
|
|
| 203 |
return "unknown"
|
| 204 |
|
| 205 |
|
|
|
|
| 258 |
key = "banks" if "banks" in blob else "monos"
|
| 259 |
ck = _flatten_stack(list(blob[key]), kind=_sd_kind(blob[key][0]),
|
| 260 |
imported_from=f"legacy_{key}")
|
| 261 |
+
elif _sd_kind(blob) != "unknown":
|
| 262 |
+
# a FLAT single-module state dict (exp002 AlephCondAdapter):
|
| 263 |
+
# one "site", stored under index 0 so per_site(0) round-trips
|
| 264 |
+
ck = _flatten_stack([blob], kind=_sd_kind(blob),
|
| 265 |
+
imported_from="legacy_flat_module")
|
| 266 |
else:
|
| 267 |
raise ValueError(f"unrecognized diffusion anchor at {path}")
|
| 268 |
+
if ck.kind == "bank":
|
| 269 |
+
ck.meta.setdefault("routing_negative", (
|
| 270 |
+
"exp007/014/015: comparative dispatch on diffusion is a 2-seed "
|
| 271 |
+
"falsified line — load for analysis, not deployment"))
|
| 272 |
+
if ck.kind == "cond":
|
| 273 |
+
ck.meta.setdefault("law2_negative", (
|
| 274 |
+
"exp002: the frozen address beside full text conditioning is "
|
| 275 |
+
"redundant-in-context (real vs deranged -0.0009); the address "
|
| 276 |
+
"ALONE steers (+0.0287). Redesign target = complementarity"))
|
| 277 |
if ck.kind == "relay":
|
| 278 |
blocks = {k.split(".", 1)[0] for k in ck.adapters}
|
| 279 |
for b in blocks:
|
src/amoe/testing/diffusion_invariants.py
CHANGED
|
@@ -151,6 +151,64 @@ def assert_band_windows():
|
|
| 151 |
assert d < 0.03, f"windows not smooth (max step {d} > theory ~0.0262)"
|
| 152 |
|
| 153 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 154 |
def assert_flow_x0_recovery():
|
| 155 |
from ..diffusion.train.objectives import flow_pieces
|
| 156 |
g = torch.Generator().manual_seed(5)
|
|
@@ -316,6 +374,8 @@ def assert_align_negative():
|
|
| 316 |
|
| 317 |
def run_all() -> None:
|
| 318 |
assert_band_windows()
|
|
|
|
|
|
|
| 319 |
assert_flow_x0_recovery()
|
| 320 |
assert_toggle_law_diffusion(tuple_site=False)
|
| 321 |
assert_toggle_law_diffusion(tuple_site=True)
|
|
@@ -324,8 +384,9 @@ def run_all() -> None:
|
|
| 324 |
assert_checkpoint_roundtrip()
|
| 325 |
assert_blob_on_eps_refusal()
|
| 326 |
assert_align_negative()
|
| 327 |
-
print("amoe.diffusion invariants: band windows,
|
| 328 |
-
"
|
|
|
|
| 329 |
"dtype law, checkpoint round trips (.pt/.safetensors/legacy/"
|
| 330 |
"fork), blob-on-eps refusal, align negative — ALL GREEN")
|
| 331 |
|
|
|
|
| 151 |
assert d < 0.03, f"windows not smooth (max step {d} > theory ~0.0262)"
|
| 152 |
|
| 153 |
|
| 154 |
+
def assert_band_coordinate():
|
| 155 |
+
"""The eps band coordinate is s01 = t/1000 — the normalized DISCRETE
|
| 156 |
+
timestep. Every certified bed trained on it (dexp008/011/012) and the
|
| 157 |
+
proven controller gates on it (dexp010). The tempting substitute,
|
| 158 |
+
1 - alphas_cumprod[t], puts t=500 in a different band entirely, which
|
| 159 |
+
silently decouples training from inference. This test pins the
|
| 160 |
+
convention and documents the divergence it guards against."""
|
| 161 |
+
from ..diffusion.core.multiband import band_of
|
| 162 |
+
|
| 163 |
+
# the REAL SD1.5 schedule (scaled_linear betas), not a stand-in: this
|
| 164 |
+
# is the schedule the shipped eps stacks were trained against
|
| 165 |
+
betas = torch.linspace(0.00085 ** 0.5, 0.012 ** 0.5, 1000,
|
| 166 |
+
dtype=torch.float64) ** 2
|
| 167 |
+
acp = torch.cumprod(1.0 - betas, dim=0)
|
| 168 |
+
|
| 169 |
+
disagreements = [t for t in range(1000)
|
| 170 |
+
if band_of(t / 1000.0) != band_of(float(1 - acp[t]))]
|
| 171 |
+
assert len(disagreements) > 300, (
|
| 172 |
+
f"only {len(disagreements)} timesteps disagree — the guard is not "
|
| 173 |
+
"sharp enough to catch a regression; revisit it")
|
| 174 |
+
# measured on this schedule: 316/1000 timesteps train a DIFFERENT band
|
| 175 |
+
# under the proxy. t=300 -> trained LOW (0.300) vs proxy MID (0.410);
|
| 176 |
+
# t=700 -> trained MID (0.700) vs proxy HIGH (0.918).
|
| 177 |
+
assert band_of(0.300) == 0 and band_of(float(1 - acp[300])) == 1
|
| 178 |
+
assert band_of(0.700) == 1 and band_of(float(1 - acp[700])) == 2
|
| 179 |
+
|
| 180 |
+
# the trainer must use the trained coordinate
|
| 181 |
+
import inspect
|
| 182 |
+
from ..diffusion.train import trainer as _tr
|
| 183 |
+
src = inspect.getsource(_tr.train)
|
| 184 |
+
assert "t.float() / 1000.0" in src, \
|
| 185 |
+
"trainer eps band coordinate is not t/1000 (band coordinate law)"
|
| 186 |
+
assert "1 - acp[t]" not in src, \
|
| 187 |
+
"trainer still uses the alphas_cumprod proxy for band windows"
|
| 188 |
+
|
| 189 |
+
# and the sampler must agree with the trainer
|
| 190 |
+
src_s = inspect.getsource(
|
| 191 |
+
__import__("amoe.diffusion.runtime.sampler", fromlist=["x"]))
|
| 192 |
+
assert "float(t) / 1000.0" in src_s, \
|
| 193 |
+
"StepGatedSampler eps gate diverged from the training coordinate"
|
| 194 |
+
|
| 195 |
+
|
| 196 |
+
def assert_kind_classification():
|
| 197 |
+
"""Dispatch banks carry BOTH key_proj and down.0.weight; the bank test
|
| 198 |
+
must win, or exp007/014/015 stacks load as multiband3 and blow up at
|
| 199 |
+
module construction. Flat single-module cond adapters (exp002) must
|
| 200 |
+
classify too."""
|
| 201 |
+
from ..io.checkpoint import _sd_kind
|
| 202 |
+
bank = {"key_proj.weight": None, "down.0.weight": None,
|
| 203 |
+
"gates": None, "codebook": None}
|
| 204 |
+
assert _sd_kind(bank) == "bank", "bank misclassified (ordering bug)"
|
| 205 |
+
assert _sd_kind({"down.0.weight": None, "gates": None}) == "multiband3"
|
| 206 |
+
assert _sd_kind({"down.weight": None, "gate": None}) == "mono"
|
| 207 |
+
assert _sd_kind({"pos_table": None, "addr_proj.weight": None}) == "cond"
|
| 208 |
+
assert _sd_kind({"addr.codebook": None, "proj.weight": None}) == "relay"
|
| 209 |
+
assert _sd_kind({"nonsense": None}) == "unknown"
|
| 210 |
+
|
| 211 |
+
|
| 212 |
def assert_flow_x0_recovery():
|
| 213 |
from ..diffusion.train.objectives import flow_pieces
|
| 214 |
g = torch.Generator().manual_seed(5)
|
|
|
|
| 374 |
|
| 375 |
def run_all() -> None:
|
| 376 |
assert_band_windows()
|
| 377 |
+
assert_band_coordinate()
|
| 378 |
+
assert_kind_classification()
|
| 379 |
assert_flow_x0_recovery()
|
| 380 |
assert_toggle_law_diffusion(tuple_site=False)
|
| 381 |
assert_toggle_law_diffusion(tuple_site=True)
|
|
|
|
| 384 |
assert_checkpoint_roundtrip()
|
| 385 |
assert_blob_on_eps_refusal()
|
| 386 |
assert_align_negative()
|
| 387 |
+
print("amoe.diffusion invariants: band windows, BAND COORDINATE "
|
| 388 |
+
"(t/1000), kind classification, exact flow x0, toggle law "
|
| 389 |
+
"(relay+multiband, tensor+tuple), P-INIT, bit-exact detach, "
|
| 390 |
"dtype law, checkpoint round trips (.pt/.safetensors/legacy/"
|
| 391 |
"fork), blob-on-eps refusal, align negative — ALL GREEN")
|
| 392 |
|