dinac3_96 v1.0

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

ATTRIBUTION.md

@@ -0,0 +1,50 @@
+# Attribution
+
+Licensing (code and weights) is stated once, in the License section of the model
+card (README.md).
+
+- **DINOv3 pretrained encoder:** Meta,
+  [official repository](https://github.com/facebookresearch/dinov3),
+  [paper](https://arxiv.org/abs/2508.10104),
+  [timm model](https://huggingface.co/timm/vit_base_patch16_dinov3.lvd_1689m).
+  The transformer blocks are used frozen; the input patch embedding is adapted
+  for this autoencoder. The weights are therefore distributed under the included
+  [DINOv3 License](LICENSE-DINOV3.md); see the model card's License section.
+- **Encoder implementation:** provided by the separately installed
+  [timm](https://github.com/huggingface/pytorch-image-models) dependency, pinned
+  to `timm==1.0.26`.
+- **semantic_vae** ([well9472/semantic_vae](https://huggingface.co/well9472/semantic_vae)):
+  the design follows its key ideas. These are a pretrained DINO encoder adapted
+  at its edges and a latent split into a reconstruction part and a DINO-derived
+  part. Its encoder has two DINOv2-B branches: one with a trainable patch
+  embedding, whose last six blocks are each LayerNormed and projected to 32
+  channels, and one fully frozen, whose final block is projected to 32 channels.
+  We also follow it in training the decoder mainly with a DINO feature loss and
+  in using VISReg as the latent regularizer.¹ No code or weights from it are included.
+- **Decoder up-path:** the VQGAN f16 decoder layout of Esser, Rombach and Ommer,
+  [Taming Transformers for High-Resolution Image Synthesis](https://arxiv.org/abs/2012.09841),
+  slightly modernized (per-position RMSNorm, folded upsampling) and behind a ViT
+  trunk.
+- **Earlier decoder head (not in the release):** a pixel-level head after
+  [PixelDiT](https://arxiv.org/abs/2511.20645)
+  ([NVlabs/PixelDiT](https://github.com/NVlabs/PixelDiT)), with the FCDM block of
+  Kwon et al.,
+  [Reviving ConvNeXt for Efficient Convolutional Diffusion Models](https://arxiv.org/abs/2603.09408),
+  as its token mixer. The released encoder and decoder trunk were trained with it
+  before the convolutional head replaced it.
+- **Latent regularizer:** [VISReg](https://arxiv.org/abs/2606.02572)
+  ([project page](https://haiyuwu.github.io/visreg/)), used during training only.
+- **Optimizer (training only, not a runtime dependency):**
+  [dionw](https://github.com/JTriggerFish/dionw), a single-GPU implementation of
+  [Dion](https://arxiv.org/abs/2504.05295) /
+  [Dion3](https://arxiv.org/abs/2608.11612) with
+  [NorMuon](https://arxiv.org/abs/2510.05491) normalization, using kernels from
+  [microsoft/dion](https://github.com/microsoft/dion) and
+  [Dao-AILab/gram-newton-schulz](https://github.com/Dao-AILab/gram-newton-schulz).
+- **Training data:** about 14 million images from a mix of public and licensed
+  or curated datasets: mostly photographs, plus book covers, a few text-heavy
+  datasets and 1% synthetic rendered text. No training images are redistributed.
+
+¹ semantic_vae's public model card names SIGReg (from LeJEPA) at 1e-3 as its
+regularizer. We verified in its released training code that the regularizer is
+VISReg, at 1e-3.

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

LICENSE-DINOV3.md

@@ -0,0 +1,66 @@
+# DINOv3 License
+
+*Last Updated: August 19, 2025*
+
+**“Agreement”** means the terms and conditions for use, reproduction, distribution and modification of the DINO Materials set forth herein.
+
+**“DINO Materials”** means, collectively, Documentation and the models, software and algorithms, including machine-learning model code, trained model weights, inference-enabling code, training-enabling code, fine-tuning enabling code, and other elements of the foregoing distributed by Meta and made available under this Agreement.
+
+**“Documentation”** means the specifications, manuals and documentation accompanying
+DINO Materials distributed by Meta.
+
+**“Licensee”** or **“you”** means you, or your employer or any other person or entity (if you are entering into this Agreement on such person or entity’s behalf), of the age required under applicable laws, rules or regulations to provide legal consent and that has legal authority to bind your employer or such other person or entity if you are entering in this Agreement on their behalf.
+
+**“Meta”** or **“we”** means Meta Platforms Ireland Limited (if you are located in or, if you are an entity, your principal place of business is in the EEA or Switzerland) or Meta Platforms, Inc. (if you are located outside of the EEA or Switzerland).
+
+**“Sanctions”** means any economic or trade sanctions or restrictions administered or enforced by the United States (including the Office of Foreign Assets Control of the U.S. Department of the Treasury (“OFAC”), the U.S. Department of State and the U.S. Department of Commerce), the United Nations, the European Union, or the United Kingdom.
+
+**“Trade Controls”** means any of the following: Sanctions and applicable export and import controls.
+
+By clicking “I Accept” below or by using or distributing any portion or element of the DINO Materials, you agree to be bound by this Agreement.
+
+## 1. License Rights and Redistribution.
+
+a. <ins>Grant of Rights</ins>. You are granted a non-exclusive, worldwide, non-transferable and royalty-free limited license under Meta’s intellectual property or other rights owned by Meta embodied in the DINO Materials to use, reproduce, distribute, copy, create derivative works of, and make modifications to the DINO Materials.
+
+b. <ins>Redistribution and Use</ins>.
+
+i. Distribution of DINO Materials, and any derivative works thereof, are subject to the terms of this Agreement. If you distribute or make the DINO Materials, or any derivative works thereof, available to a third party, you may only do so under the terms of this Agreement and you shall provide a copy of this Agreement with any such DINO Materials.
+
+ii.  If you submit for publication the results of research you perform on, using, or otherwise in connection with DINO Materials, you must acknowledge the use of DINO Materials in your publication.
+
+iii. Your use of the DINO Materials must comply with applicable laws and regulations, including Trade Control Laws and applicable privacy and data protection laws.
+
+iv. Your use of the DINO Materials will not involve or encourage others to reverse engineer, decompile or discover the underlying components of the DINO Materials.
+
+v. You are not the target of Trade Controls and your use of DINO Materials must comply with Trade Controls. You agree not to use, or permit others to use, DINO Materials for any activities subject to the International Traffic in Arms Regulations (ITAR) or end uses prohibited by Trade Controls, including those related to military or warfare purposes, nuclear industries or applications, espionage, or the development or use of guns or illegal weapons.
+
+## 2. User Support.
+
+Your use of the DINO Materials is done at your own discretion; Meta does not process any information nor provide any service in relation to such use.  Meta is under no obligation to provide any support services for the DINO Materials. Any support provided is “as is”, “with all faults”, and without warranty of any kind.
+
+## 3. Disclaimer of Warranty.
+
+UNLESS REQUIRED BY APPLICABLE LAW, THE DINO MATERIALS AND ANY OUTPUT AND RESULTS THEREFROM ARE PROVIDED ON AN “AS IS” BASIS, WITHOUT WARRANTIES OF ANY KIND, AND META DISCLAIMS ALL WARRANTIES OF ANY KIND, BOTH EXPRESS AND IMPLIED, INCLUDING, WITHOUT LIMITATION, ANY WARRANTIES OF TITLE, NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. YOU ARE SOLELY RESPONSIBLE FOR DETERMINING THE APPROPRIATENESS OF USING OR REDISTRIBUTING THE DINO MATERIALS AND ASSUME ANY RISKS ASSOCIATED WITH YOUR USE OF THE DINO MATERIALS AND ANY OUTPUT AND RESULTS.
+
+## 4. Limitation of Liability.
+
+IN NO EVENT WILL META OR ITS AFFILIATES BE LIABLE UNDER ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, TORT, NEGLIGENCE, PRODUCTS LIABILITY, OR OTHERWISE, ARISING OUT OF THIS AGREEMENT, FOR ANY LOST PROFITS OR ANY DIRECT OR INDIRECT, SPECIAL, CONSEQUENTIAL, INCIDENTAL, EXEMPLARY OR PUNITIVE DAMAGES, EVEN IF META OR ITS AFFILIATES HAVE BEEN ADVISED OF THE POSSIBILITY OF ANY OF THE FOREGOING.
+
+## 5. Intellectual Property.
+
+a. Subject to Meta’s ownership of DINO Materials and derivatives made by or for Meta, with respect to any derivative works and modifications of the DINO Materials that are made by you, as between you and Meta, you are and will be the owner of such derivative works and modifications.
+
+b. If you institute litigation or other proceedings against Meta or any entity (including a cross-claim or counterclaim in a lawsuit) alleging that the DINO Materials, outputs or results, or any portion of any of the foregoing, constitutes infringement of intellectual property or other rights owned or licensable by you, then any licenses granted to you under this Agreement shall terminate as of the date such litigation or claim is filed or instituted. You will indemnify and hold harmless Meta from and against any claim by any third party arising out of or related to your use or distribution of the DINO Materials.
+
+## 6. Term and Termination.
+
+The term of this Agreement will commence upon your acceptance of this Agreement or access to the DINO Materials and will continue in full force and effect until terminated in accordance with the terms and conditions herein. Meta may terminate this Agreement if you are in breach of any term or condition of this Agreement. Upon termination of this Agreement, you shall delete and cease use of the DINO Materials. Sections 3, 4 and 7 shall survive the termination of this Agreement.
+
+## 7. Governing Law and Jurisdiction.
+
+This Agreement will be governed and construed under the laws of the State of California without regard to choice of law principles, and the UN Convention on Contracts for the International Sale of Goods does not apply to this Agreement. The courts of California shall have exclusive jurisdiction of any dispute arising out of this Agreement.
+
+## 8. Modifications and Amendments.
+
+Meta may modify this Agreement from time to time; provided that they are similar in spirit to the current version of the Agreement, but may differ in detail to address new problems or concerns. All such changes will be effective immediately. Your continued use of the DINO Materials after any modification to this Agreement constitutes your agreement to such modification. Except as provided in this Agreement, no modification or addition to any provision of this Agreement will be binding unless it is in writing and signed by an authorized representative of both you and Meta.

README.md

@@ -0,0 +1,244 @@
+---
+license: other
+license_name: dinov3-license
+license_link: https://huggingface.co/data-archetype/dinac3_96/blob/main/LICENSE-DINOV3.md
+base_model: timm/vit_base_patch16_dinov3.lvd_1689m
+tags:
+  - autoencoder
+  - image-reconstruction
+  - latent-space
+  - latent-diffusion
+  - dinov3
+  - pytorch
+---
+
+# data-archetype/dinac3_96
+
+**dinac3_96** is a deterministic semantic autoencoder for latent diffusion. It
+maps an RGB image to a 96-channel latent grid at stride 16 and decodes it back in
+a single forward pass.
+
+It is strongly inspired by [semantic_vae][semantic-vae], the latent with the
+fastest downstream DiT convergence we know of. dinac3_96 aims to keep that
+convergence while improving reconstruction and simplifying the encoder to a
+single DINOv3 pass instead of two.
+
+- **Encoder:** a frozen [DINOv3][dinov3] ViT-B/16 behind a trainable patch
+  embedding. All twelve blocks are read, standardized per channel and mapped to
+  the latent by one trainable linear projection.
+- **Latent:** 32 semantic channels, held during training to a fixed random
+  projection of DINOv3-B's blocks summed with fitted block weights, plus 64 reconstruction channels (reconstruction
+  losses and VISReg). No posterior noise.
+- **Decoder:** a 4-block transformer trunk at width 1152 followed by a dense
+  convolutional up-path, a slightly modernized VQGAN-like decoder (stride 16 to
+  full resolution, three residual blocks per level). One pass, no diffusion
+  sampling.
+- **Training losses:** mainly a DINOv3-B feature loss on the reconstruction, with
+  pixel and blurred-image MSE terms each at about a tenth of its gradient,
+  following semantic_vae's DINO-loss-trained decoder.
+
+**[Technical report](TECHNICAL_REPORT.md)** ·
+**[Reconstruction gallery](https://huggingface.co/spaces/data-archetype/dinac3-results)**
+(39 images: original, reconstruction, RGB difference, latent PCA, per-image PSNR)
+
+## Generation benchmark
+
+We evaluate a latent by training a class-conditional DiT on it (about 190M
+parameters, 100k steps at batch 256 on ImageNet-1k, flow matching with SPRINT
+routing). We score 10,000 samples against 810,000 real images on our
+mixed-resolution ImageNet benchmark: four aspect ratios at 256- and 384-px areas.
+Lower is better.[^fid] For each model: 50 NFEs, PDG 2.5–4.0, best settings kept.
+Bold marks the better value in each column. The DiT was trained on dinac3_96's latent (the released latent
+space, frozen since the joint phase), and its samples were decoded by the
+released decoder.
+
+| model | FID | MIND | Monge-DINO |
+|---|---:|---:|---:|
+| dinac3_96 | **`9.04`** | `8.29` | **`10.92`** |
+| semantic_vae | `10.02` | **`7.42`** | `15.94` |
+
+Monge-DINO uses DINOv3-B features, the same network whose features the decoder is
+trained to match, so it is not independent of the training loss; FID and MIND, on
+Inception features, are. The same holds for rMonge-DINOv3-B below.
+
+Our blind pairwise comparison of the two models' generations put them on equal
+footing.
+
+Protocol, DiT settings and caveats:
+[technical report, section 9.1](TECHNICAL_REPORT.md#91-generation-benchmark).
+
+## Reconstruction
+
+**10,000 ImageNet images:** 10 real images per class with the generation
+benchmark's aspect-ratio quotas (9 at 256-px area and 1 at 384 px per class).
+Metrics compare the reconstructions with the same 10,000 originals.[^rfid] semantic_vae runs at its native bf16.
+
+| model | rFID | rMIND | rMonge-DINOv3-B | PSNR mean | SSIM | LPIPS-VGG | LPIPS-Alex |
+|---|---:|---:|---:|---:|---:|---:|---:|
+| dinac3_96 | **`0.809`** | **`0.166`** | **`1.205`** | **`28.16`** | **`0.8195`** | **`0.0970`** | **`0.0362`** |
+| semantic_vae | `2.149` | `0.820` | `5.994` | `25.19` | `0.7103` | `0.1715` | `0.0669` |
+
+PSNR here: reconstructions clamped and rounded to uint8, peak 255, mean of
+per-image values.
+
+**2k PSNR benchmark** (the image set of our earlier releases):
+
+| Model | Mean PSNR (dB) | Std (dB) | Median (dB) | P5 (dB) | P95 (dB) |
+|---|---:|---:|---:|---:|---:|
+| dinac3_96 | `32.15` | `5.11` | `31.64` | `24.27` | `40.70` |
+| semantic_vae | `28.03` | `4.58` | `27.79` | `20.98` | `35.44` |
+| dinac_ae_d2 | `35.59` | `4.87` | `35.40` | `27.89` | `43.51` |
+| FLUX.2 VAE | `36.28` | `4.53` | `36.07` | `28.89` | `43.63` |
+
+## Latent interface
+
+- 96 channels at stride 16; image height and width must be multiples of 16.
+- Channels 0–63 are reconstruction channels, 64–95 semantic channels
+  (`free_channels(z)`, `semantic_channels(z)`).
+- `encode(images)`: RGB in [-1, 1] → deterministic FP32 latents, whitened per
+  channel with the shipped statistics (what our DiTs were trained on).
+- `decode(latents, height, width)`: whitened latents → FP32 RGB in about
+  [-1, 1], unclamped. Any multiple of 16 works, including sizes above 1024 px.
+- `encode_raw` / `decode_raw` use the unwhitened latent; `whiten` / `dewhiten`
+  convert.
+
+## Precision
+
+- Set the dtype only through `from_pretrained(dtype=...)`; dtype casts after
+  loading raise a `TypeError` (`.to(device)` works).
+- `torch.bfloat16` (default) runs under bf16 autocast, bit-identical to the
+  training checkpoint under bf16 autocast. `torch.float32` runs without autocast.
+- Encoder and decoder are compiled by default (about 20 s per module on the
+  first call). Pass `compile_encoder=False, compile_decoder=False` for eager.
+
+## Speed
+
+RTX 5090, bf16, compiled (the default); mean of 20 timed batches after 3
+warm-up batches. Peak VRAM includes the weights.
+
+| Resolution | Batch | Encode ms/image | Decode ms/image | Peak VRAM encode / decode (GiB) |
+|---:|---:|---:|---:|---:|
+| `256x256` | `128` | `0.41` | `1.51` | `3.3` / `8.5` |
+| `512x512` | `32` | `1.79` | `6.10` | `3.3` / `8.5` |
+| `1024x1024` | `8` | `9.03` | `25.4` | `3.3` / `8.5` |
+| `2048x2048` | `1` | `69.4` | `157` | `1.8` / `6.4` |
+
+Eager mode decodes about 2.4× slower (1.8× at 2048²) and peaks at 12.7 GiB on
+the batched rows.
+
+**Size limit.** No resolution limit is coded; GPU memory is the practical one. On
+a 32 GB RTX 5090, compiled decoding handles a single 4096 × 4096 image (15.8 s,
+24.5 GiB peak); eager decoding runs out of memory at that size.
+
+## Usage
+
+Download the repository and install the package, in an environment with a
+CUDA-enabled PyTorch (dependencies: `torch>=2.13`, `timm==1.0.26`, `safetensors`,
+`huggingface-hub`):
+
+```bash
+hf download data-archetype/dinac3_96 --local-dir dinac3_96
+pip install ./dinac3_96
+```
+
+```python
+import torch
+from dinac3 import Dinac3
+
+model = Dinac3.from_pretrained(
+    "data-archetype/dinac3_96",  # or a local directory
+    device="cuda",
+    dtype=torch.bfloat16,
+)
+
+images = ...  # [B, 3, H, W] on CUDA, RGB in [-1, 1], H and W multiples of 16
+
+latents = model.encode(images)  # [B, 96, H/16, W/16], whitened
+semantic = model.semantic_channels(latents)  # [B, 32, H/16, W/16]
+recon = model.decode(latents, images.shape[-2], images.shape[-1])
+recon = recon.clamp(-1, 1)  # only for display or saving
+```
+
+`from_pretrained(path_or_repo, *, device, dtype=torch.bfloat16,
+compile_encoder=True, compile_decoder=True, revision=None)`:
+- A `Path` is always a local directory.
+- A `str` is a local directory if one exists.
+- A `str` that looks like a path (it starts with `.`, `/` or `~`, or has more
+  than one `/`) but names no directory raises an error.
+- Any other `str` is a Hub repository id; only `config.json` and
+  `model.safetensors` are downloaded.
+- The DINOv3-B weights ship inside `model.safetensors`, so a local copy loads
+  offline.
+- Inference is CUDA only.
+
+## Details
+
+- **Network:** 164M parameters, 79M of them trained.
+  - Encoder (87M): frozen DINOv3 ViT-B/16 behind a trainable patch embedding;
+    all 12 blocks standardized, concatenated and projected linearly to 96
+    channels at stride 16.
+  - Decoder (78M): 1×1 projection to width 1152, 4 transformer blocks, then a
+    convolutional up-path from stride 16 to full resolution (512-channel
+    handoff, levels of 256, 256, 128 and 128 channels, three residual 3 × 3
+    blocks per level).
+- **Training data:** about 14 million images from a mix of public and licensed
+  or curated datasets: mostly photographs, plus book covers, a few text-heavy
+  datasets and 1% synthetic rendered text. Aspect-ratio buckets, downsampled
+  only. No training images are redistributed with the model.
+- **Training:**
+  - Losses: mainly a DINOv3-B all-block feature MSE on the reconstruction, plus
+    pixel and blurred-image MSE at about a tenth of its gradient; while the
+    encoder trained, also a semantic alignment loss and [VISReg][visreg].
+  - Steps: about 150k at 256/384 px (batch 128) plus about 25k fine-tuning at five
+    resolutions from 256 to 1024 px (batch 32). The free encoder weights (patch
+    embedding and output projection) were trained jointly with the decoder for
+    the first 52k steps and then frozen.
+  - Matrix weights trained with Dion ([dionw][dionw]), the rest with AdamW. The
+    released weights are EMA weights.
+- **More:** the [technical report](TECHNICAL_REPORT.md) covers the semantic
+  target, the decoder heads tried, loss weights, training phases, ablations and
+  the benchmark protocol.
+- **Related:** [semantic_vae][semantic-vae],
+  [DINAC-AE-D2](https://huggingface.co/data-archetype/dinac_ae_d2).
+
+## License
+
+- **Code** (the `dinac3` package and scripts): Copyright 2026 data-archetype,
+  licensed under the Apache License 2.0, see [LICENSE](LICENSE).
+- **Weights:** the DINOv3 License, see [LICENSE-DINOV3.md](LICENSE-DINOV3.md).
+  The weights contain DINOv3-B's weights and a trained derivative of its patch
+  embedding, so under the DINOv3 License (section 1.b.i) they may only be
+  distributed under that licence, which must accompany them. Its terms include:
+  - acknowledging the use of DINO Materials in publications (section 1.b.ii);
+  - no reverse engineering (section 1.b.iv);
+  - compliance with trade controls, which excludes ITAR-regulated uses and
+    military, nuclear, espionage and weapons applications (sections 1.b.iii and
+    1.b.v).
+
+The `license: other` tag above refers to the weights' DINOv3 License. See
+[ATTRIBUTION.md](ATTRIBUTION.md) for sources and design credits.
+
+## Citation
+
+```bibtex
+@misc{dinac3_96,
+  title   = {dinac3_96: a DINOv3-based semantic autoencoder with a one-pass DINO-loss decoder},
+  author  = {data-archetype},
+  email   = {data-archetype@proton.me},
+  year    = {2026},
+  month   = oct,
+  url     = {https://huggingface.co/data-archetype/dinac3_96},
+}
+```
+
+[dinov3]: https://arxiv.org/abs/2508.10104
+[^fid]: Papers usually report FID on 50,000 square samples at a single
+  resolution, so our numbers are not directly comparable with published ones.
+
+[^rfid]: Papers usually report rFID on the 50,000 ImageNet validation images at
+  a single square resolution, so our numbers are not directly comparable with
+  published ones.
+
+[semantic-vae]: https://huggingface.co/well9472/semantic_vae
+[visreg]: https://arxiv.org/abs/2606.02572
+[dionw]: https://github.com/JTriggerFish/dionw

TECHNICAL_REPORT.md

@@ -0,0 +1,608 @@
+# dinac3_96 Technical Report
+
+`dinac3_96` is a deterministic semantic autoencoder for latent diffusion. A
+frozen DINOv3 ViT-B/16 behind a trainable patch embedding is read at all twelve
+blocks through one trainable linear projection, giving a 96-channel latent at
+stride 16: 32 semantic channels held to a fixed DINOv3 target and 64 free
+channels for reconstruction. A one-pass decoder (ViT trunk and dense
+convolutional up-path) is trained mainly by a DINOv3 feature loss.
+
+**Results in brief:** see the [README](README.md); protocols and full numbers in
+[section 9](#9-results).
+
+## Contents
+
+1. [Overview](#1-overview)
+2. [Encoder](#2-encoder)
+3. [The semantic target](#3-the-semantic-target)
+4. [Decoder](#4-decoder)
+5. [Why a pure DINO loss](#5-why-a-pure-dino-loss)
+6. [Losses](#6-losses)
+7. [Data](#7-data)
+8. [Training](#8-training)
+9. [Results](#9-results)
+10. [References](#references)
+
+## 1. Overview
+
+```text
+RGB image [B, 3, H, W], H and W divisible by 16
+  → trainable copy of DINOv3-B's patch-16 embedding
+  → 12 frozen DINOv3-B blocks
+  → every block's patch tokens through DINOv3's final LayerNorm,
+    standardized per channel, concatenated            [B, 12 × 768, H/16, W/16]
+  → one trainable linear projection                    [B, 96, H/16, W/16]
+      channels 0–63:  free (trained by reconstruction and VISReg)
+      channels 64–95: semantic (held to a fixed projected sum of DINOv3 blocks)
+
+latent [B, 96, H/16, W/16]
+  → 1×1 projection to width 1152
+  → 4 ViT blocks at width 1152
+  → convolutional up-path: /16 → /8 → /4 → /2 → /1, three residual
+    blocks per level
+  → RGB [B, 3, H, W]
+```
+
+Inputs are RGB in [-1, 1] of any size with height and width divisible by 16;
+DINOv3's ImageNet normalization is applied internally.
+
+The design keeps three ideas from [semantic_vae][semantic-vae]:
+
+- **A pretrained DINO encoder adapted at its edges.** Its transformer blocks are
+  frozen; the patch embedding and the projection to the latent train.
+- **A latent split into a DINO-derived part and a reconstruction part,**
+  similarly to our earlier [iRDiffAE][irdiffae] (half of its channels aligned to
+  DINOv2-S, half free) and `semantic_vae`.
+- **A decoder trained by a DINO feature loss.** Like `semantic_vae`, we use an
+  (almost) exclusively DINO-based loss. We found this critical to the best
+  compromise between fast convergence of downstream diffusion models and good
+  reconstruction (sections 5 and 6).
+
+The decoder head is a slightly modernized VQGAN-like convolutional up-path
+behind a ViT trunk (section 4.2). Differences from `semantic_vae`:
+
+- one encoder network instead of two, halving the encoder's compute;
+- DINOv3-B instead of DINOv2-B;
+- all twelve blocks read through one projection;
+- semantic channels defined by a fixed projected sum of DINOv3-B's blocks with
+  fitted block weights (`semantic_vae` projects final-block features).
+
+## 2. Encoder
+
+### 2.1 Frozen DINOv3-B with a trainable patch embedding
+
+[DINOv3][dinov3] ViT-B/16 (LVD-1689M): 12 blocks of width 768. Blocks, class and
+register tokens and final LayerNorm are frozen. A trainable copy of the
+pretrained patch embedding replaces the original; gradients reach it through the
+frozen blocks. A second, unmodified DINOv3-B computes the targets of the semantic
+alignment and of the DINO reconstruction loss.
+
+**Why the patch embedding trains.** Self-supervised encoders discard colour,
+intensity and fine-detail information that a decoder needs to recover. Training
+the patch embedding lets the model recover some of it. Because DINOv3 is a
+pre-norm residual transformer, this amounts to mixing pixel channels into DINO's
+residual stream while the frozen blocks' features mostly stay as they were (the
+alignment loss also pushes towards this).
+
+### 2.2 All twelve blocks through one projection
+
+Patch tokens \\(h_l \in \mathbb{R}^{768}\\) of each block \\(l = 0, \dots, 11\\),
+taken after DINOv3's final LayerNorm, are standardized with fixed statistics and
+projected:
+
+$$
+\tilde h_l = \frac{h_l - \mu_l}{\sigma_l}, \qquad
+x = [\tilde h_0; \dots; \tilde h_{11}] \in \mathbb{R}^{9216}, \qquad
+z = W x + b \in \mathbb{R}^{96}.
+$$
+
+\\(\mu_l, \sigma_l\\) are fixed, measured once on training images with the
+pretrained embedding.
+\\(W\\), \\(b\\) and the patch embedding are the only trainable encoder parameters.
+
+- **Per-channel standardization.** DINOv3-B's variance sits in a few
+  massive-activation channels (block 2: the top 8 of 768 hold 88%, participation
+  ratio 5.6; block 11: 11%, ratio 286). Unstandardized, the projection would
+  mostly copy them.
+- **All twelve blocks, trainable projection.** Deep blocks carry semantics and
+  are smooth; shallow blocks carry detail. Our previous latents, fixed fusions of
+  all twelve blocks, gave good DiTs (FID 12.80) but capped
+  reconstruction, since only the patch embedding could add pixel information. A
+  trainable projection (32 semantic + 64 free) lifted reconstruction to 30–32 dB
+  with a diffusion decoder, while the semantic channels kept their class content
+  (probe 0.24, participation rank 29 of 32).
+
+**Participation rank.** We measure how many channels of a group are effectively
+used by the participation ratio \\((\sum_i \lambda_i)^2 / \sum_i \lambda_i^2\\)
+of the eigenvalues \\(\lambda_i\\) of the group's token covariance. It equals
+the number of channels when all eigenvalues are equal and 1 when a single one
+dominates.
+
+### 2.3 Semantic and free channels
+
+- **Semantic (last 32 rows):** initialized to the target map of section 3, so
+  equal to their target at step 0; held to it by an alignment loss during joint
+  training.
+- **Free (64 rows):** random initialization with unit output variance; trained
+  only by the reconstruction losses and VISReg.
+
+We chose 64 free channels over 32 for reconstruction. In our experiments a
+32 + 32 latent generated only marginally better (FID 8.85 against 9.04) but
+reconstructed worse (ImageNet PSNR 26.24 against 28.16 dB, rFID 1.87 against
+0.81).
+
+### 2.4 Latent statistics and whitening
+
+The latent is deterministic. Per-channel running means and variances, tracked
+during joint training and frozen with the encoder, ship with the model. DiTs train
+on whitened latents \\((z - \mu) / \sqrt{\sigma^2 + 10^{-4}}\\); the decoder
+consumes the unwhitened latent (package interface in the [README](README.md)).
+
+## 3. The semantic target
+
+### 3.1 Why a fixed target
+
+In our initial experiments we aligned the semantic channels to DINO through a
+learned objective. These latents scored FID 14.1 to 19.1, far from `semantic_vae`'s 10.02, for a
+common reason: each objective is best satisfied by DINO's highest-variance
+directions, which are also its smoothest and carry relatively little class
+information.
+
+- **Learned up-projection with MSE or cosine** (latent → 768 → DINO features, as
+  in [DINAC-AE-D2][dinac-d2]): at fixed rank, MSE is minimized by the target's
+  top principal subspace (Eckart–Young); cosine on the full features behaves the
+  same.
+- **CCA** between a channel block and DINO features: scale-invariant, so the
+  encoder picks the directions easiest to produce, in practice the same smooth
+  modes (60% subspace overlap with the top 32 principal directions).
+- **CKA** on token similarities: dominated by the leading principal components of
+  both representations.
+- **MSE to the class token:** constrains image-level content only, with no
+  per-token structure.
+
+On DINOv3-B's final block, the top 32 principal directions have spectral slope
+2.25 (section 3.2) and class probe 0.167, against 1.84 and 0.291 for 32
+class-discriminant directions.
+
+A fixed target of the same width removes this choice. Each of the 32 semantic
+channels regresses by per-channel MSE onto one channel of a fixed 32-channel map
+of frozen DINO features, with no learned projection in between, so the map fully
+specifies which directions the latent encodes. Every target channel has unit
+variance, so each weighs equally in the loss. Section 3.2 gives the map.
+
+### 3.2 From a random projection to fitted layer weights
+
+The target is built in two steps; only the second has free parameters.
+
+**Step 1: one channel projection, shared by all blocks.** A single Gaussian
+random matrix \\(P \in \mathbb{R}^{32 \times 768}\\), drawn once (seed 0) and kept
+fixed, maps each of the 12 standardized blocks to 32 channels, which are
+standardized again:
+
+$$
+y_l = \frac{P \tilde h_l - m_l}{s_l} \in \mathbb{R}^{32}, \qquad l = 0, \dots, 11.
+$$
+
+**Step 2: scalar block weights.** The 12 projected blocks are mixed channel by
+channel with one scalar weight \\(w_l\\) per block (summing to 1) and rescaled to
+unit variance, with \\(C_c\\) the 12 × 12 cross-block correlation of channel \\(c\\):
+
+$$
+t_c = \frac{\sum_l w_l\, y_{l,c}}{\sqrt{w^\top C_c\, w}}, \qquad c = 1, \dots, 32.
+$$
+
+The spectral slope \\(\alpha\\) is the exponent of a channel's radially
+averaged power spectrum, \\(S(f) \propto f^{-\alpha}\\) (larger is smoother).
+
+**Why a Gaussian random map.** Write a standardized block's token covariance as
+\\(\Sigma = \sum_k \lambda_k u_k u_k^\top\\). \\(P\\) has i.i.d. \\(\mathcal{N}(0, 1)\\)
+entries, so its distribution is rotation-invariant (\\(PU\\) is distributed as \\(P\\)
+for any orthogonal \\(U\\)) and it favours no direction. Each output channel
+\\(p_i^\top \tilde h\\) has expected variance
+\\(\mathbb{E}[p_i^\top \Sigma p_i] = \sum_k \lambda_k\\): every principal mode
+contributes in proportion to its variance, whereas a top-32 PCA projection keeps
+modes 1 to 32 whole and drops all others.
+
+**Choosing the block weights.** Equal block weights (\\(w_l = 1/12\\)) give a
+latent that is too rough, with little class information (probe 0.123). We fitted
+the 12 block weights to maximize an ImageNet class probe (ridge
+classifier on image-mean tokens) subject to median \\(\alpha \le 1.10\\) and 90th
+percentile \\(\le 1.20\\) at 256 px and median \\(\le 1.40\\) at 512 px, keeping the
+spectrum centred near 1.0 at the 256–384 px training resolutions (`semantic_vae`:
+median 1.09 at 256 px, probe 0.207). Search: an evolution strategy over
+log-weights, selected on one half of the held-out images and reported on the
+other (0.249 / 0.245).
+
+The fit puts **0.53 on block 11** and **0.26 and 0.15 on blocks 0 and 1**; every
+other block gets 0.02 or less. Block 11 carries the class content but alone is
+too smooth (\\(\alpha\\) 1.64); the shallow blocks roughen it, while middle blocks
+would add smoothness without class. Seven of eight random restarts converged to
+the same pattern (block 11 at 0.45–0.56 plus blocks 0–2).
+
+| 32-channel candidate (DINOv3-B unless noted) | \\(\alpha\\) at 256 px, p10 / median / p90 | \\(\alpha\\) at 512 px, median | class probe | participation rank |
+| --- | --- | ---: | ---: | ---: |
+| `semantic_vae` semantic block (DINOv2-B) | 0.99 / 1.09 / 1.19 | 1.39 | 0.207 | 24.0 |
+| equal block weights, all 12 blocks | 0.78 / 0.88 / 0.96 | 1.20 | 0.123 | 27.8 |
+| equal block weights, last 6 blocks | 1.03 / 1.07 / 1.15 | 1.43 | 0.154 | 27.2 |
+| fitted block weights (probe projection) | 1.00 / 1.09 / 1.16 | 1.40 | 0.247 | 29.2 |
+| **fitted block weights, as shipped (seed-0 Gaussian \\(P\\))** | **1.00 / 1.04 / 1.14** | **1.39** | **0.241** | **29.0** |
+
+### 3.3 Training the target in
+
+The target is computed without gradient by the unmodified DINOv3-B on the same
+image, projected by \\(P\\) and summed with the fitted block weights (section 3.2).
+The encoder's trainable projection \\(W\\) is initialized with the same weights: its
+32 semantic rows are this composed linear map of the standardized blocks. With
+the trainable patch embedding starting from its pretrained values, the semantic
+channels equal their target at step 0. Training then updates the patch embedding
+and \\(W\\), and the alignment loss, a per-token, per-channel MSE between the 32
+semantic channels and the target, holds them close to it.
+
+We weighted the alignment loss so that, on the semantic rows of \\(W\\), its gradient
+is at least ten times that of the reconstruction losses, making sure the semantic
+target dominates the pull of reconstruction.
+
+## 4. Decoder
+
+Deterministic: one forward pass, no noisy-image input.
+
+### 4.1 Trunk
+
+A 1×1 projection maps the 96 channels to width 1152, one token per 16 × 16
+patch, followed by four ViT blocks: RMSNorm before and after the attention and
+MLP branches, per-head query/key RMSNorm, axial 2D RoPE on patch indices, 18
+heads of width 64, GELU MLP of ratio 4, global attention. In our
+experiments an 8-block ViT trunk reconstructed 0.5–0.6 dB better than our
+previous [FCDM][fcdm] trunk at matched steps, and extrapolated better (27.9 dB at 256 px
+to 31.2 dB at 1024 px on whole photos).
+
+### 4.2 Convolutional up-path head
+
+**Motivation.** Decoding each fixed 16 × 16 patch from its token, as in our
+previous VAEs, did not give sharp enough detail, nor the right prior for good
+reconstruction from noisy DiT latents. We therefore use a slightly modernized
+[VQGAN][vqgan]-like decoder: the f16 up-path (levels, widths, three residual
+blocks per level) with per-position RMSNorm and folded upsampling (below), and
+our ViT trunk in place of its /16 section (input convolution, mid block and three
+residual-plus-attention blocks, about 5% of its compute at 256 px).
+
+```text
+trunk tokens [1152, H/16, W/16]
+  → RMSNorm → 1×1 conv 1152 → 512 → upsample ×2 (512)              to /8
+  /8   ResBlock 512 → 256, ResBlock 256, ResBlock 256, upsample ×2  to /4
+  /4   ResBlock 256 × 3, upsample ×2                               to /2
+  /2   ResBlock 256 → 128, ResBlock 128, ResBlock 128, upsample ×2  to /1
+  /1   ResBlock 128 × 3
+  → RMSNorm (affine) → SiLU → 3×3 conv 128 → 3 (zero-initialized) → RGB
+```
+
+**Residual block** (VQGAN ResnetBlock, normalization replaced):
+
+$$
+h = \operatorname{conv}_{3\times3}\bigl(\operatorname{SiLU}(\operatorname{rms}(x))\bigr),\quad
+h = \operatorname{conv}_{3\times3}\bigl(\operatorname{SiLU}(\operatorname{rms}(h))\bigr),\quad
+\text{out} = \operatorname{shortcut}(x) + h,
+$$
+
+with a 1×1 convolution shortcut when the width changes and the identity
+otherwise. Each block's second convolution and the output convolution are
+zero-initialized, so a fresh block equals its shortcut and a fresh head outputs
+zeros.
+
+- **Per-position RMSNorm instead of GroupNorm.** Each `rms` normalizes one
+  position's channel vector (learned per-channel scale and shift) and never reads
+  other positions or images. GroupNorm's image-wide statistics make the output
+  depend on image size and distant regions, break tiled decoding, need a
+  reduction that does not fuse, and run in fp32 under autocast. Wan 2.1, KVAE and
+  DC-AE also use RMSNorm in some or all stages.
+- **Upsampling.** Nearest ×2 then a 3 × 3 convolution, as in the FLUX and
+  `semantic_vae` decoders (no transposed-convolution checkerboards). Computed
+  exactly as a 4 × 4, stride-2 transposed convolution with a kernel folded from
+  the 3 × 3 weights: 16 instead of 36 multiply-adds per channel pair and no
+  4×-size intermediate.
+- **Full resolution kept.** In a diffusion-decoder comparison, ending at /2 with
+  a pixel-shuffle readout trained about twice as fast but lost fine detail
+  (lowest-noise validation loss 36% higher, visibly worse samples).
+
+**Parameters.**
+
+| part | parameters | trained |
+| --- | ---: | --- |
+| DINOv3-B backbone (12 blocks, class and 4 register tokens, final norm) | 85.05M | frozen |
+| patch embedding (trainable copy) | 0.59M | yes |
+| layer projection (9216 → 96, with bias) | 0.88M | yes |
+| decoder trunk (1×1 latent projection and 4 ViT blocks at width 1152) | 63.85M | yes |
+| convolutional up-path head | 14.05M | yes |
+| **total** | **164.43M** (79.38M trained) | |
+
+Encoder 86.53M (1.48M trained), decoder 77.90M.
+
+**Resolution extrapolation test.** The decoder accepts any multiple of 16, including sizes
+above its 1024-px training resolution. On 12 Pexels photos it scores 28.55 dB
+PSNR at 512-px area, 29.72 dB at 1024-px area and 31.54 dB at stored size (3.9–4.2
+megapixels). The 4-megapixel
+reconstruction downscaled 2× scores 35.42 dB against the 1024-px input, above
+the native 1024-px decode on all 12 images. 2× crops show no seams, swirls or
+tiling. Rows and columns on 16-px patch boundaries carry 6–12% more error; a
+32-px border ring is 2–3 dB worse.
+
+## 5. Why a pure DINO loss
+
+Compared with diffusion decoders and diffusion + DINO decoders:
+
+1. **Smoother latents.** Diffusion decoders push detail into the latent, which
+   the DiT then has to generate. From the same encoder at matched steps (10k),
+   the DINO-loss latent had median spectral slope 1.04 against
+   0.97 (5th percentile 0.32 against 0.22) and 2.3 dB lower PSNR, but its DINO
+   loss at 10k (0.086) was already below the diffusion decoder's at 85k (0.104).
+   Every DINO-loss latent we benchmarked beat every diffusion-decoder latent on
+   FID, MIND and Monge-DINO.
+2. **Robust to the DiT's high-frequency errors.** Changing the DiT's sampler
+   moved Monge-DINO by 1.5–6.3 with our diffusion decoders and by 0.4, within
+   one standard error, with a DINO-loss decoder.
+3. **One step instead of about four** (our diffusion decoders used 4 Euler
+   steps), with more visible detail in our comparisons, including at 1024 px.
+
+**Trade-off.** At the same latent width, diffusion decoders reach higher PSNR
+(2.3 dB in the matched 10k-step comparison of point 1) and have fewer deterministic artifacts:
+DINO-loss textures can swirl, and text is sometimes distorted.
+
+## 6. Losses
+
+Weights are normalized to DINO = 1. Dion and AdamW are invariant to the overall
+loss scale, so only the ratios matter.
+
+| term | definition | weight | phases |
+| --- | --- | ---: | --- |
+| DINO reconstruction | MSE between frozen DINOv3-B features of the reconstruction and of the input, at all 12 blocks (each through DINOv3's final LayerNorm), averaged over blocks, tokens (patch and class) and channels | **1** | all |
+| pixel MSE | MSE of RGB in [-1, 1] | **0.380** | all |
+| blurred MSE | MSE between both images after a Gaussian blur with σ = 5 px | **0.505** | all |
+| semantic alignment | per-token MSE of the 32 semantic channels to the fixed target (section 3) | **57.7** | joint |
+| VISReg | [VISReg][visreg] on all 96 channels: centring, unit scale and a Gaussian-quantile shape term on 256 random projections, equal weights | **0.0202** | joint |
+
+$$
+\mathcal{L} = \mathcal{L}_{\text{DINO}}
++ 0.380\,\mathcal{L}_{\text{pix}} + 0.505\,\mathcal{L}_{\text{blur}}
++ 57.7\,\mathcal{L}_{\text{align}} + 0.0202\,\mathcal{L}_{\text{VISReg}}
+$$
+
+No diffusion, KL, adversarial or LPIPS loss. The
+DINO loss is not standardized per block: in our experiments, blocks 2–6 carried 70% of its value and block 11 under 1%, following feature
+scale.
+
+### 6.1 Why VISReg
+
+[VISReg][visreg] has a centring term on channel means, a scale term on channel
+standard deviations and a shape term matching the sorted values of 256 random
+1-D projections to standard-normal quantiles, all computed across the batch at
+each spatial position, as in `semantic_vae` (its model card names SIGReg, but its
+released training code uses VISReg).
+
+- **Normalization.** Reconstruction is blind to an invertible per-channel affine
+  change of a deterministic latent, so nothing else fixes the free channels'
+  scale and offset. The centring term is their only mean anchor: in our experiments,
+  when it was outweighed by an alignment loss on a shared encoder, the detail
+  channels' means drifted by up to 3 standard deviations within a few
+  thousand steps.
+- **Anti-collapse.** Projections are taken after per-channel standardization
+  (detached std), so the projection term acts on the correlation matrix \\(R\\):
+  unit-variance random projections ask for \\(u^\top R u \approx 1\\) for random
+  unit \\(u\\), pushing \\(R\\) weakly toward \\(I\\). This flattens the spectrum and raises
+  participation rank, penalizing a free block folded onto a few modes. Quantile
+  matching also constrains higher moments, but in our experiments 91% of the
+  quantile term was this projected-scale mismatch.
+
+In our experiments, weakening VISReg or removing its projection term let the
+free channels collapse: at 1/50 of the usual weight their participation rank fell
+to 11 of 64 (33 at the usual weight), and a latent trained without the projection
+term collapsed to 13 of 96.
+
+We also tested a variant that keeps the centring and scale terms but replaces
+quantile matching by matching only each projection's standard deviation to 1.
+With \\(\hat z\\) the per-channel-standardized latent and \\(u_1, \dots, u_K\\) random
+unit directions,
+
+$$
+\mathcal{L}_{\text{proj}} = \frac{1}{K} \sum_{k=1}^{K}
+\bigl(\operatorname{std}(u_k^\top \hat z) - 1\bigr)^2
+= \frac{1}{K} \sum_{k=1}^{K} \bigl(\sqrt{u_k^\top R\, u_k} - 1\bigr)^2,
+$$
+
+which acts on the correlation matrix \\(R\\) alone, without higher moments. It
+performed similarly to VISReg, confirming the mechanism above. We nevertheless
+kept VISReg, which is cheap and more standard.
+
+### 6.2 How the weights were set
+
+Weights were set from per-term gradient norms, each term back-propagated
+separately at a checkpoint over several image draws. Loss values are a poor
+guide: in our experiments the blurred MSE had a weighted loss of only 2.4e-4 but
+the largest decoder gradient, 3–4 times DINO's.
+
+- **Pixel and blurred MSE:** each a tenth of DINO's gradient norm on the
+  decoder (both had cosine +0.15 to +0.33 with DINO's gradient); on this
+  model's convolutional head at 2k steps they measured 0.112 and 0.116. The pixel
+  MSE acts as a light regulariser on top of the DINO loss; the blurred term is
+  there because DINO losses are nearly blind to colour.
+- **Alignment (57.7):** the ten-times rule of section 3.3.
+- **VISReg (0.0202):** set mostly empirically from the latent's statistics and
+  not extensively ablated. We also tested a roughness penalty on the free
+  channels, a hinge keeping each channel's neighbour R² (how well a token is
+  predicted from its neighbours) above a threshold, and found it unnecessary
+  with the DINO loss.
+
+## 7. Data
+
+- **Sources.** About 14 million images, mostly photographs,
+  plus book covers, a few text-heavy datasets and 1% synthetic rendered text
+  samples (single characters and simple words). A mix of public and licensed or
+  curated datasets; no training images are redistributed with the model.
+- **Aspect-ratio buckets** from 1:4 to 4:1; one aspect ratio and resolution per
+  batch.
+- **Downsampling only.** Antialiased bicubic resizing to the bucket's shape at the
+  target area (e.g. 896 × 1152 → 224 × 288 at 256 px); images smaller than a
+  target are never upsampled.
+- **Resolution mixes.** Base stage: 256-px area for 90% of batches, 384 px for
+  10%. Five-resolution stage: 256 / 384 / 512 / 768 / 1024 px at 20% each (256 as
+  square crops).
+
+## 8. Training
+
+**Optimizer.** Weight matrices use [Dion][dion] / [Dion3][dion3] orthogonalized
+updates with [NorMuon][normuon] normalization, via our single-GPU
+[dionw][dionw], with the update RMS matched to AdamW's so one learning rate
+serves both. Everything else uses AdamW (β = (0.9, 0.98)); no weight decay.
+bf16 AMP, gradient clipping at 5, EMA with decay
+0.9995; the released weights are the EMA.
+
+| stage | resolutions | batch | learning rate | steps |
+| --- | --- | ---: | ---: | ---: |
+| base | 256-px area (90%), 384 px (10%) | 128 | 1e-4 | about 150k |
+| five resolutions | 256 to 1024 px, 20% each | 32 | 2.5e-5 | about 25k |
+
+The free encoder weights (patch embedding and output projection) trained
+jointly with the decoder for the first 52k steps and were then frozen: the
+latent had stabilised, and freezing saves the encoder's backward pass.
+
+## 9. Results
+
+### 9.1 Generation benchmark
+
+We judge a latent by a DiT trained on it: in our experiments, static latent
+statistics (spectral slope, neighbour R², class probes, participation rank, PSNR)
+did not predict DiT quality.
+
+**Protocol.** Our class-conditional, mixed-resolution ImageNet benchmark.[^fid]
+
+- **Generated set:** 10,000 images, 10 per class: 9 at 256-px area and 1 at
+  384 px, in four aspect-ratio families (288 × 224, 288 × 192, 224 × 288 and
+  192 × 288 at 256 px; 416 × 320, 448 × 288, 320 × 416 and 288 × 448 at 384 px)
+  with quotas from their training-data frequency. Every model uses the same
+  class, shape and seed manifest.
+- **Reference set:** 810,000 real ImageNet training images (729,000 at 256 px and 81,000 at 384 px, 81 times each shape's generated
+  count), preprocessed as in training.
+- **Metrics** (lower is better): [FID][fid] on Inception pool3 features; MIND
+  ([Monge Inception Distance][mind]), a sliced optimal-transport distance on the
+  same features (1,024 fixed random projections); Monge-DINO, the same
+  estimator on final-normalized DINOv3-B class tokens of 224-px centre crops. The
+  Monge ± values reflect only the randomness of the projections.
+- **Guidance.** Path-drop guidance (PDG) contrasts the full model with the same model skipping its middle blocks,
+  both class-conditioned; no classifier-free guidance.
+
+**The DiT.** A class-conditional flow-matching transformer with
+[SPRINT][sprint] token routing, 190.43M parameters:
+
+| setting | value |
+| --- | --- |
+| layers | 16 at width 896: 2 dense prefix, 12 routed middle, 2 dense suffix |
+| attention | 14 heads of width 64; additive 2D sin-cos positions and axial 2D RoPE |
+| MLP | GELU, ratio 4 |
+| class conditioning | 4 learned tokens of width 256 per class, cross-attention (4 heads of width 64) in every layer; class dropped for 10% of samples |
+| time conditioning | AdaLN per layer from a 256-wide time embedding |
+| input / output | one token per latent cell (96 channels), latents whitened per channel with the VAE's statistics; velocity prediction |
+| SPRINT mix, per sample | 80%: middle layers see one random token per 2 × 2 group (75% dropped); 10%: full grid; 10%: middle layers skipped. Middle outputs are scattered back, concatenated with the prefix output, scaled per channel by the timestep and projected back to width 896 |
+| objective | flow matching, MSE on velocity, Beta(2, 2) timesteps with a resolution-dependent log-SNR shift relative to 256 px |
+| data | ImageNet-1k (1.24M training images), aspect-ratio buckets, 90% at 256-px area, 10% at 384 px |
+| optimizer | AdamW, learning rate 1e-4, β = (0.9, 0.98), weight decay 0, constant after 4,000 warmup steps, gradient clipping at 1 |
+| batch, steps | 256, 100,000 |
+| precision | bf16 AMP |
+| EMA | 2,000-step half-life (decay ≈ 0.99965), from step 20,000; EMA weights used |
+
+The latent has been fixed since the joint phase, so the released decoder
+decodes exactly the latent the DiT learned.
+
+`semantic_vae`'s DiT uses the same recipe, up to minor architectural details.
+
+**Scores.** For each model: 50 NFEs, PDG 2.5–4.0, best settings kept. The DiT's
+samples were decoded by the released convolutional decoder. Bold marks the
+better value in each column.
+
+| model | FID ↓ | MIND ↓ | Monge-DINOv3-B ↓ |
+| --- | ---: | ---: | ---: |
+| `dinac3_96` | **9.04** | 8.29 ± 0.19 | **10.92** ± 0.46 |
+| `semantic_vae` | 10.02 | **7.42** ± 0.20 | 15.94 ± 0.66 |
+
+Monge-DINO uses DINOv3-B features, as does the decoder's training loss, so it is
+not independent of it; FID and MIND, on Inception features, are.
+
+**Blind comparison.** Our blind pairwise comparison of the two models'
+generations put them on equal footing.
+
+**Limits.** One DiT seed per row; the benchmark's own noise has not been
+measured. We treat FID differences of 0.1 to 0.2 as noise.
+
+### 9.2 Reconstruction
+
+The main tables are in the [README](README.md#reconstruction); this section adds
+detail.
+
+**10,000 ImageNet images.** Originals preprocessed like the generation
+benchmark's reference images. PSNR spread and the Monge uncertainties (from the
+projections):
+
+| model | rMIND | rMonge-DINOv3-B | PSNR median | P5 | P95 |
+| --- | ---: | ---: | ---: | ---: | ---: |
+| `dinac3_96` | **0.166** ± 0.002 | **1.205** ± 0.044 | **28.07** | **21.38** | **35.58** |
+| `semantic_vae` | 0.820 ± 0.028 | 5.994 ± 0.216 | 25.09 | 18.83 | 31.82 |
+
+**2k PSNR set.** 1,330 Pexels photos, 669 book covers and one printer calibration test image,
+centre-cropped to a multiple of 16 at stored size (mostly about 1 megapixel).
+PSNR between the float [-1, 1] input and the unclamped reconstruction, peak 2,
+the DINAC-AE-D2 card's protocol. The DINAC-AE-D2 and FLUX.2 rows are quoted from
+the [DINAC-AE-D2 card](https://huggingface.co/data-archetype/dinac_ae_d2); our
+re-run of DINAC-AE-D2 reproduced it within 0.02 dB. The FLUX.2 row most likely
+used a clamped reconstruction (inferred from its value); clamped, `dinac3_96`
+scores 32.17 dB.
+
+**By resolution.** Each 2k image resized down to an aspect-ratio bucket of the
+given area; images too small are skipped, so subsets differ between rows
+(compare within a row). Mean PSNR in dB:
+
+| bucket | images | `dinac3_96` | `semantic_vae` | DINAC-AE-D2 |
+| --- | ---: | ---: | ---: | ---: |
+| 256 | 2,000 | 29.27 | 25.48 | 32.63 |
+| 512 | 1,331 | 32.48 | 28.41 | 35.72 |
+| 768 | 1,331 | 33.80 | 29.40 | 37.20 |
+| 1024 | 1,124 | 34.06 | 29.81 | 37.45 |
+
+**PSNR.** A decoder trained mainly by a DINO feature loss trades some pixel
+accuracy for perceptual fidelity: on the 2k set `dinac3_96` is 3.4 dB below the
+pixel-trained DINAC-AE-D2 and about 4 dB below FLUX.2 (the trade-off of section
+5). Text is the visible weakness: glyphs are sometimes distorted, plausibly
+because the loss encodes text style more than glyph identity.
+
+## References
+
+1. well9472. [semantic_vae][semantic-vae]. Model and training notes.
+2. O. Siméoni et al. [DINOv3][dinov3]. arXiv:2508.10104, 2025.
+3. Y. Yu, W. Xiong, W. Nie et al. [PixelDiT: Pixel Diffusion Transformers for Image Generation][pixeldit]. arXiv:2511.20645, 2025.
+4. Kwon et al. [Reviving ConvNeXt for Efficient Convolutional Diffusion Models][fcdm] (FCDM). arXiv:2603.09408, 2026.
+5. K. Ahn, B. Xu, N. Abreu, Y. Fan et al. [Dion: Distributed Orthonormalized Updates][dion]. arXiv:2504.05295, 2025.
+6. N. Amsel, J. Zhang, K. Ahn, A. Naeimi et al. [Dion3: Full-Stack Orthogonal Updates][dion3]. arXiv:2608.11612, 2026.
+7. Z. Li, L. Liu, C. Liang, W. Chen et al. [NorMuon: Making Muon more efficient and scalable][normuon]. arXiv:2510.05491, 2025.
+8. [dionw][dionw]: single-GPU Dion with fused AdamW.
+9. Wu, Balestriero and Levine. [VISReg: Variance-Invariance-Sketching Regularization for JEPA training][visreg]. arXiv:2606.02572, 2026.
+10. Park et al. [Sprint: Sparse-Dense Residual Fusion for Efficient Diffusion Transformers][sprint]. arXiv:2510.21986, 2025.
+11. Berthet et al. [MIND: Monge Inception Distance for Generative Models Evaluation][mind]. arXiv:2605.06797, 2026.
+12. M. Heusel et al. [GANs Trained by a Two Time-Scale Update Rule Converge to a Local Nash Equilibrium][fid] (FID). arXiv:1706.08500, 2017.
+13. P. Esser, R. Rombach, B. Ommer. [Taming Transformers for High-Resolution Image Synthesis][vqgan] (VQGAN). arXiv:2012.09841, 2020.
+14. Data Archetype. [DINAC-AE-D2][dinac-d2].
+
+[^fid]: Papers usually report FID on 50,000 square samples at a single
+  resolution, so our numbers are not directly comparable with published ones.
+
+[semantic-vae]: https://huggingface.co/well9472/semantic_vae
+[irdiffae]: https://huggingface.co/data-archetype/irdiffae-v1/blob/main/technical_report.md
+[dinov3]: https://arxiv.org/abs/2508.10104
+[pixeldit]: https://arxiv.org/abs/2511.20645
+[fcdm]: https://arxiv.org/abs/2603.09408
+[dion]: https://arxiv.org/abs/2504.05295
+[dion3]: https://arxiv.org/abs/2608.11612
+[normuon]: https://arxiv.org/abs/2510.05491
+[dionw]: https://github.com/JTriggerFish/dionw
+[visreg]: https://arxiv.org/abs/2606.02572
+[sprint]: https://arxiv.org/abs/2510.21986
+[mind]: https://arxiv.org/abs/2605.06797
+[fid]: https://arxiv.org/abs/1706.08500
+[vqgan]: https://arxiv.org/abs/2012.09841
+[dinac-d2]: https://huggingface.co/data-archetype/dinac_ae_d2

config.json

@@ -0,0 +1,19 @@
+{
+  "format_version": 1,
+  "backbone": "vit_base_patch16_dinov3",
+  "latent_channels": 96,
+  "semantic_channels": 32,
+  "decoder_width": 1152,
+  "decoder_depth": 4,
+  "decoder_head_dim": 64,
+  "decoder_mlp_ratio": 4.0,
+  "conv_up_handoff_channels": 512,
+  "conv_up_channels": [
+    256,
+    256,
+    128,
+    128
+  ],
+  "conv_up_blocks_per_level": 3,
+  "latent_stats_eps": 0.0001
+}

dinac3/__init__.py

@@ -0,0 +1,6 @@
+"""dinac3: deterministic image autoencoders with a DINOv3-aligned latent."""
+
+from .config import CONFIG_FORMAT_VERSION, PATCH, Backbone, Dinac3Config
+from .model import Dinac3
+
+__all__ = ["CONFIG_FORMAT_VERSION", "PATCH", "Backbone", "Dinac3", "Dinac3Config"]

dinac3/config.py

@@ -0,0 +1,195 @@
+"""Explicit architecture configuration for dinac3 autoencoders."""
+
+from __future__ import annotations
+
+import json
+import math
+from dataclasses import asdict, dataclass, fields
+from enum import Enum
+from typing import TYPE_CHECKING, cast
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+# Pixels per latent token along each side (DINOv3-B patch size).
+PATCH = 16
+# Version of the config.json schema this package reads and writes.
+CONFIG_FORMAT_VERSION = 1
+
+
+class Backbone(str, Enum):
+    """Supported frozen encoder backbones (timm architecture names)."""
+
+    DINO_V3_B = "vit_base_patch16_dinov3"
+
+
+def _require_int(data: dict[str, object], name: str) -> int:
+    """An integer JSON field (not a float such as ``64.0``, not a boolean).
+
+    Raises:
+        TypeError: If the value is not a JSON integer.
+    """
+
+    value = data[name]
+    if type(value) is not int:
+        raise TypeError(f"Config field {name} must be an integer, got {value!r}")
+    return value
+
+
+def _require_int_tuple(data: dict[str, object], name: str) -> tuple[int, ...]:
+    """A JSON list of integers, as a tuple.
+
+    Raises:
+        TypeError: If the value is not a list of JSON integers.
+    """
+
+    value = data[name]
+    if not isinstance(value, list) or any(
+        type(item) is not int for item in cast("list[object]", value)
+    ):
+        raise TypeError(f"Config field {name} must be a list of integers: {value!r}")
+    return tuple(cast("list[int]", value))
+
+
+def _require_float(data: dict[str, object], name: str) -> float:
+    """A real-number JSON field (an integer literal is accepted, a boolean not).
+
+    Raises:
+        TypeError: If the value is not a JSON number.
+    """
+
+    value = data[name]
+    if type(value) not in (int, float):
+        raise TypeError(f"Config field {name} must be a number, got {value!r}")
+    return float(cast("float", value))
+
+
+def _require_str(data: dict[str, object], name: str) -> str:
+    """A string JSON field.
+
+    Raises:
+        TypeError: If the value is not a JSON string.
+    """
+
+    value = data[name]
+    if not isinstance(value, str):
+        raise TypeError(f"Config field {name} must be a string, got {value!r}")
+    return value
+
+
+@dataclass(frozen=True)
+class Dinac3Config:
+    """Architecture of one dinac3 model; every persisted field is required.
+
+    The latent has ``latent_channels`` channels at stride 16: the first
+    ``latent_channels - semantic_channels`` are free (detail) channels, the last
+    ``semantic_channels`` the DINOv3-aligned semantic channels.
+    """
+
+    format_version: int
+    backbone: Backbone
+    latent_channels: int
+    semantic_channels: int
+    decoder_width: int
+    decoder_depth: int
+    decoder_head_dim: int
+    decoder_mlp_ratio: float
+    conv_up_handoff_channels: int
+    conv_up_channels: tuple[int, ...]
+    conv_up_blocks_per_level: int
+    latent_stats_eps: float
+
+    def __post_init__(self) -> None:
+        """Reject unsupported topologies and invalid scalars.
+
+        Raises:
+            ValueError: On another format version, a non-positive width, depth
+                or count, a semantic block not smaller than the latent, a width
+                not divisible into heads, a head width not a multiple of 4, a
+                number of up-path levels other than log2(16) = 4, or a
+                non-finite or non-positive ratio or epsilon.
+            TypeError: If ``backbone`` is not a :class:`Backbone` or
+                ``conv_up_channels`` is not a tuple.
+        """
+
+        if self.format_version != CONFIG_FORMAT_VERSION:
+            raise ValueError(
+                f"config format_version {self.format_version} is not supported; "
+                f"this package reads version {CONFIG_FORMAT_VERSION}"
+            )
+        if not isinstance(self.backbone, Backbone):
+            raise TypeError("backbone must be a Backbone enum")
+        counts = (
+            self.latent_channels,
+            self.semantic_channels,
+            self.decoder_width,
+            self.decoder_depth,
+            self.decoder_head_dim,
+            self.conv_up_handoff_channels,
+            self.conv_up_blocks_per_level,
+            *self.conv_up_channels,
+        )
+        if any(value <= 0 for value in counts):
+            raise ValueError("Every width, depth and count must be positive")
+        if self.semantic_channels >= self.latent_channels:
+            raise ValueError("semantic_channels must be smaller than latent_channels")
+        if self.decoder_width % self.decoder_head_dim:
+            raise ValueError("decoder_width must be a multiple of decoder_head_dim")
+        if self.decoder_head_dim % 4:
+            raise ValueError("decoder_head_dim must be a multiple of 4 (2-D RoPE)")
+        if not isinstance(self.conv_up_channels, tuple):
+            raise TypeError("conv_up_channels must be a tuple of level widths")
+        if 2 ** len(self.conv_up_channels) != PATCH:
+            raise ValueError(
+                f"conv_up_channels needs one width per level from /{PATCH // 2} to "
+                f"/1 ({PATCH.bit_length() - 1} levels)"
+            )
+        scalars = (self.decoder_mlp_ratio, self.latent_stats_eps)
+        if not all(math.isfinite(v) and v > 0 for v in scalars):
+            raise ValueError(
+                "decoder_mlp_ratio and latent_stats_eps must be finite and positive"
+            )
+
+    @property
+    def free_channels(self) -> int:
+        """Number of leading free (detail) latent channels."""
+
+        return self.latent_channels - self.semantic_channels
+
+    def save(self, path: Path) -> None:
+        """Write every field as JSON."""
+
+        path.write_text(json.dumps(asdict(self), indent=2) + "\n")
+
+    @classmethod
+    def load(cls, path: Path) -> Dinac3Config:
+        """Load strictly: missing or unknown fields and wrong JSON types are errors.
+
+        Raises:
+            ValueError: If the JSON fields differ from the dataclass fields or
+                the format version is not this package's.
+            TypeError: If a field has the wrong JSON type (e.g. ``64.0`` or
+                ``true`` for an integer).
+        """
+
+        data = json.loads(path.read_text())
+        if not isinstance(data, dict):
+            raise TypeError(f"{path} must hold a JSON object")
+        values = cast("dict[str, object]", data)
+        expected = {field.name for field in fields(cls)}
+        if set(values) != expected:
+            raise ValueError(f"Config fields must be exactly {sorted(expected)}")
+        return cls(
+            format_version=_require_int(values, "format_version"),
+            backbone=Backbone(_require_str(values, "backbone")),
+            latent_channels=_require_int(values, "latent_channels"),
+            semantic_channels=_require_int(values, "semantic_channels"),
+            decoder_width=_require_int(values, "decoder_width"),
+            decoder_depth=_require_int(values, "decoder_depth"),
+            decoder_head_dim=_require_int(values, "decoder_head_dim"),
+            decoder_mlp_ratio=_require_float(values, "decoder_mlp_ratio"),
+            conv_up_handoff_channels=_require_int(values, "conv_up_handoff_channels"),
+            conv_up_channels=_require_int_tuple(values, "conv_up_channels"),
+            conv_up_blocks_per_level=_require_int(values, "conv_up_blocks_per_level"),
+            latent_stats_eps=_require_float(values, "latent_stats_eps"),
+        )

dinac3/conv_up_head.py

@@ -0,0 +1,180 @@
+"""Deterministic convolutional up-path head: trunk tokens at /16 to RGB.
+
+A handoff (per-position RMSNorm, 1x1 projection, upsampling to /8), then one
+level per stride (/8, /4, /2, /1) of residual blocks, upsampling between levels,
+and a readout (per-position RMSNorm, SiLU, 3x3 convolution to RGB). Every
+upsampling is nearest x2 followed by a 3x3 convolution, computed exactly as a
+stride-2 transposed convolution whose 4x4 kernel folds the 3x3 weights. No
+attention, no time conditioning, no noise input: one deterministic pass.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+from torch import Tensor, nn
+from torch.nn import functional as F
+
+from .layers import ChannelRMSNorm, pointwise_conv, same_conv
+
+if TYPE_CHECKING:
+    from .config import Dinac3Config
+
+# Row k maps a 3x3 kernel's taps (offsets -1, 0, +1) onto tap k of a 4-tap,
+# stride-2 transposed convolution: nearest x2 followed by the 3x3 convolution.
+NEAREST_CONV_FOLD = (
+    (0.0, 0.0, 1.0),
+    (0.0, 1.0, 1.0),
+    (1.0, 1.0, 0.0),
+    (1.0, 0.0, 0.0),
+)
+
+
+class UpsampleConv(nn.Module):
+    """Nearest x2 then a 3x3 convolution, as one folded transposed convolution.
+
+    The 4x4 transposed kernel is folded once, at load time (:meth:`fold_kernel`),
+    from the FP32 3x3 weight in IEEE FP32 arithmetic, so a global TF32 matmul
+    setting cannot change it; the transposed convolution then runs in the
+    activation precision. Loading a state dict into a folded module folds again,
+    so the kernel never goes stale.
+    """
+
+    fold: Tensor
+    kernel: Tensor | None
+
+    def __init__(self, channels: int) -> None:
+        """Allocate the 3x3 convolution, the constant fold matrix and the
+        (not yet folded) kernel buffer."""
+
+        super().__init__()
+        self.conv = nn.Conv2d(channels, channels, 3, padding=1)
+        self.register_buffer("fold", torch.tensor(NEAREST_CONV_FOLD))
+        self.register_buffer("kernel", None, persistent=False)
+        self.register_load_state_dict_post_hook(_refold_after_load)
+
+    @torch.no_grad()
+    def fold_kernel(self) -> None:
+        """Fold the 3x3 weight into the ``[C_in, C_out, 4, 4]`` transposed kernel
+        (IEEE FP32 matmuls for the fold, restored afterwards).
+
+        Raises:
+            ValueError: If the stored fold matrix is not ``NEAREST_CONV_FOLD``.
+        """
+
+        expected = torch.tensor(NEAREST_CONV_FOLD, device=self.fold.device)
+        if not torch.equal(self.fold.float(), expected):
+            raise ValueError("UpsampleConv.fold is not the nearest x2 fold matrix")
+        matmul = torch.backends.cuda.matmul
+        previous = matmul.fp32_precision
+        matmul.fp32_precision = "ieee"
+        try:
+            with torch.autocast(device_type=self.fold.device.type, enabled=False):
+                fold = self.fold.to(dtype=self.conv.weight.dtype)
+                kernel = torch.einsum("ka,ocab,lb->ockl", fold, self.conv.weight, fold)
+        finally:
+            matmul.fp32_precision = previous
+        self.kernel = kernel.transpose(0, 1)
+
+    def forward(self, x: Tensor) -> Tensor:
+        """``[B, C, H, W]`` to ``[B, C, 2H, 2W]``.
+
+        Raises:
+            RuntimeError: Before :meth:`fold_kernel` (done by ``Dinac3.prepare``).
+        """
+
+        kernel = self.kernel
+        if kernel is None:
+            raise RuntimeError("UpsampleConv.fold_kernel() has not run")
+        return F.conv_transpose2d(x, kernel, self.conv.bias, stride=2, padding=1)
+
+
+def _refold_after_load(module: nn.Module, _incompatible: object) -> None:
+    """After a state-dict load (children included), fold an already folded
+    upsampler again so its kernel never goes stale.
+
+    Raises:
+        TypeError: If registered on another module.
+    """
+
+    if not isinstance(module, UpsampleConv):
+        raise TypeError(f"Fold hook registered on {type(module)}")
+    if module.kernel is not None:
+        module.fold_kernel()
+
+
+class ResBlock(nn.Module):
+    """``shortcut(x) + conv2(SiLU(norm2(conv1(SiLU(norm1(x))))))``, with
+    per-position RMSNorms carrying per-channel gains and biases, 3x3
+    convolutions and a 1x1 shortcut when the width changes."""
+
+    def __init__(self, in_channels: int, out_channels: int) -> None:
+        """Allocate norms, convolutions and the optional shortcut."""
+
+        super().__init__()
+        self.norm1 = ChannelRMSNorm(in_channels, affine=True)
+        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, padding=1)
+        self.norm2 = ChannelRMSNorm(out_channels, affine=True)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, 3, padding=1)
+        self.shortcut = (
+            nn.Conv2d(in_channels, out_channels, 1)
+            if in_channels != out_channels
+            else None
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        """Residual update of a ``[B, C, H, W]`` map."""
+
+        h = same_conv(self.conv1, F.silu(self.norm1(x)))
+        h = same_conv(self.conv2, F.silu(self.norm2(h)))
+        skip = x if self.shortcut is None else pointwise_conv(self.shortcut, x)
+        return skip + h
+
+
+class ConvUpHead(nn.Module):
+    """Trunk tokens ``[B, width, h, w]`` to RGB ``[B, 3, 16 h, 16 w]``."""
+
+    def __init__(self, config: Dinac3Config) -> None:
+        """Allocate handoff, residual blocks, upsamplers and readout."""
+
+        super().__init__()
+        handoff = config.conv_up_handoff_channels
+        self.blocks_per_level = config.conv_up_blocks_per_level
+        self.handoff_norm = ChannelRMSNorm(config.decoder_width, affine=False)
+        self.handoff_proj = nn.Conv2d(config.decoder_width, handoff, 1)
+        self.handoff_upsample = UpsampleConv(handoff)
+        blocks: list[ResBlock] = []
+        upsamples: list[UpsampleConv] = []
+        width = handoff
+        for index, channels in enumerate(config.conv_up_channels):
+            for _ in range(self.blocks_per_level):
+                blocks.append(ResBlock(width, channels))
+                width = channels
+            if index + 1 < len(config.conv_up_channels):
+                upsamples.append(UpsampleConv(width))
+        self.blocks = nn.ModuleList(blocks)
+        self.upsamples = nn.ModuleList(upsamples)
+        self.out_norm = ChannelRMSNorm(width, affine=True)
+        self.out_proj = nn.Conv2d(width, 3, 3, padding=1)
+
+    def fold_kernels(self) -> None:
+        """Fold every upsampler's kernel (after loading and any device move)."""
+
+        self.handoff_upsample.fold_kernel()
+        for upsampler in self.upsamples:
+            if not isinstance(upsampler, UpsampleConv):
+                raise TypeError(f"Unexpected upsampler {type(upsampler)}")
+            upsampler.fold_kernel()
+
+    def forward(self, features: Tensor) -> Tensor:
+        """Decode a contiguous NCHW token map."""
+
+        h = pointwise_conv(self.handoff_proj, self.handoff_norm(features))
+        h = self.handoff_upsample(h)
+        for index, block in enumerate(self.blocks):
+            h = block(h)
+            level, position = divmod(index, self.blocks_per_level)
+            if position == self.blocks_per_level - 1 and level < len(self.upsamples):
+                h = self.upsamples[level](h)
+        return same_conv(self.out_proj, F.silu(self.out_norm(h)))

dinac3/decoder.py

@@ -0,0 +1,45 @@
+"""Deterministic one-pass decoder: latent projection, ViT trunk and conv up-path."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from torch import Tensor, nn
+
+from .conv_up_head import ConvUpHead
+from .trunk import DitBlock, rope_tables
+
+if TYPE_CHECKING:
+    from .config import Dinac3Config
+
+
+class Decoder(nn.Module):
+    """Raw latents ``[B, C, h, w]`` to RGB ``[B, 3, 16 h, 16 w]``."""
+
+    def __init__(self, config: Dinac3Config) -> None:
+        """Allocate the latent projection, trunk blocks and conv up-path head."""
+
+        super().__init__()
+        width = config.decoder_width
+        self.head_dim = config.decoder_head_dim
+        self.latent_up = nn.Conv2d(config.latent_channels, width, 1)
+        self.trunk = nn.ModuleList(
+            [
+                DitBlock(width, config.decoder_head_dim, config.decoder_mlp_ratio)
+                for _ in range(config.decoder_depth)
+            ]
+        )
+        self.conv_up_head = ConvUpHead(config)
+
+    def forward(self, latents: Tensor) -> Tensor:
+        """Decode in one pass; latents are validated by the public API."""
+
+        x = self.latent_up(latents)
+        b, c, h, w = x.shape
+        sin, cos = rope_tables(h, w, head_dim=self.head_dim, device=x.device)
+        tokens = x.permute(0, 2, 3, 1).reshape(b, h * w, c)
+        for block in self.trunk:
+            tokens = block(tokens, sin, cos)
+        # The head reads the trunk's tokens as a contiguous NCHW map.
+        features = tokens.transpose(1, 2).reshape(b, c, h, w).contiguous()
+        return self.conv_up_head(features)

dinac3/encoder.py

@@ -0,0 +1,137 @@
+"""Frozen DINOv3-B behind a trained patch embedding, and the layer projection.
+
+Every block's patch tokens (under the final LayerNorm) are standardized per
+channel with fixed statistics, concatenated (12 x 768) and mapped by one linear
+projection to the latent channels.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, cast
+
+import timm
+import torch
+from timm.layers.pos_embed_sincos import RotaryEmbeddingDinoV3
+from torch import Tensor, nn
+
+from .config import PATCH
+
+if TYPE_CHECKING:
+    from timm.models.eva import Eva
+
+    from .config import Dinac3Config
+
+IMAGENET_MEAN = (0.485, 0.456, 0.406)
+IMAGENET_STD = (0.229, 0.224, 0.225)
+
+
+def _require_supported_backbone(backbone: Eva) -> None:
+    """Fail fast if timm builds a backbone the explicit forward does not cover.
+
+    Raises:
+        ValueError: On absolute position embeddings, patch dropout, a pre-norm,
+            missing CLS/register tokens or a RoPE other than DINOv3's without
+            coordinate augmentation.
+    """
+
+    supported = (
+        backbone.pos_embed is None
+        and backbone.patch_drop is None
+        and isinstance(backbone.norm_pre, nn.Identity)
+        and backbone.cls_token is not None
+        and backbone.reg_token is not None
+        and isinstance(backbone.rope, RotaryEmbeddingDinoV3)
+        and not backbone.rope.aug_active
+    )
+    if not supported:
+        raise ValueError("Unsupported timm DINOv3 backbone structure for dinac3")
+
+
+class Encoder(nn.Module):
+    """Images in [-1, 1] to raw latents ``[B, C, H / 16, W / 16]``."""
+
+    pixel_mean: Tensor
+    pixel_std: Tensor
+    layer_mean: Tensor
+    layer_std: Tensor
+
+    def __init__(self, config: Dinac3Config) -> None:
+        """Build the architecture only; every tensor comes from the artifact."""
+
+        super().__init__()
+        self.backbone = cast(
+            "Eva",
+            timm.create_model(
+                config.backbone.value,
+                pretrained=False,
+                num_classes=0,
+                dynamic_img_size=True,
+                dynamic_img_pad=True,
+            ),
+        )
+        _require_supported_backbone(self.backbone)
+        # timm keeps RoPE periods non-persistent; the artifact stores the exact
+        # values the training run used.
+        rope = cast("nn.Module", self.backbone.rope)
+        rope.register_buffer("periods", cast("Tensor", rope.periods), persistent=True)
+        self.prefix_tokens = self.backbone.num_prefix_tokens
+        width = self.backbone.embed_dim * len(self.backbone.blocks)
+        self.register_buffer("pixel_mean", torch.tensor(IMAGENET_MEAN).view(1, 3, 1, 1))
+        self.register_buffer("pixel_std", torch.tensor(IMAGENET_STD).view(1, 3, 1, 1))
+        self.register_buffer("layer_mean", torch.zeros(width))
+        self.register_buffer("layer_std", torch.ones(width))
+        self.projection = nn.Linear(width, config.latent_channels)
+        self._rope_cache: dict[tuple[int, int, torch.device, torch.dtype], Tensor] = {}
+
+    def rope_embed(self, height: int, width: int) -> Tensor:
+        """The DINOv3 RoPE table of an image size, built once per token grid
+        (outside any compiled graph: timm builds coordinates on the CPU)."""
+
+        rope = cast("RotaryEmbeddingDinoV3", self.backbone.rope)
+        periods = cast("Tensor", rope.periods)
+        key = (height // PATCH, width // PATCH, periods.device, periods.dtype)
+        cached = self._rope_cache.get(key)
+        if cached is None:
+            with torch.inference_mode(False), torch.no_grad():
+                cached = rope.get_embed(shape=[key[0], key[1]])
+            self._rope_cache[key] = cached
+        return cached
+
+    def forward(self, images: Tensor, rope: Tensor) -> Tensor:
+        """Raw latents in the activation dtype; ``rope`` from :meth:`rope_embed`."""
+
+        dtype = self._prefix_tokens()[0].dtype
+        x = ((images.float().add(1.0).mul(0.5) - self.pixel_mean) / self.pixel_std).to(
+            dtype=dtype
+        )
+        features = self._block_features(x, rope)
+        tokens = torch.cat([feature.to(dtype=dtype) for feature in features], dim=-1)
+        standardized = (tokens.float() - self.layer_mean) / self.layer_std
+        latents = self.projection(standardized)
+        b, _, height, width = images.shape
+        return latents.transpose(1, 2).reshape(b, -1, height // PATCH, width // PATCH)
+
+    def _block_features(self, x: Tensor, rope: Tensor) -> list[Tensor]:
+        """Final-normed patch tokens of every block (timm's intermediates path)."""
+
+        backbone = self.backbone
+        tokens = backbone.patch_embed(x)
+        b, _, _, c = tokens.shape
+        tokens = tokens.view(b, -1, c)
+        cls_token, reg_token = self._prefix_tokens()
+        tokens = torch.cat(
+            [cls_token.expand(b, -1, -1), reg_token.expand(b, -1, -1), tokens], dim=1
+        )
+        features: list[Tensor] = []
+        for block in backbone.blocks:
+            tokens = block(tokens, rope=rope)
+            features.append(backbone.norm(tokens)[:, self.prefix_tokens :])
+        return features
+
+    def _prefix_tokens(self) -> tuple[Tensor, Tensor]:
+        """The CLS and register tokens (present: checked at construction)."""
+
+        cls_token, reg_token = self.backbone.cls_token, self.backbone.reg_token
+        if cls_token is None or reg_token is None:
+            raise RuntimeError("The DINOv3 backbone lost its CLS or register tokens")
+        return cls_token, reg_token

dinac3/layers.py

@@ -0,0 +1,97 @@
+"""Normalization, MLP and convolution primitives with the training arithmetic.
+
+Each op reproduces the training model's dtype handling under BF16 autocast
+(reductions in FP32, results in the activation dtype), so BF16 inference matches
+the training precision bit for bit.
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import Tensor, nn
+from torch.nn import functional as F
+
+NORM_EPS = 1e-6
+
+
+class RMSNorm(nn.Module):
+    """RMS normalization over the last dimension, optional per-channel gain.
+
+    The gain is cast to the input dtype and the result returned in it.
+    """
+
+    def __init__(self, dim: int, *, affine: bool) -> None:
+        """Build the norm; ``affine`` adds a gain initialized to one."""
+
+        super().__init__()
+        self.dim = dim
+        self.weight = nn.Parameter(torch.ones(dim)) if affine else None
+
+    def forward(self, x: Tensor) -> Tensor:
+        """Normalize ``x`` ``[..., dim]``."""
+
+        weight = None if self.weight is None else self.weight.to(dtype=x.dtype)
+        return F.rms_norm(x, (self.dim,), weight, NORM_EPS).to(dtype=x.dtype)
+
+
+class ChannelRMSNorm(nn.Module):
+    """Per-position RMS normalization over the channels of ``[B, C, H, W]``,
+    reduced in FP32 and applied in the input dtype, with an optional
+    per-channel gain and bias (cast to the input dtype)."""
+
+    def __init__(self, channels: int, *, affine: bool) -> None:
+        """Build the norm; ``affine`` adds a gain (one) and a bias (zero)."""
+
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(channels)) if affine else None
+        self.bias = nn.Parameter(torch.zeros(channels)) if affine else None
+
+    def forward(self, x: Tensor) -> Tensor:
+        """Normalize each position's channel vector."""
+
+        mean_square = torch.mean(
+            torch.square(x), dim=1, keepdim=True, dtype=torch.float32
+        )
+        y = x * torch.rsqrt(mean_square + NORM_EPS).to(dtype=x.dtype)
+        if self.weight is not None and self.bias is not None:
+            y = y * self.weight.view(1, -1, 1, 1).to(dtype=x.dtype)
+            y = y + self.bias.view(1, -1, 1, 1).to(dtype=x.dtype)
+        return y
+
+
+class Mlp(nn.Module):
+    """Two linear layers with a GELU between them, both with biases."""
+
+    def __init__(self, dim: int, hidden: int) -> None:
+        """Allocate the up and down projections."""
+
+        super().__init__()
+        self.up = nn.Linear(dim, hidden)
+        self.down = nn.Linear(hidden, dim)
+
+    def forward(self, x: Tensor) -> Tensor:
+        """``down(gelu(up(x)))``."""
+
+        return self.down(F.gelu(self.up(x)))
+
+
+def pointwise_conv(module: nn.Conv2d, x: Tensor) -> Tensor:
+    """A 1x1 convolution as a channel matmul (as the training graph runs it)."""
+
+    y = torch.einsum("bchw,oc->bohw", x, module.weight.flatten(1))
+    if module.bias is not None:
+        y = y + module.bias.view(1, -1, 1, 1).to(dtype=y.dtype)
+    return y
+
+
+def same_conv(module: nn.Conv2d, x: Tensor) -> Tensor:
+    """A stride-1, same-padded convolution with ``module``'s parameters."""
+
+    return F.conv2d(
+        x,
+        module.weight,
+        module.bias,
+        stride=1,
+        padding=module.kernel_size[0] // 2,
+        groups=module.groups,
+    )

dinac3/model.py

@@ -0,0 +1,370 @@
+"""Loading and inference API of dinac3 autoencoders."""
+
+from __future__ import annotations
+
+import os
+from collections.abc import Callable
+from pathlib import Path
+from typing import NoReturn
+
+import torch
+from safetensors.torch import load_file
+from torch import Tensor, nn
+from torch.fx.experimental import _config as fx_config
+
+from .config import PATCH, Dinac3Config
+from .decoder import Decoder
+from .encoder import Encoder
+from .precision import require_storage_policy
+
+CONFIG_FILENAME = "config.json"
+WEIGHTS_FILENAME = "model.safetensors"
+SUPPORTED_DTYPES = (torch.bfloat16, torch.float32)
+
+
+def _looks_like_local_path(name: str) -> bool:
+    """Whether a string names a filesystem path rather than a Hub ``org/name``
+    id: it starts with ``.``, ``/`` or ``~``, uses a path separator other than
+    the id's single ``/``, or has more than one ``/``."""
+
+    other_separator = os.sep != "/" and os.sep in name
+    return name.startswith((".", "/", "~")) or other_separator or name.count("/") > 1
+
+
+def resolve_model_dir(path_or_repo: str | Path, *, revision: str | None) -> Path:
+    """A local artifact directory, or a Hugging Face Hub snapshot of a repo id.
+
+    A ``Path`` is always local. A ``str`` is local when it names an existing
+    directory; a string that looks like a path but names none is an error;
+    any other string is a Hub repository id, of which only ``config.json`` and
+    ``model.safetensors`` are downloaded.
+
+    Raises:
+        FileNotFoundError: If a local path does not name a directory.
+        ValueError: If ``revision`` is given for a local directory.
+    """
+
+    match path_or_repo:
+        case Path() as directory:
+            local = directory.expanduser()
+        case str() as name if Path(name).expanduser().is_dir():
+            local = Path(name).expanduser()
+        case str() as name if _looks_like_local_path(name):
+            raise FileNotFoundError(f"Local model path not found: {name}")
+        case str() as repo_id:
+            from huggingface_hub import snapshot_download
+
+            return Path(
+                snapshot_download(
+                    repo_id,
+                    revision=revision,
+                    allow_patterns=[CONFIG_FILENAME, WEIGHTS_FILENAME],
+                )
+            )
+        case other:
+            raise TypeError(f"Expected a str or Path, got {type(other)}")
+    if not local.is_dir():
+        raise FileNotFoundError(f"Model directory does not exist: {local}")
+    if revision is not None:
+        raise ValueError(f"revision applies to Hub repositories, not to {local}")
+    return local
+
+
+def compile_dynamic(call: Callable[..., Tensor]) -> Callable[..., Tensor]:
+    """``torch.compile`` with dynamic shapes and without duck sizing.
+
+    Dynamo otherwise assumes sizes that happen to be equal on the first call
+    (an image width and a token count, a batch and a grid side) stay equal, and
+    recompiles at the first shape that breaks the coincidence. Dimensions of
+    size 1 are still specialized: a batch of 1, or a 16-pixel image side (a
+    latent grid side of 1), compiles one more graph per such pattern.
+
+    Relies on PyTorch's private ``torch.fx.experimental._config.use_duck_shape``
+    (tested with PyTorch 2.13).
+    """
+
+    compiled = torch.compile(call, dynamic=True, fullgraph=True)
+
+    def traced(*args: Tensor) -> Tensor:
+        """Run the compiled call with duck sizing off (a per-thread setting)."""
+
+        with fx_config.patch(use_duck_shape=False):  # ty: ignore[unresolved-attribute]
+            return compiled(*args)
+
+    return traced
+
+
+class Dinac3(nn.Module):
+    """Deterministic image autoencoder with a DINOv3-aligned latent.
+
+    Images are RGB in [-1, 1] with height and width multiples of 16. The latent
+    has ``config.latent_channels`` channels at stride 16: free (detail)
+    channels first, the ``config.semantic_channels`` semantic channels last.
+    ``encode``/``decode`` use the whitened latent (zero mean, unit variance per
+    channel under the training distribution); ``encode_raw``/``decode_raw`` the
+    raw one.
+    """
+
+    latent_mean: Tensor
+    latent_var: Tensor
+
+    def __init__(self, config: Dinac3Config) -> None:
+        """Build the architecture without weights (no network access)."""
+
+        super().__init__()
+        self.config = config
+        self.encoder = Encoder(config)
+        self.decoder = Decoder(config)
+        self.register_buffer("latent_mean", torch.zeros(config.latent_channels))
+        self.register_buffer("latent_var", torch.ones(config.latent_channels))
+        self.compute_dtype = torch.bfloat16
+        self._encode_call: Callable[[Tensor, Tensor], Tensor] = self.encoder.forward
+        self._decode_call: Callable[[Tensor], Tensor] = self.decoder.forward
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        path_or_repo: str | Path,
+        *,
+        device: torch.device | str,
+        dtype: torch.dtype = torch.bfloat16,
+        compile_encoder: bool = True,
+        compile_decoder: bool = True,
+        revision: str | None = None,
+    ) -> Dinac3:
+        """Load an artifact strictly and prepare it for CUDA inference.
+
+        Args:
+            path_or_repo: Local directory (``Path``, or ``str`` naming one) or
+                Hugging Face Hub repository id.
+            device: CUDA device.
+            dtype: ``torch.bfloat16`` (the stored mixed precision, BF16
+                autocast: the training precision) or ``torch.float32`` (every
+                tensor upcast, no autocast).
+            compile_encoder: ``torch.compile`` the encoder with dynamic shapes.
+            compile_decoder: ``torch.compile`` the decoder with dynamic shapes.
+            revision: Hub revision, for repository ids only.
+
+        Raises:
+            ValueError: On an unsupported dtype or device, or tensors stored in
+                another dtype than the storage policy.
+        """
+
+        directory = resolve_model_dir(path_or_repo, revision=revision)
+        config = Dinac3Config.load(directory / CONFIG_FILENAME)
+        state = load_file(str(directory / WEIGHTS_FILENAME))
+        require_storage_policy(state)
+        # Every tensor comes from the artifact: build on the meta device and
+        # let the loaded tensors become the parameters and buffers.
+        with torch.device("meta"):
+            model = cls(config)
+        model.load_state_dict(state, strict=True, assign=True)
+        model.prepare(
+            device=torch.device(device),
+            dtype=dtype,
+            compile_encoder=compile_encoder,
+            compile_decoder=compile_decoder,
+        )
+        return model
+
+    def prepare(
+        self,
+        *,
+        device: torch.device,
+        dtype: torch.dtype,
+        compile_encoder: bool,
+        compile_decoder: bool,
+    ) -> None:
+        """Move to ``device``, set the compute precision and compile once.
+
+        Raises:
+            ValueError: On a non-CUDA device or an unsupported dtype.
+        """
+
+        if device.type != "cuda":
+            raise ValueError("dinac3 inference requires a CUDA device")
+        if dtype not in SUPPORTED_DTYPES:
+            raise ValueError(f"dtype must be one of {SUPPORTED_DTYPES}, got {dtype}")
+        self.to(device=device)
+        if dtype == torch.float32:
+            for tensor in (*self.parameters(), *self.buffers()):
+                tensor.data = tensor.data.float()
+        self.compute_dtype = dtype
+        self.eval().requires_grad_(False)
+        self.decoder.conv_up_head.fold_kernels()
+        self._encode_call = (
+            compile_dynamic(self.encoder.forward)
+            if compile_encoder
+            else self.encoder.forward
+        )
+        self._decode_call = (
+            compile_dynamic(self.decoder.forward)
+            if compile_decoder
+            else self.decoder.forward
+        )
+
+    @torch.inference_mode()
+    def encode(self, images: Tensor) -> Tensor:
+        """Whitened FP32 latents ``[B, C, H / 16, W / 16]`` of [-1, 1] images."""
+
+        return self.whiten(self.encode_raw(images))
+
+    @torch.inference_mode()
+    def encode_raw(self, images: Tensor) -> Tensor:
+        """Raw (unwhitened) FP32 latents of [-1, 1] images."""
+
+        self._require_images(images)
+        rope = self.encoder.rope_embed(images.shape[-2], images.shape[-1])
+        with self._autocast():
+            return self._encode_call(images, rope).float()
+
+    @torch.inference_mode()
+    def decode(self, latents: Tensor, height: int, width: int) -> Tensor:
+        """FP32 RGB ``[B, 3, height, width]`` (unclamped, about [-1, 1]) from
+        whitened latents, in one deterministic pass."""
+
+        if (height, width) != (latents.shape[-2] * PATCH, latents.shape[-1] * PATCH):
+            raise ValueError(
+                f"{height}x{width} is not the 16x image of a "
+                f"{latents.shape[-2]}x{latents.shape[-1]} latent grid"
+            )
+        return self.decode_raw(self.dewhiten(latents))
+
+    @torch.inference_mode()
+    def decode_raw(self, latents: Tensor) -> Tensor:
+        """FP32 RGB from raw (unwhitened) latents, in one deterministic pass."""
+
+        self._require_latents(latents)
+        with self._autocast():
+            images = self._decode_call(latents.float())
+        return images.float().contiguous()
+
+    def whiten(self, latents: Tensor) -> Tensor:
+        """Raw to whitened latents, in FP32."""
+
+        mean, std = self._latent_stats()
+        return (latents.float() - mean) / std
+
+    def dewhiten(self, latents: Tensor) -> Tensor:
+        """Whitened to raw latents, in FP32."""
+
+        mean, std = self._latent_stats()
+        return latents.float() * std + mean
+
+    def semantic_channels(self, latents: Tensor) -> Tensor:
+        """The DINOv3-aligned channels: the last ``semantic_channels``."""
+
+        return latents[:, self.config.free_channels :]
+
+    def free_channels(self, latents: Tensor) -> Tensor:
+        """The free (detail) channels: all but the semantic ones."""
+
+        return latents[:, : self.config.free_channels]
+
+    def to(self, *args: object, **kwargs: object) -> Dinac3:
+        """Move to a device; dtype casts are rejected (see :meth:`_reject_cast`)."""
+
+        casts = (
+            "dtype" in kwargs
+            or "tensor" in kwargs
+            or any(isinstance(arg, torch.dtype | Tensor) for arg in args)
+        )
+        if casts:
+            self._reject_cast()
+        return super().to(*args, **kwargs)  # ty: ignore[no-matching-overload]
+
+    def half(self) -> NoReturn:
+        """Rejected: see :meth:`_reject_cast`."""
+
+        self._reject_cast()
+
+    def bfloat16(self) -> NoReturn:
+        """Rejected: see :meth:`_reject_cast`."""
+
+        self._reject_cast()
+
+    def float(self) -> NoReturn:
+        """Rejected: see :meth:`_reject_cast`."""
+
+        self._reject_cast()
+
+    def double(self) -> NoReturn:
+        """Rejected: see :meth:`_reject_cast`."""
+
+        self._reject_cast()
+
+    def type(self, dst_type: object) -> NoReturn:
+        """Rejected: see :meth:`_reject_cast`."""
+
+        del dst_type
+        self._reject_cast()
+
+    def _reject_cast(self) -> NoReturn:
+        """Module-wide dtype casts would break the mixed storage policy.
+
+        Raises:
+            TypeError: Always, pointing to ``from_pretrained(dtype=...)``.
+        """
+
+        raise TypeError(
+            "dinac3 keeps a mixed storage policy (BF16 weights, FP32 residual-path, "
+            "statistics and RoPE tensors); a module-wide dtype cast would break it. "
+            "Choose the precision with Dinac3.from_pretrained(..., "
+            "dtype=torch.bfloat16 | torch.float32)."
+        )
+
+    def _latent_stats(self) -> tuple[Tensor, Tensor]:
+        """Per-channel ``(mean, std)`` ``[1, C, 1, 1]`` in FP32."""
+
+        mean = self.latent_mean.float().view(1, -1, 1, 1)
+        var = self.latent_var.float().view(1, -1, 1, 1)
+        return mean, torch.sqrt(var + self.config.latent_stats_eps)
+
+    def _autocast(self) -> torch.autocast:
+        """BF16 autocast for BF16 inference, disabled for FP32."""
+
+        return torch.autocast(
+            "cuda", dtype=torch.bfloat16, enabled=self.compute_dtype == torch.bfloat16
+        )
+
+    def _require_images(self, images: Tensor) -> None:
+        """Validate images at the API boundary, outside compiled graphs.
+
+        Raises:
+            ValueError: On a wrong shape, size, dtype or device.
+        """
+
+        if images.ndim != 4 or images.shape[1] != 3 or images.shape[0] < 1:
+            raise ValueError("images must have shape [B, 3, H, W] with B >= 1")
+        height, width = images.shape[-2:]
+        if min(height, width) < PATCH or height % PATCH or width % PATCH:
+            raise ValueError(f"Image sides must be positive multiples of {PATCH}")
+        self._require_tensor(images)
+
+    def _require_latents(self, latents: Tensor) -> None:
+        """Validate latents at the API boundary.
+
+        Raises:
+            ValueError: On a wrong shape, dtype or device.
+        """
+
+        if latents.ndim != 4 or latents.shape[1] != self.config.latent_channels:
+            raise ValueError(
+                f"latents must have shape [B, {self.config.latent_channels}, h, w]"
+            )
+        if min(latents.shape[0], *latents.shape[-2:]) < 1:
+            raise ValueError("Latent batch and grid sides must be positive")
+        self._require_tensor(latents)
+
+    def _require_tensor(self, tensor: Tensor) -> None:
+        """Require a floating tensor on the model's CUDA device.
+
+        Raises:
+            ValueError: If the tensor is elsewhere or not floating point.
+        """
+
+        device = self.latent_mean.device
+        if tensor.device != device:
+            raise ValueError(f"Inputs must be on the model device {device}")
+        if not tensor.is_floating_point():
+            raise ValueError("Inputs must be floating-point tensors")

dinac3/precision.py

@@ -0,0 +1,79 @@
+"""Storage dtype of every exported tensor.
+
+BF16 for every weight that BF16 autocast consumes in BF16 anyway (matrices,
+convolutions, biases, and norm gains and biases applied to BF16 activations):
+storing them in BF16 gives bit-identical outputs. FP32 where the value takes
+part in FP32 arithmetic: the frozen DINOv3 residual stream (CLS and register
+tokens, layer scales, LayerNorms), its RoPE periods, the input and
+layer-standardization statistics, the latent statistics used for whitening,
+and the upsamplers' 3x3 weights and fold matrices (the 4x4 transposed kernel is
+folded from them outside autocast, in their storage dtype).
+"""
+
+from __future__ import annotations
+
+import re
+from typing import TYPE_CHECKING
+
+import torch
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+
+    from torch import Tensor
+
+_FP32_NAMES = frozenset(
+    {
+        "latent_mean",
+        "latent_var",
+        "encoder.pixel_mean",
+        "encoder.pixel_std",
+        "encoder.layer_mean",
+        "encoder.layer_std",
+        "encoder.backbone.cls_token",
+        "encoder.backbone.reg_token",
+        "encoder.backbone.rope.periods",
+        "encoder.backbone.norm.weight",
+        "encoder.backbone.norm.bias",
+    }
+)
+_UPSAMPLER = re.compile(
+    r"decoder\.conv_up_head\.(handoff_upsample|upsamples\.\d+)\.(conv\.weight|fold)"
+)
+_BACKBONE_BLOCK_PREFIX = "encoder.backbone.blocks."
+_FP32_BLOCK_SUFFIXES = (
+    ".gamma_1",
+    ".gamma_2",
+    ".norm1.weight",
+    ".norm1.bias",
+    ".norm2.weight",
+    ".norm2.bias",
+)
+
+
+def storage_dtype(name: str) -> torch.dtype:
+    """BF16 or FP32 storage of one exported tensor, by its state-dict name."""
+
+    fp32 = (
+        name in _FP32_NAMES
+        or (
+            name.startswith(_BACKBONE_BLOCK_PREFIX)
+            and name.endswith(_FP32_BLOCK_SUFFIXES)
+        )
+        or _UPSAMPLER.fullmatch(name) is not None
+    )
+    return torch.float32 if fp32 else torch.bfloat16
+
+
+def require_storage_policy(state: Mapping[str, Tensor]) -> None:
+    """Require every tensor of an artifact to be stored in its policy dtype.
+
+    Raises:
+        ValueError: Naming every tensor stored in another dtype.
+    """
+
+    wrong = sorted(
+        name for name, tensor in state.items() if tensor.dtype != storage_dtype(name)
+    )
+    if wrong:
+        raise ValueError(f"Tensors not stored in their policy dtype: {wrong}")

dinac3/trunk.py

@@ -0,0 +1,107 @@
+"""The decoder's ViT trunk: RMS-sandwich blocks with 2-D RoPE."""
+
+from __future__ import annotations
+
+import torch
+from torch import Tensor, nn
+from torch.nn import functional as F
+
+from .layers import Mlp, RMSNorm
+
+ROPE_BASE = 10_000.0
+
+
+def rope_periods(head_dim: int, device: torch.device) -> Tensor:
+    """FP32 periods ``ROPE_BASE ** (2 i / (head_dim / 2))`` of the
+    ``head_dim / 4`` rotation frequencies per axis."""
+
+    exponents = (
+        2.0
+        * torch.arange(head_dim // 4, device=device, dtype=torch.float32)
+        / (head_dim // 2)
+    )
+    return torch.tensor(ROPE_BASE, device=device, dtype=torch.float32) ** exponents
+
+
+def rope_tables(
+    height: int, width: int, *, head_dim: int, device: torch.device
+) -> tuple[Tensor, Tensor]:
+    """FP32 ``(sin, cos)`` ``[H * W, head_dim]`` of axial 2-D RoPE on the token
+    grid: unnormalized patch-index coordinates (row, column), the
+    :func:`rope_periods`, adjacent-pair layout."""
+
+    periods = rope_periods(head_dim, device)
+    rows = torch.arange(0.0, float(height), device=device, dtype=torch.float32)
+    cols = torch.arange(0.0, float(width), device=device, dtype=torch.float32)
+    coords = torch.stack(torch.meshgrid(rows, cols, indexing="ij"), dim=-1).flatten(
+        0, 1
+    )
+    angles = 1.0 * coords[:, :, None] / periods[None, None, :].expand(1, 2, -1)
+    angles = angles.repeat_interleave(2, dim=-1).flatten(1, 2)
+    return torch.sin(angles), torch.cos(angles)
+
+
+def _rotate_pairs(x: Tensor) -> Tensor:
+    """Rotate consecutive channel pairs: ``(a, b) -> (-b, a)``."""
+
+    pairs = x.reshape(*x.shape[:-1], x.shape[-1] // 2, 2)
+    return torch.stack((-pairs[..., 1], pairs[..., 0]), dim=-1).reshape_as(x)
+
+
+def apply_rope(x: Tensor, sin: Tensor, cos: Tensor) -> Tensor:
+    """Rotate ``[B, heads, N, head_dim]`` queries or keys in the tables' dtype."""
+
+    rotated = x.to(dtype=sin.dtype)
+    sin_b = sin.view(1, 1, *sin.shape)
+    cos_b = cos.view(1, 1, *cos.shape)
+    return ((rotated * cos_b) + (_rotate_pairs(rotated) * sin_b)).to(dtype=x.dtype)
+
+
+class Attention(nn.Module):
+    """Bias-free multi-head self-attention with RMS-normalized queries and keys."""
+
+    def __init__(self, width: int, head_dim: int) -> None:
+        """Allocate the fused QKV and output projections and the q/k norms."""
+
+        super().__init__()
+        self.heads = width // head_dim
+        self.head_dim = head_dim
+        self.qkv = nn.Linear(width, 3 * width, bias=False)
+        self.proj_out = nn.Linear(width, width, bias=False)
+        self.q_norm = RMSNorm(head_dim, affine=True)
+        self.k_norm = RMSNorm(head_dim, affine=True)
+
+    def forward(self, x: Tensor, sin: Tensor, cos: Tensor) -> Tensor:
+        """Attend over all tokens of ``[B, N, width]``."""
+
+        b, n, width = x.shape
+        q, k, v = self.qkv(x).chunk(3, dim=-1)
+        shape = (b, n, self.heads, self.head_dim)
+        q = q.view(shape).transpose(1, 2).contiguous()
+        k = k.view(shape).transpose(1, 2).contiguous()
+        v = v.view(shape).transpose(1, 2).contiguous()
+        q = apply_rope(self.q_norm(q), sin, cos)
+        k = apply_rope(self.k_norm(k), sin, cos)
+        out = F.scaled_dot_product_attention(q, k, v, dropout_p=0.0, is_causal=False)
+        return self.proj_out(out.transpose(1, 2).contiguous().view(b, n, width))
+
+
+class DitBlock(nn.Module):
+    """ViT block: RMSNorm before and after attention and MLP."""
+
+    def __init__(self, width: int, head_dim: int, mlp_ratio: float) -> None:
+        """Allocate attention, MLP and the four norms."""
+
+        super().__init__()
+        self.attn_norm1 = RMSNorm(width, affine=True)
+        self.attn_norm2 = RMSNorm(width, affine=True)
+        self.mlp_norm1 = RMSNorm(width, affine=True)
+        self.mlp_norm2 = RMSNorm(width, affine=True)
+        self.attn = Attention(width, head_dim)
+        self.mlp = Mlp(width, int(mlp_ratio * width))
+
+    def forward(self, x: Tensor, sin: Tensor, cos: Tensor) -> Tensor:
+        """Residual attention then residual MLP on ``[B, N, width]`` tokens."""
+
+        x = x + self.attn_norm2(self.attn(self.attn_norm1(x), sin, cos))
+        return x + self.mlp_norm2(self.mlp(self.mlp_norm1(x)))

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2696c98a87d595afc386f068ab5f9fa99d6ff1f71e521145bc31371612b62856
+size 336457454

pyproject.toml

@@ -0,0 +1,13 @@
+[build-system]
+requires = ["setuptools>=68"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "dinac3"
+version = "1.0"
+description = "dinac3: deterministic image autoencoders with a DINOv3-aligned latent"
+requires-python = ">=3.10"
+dependencies = ["torch>=2.13", "timm==1.0.26", "safetensors>=0.5", "huggingface-hub>=0.28"]
+
+[tool.setuptools.packages.find]
+include = ["dinac3*"]