Release experimental-X1.1: completed 448px ONNX and FP32 safetensors

.gitattributes

@@ -1,35 +1,2 @@
-*.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
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,48 @@
+# Educational Use Only License
+
+Copyright (c) 2025
+
+## Grant of License
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to use, copy, modify, merge, and distribute the Software solely for **educational purposes**, subject to the following conditions:
+
+## Educational Purposes Definition
+
+"Educational purposes" means use of the Software for:
+- Personal learning and skill development
+- Academic research and study
+- Teaching in educational institutions (schools, universities, colleges)
+- Non-commercial educational workshops and tutorials
+- Student projects and assignments
+- Open-source educational content creation
+
+## Restrictions
+
+The Software may **NOT** be used for:
+- Commercial purposes of any kind
+- Revenue-generating activities
+- Production environments in for-profit organizations
+- Consulting or professional services
+- Any purpose that directly or indirectly generates monetary compensation
+
+## Attribution
+
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+
+Any public use or distribution of work derived from this Software must include clear attribution to the original author and a link to the original repository.
+
+## Disclaimer
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
+
+## Additional Terms
+
+1. **Derivative Works**: Any modifications or derivative works must also be licensed under these same terms and clearly marked as modified versions.
+
+2. **No Endorsement**: Use of this Software does not grant permission to use the trade names, trademarks, service marks, or product names of the Licensor.
+
+3. **Termination**: This license automatically terminates if you violate any of these restrictions and may be terminated by the copyright holder at any time.
+
+## Contact
+
+For commercial licensing inquiries or permissions beyond the scope of this license, please contact: on telegram @grio43

README.md

@@ -0,0 +1,277 @@
+---
+pipeline_tag: image-classification
+license: other
+license_name: educational-use-only
+license_link: LICENSE
+language:
+- en
+base_model: Grio43/OppaiOracle-experimental-X1
+base_model_relation: finetune
+tags:
+- pytorch
+- onnx
+- safetensors
+- anime
+- image-tagging
+- multi-label
+- vision-transformer
+- experimental
+---
+
+# OppaiOracle experimental-X1.1
+
+**experimental-X1.1 is the 448×448 continuation of
+[experimental-X1](https://huggingface.co/Grio43/OppaiOracle-experimental-X1).**
+It is a multi-label anime and illustration tagger, released in **FP32 ONNX and
+FP32 safetensors** formats. Both contain the selected **448px epoch 9, update
+55,719** weights. This continuation completed its nine-epoch schedule and WSD
+cooldown to zero learning rate on September 17, 2026.
+
+This remains an experimental release. The shared-tag comparison below shows higher
+micro-F1 than V1.1 at development-tuned thresholds, with slightly lower overall
+exact mAP. Performance varies by tag and frequency; there is no blanket claim that
+X1.1 replaces V1.1 for every use.
+
+## Download and run
+
+The complete download includes both weight formats, about **1.18 GB** in total.
+Use a fresh Python environment, download the repository, then install the
+dependencies for your chosen format:
+
+```bash
+python -m pip install huggingface_hub
+python -c "from huggingface_hub import snapshot_download; snapshot_download('Grio43/OppaiOracle-experimental-X1.1', local_dir='OppaiOracle-experimental-X1.1')"
+cd OppaiOracle-experimental-X1.1
+```
+
+For ONNX Runtime on CPU:
+
+```bash
+python -m pip install -r requirements.txt
+python infer.py image.jpg --format onnx --device cpu --output predictions.json
+```
+
+For safetensors with PyTorch:
+
+```bash
+python -m pip install -r requirements-torch.txt
+python infer.py image.jpg --format safetensors --device cpu --output predictions.json
+```
+
+For a smaller download, obtain just [download.py](download.py), install
+`huggingface_hub`, and run one of the following from the directory containing it:
+
+```bash
+python download.py OppaiOracle-experimental-X1.1 --format onnx
+python download.py OppaiOracle-experimental-X1.1 --format safetensors
+```
+
+Each command downloads the selected weights and required support files. The
+helper resolves `--revision` (default `main`) to one commit before downloading,
+keeping the configuration, vocabulary and weights aligned. Use `--format all`
+to obtain both formats. The resulting directory supports the same install and
+inference commands above.
+
+The supplied PyTorch model code loads `model.safetensors` and the real architecture
+configuration in `config.json`. These are inference weights; the safetensors file
+does not contain a pickled training checkpoint or optimizer state. The release
+preserves FP32 values without quantization or conversion to BF16/FP16.
+
+Both inference paths apply the same preprocessing and return tag names with sigmoid
+scores, excluding PAD and UNK. The default threshold is **0.7085427045822144**, the
+best micro-F1 threshold on the shared-tag development comparison's 200-point grid.
+It is a useful starting point, **not final deployment calibration**, and was not
+optimized independently for every tag or for the complete vocabulary. Set
+`--threshold` for your application. Scores are not calibrated probabilities that a
+label is correct.
+
+To inspect the 30 highest scores regardless of the default threshold:
+
+```bash
+python infer.py image.jpg --format onnx --threshold 0 --top-k 30
+```
+
+Multiple image paths and `--batch-size` are supported. `--device cuda` requires
+the corresponding CUDA-capable PyTorch or ONNX Runtime environment. Local
+inference does not make network requests.
+
+## Training
+
+The model lineage was trained **from scratch**, without an externally pretrained
+backbone. X1.1 starts from the released 320px X1 weights and adapts **all model
+weights**, rather than freezing the backbone.
+
+| Stage | Recorded recipe |
+|---|---|
+| 320px source | V2 Phase 1, epoch 21, update 127,974; the checkpoint released as experimental-X1 |
+| 320px optimization | Effective batch 1,024; AdamW8bit with FP32 head/position optimizer states; BF16 AMP; weight decay 0.05; 10,000-update warmup; stable LR 5.4e-4 |
+| Source loss history | ASL negative gamma 4 initially, explicitly changed to 3 at update 41,223; positive gamma 0, probability clip 0.05, alpha 1, no label smoothing, detached focal weights |
+| 448px continuation | Nine phase-local epochs, 55,719 optimizer updates, September 13–17, 2026 |
+| Initialization | One-time bicubic positional-grid interpolation from 20×20 to 28×28; fresh optimizer, scheduler and phase-local history; learned weights and frozen vocabulary retained |
+| 448px batch and precision | Microbatch 48 × 21 accumulation = 1,008; BF16 AMP; no activation checkpointing |
+| 448px optimizer | AdamW8bit; FP32 head/position optimizer states; weight decay 0.05; gradient norm clipping 1.0 |
+| 448px schedule | Base LR 1.5e-5, square-root scaling from batch 256 gives peak LR 2.97647022495e-5; one warmup epoch; budget-driven WSD cooldown begins at update 48,451, lasts 7,268 updates and ends at zero LR |
+| 448px loss | Fixed ASL negative/positive gamma 3/0, probability clip 0.05, alpha 1, no label smoothing, detached focal weights; PAD/UNK excluded |
+| Selection | Highest recorded cooldown validation binned mAP; epoch 9 is both the selected best and final checkpoint |
+
+The 320px source did not complete its original planned cooldown. The separate
+448px experiment did complete its own cooldown; it does not imply completion of a
+later 512px phase. Plateau-triggered cooldown was disabled for the nine-epoch
+continuation; its total budget included warmup and cooldown.
+
+At 448px, drop-path was 0.15 and ordinary dropout was zero. Augmentation used
+brightness/contrast/saturation jitter of 0.22/0.15/0.06, each at probability 0.5;
+2–5-degree rotation at probability 0.30; Gaussian blur sigma 0.1–1.0 at probability
+0.15; and horizontal flipping at probability 0.5. Images carrying any of
+`gen:left-handed`, `gen:left-to-right_manga`, `gen:right-over-left_kimono` or
+`gen:right-to-left_comic` were excluded from horizontal flips. Random erasing was
+disabled.
+
+## Dataset and tag updates
+
+The frozen prepared corpus contains **6,269,823 image/metadata pairs**, split into
+**6,239,823 TRAIN** and **30,000 EVAL** images. The 448px continuation retained the
+same split, labels and vocabulary as X1.
+
+| Source | Pairs |
+|---|---:|
+| Danbooru | 6,142,955 |
+| Real-life photos without people | 123,335 |
+| Requested images and tags | 3,533 |
+
+**Dataset-maintainer provenance:** the author reports that the Danbooru metadata
+was synchronized in **August 2026**, deprecated tags were removed, and renamed
+tags were updated to their correct canonical names. This describes the metadata
+preparation, not an assertion that every image was uploaded in August. These
+cleanup statements are supplied by the dataset maintainer; the release checks
+verify the frozen artifacts and model alignment.
+
+General tags have a **450-image minimum**, down from V1.1's **500**. Counts are
+distinct images in the combined prepared corpus. Character, copyright, artist and
+metadata tags use a 550-image minimum; the four rating tags are mandatory.
+Despite the lower general-tag floor, the cleaned vocabulary has **14,994 tags**
+versus **19,292 in V1.1**. The overall reduction accompanies deprecated-tag removal,
+canonical-name updates and the changed vocabulary policy; it is not a count of
+deprecated tags alone.
+
+| Category | Tags |
+|---|---:|
+| General | 9,920 |
+| Character | 3,069 |
+| Copyright | 1,139 |
+| Artist | 745 |
+| Metadata | 117 |
+| Rating | 4 |
+| **Total real tags** | **14,994** |
+
+The model has **14,996 output positions**, including PAD at index 0 and UNK at
+index 1. Keep the supplied vocabulary ordering: V1.1 indices and tag counts are
+not compatible with this model.
+
+## Evaluation against V1.1
+
+Both models were evaluated at **448px in FP32 ONNX**, on the same frozen **30,000
+EVAL images** and recorded labels, with identical preprocessing and TF32 disabled.
+All source-sidecar hashes were verified. One-to-one shared names yielded **13,569
+tags**, with **13,394 supported tags** contributing to macro AP. V1.1 results were
+reused from a verified archived evaluation of the same images and support.
+
+| Shared-tag measurement | experimental-X1.1 | V1.1 |
+|---|---:|---:|
+| Exact mAP | 58.7513% | 59.2716% |
+| 200-bin mAP | 58.5193% | 58.8611% |
+| Best grid micro-F1 | 68.6410% | 66.5864% |
+| Threshold for that micro-F1 | 0.708543 | 0.798995 |
+| General-tag exact mAP, 8,676 supported shared tags | 43.7549% | 43.9700% |
+
+X1.1's overall shared-tag exact mAP is **0.5203 percentage points lower**; its
+micro-F1 is **2.0546 points higher** at independently selected grid thresholds.
+Across supported shared tags, X1.1 has higher AP on 5,425, V1.1 on 6,271, and 1,698
+tie. For general tags with at least 50 EVAL positives, X1.1's exact mAP is 54.893%
+versus 53.134%; V1.1 retains an advantage among the sparsest general tags.
+
+These are **development-set measurements**, not an independent clean-label test.
+This EVAL set has informed training decisions and threshold selection. V1.1's
+training overlap with these images is unverified. Recorded labels may be missing
+or incorrect, and aggregate gains do not establish a gain for every tag.
+See [comparison details](comparison_v1_1.json) and
+[shared-tag AP results](comparison_v1_1_per_tag.csv).
+
+For the complete X1.1 supported vocabulary, the FP32 export's exact mAP was
+**57.7464%** over **14,767 tags**. The final BF16 training validation recorded exact
+mAP **57.7166%**, binned mAP **57.5063%**, and micro-F1 **66.2121%** at the inherited
+training diagnostic threshold 0.7927. These differ in numerical precision and/or
+tag scope from the shared-tag table; they should not be substituted for one
+another.
+The archived full FP32 measurement is in [validation.json](validation.json);
+[training_metadata.json](training_metadata.json) records the training provenance.
+
+## Architecture and preprocessing
+
+The **147,706,932-parameter** ViT uses patch size 16, width 896, 18 transformer
+layers, 14 attention heads, MLP width 2,320, QK normalization and LayerScale
+initialization 0.1. Its positional grid is 28×28 at 448px.
+
+| ONNX input/output | Type | Shape |
+|---|---|---|
+| `pixel_values` | float32 | `[batch, 3, 448, 448]` |
+| `padding_mask` | bool | `[batch, 448, 448]` |
+| `probabilities` | float32 | `[batch, 14996]` |
+
+Batch size is dynamic; spatial dimensions are fixed. Apply EXIF orientation,
+convert to RGB, composite transparency onto `[114,114,114]`, and downscale with
+LANCZOS while preserving aspect ratio, **without upscaling**. Center the result on
+a 448×448 canvas with the same pad color. Padding-mask values are **True for
+padding**. Divide pixels by 255, normalize using mean and standard deviation
+`[0.5,0.5,0.5]`, and arrange as NCHW. The supplied script implements this contract.
+
+The ONNX output already includes sigmoid. The safetensors inference helper
+applies sigmoid to the PyTorch logits. Four rating scores are independent sigmoid
+outputs, not a softmax. A frontend must support the two inputs and this
+preprocessing; a raw-image, single-input tagger interface needs an adapter.
+
+## Files, integrity and download statistics
+
+| File | Purpose |
+|---|---|
+| `model.onnx` | FP32 inference graph, about 591 MB, with embedded vocabulary/preprocessing metadata |
+| `model.safetensors` | FP32 tensor-only weights, about 591 MB |
+| `config.json` | Real architecture and release configuration consumed by the supplied loader |
+| `preprocessing.json` | Image transform and input contract |
+| `vocabulary.json`, `selected_tags.csv` | Exact output ordering and tag categories |
+| `infer.py`, `download.py` | Standalone inference and selective release download helpers |
+| `requirements.txt`, `requirements-torch.txt` | Dependencies for CPU ONNX and PyTorch/safetensors inference |
+| `training_metadata.json`, `validation.json` | Training provenance and archived full FP32 evaluation |
+| `inference_checks.json` | Release-package checks using a bounded set of fixtures; not another full evaluation |
+| `release_manifest.json`, `SHA256SUMS` | Provenance and file-integrity records |
+
+The normal repository download includes `config.json`, which the supplied loader
+actually uses. [Hugging Face download statistics](https://huggingface.co/docs/hub/models-download-stats)
+count requests to supported query files such as `config.json`. The public counter
+is not a unique-user count; direct downloads of weights alone may not be reflected.
+There are no extra telemetry calls or artificial counting requests in the
+inference helper.
+
+The selected source checkpoint SHA256 is:
+
+```text
+d400c2e460e7287357f86083f84c187790a45b1d47985e908c52825373f888c3
+```
+
+This identifies the preserved local training source. The release manifest and
+checksums identify the distributed files; optimizer state and the pickled training
+checkpoint are not included. The release ONNX has refreshed descriptive metadata;
+its inference graph and learned weights are unchanged from the evaluated artifact.
+The manifest records both artifact identities.
+
+## Scope and license
+
+X1.1 is intended for experimental tagging, comparison and feedback within the
+license terms. Training annotations are noisy and incomplete; rare concepts and
+visually similar tags can be unreliable. Review suggestions before applying bulk
+dataset edits. The vocabulary includes sexual-content labels and independent
+rating tags; this model has not been validated as a safety classifier.
+
+The release retains the **[Educational Use Only License](LICENSE)**. It permits
+the educational uses described in that file and prohibits commercial use. Read
+the full license for attribution, redistribution and other conditions.

SHA256SUMS

@@ -0,0 +1,29 @@
+5a4c917343b3f76df936dbd7f23ad4693aea7f96ad4c1029cdfd70ac3c78cb9b  .gitattributes
+322447c97c7fc1b266286efd31b7a9aefd21a111fdbde60090ebfa0601cca497  LICENSE
+4023de4f067e0cf475dfc80d0d038477bb1d4bbea4c5edf9f23d24913e575206  README.md
+dfc682758528177f0d96a0a1230add733f537dd40932239095893c6c585eef5e  checkpoint_evidence.json
+565667f7bd86c3b61c4fda46cec38e8838ad0719401b991513fedaba9961e3f8  comparison_groups.csv
+ad522ba09cd16227f88215a62c4f71f64f0af9fe930ea39f93a0c406f6d9a279  comparison_v1_1.json
+bf6ce722e843f4246d42fa86bb17a665eb0de514a378bd7aed9ba36b1f9c47e3  comparison_v1_1_per_tag.csv
+39e218e1674f5fb4e00023aa4bcea778c09704d856f9f833ef01bb251d1415fd  config.json
+c83f832a32b6a8b8802dd9207f2f9be2ebef2a6d169d6daf1cdb2530cfe287ab  custom_drop_path.py
+ab4c1c04c03f5cefe3153bede72878a0d2df9e8378724922ed1571338b38d3c6  download.py
+23f6dd10da65f031e5aae31fe88aefa35c6aa34e33d31986d45c9f9d57c705b8  frozen_evaluation_contract.json
+79b15ad63de33354bc0a39176cd9eddd915361bd48b6c27b885009aac77185ea  infer.py
+f828ef9205af946e5633e416044fd84d6568292bfc1e02cb0852f484368426da  inference_checks.json
+373cc75ed9975e58864e4ed73a34316a696b6091b83971207677a6dae924fcd2  mask_utils.py
+49ba0f66f83ba88a101b6edac54ec39f0a12eab66743e13ae42efc3570745bf7  model.onnx
+ec07d0034f48ba3ceff8624993baf3ea30fa3a2e3cc8908682eb3e4f84eda766  model.safetensors
+bce3773157a4c4dfc153c633d489c1e7c18f544b7fac6f51bd318ebe0fb7c6aa  model_architecture.py
+1c9af089a8c150b5270bc933fe4053a913f284a4242b53cdb94d108c529b2ef9  preprocessing.json
+b058043b832bbb02c4139d3766dc8f1aac1fb781f90c897cdcf9230694cc2126  release_manifest.json
+d604ef2e5271edac5d5b4bba2eeccf11102e69dd82667f40ca0be417b806ea62  requirements-gpu.txt
+9d4beb51ff7c74d2ff68981496e6c1b37f1ab9cc85a53e692cb22c345c1ee3d7  requirements-torch.txt
+abfcf5a28c6c88fa86822c1d7bcbf7afca92a893b8a6a4a65d3267a8786c3b28  requirements.txt
+ec2b028ae5c8355bc7ccbab51ef967ef4bec940b686fe3e564e0177a865c9799  selected_tags.csv
+3b20a011b5acf7732c1546dab99d6a56f48730eb54cadf30434f1dba62053436  thresholds.json
+1a9ba147bc89a2eef54f2e80ae406f8652332fb91c9f2e3cf7fb05556388ad44  training_buckets.csv
+54ef172f7c21c8cb414bef220e6ca45a33871134c99a4b3a1532c425f8ebc989  training_epochs.csv
+6d4af05a7ee8d54fd09aae488b0fd470a39c2cb16430e0f31597a3b9a867800a  training_metadata.json
+fd2475144e7583842f4c733ad77d770749964b147a41e732a58c868516fc4e6f  validation.json
+e3e6ccc137d9152d8a17fbd19e0c1ae595b2317a6c55893fd25779620b4b8087  vocabulary.json

checkpoint_evidence.json

@@ -0,0 +1,45 @@
+{
+  "checkpoint": "checkpoint_epoch_9_step_55719.pt",
+  "checkpoint_sha256": "d400c2e460e7287357f86083f84c187790a45b1d47985e908c52825373f888c3",
+  "epoch": 9,
+  "step": 55719,
+  "scheduler": {
+    "warmup_steps": 6191,
+    "total_steps": 55719,
+    "cooldown_fraction": 0.15,
+    "budget_cooldown_start": 48451,
+    "min_lr": 0.0,
+    "base_max_lr": 2.9764702249476646e-05,
+    "cooldown_start": 48451,
+    "cooldown_steps": 7268,
+    "cooldown_reason": "budget",
+    "cooldown_best_metric": 0.5750626921653748,
+    "base_lrs": [
+      2.9764702249476646e-05,
+      2.9764702249476646e-05
+    ],
+    "last_epoch": 55719,
+    "_step_count": 55720,
+    "_is_initial": false,
+    "_get_lr_called_within_step": false,
+    "_last_lr": [
+      0.0,
+      0.0
+    ]
+  },
+  "preprocessing": {
+    "normalize_mean": [
+      0.5,
+      0.5,
+      0.5
+    ],
+    "normalize_std": [
+      0.5,
+      0.5,
+      0.5
+    ],
+    "image_size": 448,
+    "patch_size": 16,
+    "color_order": "RGB"
+  }
+}

comparison_groups.csv

@@ -0,0 +1,14 @@
+group,tags,x1_ap,v11_ap,delta_points,x1_wins,v11_wins,ties
+all shared supported,13394,58.75134140920024,59.27161263163755,-0.5202712200966514,5425,6271,1698
+shared with >=50 EVAL positives,2028,57.690053587344046,55.98307259190696,1.7069809948860188,1161,867,0
+category gen,8676,43.75492858232571,43.96996116217886,-0.21503257453254906,4069,4420,187
+category char,2870,91.87811929600326,91.97710045881874,-0.0989811631562091,815,965,1090
+category copyright,1050,78.31961239122653,79.32927887092089,-1.0096664829623132,378,524,148
+category artist,695,83.73359510134034,89.3767281228395,-5.643133035251241,109,314,272
+category meta,99,28.44267269873385,27.431288556721427,1.0113841334870555,52,46,1
+category rating,4,90.30946791172028,74.51295033097267,15.796517580747604,2,2,0
+frequency 300-499,63,45.1914965982608,51.18916831404295,-5.997671779348618,14,36,13
+frequency 500-999,3962,63.82945302063262,66.1829312171874,-2.3534781905738833,1093,1742,1127
+frequency 1000-4999,6049,57.39406408068191,57.68059412367721,-0.286530040883122,2509,2986,554
+frequency 5000-9999,1249,52.204597668254,51.39728918997086,0.8073084776866152,638,607,4
+frequency 10000-plus,2071,57.361599985368,55.69153087854903,1.6700691066390936,1171,900,0

comparison_v1_1.json

@@ -0,0 +1,222 @@
+{
+  "x1": {
+    "model": "X1 448 final epoch 9",
+    "checkpoint": {
+      "checkpoint": "checkpoint_epoch_9_step_55719.pt",
+      "checkpoint_sha256": "d400c2e460e7287357f86083f84c187790a45b1d47985e908c52825373f888c3",
+      "epoch": 9,
+      "step": 55719,
+      "scheduler": {
+        "warmup_steps": 6191,
+        "total_steps": 55719,
+        "cooldown_fraction": 0.15,
+        "budget_cooldown_start": 48451,
+        "min_lr": 0.0,
+        "base_max_lr": 2.9764702249476646e-05,
+        "cooldown_start": 48451,
+        "cooldown_steps": 7268,
+        "cooldown_reason": "budget",
+        "cooldown_best_metric": 0.5750626921653748,
+        "base_lrs": [
+          2.9764702249476646e-05,
+          2.9764702249476646e-05
+        ],
+        "last_epoch": 55719,
+        "_step_count": 55720,
+        "_is_initial": false,
+        "_get_lr_called_within_step": false,
+        "_last_lr": [
+          0.0,
+          0.0
+        ]
+      },
+      "preprocessing": {
+        "normalize_mean": [
+          0.5,
+          0.5,
+          0.5
+        ],
+        "normalize_std": [
+          0.5,
+          0.5,
+          0.5
+        ],
+        "image_size": 448,
+        "patch_size": 16,
+        "color_order": "RGB"
+      }
+    },
+    "model_sha256": "c5fcab0fb53e07aee437f3ef29e19a1e60b3dc18380048fa6d99baf35fa6c8f3",
+    "native_image_size": 448,
+    "samples": 30000,
+    "sample_sha256": "20a60e7a747cdf9a698a464c9d324612aad2ca6b43a541db0813554eb9af8694",
+    "shared_tags": 13569,
+    "supported_shared_tags": 13394,
+    "binned_mAP": 0.5851927682386426,
+    "exact_mAP": 0.5875134140920024,
+    "best_grid_micro_f1": 0.6864101262842809,
+    "best_grid_threshold": 0.7085427045822144,
+    "precision_at_best_grid": 0.7065020665350971,
+    "recall_at_best_grid": 0.6674293596992867,
+    "all_supported_tags": 14767,
+    "all_tag_exact_mAP": 0.5774638913821051,
+    "all_tag_binned_mAP": 0.5750950549967043,
+    "inference_seconds": 347.91008079855237,
+    "total_seconds": 487.7315649986267,
+    "preprocessing": "FP32 RGB, native resolution, downscale-only centered LANCZOS letterbox, mean/std 0.5, pad 114, True=padding",
+    "verified_sidecars": 30000
+  },
+  "v11": {
+    "model": "V1.1",
+    "model_sha256": "8567852deb135eccfe4b8445d48e4476ee8846436486679adc0642cfeda07d13",
+    "native_image_size": 448,
+    "samples": 30000,
+    "sample_sha256": "20a60e7a747cdf9a698a464c9d324612aad2ca6b43a541db0813554eb9af8694",
+    "shared_tags": 13569,
+    "supported_shared_tags": 13394,
+    "binned_mAP": 0.5886105397423468,
+    "exact_mAP": 0.5927161263163755,
+    "best_grid_micro_f1": 0.6658638534826367,
+    "best_grid_threshold": 0.7989949584007263,
+    "precision_at_best_grid": 0.6852139862873703,
+    "recall_at_best_grid": 0.6475765838848551,
+    "inference_seconds": 416.6925043994561,
+    "total_seconds": 522.6292808055878,
+    "preprocessing": "FP32 RGB, native resolution, downscale-only centered LANCZOS letterbox, mean/std 0.5, pad 114, True=padding"
+  },
+  "shared_tag_results": [
+    {
+      "group": "all shared supported",
+      "tags": 13394,
+      "x1_ap": 58.75134140920024,
+      "v11_ap": 59.27161263163755,
+      "delta_points": -0.5202712200966514,
+      "x1_wins": 5425,
+      "v11_wins": 6271,
+      "ties": 1698
+    },
+    {
+      "group": "shared with >=50 EVAL positives",
+      "tags": 2028,
+      "x1_ap": 57.690053587344046,
+      "v11_ap": 55.98307259190696,
+      "delta_points": 1.7069809948860188,
+      "x1_wins": 1161,
+      "v11_wins": 867,
+      "ties": 0
+    },
+    {
+      "group": "category gen",
+      "tags": 8676,
+      "x1_ap": 43.75492858232571,
+      "v11_ap": 43.96996116217886,
+      "delta_points": -0.21503257453254906,
+      "x1_wins": 4069,
+      "v11_wins": 4420,
+      "ties": 187
+    },
+    {
+      "group": "category char",
+      "tags": 2870,
+      "x1_ap": 91.87811929600326,
+      "v11_ap": 91.97710045881874,
+      "delta_points": -0.0989811631562091,
+      "x1_wins": 815,
+      "v11_wins": 965,
+      "ties": 1090
+    },
+    {
+      "group": "category copyright",
+      "tags": 1050,
+      "x1_ap": 78.31961239122653,
+      "v11_ap": 79.32927887092089,
+      "delta_points": -1.0096664829623132,
+      "x1_wins": 378,
+      "v11_wins": 524,
+      "ties": 148
+    },
+    {
+      "group": "category artist",
+      "tags": 695,
+      "x1_ap": 83.73359510134034,
+      "v11_ap": 89.3767281228395,
+      "delta_points": -5.643133035251241,
+      "x1_wins": 109,
+      "v11_wins": 314,
+      "ties": 272
+    },
+    {
+      "group": "category meta",
+      "tags": 99,
+      "x1_ap": 28.44267269873385,
+      "v11_ap": 27.431288556721427,
+      "delta_points": 1.0113841334870555,
+      "x1_wins": 52,
+      "v11_wins": 46,
+      "ties": 1
+    },
+    {
+      "group": "category rating",
+      "tags": 4,
+      "x1_ap": 90.30946791172028,
+      "v11_ap": 74.51295033097267,
+      "delta_points": 15.796517580747604,
+      "x1_wins": 2,
+      "v11_wins": 2,
+      "ties": 0
+    },
+    {
+      "group": "frequency 300-499",
+      "tags": 63,
+      "x1_ap": 45.1914965982608,
+      "v11_ap": 51.18916831404295,
+      "delta_points": -5.997671779348618,
+      "x1_wins": 14,
+      "v11_wins": 36,
+      "ties": 13
+    },
+    {
+      "group": "frequency 500-999",
+      "tags": 3962,
+      "x1_ap": 63.82945302063262,
+      "v11_ap": 66.1829312171874,
+      "delta_points": -2.3534781905738833,
+      "x1_wins": 1093,
+      "v11_wins": 1742,
+      "ties": 1127
+    },
+    {
+      "group": "frequency 1000-4999",
+      "tags": 6049,
+      "x1_ap": 57.39406408068191,
+      "v11_ap": 57.68059412367721,
+      "delta_points": -0.286530040883122,
+      "x1_wins": 2509,
+      "v11_wins": 2986,
+      "ties": 554
+    },
+    {
+      "group": "frequency 5000-9999",
+      "tags": 1249,
+      "x1_ap": 52.204597668254,
+      "v11_ap": 51.39728918997086,
+      "delta_points": 0.8073084776866152,
+      "x1_wins": 638,
+      "v11_wins": 607,
+      "ties": 4
+    },
+    {
+      "group": "frequency 10000-plus",
+      "tags": 2071,
+      "x1_ap": 57.361599985368,
+      "v11_ap": 55.69153087854903,
+      "delta_points": 1.6700691066390936,
+      "x1_wins": 1171,
+      "v11_wins": 900,
+      "ties": 0
+    }
+  ],
+  "exact_map_delta_points": -0.5202712224373074,
+  "micro_f1_delta_points": 2.0546272801644228,
+  "comparability": "Same frozen 30,000 EVAL identities and labels; every source sidecar hash verified; identical one-to-one shared tags and support; both 448 FP32 ONNX CUDA with TF32 disabled and identical preprocessing. V1.1 metrics reused from verified archived evaluation. Development-set comparison; V1.1 training overlap is unverified. F1 thresholds independently optimized on the same 200-point grid; not deployment calibration."
+}

config.json

@@ -0,0 +1,77 @@
+{
+  "model_type": "oppaioracle_vit",
+  "architectures": [
+    "SimplifiedTagger"
+  ],
+  "release": "experimental-X1.1",
+  "repo_id": "Grio43/OppaiOracle-experimental-X1.1",
+  "image_size": 448,
+  "num_labels": 14996,
+  "parameter_count": 147706932,
+  "torch_dtype": "float32",
+  "architecture": {
+    "image_size": 448,
+    "patch_size": 16,
+    "num_channels": 3,
+    "hidden_size": 896,
+    "num_hidden_layers": 18,
+    "num_attention_heads": 14,
+    "intermediate_size": 2320,
+    "num_tags": 14996,
+    "dropout": 0.0,
+    "pos_dropout": 0.0,
+    "attention_dropout": 0.0,
+    "qk_norm": true,
+    "layer_scale_init": 0.1,
+    "layer_norm_eps": 1e-06,
+    "use_flex_attention": false,
+    "flex_block_size": 128,
+    "attention_bias": true,
+    "token_ignore_threshold": 0.9,
+    "gradient_checkpointing": false,
+    "checkpoint_every_n_layers": 1,
+    "drop_path_rate": 0.15,
+    "check_numerical_stability": false,
+    "logit_clamp_value": 15.0,
+    "use_fp32_layernorm": false
+  },
+  "preprocessing": {
+    "normalize_mean": [
+      0.5,
+      0.5,
+      0.5
+    ],
+    "normalize_std": [
+      0.5,
+      0.5,
+      0.5
+    ],
+    "image_size": 448,
+    "patch_size": 16,
+    "color_order": "RGB",
+    "pad_color": [
+      114,
+      114,
+      114
+    ],
+    "resize_filter": "LANCZOS",
+    "resize_mode": "letterbox_downscale_only",
+    "upscale": false,
+    "exif_transpose": true,
+    "alpha_composite": true,
+    "input_layout": "NCHW",
+    "padding_mask_true_means": "padding",
+    "output_activation": "sigmoid",
+    "skip_indices": [
+      0,
+      1
+    ]
+  },
+  "vocab_file": "vocabulary.json",
+  "vocab_sha256": "e3e6ccc137d9152d8a17fbd19e0c1ae595b2317a6c55893fd25779620b4b8087",
+  "model_file_onnx": "model.onnx",
+  "model_file_safetensors": "model.safetensors",
+  "threshold": 0.7085427045822144,
+  "threshold_calibrated": false,
+  "threshold_source": "200-point micro-F1 grid on 13,569 shared development tags; not final deployment calibration"
+}

custom_drop_path.py

@@ -0,0 +1,22 @@
+import torch
+import torch.nn as nn
+
+class SafeDropPath(nn.Module):
+    def __init__(self, drop_prob: float = 0.0):
+        super().__init__()
+        drop_prob = float(drop_prob)
+        if not (0.0 <= drop_prob < 1.0):
+            raise ValueError(f"drop_prob must be in [0, 1), got {drop_prob}")
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        if not self.training or self.drop_prob == 0.0:
+            return x
+        keep_prob = 1.0 - self.drop_prob
+        # keep_prob is always > 0 due to __init__ validation
+        shape = (x.shape[0],) + (1,) * (x.ndim - 1)
+        # Draw in float32: drawing in x.dtype (e.g. bf16) quantizes keep_prob
+        # and skews the realized keep rate. The 0/1 mask casts back exactly.
+        random_tensor = keep_prob + torch.rand(shape, dtype=torch.float32, device=x.device)
+        random_tensor.floor_()
+        return x.div(keep_prob) * random_tensor.to(x.dtype)

download.py

@@ -0,0 +1,66 @@
+#!/usr/bin/env python3
+"""Download an immutable OppaiOracle release snapshot using Hugging Face Hub."""
+import argparse
+import json
+from pathlib import Path, PurePosixPath
+
+
+DEFAULT_REPO_ID = 'Grio43/OppaiOracle-experimental-X1.1'
+SUPPORT_FILES = [
+    'README.md', 'LICENSE', 'config.json', 'selected_tags.csv', 'thresholds.json',
+    'infer.py', 'download.py', 'model_architecture.py', 'custom_drop_path.py',
+    'mask_utils.py', 'requirements*.txt', 'release_manifest.json',
+    'preprocessing.json', 'training_metadata.json', 'validation.json',
+    'comparison_v1_1.json', 'comparison_v1_1_per_tag.csv', 'comparison_groups.csv',
+    'checkpoint_evidence.json', 'frozen_evaluation_contract.json',
+    'training_epochs.csv', 'training_buckets.csv', 'inference_checks.json', 'SHA256SUMS',
+]
+
+
+def release_filename(name):
+    """Require ordinary relative Hub filenames, without glob expansion."""
+    path = PurePosixPath(name)
+    if (not name or path.is_absolute() or '..' in path.parts
+            or any(char in name for char in ('\\', ':', '*', '?', '[', ']'))):
+        raise ValueError(f'Invalid release filename: {name!r}')
+    return name
+
+
+def download_release(destination, repo_id=DEFAULT_REPO_ID, model_format='onnx', revision='main'):
+    from huggingface_hub import HfApi, hf_hub_download, snapshot_download
+
+    # Resolve once, so config and model files cannot straddle a moving main.
+    commit = HfApi().model_info(repo_id=repo_id, revision=revision).sha
+    if not commit:
+        raise RuntimeError('Hugging Face did not return an immutable commit SHA')
+    destination = Path(destination)
+    config_path = hf_hub_download(repo_id=repo_id, filename='config.json',
+                                  revision=commit, local_dir=destination)
+    config = json.loads(Path(config_path).read_text(encoding='utf-8-sig'))
+    formats = ('onnx', 'safetensors') if model_format == 'all' else (model_format,)
+    models = [release_filename(config[f'model_file_{kind}']) for kind in formats]
+    vocab = release_filename(config.get('vocab_file', 'vocabulary.json'))
+    patterns = SUPPORT_FILES + [vocab] + models
+    snapshot_download(repo_id=repo_id, revision=commit, local_dir=destination,
+                      allow_patterns=patterns)
+    record = {'repo_id': repo_id, 'commit': commit, 'format': model_format,
+              'model_files': models}
+    (destination / 'download_info.json').write_text(
+        json.dumps(record, indent=2) + '\n', encoding='utf-8')
+    return record
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument('destination', nargs='?', type=Path,
+                        default=Path('OppaiOracle-experimental-X1.1'))
+    parser.add_argument('--repo-id', default=DEFAULT_REPO_ID)
+    parser.add_argument('--revision', default='main', help='Branch, tag, or commit; resolved to a pinned commit')
+    parser.add_argument('--format', choices=('onnx', 'safetensors', 'all'), default='onnx')
+    args = parser.parse_args()
+    record = download_release(args.destination, args.repo_id, args.format, args.revision)
+    print(json.dumps({'directory': str(args.destination.resolve()), **record}, indent=2))
+
+
+if __name__ == '__main__':
+    main()

frozen_evaluation_contract.json

@@ -0,0 +1,52 @@
+{
+  "manifest": "./logs/v2_phase1/dataset_manifest.json",
+  "frozen_ignore_hash": "2194183cb84a03b1701e2aa1833176633223e906055a8e6d031ba13a64c1b021",
+  "current_ignore_hash": "b8fa4041e90e598c8ffcde983500b849606e93d349d5810bac1762f4a812bd60",
+  "frozen_vocab_file_sha256": "a168f4cb343e75fdbacea13de0c3ec2d970829ee0df68a193b4b18cea6a1ed81",
+  "newly_ignored_existing_tags": [
+    "meta:alt_text",
+    "meta:alternate_language",
+    "meta:annotated",
+    "meta:archived_source",
+    "meta:bad_bcy_id",
+    "meta:bad_blogspot_link",
+    "meta:bad_discord_link",
+    "meta:bad_facebook_link",
+    "meta:bad_fc2_link",
+    "meta:bad_hentai_foundry_id",
+    "meta:bad_imgur_id",
+    "meta:bad_instagram_id",
+    "meta:bad_lofter_id",
+    "meta:bad_nico_seiga_link",
+    "meta:bad_nijie_id",
+    "meta:bad_poipiku_id",
+    "meta:bad_source",
+    "meta:bad_tumblr_link",
+    "meta:bad_twitter_link",
+    "meta:bad_weibo_id",
+    "meta:check_artist",
+    "meta:check_character",
+    "meta:cropped",
+    "meta:derivative_work",
+    "meta:detexted",
+    "meta:fancy_embedded_notes",
+    "meta:hard-translated",
+    "meta:has_bad_revision",
+    "meta:has_censored_revision",
+    "meta:has_downscaled_revision",
+    "meta:imageboard_desourced",
+    "meta:inactive_account",
+    "meta:large_variant_set",
+    "meta:pixel-perfect_duplicate",
+    "meta:protected_link",
+    "meta:redrawn",
+    "meta:second-party_source",
+    "meta:source_larger",
+    "meta:source_mismatch",
+    "meta:source_smaller",
+    "meta:third-party_source",
+    "meta:variant_set"
+  ],
+  "verified_sidecars": 30000,
+  "method": "Direct frozen-EVAL sidecars and original prepared vocabulary. Original ignored tags were excluded when this vocabulary was built; all original output labels remain evaluated. No preparation or training startup is performed."
+}

infer.py

@@ -0,0 +1,244 @@
+#!/usr/bin/env python3
+"""Offline ONNX or safetensors inference for OppaiOracle experimental-X1.1."""
+import argparse
+import base64
+import gzip
+import hashlib
+import json
+from pathlib import Path
+
+import numpy as np
+from PIL import Image, ImageOps
+
+
+def read_json(path):
+    return json.loads(Path(path).read_text(encoding='utf-8-sig'))
+
+
+def local_file(directory, name):
+    """Resolve a release-relative file without escaping the package directory."""
+    root = Path(directory).resolve()
+    path = (root / name).resolve()
+    if not path.is_relative_to(root):
+        raise ValueError(f'Release file must be inside {root}: {name}')
+    return path
+
+
+def vocabulary_sha256(vocabulary):
+    canonical = json.dumps(vocabulary, ensure_ascii=False, sort_keys=True,
+                           separators=(',', ':')).encode('utf-8')
+    return hashlib.sha256(canonical).hexdigest()
+
+
+def load_config(config_path):
+    config = read_json(config_path)
+    pre = config['preprocessing']
+    if (int(pre['image_size']) != 448 or pre.get('color_order', 'RGB') != 'RGB'
+            or list(pre['pad_color']) != [114, 114, 114]
+            or list(pre['normalize_mean']) != [0.5, 0.5, 0.5]
+            or list(pre['normalize_std']) != [0.5, 0.5, 0.5]
+            or pre.get('upscale', False)):
+        raise ValueError('Expected the X1.1 448px RGB/downscale-only preprocessing contract')
+    if pre.get('resize_filter', 'LANCZOS').upper() != 'LANCZOS':
+        raise ValueError('X1.1 requires LANCZOS resampling')
+    if int(config['architecture']['image_size']) != int(pre['image_size']):
+        raise ValueError('Architecture and preprocessing image sizes differ')
+    directory = Path(config_path).resolve().parent
+    vocab = read_json(local_file(directory, config.get('vocab_file', 'vocabulary.json')))
+    digest = vocabulary_sha256(vocab)
+    if digest != config['vocab_sha256']:
+        raise ValueError('Vocabulary checksum does not match config.json')
+    names = vocab['index_to_tag']
+    num_tags = int(config['architecture']['num_tags'])
+    if set(names) != {str(i) for i in range(num_tags)}:
+        raise ValueError('Vocabulary indices must cover every model output exactly')
+    if names['0'] != '<PAD>' or names['1'] != '<UNK>':
+        raise ValueError('Expected PAD=0 and UNK=1')
+    if vocab['tag_to_index'] != {tag: int(index) for index, tag in names.items()}:
+        raise ValueError('Vocabulary index maps disagree')
+    return config, vocab
+
+
+def preprocess(path, metadata):
+    """EXIF transpose, alpha composite, downscale-only LANCZOS, centered pad."""
+    size = int(metadata['image_size'])
+    pad = tuple(metadata['pad_color'])
+    with Image.open(path) as source:
+        img = ImageOps.exif_transpose(source)
+        if img.mode in ('RGBA', 'LA') or 'transparency' in img.info:
+            rgba = img.convert('RGBA')
+            background = Image.new('RGB', img.size, pad)
+            background.paste(rgba.convert('RGB'), mask=rgba.getchannel('A'))
+            img = background
+        else:
+            img = img.convert('RGB')
+        scale = min(1.0, size / img.width, size / img.height)
+        width = max(1, round(img.width * scale))
+        height = max(1, round(img.height * scale))
+        resized = img.resize((width, height), Image.Resampling.LANCZOS)
+        canvas = Image.new('RGB', (size, size), pad)
+        left, top = (size - width) // 2, (size - height) // 2
+        canvas.paste(resized, (left, top))
+        array = np.asarray(canvas, dtype=np.float32)
+    mean = np.asarray(metadata['normalize_mean'], dtype=np.float32)
+    std = np.asarray(metadata['normalize_std'], dtype=np.float32)
+    array = (array / 255.0 - mean) / std
+    mask = np.ones((size, size), dtype=bool)
+    mask[top:top + height, left:left + width] = False
+    return np.ascontiguousarray(array.transpose(2, 0, 1)), mask
+
+
+class ONNXRuntime:
+    def __init__(self, model_path, config, device):
+        import onnxruntime as ort
+
+        if device == 'cuda':
+            # Reuse CUDA DLLs from PyTorch when it is installed.
+            try:
+                import torch  # noqa: F401
+            except ImportError:
+                pass
+            if hasattr(ort, 'preload_dlls'):
+                ort.preload_dlls()
+            if 'CUDAExecutionProvider' not in ort.get_available_providers():
+                raise RuntimeError('Install compatible onnxruntime-gpu or use --device cpu')
+            providers = [('CUDAExecutionProvider', {'use_tf32': False}), 'CPUExecutionProvider']
+        else:
+            providers = ['CPUExecutionProvider']
+        options = ort.SessionOptions()
+        options.intra_op_num_threads = 8
+        self.session = ort.InferenceSession(str(model_path), sess_options=options,
+                                           providers=providers)
+        if device == 'cuda' and 'CUDAExecutionProvider' not in self.session.get_providers():
+            raise RuntimeError('CUDA provider initialization failed')
+        metadata = self.session.get_modelmeta().custom_metadata_map
+        if metadata.get('preprocessing') != 'external':
+            raise ValueError('Expected the release ONNX model with external preprocessing')
+        if metadata.get('vocab_sha256') != config['vocab_sha256']:
+            raise ValueError('ONNX vocabulary checksum does not match config.json')
+        embedded = json.loads(gzip.decompress(base64.b64decode(metadata['vocab_b64_gzip'])))
+        if vocabulary_sha256(embedded) != config['vocab_sha256']:
+            raise ValueError('ONNX embedded vocabulary checksum mismatch')
+        for key in ('image_size', 'pad_color', 'normalize_mean', 'normalize_std'):
+            if json.loads(metadata[key]) != config['preprocessing'][key]:
+                raise ValueError(f'ONNX preprocessing differs from config.json: {key}')
+        if metadata.get('color_order', 'RGB') != config['preprocessing'].get('color_order', 'RGB'):
+            raise ValueError('ONNX color order differs from config.json')
+        inputs = {item.name: item.type for item in self.session.get_inputs()}
+        if set(inputs) != {'pixel_values', 'padding_mask'} or inputs['padding_mask'] != 'tensor(bool)':
+            raise ValueError('Unexpected ONNX input contract')
+        if inputs['pixel_values'] not in ('tensor(float)', 'tensor(float16)'):
+            raise ValueError('Unsupported ONNX input precision')
+        self.dtype = np.float16 if inputs['pixel_values'] == 'tensor(float16)' else np.float32
+
+    def predict(self, pixels, masks):
+        # The exported ONNX graph already includes sigmoid.
+        return self.session.run(['probabilities'], {
+            'pixel_values': pixels.astype(self.dtype, copy=False),
+            'padding_mask': masks,
+        })[0]
+
+
+class SafetensorsRuntime:
+    def __init__(self, model_path, config, device):
+        import torch
+        from safetensors import safe_open
+        from safetensors.torch import load_file
+        from model_architecture import SimplifiedTagger, VisionTransformerConfig
+
+        if device == 'cuda' and not torch.cuda.is_available():
+            raise RuntimeError('CUDA is unavailable; use --device cpu')
+        torch.set_num_threads(8)
+        if device == 'cuda':
+            torch.backends.cuda.matmul.allow_tf32 = False
+            torch.backends.cudnn.allow_tf32 = False
+        architecture = dict(config['architecture'])
+        architecture.update(use_flex_attention=False, gradient_checkpointing=False)
+        self.model = SimplifiedTagger(VisionTransformerConfig(**architecture))
+        # This architecture selects eager SDPA through set_onnx_mode; changing
+        # use_flex_attention alone does not switch its forward implementation.
+        self.model.set_onnx_mode(True)
+        with safe_open(str(model_path), framework='pt', device='cpu') as handle:
+            metadata = handle.metadata() or {}
+        if metadata.get('vocab_sha256') != config['vocab_sha256']:
+            raise ValueError('Safetensors vocabulary checksum does not match config.json')
+        # No pickle loading or remote code downloads. Strict loading preserves
+        # every output index and catches incompatible architectures.
+        state = load_file(str(model_path), device='cpu')
+        self.model.load_state_dict(state, strict=True)
+        del state
+        self.model = self.model.float().to(device).eval()
+        self.device = device
+
+    def predict(self, pixels, masks):
+        import torch
+
+        with torch.inference_mode():
+            logits = self.model(torch.from_numpy(pixels).to(self.device),
+                                padding_mask=torch.from_numpy(masks).to(self.device))['tag_logits']
+            return torch.sigmoid(logits.float()).cpu().numpy()
+
+
+def select_tags(scores, names, threshold, top_k):
+    indices = np.flatnonzero(scores >= threshold)
+    indices = indices[indices >= 2]  # PAD and UNK are never predictions.
+    indices = indices[np.argsort(-scores[indices], kind='stable')]
+    if top_k:
+        indices = indices[:top_k]
+    return [{'name': names[str(int(i))], 'score': float(scores[i])} for i in indices]
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument('images', nargs='+', type=Path)
+    parser.add_argument('--format', choices=('onnx', 'safetensors'), default='onnx')
+    parser.add_argument('--model', type=Path, help='Override the local model file')
+    parser.add_argument('--config', type=Path, default=Path(__file__).resolve().parent / 'config.json')
+    parser.add_argument('--device', choices=('cpu', 'cuda'), default='cpu')
+    parser.add_argument('--threshold', type=float, help='Override the decision threshold in config.json')
+    parser.add_argument('--top-k', type=int, default=0, help='0 returns every tag above threshold')
+    parser.add_argument('--batch-size', type=int, default=1)
+    parser.add_argument('--output', type=Path)
+    args = parser.parse_args()
+    if not 1 <= args.batch_size <= 128:
+        parser.error('--batch-size must be between 1 and 128')
+    if args.top_k < 0 or (args.threshold is not None and not 0 <= args.threshold <= 1):
+        parser.error('Use a nonnegative --top-k and a threshold between 0 and 1')
+    config, vocab = load_config(args.config)
+    model_path = args.model or local_file(args.config.parent, config[f'model_file_{args.format}'])
+    threshold = args.threshold if args.threshold is not None else float(config['threshold'])
+    if not 0 <= threshold <= 1:
+        raise ValueError('Config threshold must be between 0 and 1')
+    runtime_type = ONNXRuntime if args.format == 'onnx' else SafetensorsRuntime
+    runtime = runtime_type(model_path, config, args.device)
+    results = []
+    for offset in range(0, len(args.images), args.batch_size):
+        paths = args.images[offset:offset + args.batch_size]
+        batch = [preprocess(path, config['preprocessing']) for path in paths]
+        pixels = np.stack([item[0] for item in batch])
+        masks = np.stack([item[1] for item in batch])
+        probabilities = runtime.predict(pixels, masks)
+        if probabilities.shape != (len(paths), int(config['architecture']['num_tags'])):
+            raise RuntimeError('Model output shape does not match vocabulary')
+        if not np.isfinite(probabilities).all() or ((probabilities < 0) | (probabilities > 1)).any():
+            raise RuntimeError('Model returned invalid probabilities')
+        for path, scores in zip(paths, probabilities):
+            results.append({'image': str(path), 'tags': select_tags(
+                scores, vocab['index_to_tag'], threshold, args.top_k)})
+    output = {
+        'model': model_path.name,
+        'format': args.format,
+        'threshold': threshold,
+        'threshold_source': 'user override' if args.threshold is not None else config.get('threshold_source', 'config.json'),
+        'vocab_sha256': config['vocab_sha256'],
+        'predictions': results,
+    }
+    rendered = json.dumps(output, indent=2, ensure_ascii=False)
+    if args.output:
+        args.output.write_text(rendered + '\n', encoding='utf-8')
+    else:
+        print(rendered)
+
+
+if __name__ == '__main__':
+    main()

inference_checks.json

@@ -0,0 +1,70 @@
+{
+  "passed": true,
+  "scope": "Four deterministic synthetic image fixtures; full development AP is archived in validation.json",
+  "atol": 0.0005,
+  "rtol": 0.001,
+  "checks": [
+    {
+      "format": "onnx",
+      "device": "cpu",
+      "batch": 1,
+      "max_absolute_error_vs_cpu_safetensors": 1.1324882507324219e-06,
+      "mean_absolute_error": 6.287191922638158e-08
+    },
+    {
+      "format": "onnx",
+      "device": "cpu",
+      "batch": 2,
+      "max_absolute_error_vs_cpu_safetensors": 1.3113021850585938e-06,
+      "mean_absolute_error": 5.217425069758974e-08
+    },
+    {
+      "format": "onnx",
+      "device": "cpu",
+      "batch": 4,
+      "max_absolute_error_vs_cpu_safetensors": 1.3113021850585938e-06,
+      "mean_absolute_error": 6.336669855500077e-08
+    },
+    {
+      "device": "cpu",
+      "padding_mask_max_effect": 0.6925351023674011
+    },
+    {
+      "format": "onnx",
+      "device": "cuda",
+      "batch": 1,
+      "max_absolute_error_vs_cpu_safetensors": 7.748603820800781e-07,
+      "mean_absolute_error": 6.658750351107301e-08
+    },
+    {
+      "format": "onnx",
+      "device": "cuda",
+      "batch": 2,
+      "max_absolute_error_vs_cpu_safetensors": 1.043081283569336e-06,
+      "mean_absolute_error": 5.9538443508699856e-08
+    },
+    {
+      "format": "onnx",
+      "device": "cuda",
+      "batch": 4,
+      "max_absolute_error_vs_cpu_safetensors": 9.238719940185547e-07,
+      "mean_absolute_error": 6.518111206332833e-08
+    },
+    {
+      "device": "cuda",
+      "padding_mask_max_effect": 0.6925351023674011
+    },
+    {
+      "format": "safetensors",
+      "device": "cuda",
+      "batch": 2,
+      "max_absolute_error_vs_cpu_safetensors": 1.1920928955078125e-06
+    }
+  ],
+  "versions": {
+    "torch": "2.14.0+cu130",
+    "numpy": "2.2.6"
+  },
+  "model_sha256": "49ba0f66f83ba88a101b6edac54ec39f0a12eab66743e13ae42efc3570745bf7",
+  "safetensors_sha256": "ec07d0034f48ba3ceff8624993baf3ea30fa3a2e3cc8908682eb3e4f84eda766"
+}

mask_utils.py

@@ -0,0 +1,96 @@
+#!/usr/bin/env python3
+"""
+Mask migration shims.
+Standardize pixel padding masks to (B,1,H,W) bool with True=PAD,
+and pool to token-level ignore masks (B,L) with True=IGNORE.
+"""
+from typing import Optional
+import torch
+import torch.nn.functional as F
+
+@torch.no_grad()
+def ensure_pixel_padding_mask(
+    mask: torch.Tensor,
+    mask_semantics: str = 'auto'
+) -> Optional[torch.Tensor]:
+    """
+    (mask) -> (B,1,H,W,bool, True=PAD)
+    Accepts (B,H,W) or (B,1,H,W) of {bool,uint8,int,float}.
+
+    Args:
+        mask: Input mask tensor
+        mask_semantics: Interpretation of True values:
+            - 'auto': Heuristic detection (legacy, unreliable)
+            - 'pad': True means padding (no inversion needed)
+            - 'valid': True means valid pixels (will be inverted)
+
+    Returns:
+        Standardized mask with True=PAD semantics
+    """
+    if mask is None:
+        return None
+
+    # Use tensor-based assertions that work during tracing
+    original_dims = mask.dim()
+    if mask.dim() == 3:
+        mask = mask.unsqueeze(1)  # (B,1,H,W)
+
+    # Trace-friendly validation
+    if not torch.jit.is_tracing():
+        assert mask.dim() == 4, \
+            f"padding_mask must be (B,H,W) or (B,1,H,W) after expansion; got {original_dims}D -> {mask.dim()}D with shape {tuple(mask.shape)}"
+
+    # Binarize + cast
+    if mask.dtype != torch.bool:
+        mask = (mask > 0.5)
+
+    # Apply semantics
+    if mask_semantics == 'valid':
+        # Invert: True=valid -> True=pad
+        mask = ~mask
+    elif mask_semantics == 'pad':
+        # Already correct semantics
+        pass
+    elif mask_semantics == 'auto':
+        # Legacy heuristic (unreliable, deprecated) - per-sample decision
+        frac_true = mask.flatten(2).to(torch.float16).mean(dim=2)  # (B, 1)
+        should_invert = frac_true > 0.5  # (B, 1)
+        mask = torch.where(should_invert.unsqueeze(-1).unsqueeze(-1), ~mask, mask)
+        # Log warning in non-tracing mode
+        if not torch.jit.is_tracing():
+            import logging
+            logging.getLogger(__name__).warning(
+                "Using 'auto' mask semantics detection is unreliable. "
+                "Please specify mask_semantics='pad' or 'valid' explicitly."
+            )
+    else:
+        raise ValueError(f"Invalid mask_semantics: {mask_semantics}. Use 'auto', 'pad', or 'valid'.")
+
+    return mask
+
+@torch.no_grad()
+def pixel_to_token_ignore(mask_b1hw: torch.Tensor, patch: int, threshold: float = 0.9) -> torch.Tensor:
+    """
+    (B,1,H,W,bool True=PAD) -> (B,L,bool True=IGNORE)
+    Pools pad-fraction per patch via avg_pool2d(kernel=stride=patch) and marks
+    tokens IGNORE when pad_fraction >= threshold.
+    """
+    B, C, H, W = mask_b1hw.shape
+
+    # Trace-friendly validation
+    if not torch.jit.is_tracing():
+        assert mask_b1hw.dim() == 4, \
+            f"Expected 4D tensor (B,1,H,W), got {mask_b1hw.dim()}D with shape {tuple(mask_b1hw.shape)}"
+        assert C == 1, \
+            f"Expected channel dimension of 1, got {C} with shape {tuple(mask_b1hw.shape)}"
+        assert mask_b1hw.dtype == torch.bool, \
+            f"mask_b1hw must be bool (True=PAD), got dtype {mask_b1hw.dtype}"
+        assert H % patch == 0, \
+            f"Image height {H} must be divisible by patch size {patch}"
+        assert W % patch == 0, \
+            f"Image width {W} must be divisible by patch size {patch}"
+
+    # pad_fraction per patch
+    pad_frac = F.avg_pool2d(mask_b1hw.to(torch.float16), kernel_size=patch, stride=patch)  # (B,1,GH,GW)
+    ignore = (pad_frac.squeeze(1) >= threshold)                                   # (B,GH,GW)
+    return ignore.flatten(1)                                                      # (B,L)

model.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:49ba0f66f83ba88a101b6edac54ec39f0a12eab66743e13ae42efc3570745bf7
+size 592413547

model.safetensors

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

model_architecture.py

@@ -0,0 +1,860 @@
+#!/usr/bin/env python3
+"""
+Model Architecture for Anime Image Tagger - Direct Training (modified)
+Vision Transformer for anime image tagging
+"""
+
+from abc import ABC, abstractmethod
+import logging
+import math
+import warnings
+from dataclasses import dataclass, fields
+from typing import Any, Dict, Optional, Tuple
+
+logger = logging.getLogger(__name__)
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from mask_utils import ensure_pixel_padding_mask, pixel_to_token_ignore
+from custom_drop_path import SafeDropPath
+
+# Import Flex Attention (PyTorch 2.5+)
+from torch.nn.attention.flex_attention import flex_attention, create_block_mask, BlockMask
+
+logger.info("Using PyTorch Flex Attention")
+
+
+def _check_triton_available() -> bool:
+    """Check if Triton is available for compiled block mask creation."""
+    try:
+        import triton
+        return True
+    except ImportError:
+        return False
+
+
+# Cache the Triton availability check at module load time
+_TRITON_AVAILABLE = _check_triton_available()
+# PyTorch 2.14 deprecates create_block_mask(_compile=True). Reuse the public
+# compiled callable; compilation remains lazy until the first forward pass.
+_CREATE_BLOCK_MASK = torch.compile(create_block_mask) if _TRITON_AVAILABLE else create_block_mask
+
+
+class BaseTagger(ABC, nn.Module):
+    """Abstract base class for image tagging models."""
+
+    @abstractmethod
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        padding_mask: Optional[torch.Tensor] = None,
+    ) -> Dict[str, torch.Tensor]:
+        """Forward pass returning tag_logits dict."""
+        pass
+
+    @abstractmethod
+    def set_onnx_mode(self, enabled: bool = True) -> None:
+        """Enable/disable ONNX-compatible attention mode."""
+        pass
+
+    @property
+    @abstractmethod
+    def config(self):
+        """Return model configuration."""
+        pass
+
+
+def initialize_tag_head_bias(
+    model: BaseTagger,
+    index_to_tag: Dict[int, str],
+    tag_frequencies: Dict[str, int],
+    total_samples: int,
+    min_prior: float = 1e-5,
+    max_prior: float = 0.99,
+    rated_samples: Optional[int] = None,
+) -> None:
+    """Initialize tag_head bias with log-prior for focal loss (RetinaNet technique).
+
+    Sets each tag's bias to log(prior / (1 - prior)) based on empirical tag frequency.
+    This makes the model start by predicting "mostly negative" for sparse labels,
+    providing meaningful gradients from the first step instead of random 0.5 predictions.
+
+    Must be called AFTER model creation (which zeros the bias) and BEFORE checkpoint
+    loading (which will overwrite the bias if resuming).
+    """
+    with torch.no_grad():
+        bias = model.tag_head.bias
+        num_tags = bias.shape[0]
+        for idx in range(num_tags):
+            tag = index_to_tag.get(idx, "")
+            freq = tag_frequencies.get(tag, 0)
+            denominator = (rated_samples if tag.startswith('rating:') and rated_samples is not None
+                           else total_samples)
+            # An entirely unrated corpus provides no rating prior evidence.
+            empirical = 0.25 if denominator == 0 and tag.startswith('rating:') else freq / max(1, denominator)
+            prior = max(min_prior, min(max_prior, empirical))
+            bias[idx] = math.log(prior / (1 - prior))
+
+        bias_vals = bias.tolist()
+        logger.info(
+            "Tag head bias initialized with log-prior: min=%.2f, max=%.2f, mean=%.2f",
+            min(bias_vals), max(bias_vals), sum(bias_vals) / len(bias_vals),
+        )
+
+
+class LayerNormFp32(nn.LayerNorm):
+    """
+    LayerNorm that optionally casts to float32 before calling the original LayerNorm.
+    This is to improve stability when using mixed precision training.
+
+    Args:
+        normalized_shape: Input shape from an expected input of size
+        eps: A value added to the denominator for numerical stability
+        elementwise_affine: A boolean value that when set to True, gives learnable parameters
+        use_fp32: If True, cast to float32 before LayerNorm (better stability).
+                  If False, use native dtype (faster but potentially less stable).
+    """
+    def __init__(self, normalized_shape, eps=1e-5, elementwise_affine=True, use_fp32=True):
+        super().__init__(normalized_shape, eps=eps, elementwise_affine=elementwise_affine)
+        self.use_fp32 = use_fp32
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.use_fp32:
+            return super().forward(x.float()).type_as(x)
+        else:
+            return super().forward(x)
+
+
+@dataclass
+class VisionTransformerConfig:
+    """Configuration for the Vision Transformer used in direct training."""
+    image_size: int = 448
+    # Patch size must divide image_size evenly to avoid losing border information (validated in __post_init__)
+    patch_size: int = 16
+    num_channels: int = 3
+    hidden_size: int = 1024
+    num_hidden_layers: int = 18
+    num_attention_heads: int = 16
+    intermediate_size: int = 4096
+    num_tags: int = 100000  # This should be overridden with actual vocab size (includes rating tags)
+    dropout: float = 0.1
+    pos_dropout: float = 0.0  # Position embedding dropout (0.0 = modern standard; drop_path handles regularization)
+    attention_dropout: float = 0.1
+    qk_norm: bool = False  # Legacy checkpoints omit this; V2 explicitly enables it.
+    layer_scale_init: float = 0.0  # Zero disables LayerScale for legacy checkpoints.
+    layer_norm_eps: float = 1e-6
+    use_flex_attention: bool = True  # Use Flex Attention (PyTorch 2.5+)
+    flex_block_size: int = 128  # Block size for Flex Attention sparse computation
+    attention_bias: bool = True  # Use bias in attention QKV and projection layers
+    # Token ignore threshold: a token is ignored if >= this fraction of its pixels are PAD
+    token_ignore_threshold: float = 0.9
+    # Enable gradient checkpointing by default to reduce memory usage
+    gradient_checkpointing: bool = True
+    # Checkpoint every N-th layer (1=all layers, 2=every 2nd, 4=every 4th, etc.)
+    checkpoint_every_n_layers: int = 1
+    drop_path_rate: float = 0.0
+    # Enable numerical stability checking (for debugging only)
+    check_numerical_stability: bool = False
+    # Logit clamping for inference/export outputs (None to disable)
+    # exp(15) ~ 3.3M which is safe for softmax in float16/bfloat16
+    # Applied only in eval mode: training logits are left unclamped because
+    # clamp has zero gradient outside the bounds (see forward()).
+    logit_clamp_value: Optional[float] = 15.0
+    # Precision configuration
+    use_fp32_layernorm: bool = False  # Use FP32 for LayerNorm (better stability, slight speed cost). Set to False for full bfloat16.
+
+    def __post_init__(self):
+        # Validate drop_path_rate
+        assert 0.0 <= self.drop_path_rate < 1.0, (
+            f"drop_path_rate must be in [0, 1), got {self.drop_path_rate}"
+        )
+        if self.drop_path_rate > 0.5:
+            warnings.warn(
+                f"drop_path_rate={self.drop_path_rate} is unusually high; "
+                "values between 0.08 and 0.3 are typical for ViT."
+            )
+
+        # Validate image_size and patch_size
+        if self.image_size <= 0:
+            raise ValueError(
+                f"image_size must be positive, got {self.image_size}"
+            )
+
+        if self.patch_size <= 0:
+            raise ValueError(
+                f"patch_size must be positive, got {self.patch_size}"
+            )
+
+        if self.image_size % self.patch_size != 0:
+            # Calculate valid alternatives
+            valid_sizes = [
+                s for s in [224, 256, 384, 448, 512, 576, 640, 768, 896, 1024]
+                if s % self.patch_size == 0 and abs(s - self.image_size) < 200
+            ]
+
+            raise ValueError(
+                f"image_size ({self.image_size}) must be evenly divisible by "
+                f"patch_size ({self.patch_size}). "
+                f"Current: {self.image_size} % {self.patch_size} = {self.image_size % self.patch_size}. "
+                f"\n\nSuggested fixes:"
+                f"\n  1. Use a standard image size: {valid_sizes if valid_sizes else 'N/A'}"
+                f"\n  2. Change patch_size to a factor of {self.image_size}: "
+                f"{[d for d in [8, 14, 16, 20, 32] if self.image_size % d == 0]}"
+                f"\n  3. Adjust image_size to nearest multiple of {self.patch_size}: "
+                f"{(self.image_size // self.patch_size) * self.patch_size} or "
+                f"{((self.image_size // self.patch_size) + 1) * self.patch_size}"
+            )
+
+        # Validate computed values make sense
+        num_patches = (self.image_size // self.patch_size) ** 2
+        if num_patches < 4:
+            warnings.warn(
+                f"Very few patches ({num_patches}) with image_size={self.image_size} "
+                f"and patch_size={self.patch_size}. Model may underperform."
+            )
+
+        if num_patches > 10000:
+            warnings.warn(
+                f"Very many patches ({num_patches}) with image_size={self.image_size} "
+                f"and patch_size={self.patch_size}. May cause memory issues."
+            )
+
+        # Validate hidden_size is divisible by num_attention_heads
+        if self.hidden_size % self.num_attention_heads != 0:
+            raise ValueError(
+                f"hidden_size ({self.hidden_size}) must be divisible by "
+                f"num_attention_heads ({self.num_attention_heads}). "
+                f"Current remainder: {self.hidden_size % self.num_attention_heads}. "
+                f"head_dim would be {self.hidden_size / self.num_attention_heads:.2f} (must be integer)."
+            )
+
+
+class TransformerBlock(nn.Module):
+    """Single transformer block using Flex Attention."""
+
+    def __init__(self, config: VisionTransformerConfig, drop_path: float = 0.):
+        super().__init__()
+        self.config = config
+        self.norm1 = LayerNormFp32(config.hidden_size, eps=config.layer_norm_eps, use_fp32=config.use_fp32_layernorm)
+
+        # Flex Attention setup
+        self.num_heads = config.num_attention_heads
+        self.head_dim = config.hidden_size // config.num_attention_heads
+        self.scale = self.head_dim ** -0.5
+        self.flex_block_size = config.flex_block_size
+
+        # QKV projection (fused linear for efficiency)
+        self.qkv = nn.Linear(config.hidden_size, 3 * config.hidden_size, bias=config.attention_bias)
+        self.proj = nn.Linear(config.hidden_size, config.hidden_size, bias=config.attention_bias)
+        self.attn_dropout = config.attention_dropout
+        self.q_norm = LayerNormFp32(self.head_dim, eps=config.layer_norm_eps, use_fp32=True) if config.qk_norm else nn.Identity()
+        self.k_norm = LayerNormFp32(self.head_dim, eps=config.layer_norm_eps, use_fp32=True) if config.qk_norm else nn.Identity()
+        self.ls1 = nn.Parameter(torch.full((config.hidden_size,), config.layer_scale_init)) if config.layer_scale_init > 0 else None
+        self.ls2 = nn.Parameter(torch.full((config.hidden_size,), config.layer_scale_init)) if config.layer_scale_init > 0 else None
+        self.capture_attention_stats = False
+        self.max_attention_logit = None
+
+        # Drop path and MLP
+        self.drop_path = SafeDropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+        self.norm2 = LayerNormFp32(config.hidden_size, eps=config.layer_norm_eps, use_fp32=config.use_fp32_layernorm)
+        self.mlp = nn.Sequential(
+            nn.Linear(config.hidden_size, config.intermediate_size),
+            nn.GELU(),
+            nn.Dropout(config.dropout),
+            nn.Linear(config.intermediate_size, config.hidden_size),
+            nn.Dropout(config.dropout)
+        )
+
+    def _normalize_qk(self, q, k):
+        q, k = self.q_norm(q), self.k_norm(k)
+        if self.capture_attention_stats:
+            # Chunked exact maximum: bounded scratch rather than retaining L x L
+            # scores for every layer. Only enabled for the diagnostic forward.
+            with torch.no_grad(), torch.autocast(device_type=q.device.type, enabled=False):
+                maximum = q.new_zeros((), dtype=torch.float32)
+                for start in range(0, q.shape[-2], 64):
+                    scores = q[:, :, start:start + 64].float() @ k.float().transpose(-2, -1)
+                    maximum = torch.maximum(maximum, scores.abs().amax() * self.scale)
+                self.max_attention_logit = maximum
+        return q, k
+
+    def _residual(self, x, attn_out):
+        attention = self.proj(attn_out)
+        x = x + self.drop_path(attention if self.ls1 is None else attention * self.ls1.to(attention.dtype))
+        mlp = self.mlp(self.norm2(x))
+        return x + self.drop_path(mlp if self.ls2 is None else mlp * self.ls2.to(mlp.dtype))
+
+    def forward(self, x: torch.Tensor, block_mask: Optional[BlockMask] = None) -> torch.Tensor:
+        """Forward pass using Flex Attention.
+
+        Args:
+            x: Input tensor of shape (batch_size, seq_len, hidden_size)
+            block_mask: Optional pre-computed BlockMask for padding-aware attention.
+                       Created once by SimplifiedTagger and shared across all layers.
+
+        Returns:
+            torch.Tensor: Output tensor of shape (batch_size, seq_len, hidden_size)
+
+        Note:
+            Uses pre-layer normalization (LayerNorm -> Attention -> Residual).
+        """
+        normed_x = self.norm1(x)
+        B, L, D = normed_x.shape
+
+        # QKV projection and reshape
+        qkv = self.qkv(normed_x).reshape(B, L, 3, self.num_heads, self.head_dim)
+        q, k, v = qkv.unbind(2)  # Each: (B, L, num_heads, head_dim)
+
+        # Transpose to (B, num_heads, L, head_dim) for flex_attention
+        # IMPORTANT: Make tensors contiguous after transpose to ensure proper memory layout.
+        # Non-contiguous tensors from reshape->unbind->transpose cause stride assertion
+        # failures when torch.compile traces the graph (inductor expects contiguous strides).
+        q = q.transpose(1, 2).contiguous()
+        k = k.transpose(1, 2).contiguous()
+        v = v.transpose(1, 2).contiguous()
+        q, k = self._normalize_qk(q, k)
+
+        # flex_attention has no native attention-weight dropout, so apply V-dropout
+        # before the call. Zeroing random value tokens removes their contribution to
+        # all queries — a documented approximation of attention dropout used in T5
+        # and several BERT variants. Replaces the previous post-attention output
+        # dropout, which did not regularize attention weights as the config name
+        # suggested.
+        if self.training and self.attn_dropout > 0:
+            v = F.dropout(v, p=self.attn_dropout, training=True)
+
+        # Flex Attention computation (block_mask is pre-computed and shared across layers)
+        attn_out = flex_attention(
+            q, k, v,
+            block_mask=block_mask,
+            scale=self.scale,
+        )
+
+        # Reshape back to (B, L, D) - ensure contiguous for torch.compile compatibility
+        attn_out = attn_out.transpose(1, 2).reshape(B, L, D).contiguous()
+
+        # Residual connections
+        x = self._residual(x, attn_out)
+
+        return x
+
+    def forward_sdpa(self, x: torch.Tensor, attn_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+        """Forward pass using scaled_dot_product_attention (ONNX-compatible).
+
+        This method provides an ONNX-exportable alternative to the flex_attention path.
+        It uses F.scaled_dot_product_attention which is supported by ONNX opset 14+.
+
+        Args:
+            x: Input tensor of shape (batch_size, seq_len, hidden_size)
+            attn_mask: Optional boolean attention mask (B, L) with True=ATTEND, False=IGNORE.
+                      Passed directly to SDPA as a boolean key-padding mask — avoids the
+                      additive float mask path that forces PyTorch off the Flash/efficient
+                      backends.
+
+        Returns:
+            torch.Tensor: Output tensor of shape (batch_size, seq_len, hidden_size)
+
+        Note:
+            Uses pre-layer normalization (LayerNorm -> Attention -> Residual).
+            For ONNX export, this path is used instead of flex_attention.
+        """
+        normed_x = self.norm1(x)
+        B, L, D = normed_x.shape
+
+        # QKV projection and reshape
+        qkv = self.qkv(normed_x).reshape(B, L, 3, self.num_heads, self.head_dim)
+        q, k, v = qkv.unbind(2)  # Each: (B, L, num_heads, head_dim)
+
+        # Transpose to (B, num_heads, L, head_dim) for scaled_dot_product_attention
+        q = q.transpose(1, 2).contiguous()
+        k = k.transpose(1, 2).contiguous()
+        v = v.transpose(1, 2).contiguous()
+        q, k = self._normalize_qk(q, k)
+
+        # Pass the boolean mask straight through. SDPA accepts a bool key-padding mask
+        # broadcastable to (B, num_heads, L_q, L_k); we broadcast (B, L) -> (B, 1, 1, L).
+        # Bool semantics for SDPA: True=ATTEND, False=IGNORE — already matches `attn_mask`.
+        sdpa_mask: Optional[torch.Tensor] = None
+        if attn_mask is not None:
+            if attn_mask.ndim != 2 or attn_mask.shape[0] != B or attn_mask.shape[1] != L:
+                raise ValueError(
+                    f"attn_mask shape {attn_mask.shape} mismatch. "
+                    f"Expected ({B}, {L}) for input shape ({B}, {L}, {D})"
+                )
+            sdpa_mask = attn_mask.unsqueeze(1).unsqueeze(2)  # (B, 1, 1, L) bool
+
+        # Scaled dot-product attention (ONNX-compatible)
+        attn_out = F.scaled_dot_product_attention(
+            q, k, v,
+            attn_mask=sdpa_mask,
+            dropout_p=self.attn_dropout if self.training else 0.0,
+            scale=self.scale,
+        )
+
+        # Reshape back to (B, L, D)
+        attn_out = attn_out.transpose(1, 2).reshape(B, L, D).contiguous()
+
+        # Residual connections
+        x = self._residual(x, attn_out)
+
+        return x
+
+
+@torch.no_grad()
+def sample_attention_logits(model, images, padding_mask=None):
+    """Exact absolute score maximum on a small diagnostic batch, without RNG drift."""
+    model = getattr(model, '_orig_mod', model)
+    old_mode = model._onnx_mode
+    devices = [images.device.index] if images.is_cuda else []
+    try:
+        model._onnx_mode = True  # Eager SDPA avoids compiling a diagnostic graph.
+        for block in model.blocks:
+            block.capture_attention_stats = True
+        with torch.random.fork_rng(devices=devices):
+            model(images, padding_mask=padding_mask)
+        return torch.stack([block.max_attention_logit for block in model.blocks]).amax().item()
+    finally:
+        model._onnx_mode = old_mode
+        for block in model.blocks:
+            block.capture_attention_stats = False
+            block.max_attention_logit = None
+
+
+class SimplifiedTagger(BaseTagger):
+    """Vision Transformer based tagger for anime images."""
+    def __init__(self, config: VisionTransformerConfig):
+        super().__init__()
+        self._config = config
+        self._onnx_mode = False  # When True, use SDPA instead of flex_attention
+        # Patch embedding layer
+        self.patch_embed = nn.Conv2d(
+            config.num_channels,
+            config.hidden_size,
+            kernel_size=config.patch_size,
+            stride=config.patch_size
+        )
+        # Position embeddings
+        num_patches = (config.image_size // config.patch_size) ** 2
+        self.pos_embed = nn.Parameter(
+            torch.zeros(1, num_patches + 1, config.hidden_size)
+        )
+        self.cls_token = nn.Parameter(
+            torch.zeros(1, 1, config.hidden_size)
+        )
+        # Initialize pos_embed and cls_token with truncated normal (std=0.02)
+        # These are nn.Parameter objects, not nn.Module, so _init_weights() doesn't handle them.
+        # Standard ViT initialization follows timm/HuggingFace convention.
+        nn.init.trunc_normal_(self.pos_embed, std=0.02)
+        nn.init.trunc_normal_(self.cls_token, std=0.02)
+        self.pos_drop = nn.Dropout(p=config.pos_dropout)
+        # Transformer blocks
+        rate = getattr(config, 'drop_path_rate', 0.0)
+        dpr = torch.linspace(0.0, rate, config.num_hidden_layers)
+        dpr = dpr.clamp_max(1.0 - 1e-6).tolist()
+        self.blocks = nn.ModuleList([
+            TransformerBlock(config, drop_path=float(p))
+            for p in dpr
+        ])
+        # Final layer norm
+        self.norm = LayerNormFp32(config.hidden_size, eps=config.layer_norm_eps, use_fp32=config.use_fp32_layernorm)
+        # Classification head (ratings are now part of the tag vocabulary)
+        self.tag_head = nn.Linear(config.hidden_size, config.num_tags)
+        # Weight initialization
+        self.apply(self._init_weights)
+        # Override patch_embed init: standard ViT uses trunc_normal(std=0.02).
+        # The Conv2d branch in _init_weights sets std = sqrt(2/(k*k*out_channels))
+        # ≈ 0.0028 for our k=16, out_channels=1024 — ~7x smaller than ViT convention,
+        # which produces vanishing patch-embed activations early in from-scratch
+        # training. Re-init here so the apply() pass above doesn't overwrite us.
+        nn.init.trunc_normal_(self.patch_embed.weight, std=0.02)
+        if self.patch_embed.bias is not None:
+            nn.init.constant_(self.patch_embed.bias, 0)
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            torch.nn.init.trunc_normal_(module.weight, std=0.02)
+            if module.bias is not None:
+                nn.init.constant_(module.bias, 0)
+        elif isinstance(module, (nn.LayerNorm, LayerNormFp32)):
+            nn.init.constant_(module.bias, 0)
+            nn.init.constant_(module.weight, 1.0)
+        elif isinstance(module, nn.Conv2d):
+            fan_out = module.kernel_size[0] * module.kernel_size[1] * module.out_channels
+            fan_out //= module.groups
+            module.weight.data.normal_(0, math.sqrt(2.0 / fan_out))
+            if module.bias is not None:
+                module.bias.data.zero_()
+
+    def set_onnx_mode(self, enabled: bool = True) -> None:
+        """Enable or disable ONNX-compatible mode.
+
+        When enabled, the model uses F.scaled_dot_product_attention instead of
+        flex_attention. This is required for ONNX export since flex_attention
+        is a PyTorch 2.5+ CUDA kernel that cannot be traced by ONNX.
+
+        Args:
+            enabled: If True, use SDPA (ONNX-compatible). If False, use flex_attention.
+
+        Note:
+            This only affects the attention computation path. Model weights are unchanged.
+            The SDPA path produces numerically equivalent results to flex_attention.
+        """
+        self._onnx_mode = enabled
+        logger.info(f"ONNX mode {'enabled' if enabled else 'disabled'} - using {'SDPA' if enabled else 'flex_attention'}")
+
+    @property
+    def config(self):
+        """Return model configuration."""
+        return self._config
+
+    def _get_pos_embed(self, grid_h: int, grid_w: int) -> torch.Tensor:
+        """Return position embeddings matched to the runtime token grid.
+
+        Fast path (the common case, and the ONNX-export resolution): the runtime
+        grid equals the stored square grid, so the stored parameter is returned
+        unchanged. Otherwise the 2D patch position embeddings are bicubically
+        interpolated to (grid_h, grid_w) — same approach as the checkpoint-resume
+        interpolation in training_utils.py — keeping the CLS embedding as-is.
+        Without this, inputs at a different resolution would silently add the
+        raster-order *prefix* of the stored grid.
+        """
+        stored_len = self.pos_embed.shape[1]  # 1 (CLS) + stored_grid**2
+        stored_grid = int(math.isqrt(stored_len - 1))
+        if grid_h == stored_grid and grid_w == stored_grid:
+            return self.pos_embed
+        cls_pos = self.pos_embed[:, :1, :]
+        patch_pos = self.pos_embed[:, 1:, :]
+        hidden_size = patch_pos.shape[-1]
+        # Reshape to 2D spatial grid, interpolate, flatten back
+        patch_pos = patch_pos.reshape(1, stored_grid, stored_grid, hidden_size).permute(0, 3, 1, 2).float()
+        patch_pos = F.interpolate(
+            patch_pos, size=(grid_h, grid_w), mode='bicubic', align_corners=False
+        )
+        patch_pos = patch_pos.permute(0, 2, 3, 1).reshape(1, grid_h * grid_w, hidden_size)
+        patch_pos = patch_pos.to(self.pos_embed.dtype)
+        return torch.cat([cls_pos, patch_pos], dim=1)
+
+    def _check_numerical_stability(
+        self,
+        tag_logits: torch.Tensor,
+    ) -> None:
+        """Check for NaN/Inf in logits and log statistics.
+
+        This method is only called when config.check_numerical_stability=True.
+        It helps diagnose numerical instability issues during training/inference.
+        """
+        tag_has_nan = torch.isnan(tag_logits).any().item()
+        tag_has_inf = torch.isinf(tag_logits).any().item()
+
+        if tag_has_nan or tag_has_inf:
+            warnings.warn(
+                f"Numerical instability in tag_logits: "
+                f"NaN={tag_has_nan}, Inf={tag_has_inf}. "
+                f"Stats: min={tag_logits.min():.2f}, max={tag_logits.max():.2f}, "
+                f"mean={tag_logits.mean():.2f}, std={tag_logits.std():.2f}"
+            )
+
+        # Diagnostic only: report how many logits exceed the *configured* clamp
+        # value (the real clamp happens later in forward() at logit_clamp_value,
+        # eval mode only — training logits are left unclamped).
+        clamp_threshold = self._config.logit_clamp_value
+        if clamp_threshold is not None:
+            tag_over_clamp = (tag_logits.abs() > clamp_threshold).any().item()
+            if tag_over_clamp:
+                num_over = (tag_logits.abs() > clamp_threshold).sum().item()
+                warnings.warn(
+                    f"{num_over} tag logits exceed logit_clamp_value={clamp_threshold} "
+                    f"(clamped downstream in eval mode only; max abs value: {tag_logits.abs().max():.2f})"
+                )
+
+    @torch.compiler.disable(recursive=False)
+    def _create_block_mask(self, key_padding_mask: torch.Tensor, seq_len: int) -> BlockMask:
+        """Create BlockMask for padding-aware attention.
+
+        This method creates the mask once per forward pass, which is then shared
+        across all transformer layers for efficiency (instead of creating per-layer).
+        Keep its host dispatch outside the model graph: PyTorch 2.14 can raise
+        Inductor CantSplit when block-mask reductions are fused with dynamic image
+        shapes. recursive=False lets the inner mask builder remain compiled.
+
+        Args:
+            key_padding_mask: (B, L) bool, True=IGNORE (padding tokens)
+            seq_len: Sequence length
+
+        Returns:
+            BlockMask for use with flex_attention
+        """
+        B = key_padding_mask.shape[0]
+        attend_mask = ~key_padding_mask  # (B, L) True=ATTEND
+
+        def mask_mod(b, h, q_idx, kv_idx):
+            # Both query and key must be non-padding for attention
+            # This prevents padding tokens from attending (wasteful) and being attended to
+            # Use bitwise & instead of logical 'and' for Triton 3.4+ compatibility
+            return attend_mask[b, q_idx] & attend_mask[b, kv_idx]
+
+        # Use first block's config for flex_block_size
+        flex_block_size = self._config.flex_block_size
+
+        return _CREATE_BLOCK_MASK(
+            mask_mod,
+            B=B,
+            H=None,  # Broadcast across heads - mask is head-independent (saves memory)
+            Q_LEN=seq_len,
+            KV_LEN=seq_len,
+            device=key_padding_mask.device,
+            BLOCK_SIZE=min(flex_block_size, seq_len),
+        )
+
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        padding_mask: Optional[torch.Tensor] = None,  # (B,H,W) or (B,1,H,W), auto-detected semantics
+    ) -> Dict[str, torch.Tensor]:
+        B = pixel_values.shape[0]
+        # Patch embedding (optionally force fp32 for numerical stability under AMP)
+        if self._config.use_fp32_layernorm and pixel_values.dtype in (torch.float16, torch.bfloat16):
+            # Detect device type dynamically to support cuda/cpu/mps
+            device_type = pixel_values.device.type
+            # Autocast is only supported on certain device types
+            supported_devices = {'cuda', 'cpu', 'mps', 'xpu'}
+
+            if device_type in supported_devices:
+                # Use autocast to disable AMP for this operation
+                with torch.autocast(device_type=device_type, enabled=False):
+                    x = self.patch_embed(pixel_values.float())
+            else:
+                # Fallback: Just convert to float32 without autocast context
+                # This works on all devices but doesn't interact with AMP
+                warnings.warn(
+                    f"Device type '{device_type}' doesn't support autocast. "
+                    f"Using fallback path for patch embedding."
+                )
+                x = self.patch_embed(pixel_values.float())
+
+            x = x.to(pixel_values.dtype)
+        else:
+            x = self.patch_embed(pixel_values)
+        # Capture the runtime token grid before flattening (for pos-embed matching)
+        grid_h, grid_w = x.shape[-2], x.shape[-1]
+        x = x.flatten(2).transpose(1, 2).contiguous()
+        # Add CLS token
+        cls_tokens = self.cls_token.expand(B, -1, -1)
+        x = torch.cat([cls_tokens, x], dim=1)
+        # Add position embeddings (bicubically interpolated to the runtime grid
+        # if it differs from the stored grid; no-op at the trained resolution)
+        x = x + self._get_pos_embed(grid_h, grid_w)
+        x = self.pos_drop(x)
+        # =======================================================================
+        # MASK SEMANTICS DOCUMENTATION
+        # =======================================================================
+        # This section handles padding masks with multiple transformations.
+        # Understanding the semantics at each step is crucial for correctness.
+        #
+        # STEP 1: Input padding_mask (pixel-level)
+        #   - Shape: (B, H, W) or (B, 1, H, W)
+        #   - Semantics: True = PADDING pixel (should be ignored)
+        #
+        # STEP 2: attn_kpm (token-level key padding mask)
+        #   - Shape: (B, 1+Lp) where Lp = number of patch tokens
+        #   - Semantics: True = IGNORE this token (it's mostly padding)
+        #   - CLS token is always False (never ignored)
+        #
+        # STEP 3a: For ONNX/SDPA mode -> sdpa_attn_mask
+        #   - We invert: sdpa_attn_mask = ~attn_kpm
+        #   - Semantics: True = ATTEND to this token
+        #   - forward_sdpa() passes this directly to F.scaled_dot_product_attention as
+        #     a boolean key-padding mask (broadcast to (B, 1, 1, L)). SDPA's bool-mask
+        #     contract matches: True=attend, False=ignore — so no additive conversion
+        #     is needed, and PyTorch can keep using the Flash/efficient backends.
+        #
+        # STEP 3b: For flex_attention mode -> block_mask
+        #   - _create_block_mask inverts internally: attend_mask = ~key_padding_mask
+        #   - BlockMask uses: True = ATTEND (allow attention between these positions)
+        #   - flex_attention natively handles this without further inversion
+        # =======================================================================
+
+        # Build key-padding mask (CLS + patch tokens) from pixel-level mask.
+        # Output: attn_kpm with semantics True=IGNORE (padding tokens to exclude)
+        attn_kpm: Optional[torch.Tensor] = None
+        if padding_mask is not None:
+            # Convert pixel mask to standard format: (B,1,H,W) bool, True=PAD
+            pm = ensure_pixel_padding_mask(padding_mask, mask_semantics='pad')
+            thr = getattr(self._config, "token_ignore_threshold", 0.9)
+            # Pool pixel mask to token-level: token is IGNORE if >= threshold fraction is padding
+            token_ignore = pixel_to_token_ignore(pm, patch=self._config.patch_size, threshold=thr)  # (B,Lp)
+            # Prepend CLS token (always attend, never ignored)
+            cls_keep = torch.zeros(B, 1, dtype=torch.bool, device=token_ignore.device)  # False = don't ignore
+            attn_kpm = torch.cat([cls_keep, token_ignore], dim=1)  # (B, 1+Lp), True=IGNORE
+
+        if attn_kpm is not None:
+            # Only validate in debug mode to avoid expensive GPU-CPU synchronization
+            # (skip inside tracing to avoid control flow issues)
+            if self._config.check_numerical_stability and not torch.jit.is_tracing() and not self._onnx_mode:
+                if attn_kpm.all(dim=1).any().item():
+                    raise RuntimeError("attn_kpm masks all keys for at least one sample.")
+
+        # Create attention mask based on mode (ONNX vs flex_attention)
+        # Each path converts attn_kpm (True=IGNORE) to the format expected by the attention mechanism
+        block_mask: Optional[BlockMask] = None
+        sdpa_attn_mask: Optional[torch.Tensor] = None
+
+        if self._onnx_mode:
+            # ONNX mode: use SDPA with simple boolean mask
+            # For ONNX export, images are padded to square by InferenceWrapper,
+            # so typically no masking is needed. But support it for completeness.
+            if attn_kpm is not None:
+                # Invert semantics: attn_kpm True=IGNORE -> sdpa_attn_mask True=ATTEND
+                # forward_sdpa() passes this boolean mask straight to SDPA (True=ATTEND,
+                # False=IGNORE); no additive -inf conversion, so the Flash/efficient backend stays in use.
+                # Always apply (no .any() guard) so the mask branch is captured during export.
+                # An all-False attn_kpm produces an all-True sdpa_attn_mask (all-attend),
+                # which is mathematically identical to no mask.
+                sdpa_attn_mask = ~attn_kpm
+        else:
+            # Normal mode: use flex_attention with BlockMask
+            # _create_block_mask handles the inversion internally:
+            #   attend_mask = ~key_padding_mask (True=ATTEND)
+            #   BlockMask allows attention where attend_mask[q] & attend_mask[kv] is True
+            # No .any() guard: it forces a GPU->CPU sync every forward and with
+            # letterboxed data the mask is almost always non-empty anyway. An
+            # all-False attn_kpm yields an all-attend BlockMask (same result).
+            if attn_kpm is not None:
+                block_mask = self._create_block_mask(attn_kpm, x.size(1))
+
+        # Transformer blocks with optional selective gradient checkpointing
+        # checkpoint_every_n_layers controls granularity: 1=all, 2=every 2nd, 4=every 4th
+        checkpoint_interval = getattr(self._config, 'checkpoint_every_n_layers', 1)
+        for idx, block in enumerate(self.blocks):
+            should_checkpoint = (
+                self._config.gradient_checkpointing
+                and self.training
+                and (idx % checkpoint_interval == 0)
+            )
+
+            if self._onnx_mode:
+                # ONNX mode: use SDPA path (no checkpointing during export)
+                x = block.forward_sdpa(x, attn_mask=sdpa_attn_mask)
+            elif should_checkpoint:
+                # use_reentrant=False is recommended for PyTorch 2.x and works better with
+                # complex inputs like attention masks. It's also more memory efficient.
+                # See: https://pytorch.org/docs/stable/checkpoint.html
+
+                def create_block_forward(b, mask):
+                    """Create a closure that captures block and mask for checkpointing."""
+                    def block_forward(hidden_states):
+                        return b(hidden_states, block_mask=mask)
+                    return block_forward
+
+                x = torch.utils.checkpoint.checkpoint(
+                    create_block_forward(block, block_mask),
+                    x,
+                    use_reentrant=False
+                )
+            else:
+                x = block(x, block_mask=block_mask)
+        # Final norm
+        x = self.norm(x)
+        # Use CLS token for classification
+        cls_output = x[:, 0]
+        # Predictions (ratings are included as tags in the vocabulary)
+        tag_logits = self.tag_head(cls_output)
+
+        # Monitor for numerical issues (optional, controlled by config)
+        if self._config.check_numerical_stability:
+            self._check_numerical_stability(tag_logits)
+
+        # Clamp logits for inference/export outputs only. During training the
+        # clamp is skipped: torch.clamp has zero gradient outside the bounds,
+        # which would kill the learning signal exactly for confidently-wrong
+        # predictions. The loss (loss_functions.py) computes in fp32 with its
+        # own log/exp clamps, so unclamped training logits are numerically safe.
+        if self._config.logit_clamp_value is not None and not self.training:
+            clamp_val = self._config.logit_clamp_value
+            tag_logits = torch.clamp(tag_logits, min=-clamp_val, max=clamp_val)
+
+        return {
+            'tag_logits': tag_logits,
+            'logits': tag_logits
+        }
+
+    def cleanup(self):
+        """Explicitly release GPU memory and clear cached tensors."""
+        import gc
+
+        # Move all parameters to CPU to release GPU memory
+        self.cpu()
+
+        # Clear any cached gradients
+        for param in self.parameters():
+            if param.grad is not None:
+                param.grad = None
+
+        # Force garbage collection
+        gc.collect()
+
+        # Clear CUDA cache if available
+        if torch.cuda.is_available():
+            # Synchronize before clearing cache to ensure all GPU operations complete
+            torch.cuda.synchronize()
+            torch.cuda.empty_cache()
+
+    def __del__(self):
+        """Cleanup on deletion."""
+        try:
+            self.cleanup()
+        except Exception:
+            pass
+
+    def __enter__(self):
+        """Context manager support."""
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Ensure cleanup on exit."""
+        self.cleanup()
+        return False
+
+
+def create_model(config: Optional[VisionTransformerConfig] = None, **kwargs) -> BaseTagger:
+    """Create the ViT tagger model (SimplifiedTagger) from configuration.
+
+    Args:
+        config: VisionTransformerConfig instance. If provided, config fields in kwargs are ignored.
+        **kwargs: Alternative to config - specify config fields as keyword arguments.
+                 See VisionTransformerConfig for details on each parameter.
+
+    Returns:
+        SimplifiedTagger instance with the specified configuration.
+
+    Examples:
+        # Using config object
+        config = VisionTransformerConfig(image_size=512, num_tags=10000)
+        model = create_model(config=config)
+
+        # Using kwargs
+        model = create_model(image_size=512, num_tags=10000)
+    """
+    # Pop the legacy architecture selector so it doesn't reach the config-field filter
+    # below. Only the ViT path (SimplifiedTagger) is supported now; the value is ignored.
+    kwargs.pop('architecture_type', None)
+
+    if config is None:
+        # Get the names of the fields in the VisionTransformerConfig dataclass
+        config_fields = {f.name for f in fields(VisionTransformerConfig)}
+
+        # Filter kwargs to only include keys that are in the config_fields
+        filtered_kwargs = {k: v for k, v in kwargs.items() if k in config_fields}
+
+        # Warn if kwargs contains invalid keys (excluding already-extracted keys)
+        invalid_keys = set(kwargs.keys()) - config_fields
+        if invalid_keys:
+            warnings.warn(
+                f"Ignoring unknown configuration parameters: {sorted(invalid_keys)}. "
+                f"Valid parameters: {sorted(config_fields)}"
+            )
+
+        config = VisionTransformerConfig(**filtered_kwargs)
+
+    return SimplifiedTagger(config)

preprocessing.json

@@ -0,0 +1,32 @@
+{
+  "normalize_mean": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "normalize_std": [
+    0.5,
+    0.5,
+    0.5
+  ],
+  "image_size": 448,
+  "patch_size": 16,
+  "color_order": "RGB",
+  "pad_color": [
+    114,
+    114,
+    114
+  ],
+  "resize_filter": "LANCZOS",
+  "resize_mode": "letterbox_downscale_only",
+  "upscale": false,
+  "exif_transpose": true,
+  "alpha_composite": true,
+  "input_layout": "NCHW",
+  "padding_mask_true_means": "padding",
+  "output_activation": "sigmoid",
+  "skip_indices": [
+    0,
+    1
+  ]
+}

release_manifest.json

@@ -0,0 +1,20 @@
+{
+  "release": "OppaiOracle experimental-X1.1",
+  "repo_id": "Grio43/OppaiOracle-experimental-X1.1",
+  "created_at": "2026-10-12T09:59:58.688810+00:00",
+  "image_size": 448,
+  "epoch": 9,
+  "step": 55719,
+  "source_checkpoint_sha256": "d400c2e460e7287357f86083f84c187790a45b1d47985e908c52825373f888c3",
+  "validated_onnx_source_sha256": "c5fcab0fb53e07aee437f3ef29e19a1e60b3dc18380048fa6d99baf35fa6c8f3",
+  "onnx_graph_sha256": "bbb2fd38ed9fdfe5977b0fd79ca2c7820a1eab5839b416aab5579109e6c0fa74",
+  "onnx_change": "descriptive metadata only; graph identical to fully evaluated artifact",
+  "safetensors": "all 332 FP32 tensors exact round-trip; no optimizer or pickle",
+  "model_sha256": "49ba0f66f83ba88a101b6edac54ec39f0a12eab66743e13ae42efc3570745bf7",
+  "safetensors_sha256": "ec07d0034f48ba3ceff8624993baf3ea30fa3a2e3cc8908682eb3e4f84eda766",
+  "vocab_sha256": "e3e6ccc137d9152d8a17fbd19e0c1ae595b2317a6c55893fd25779620b4b8087",
+  "storage": "Hugging Face Hub / Xet",
+  "license": "educational-use-only",
+  "download_count_query_file": "config.json",
+  "validation": "validation.json"
+}

requirements-gpu.txt

@@ -0,0 +1,4 @@
+numpy>=1.26
+Pillow>=10.0
+onnxruntime-gpu>=1.20
+huggingface_hub>=1.0

requirements-torch.txt

@@ -0,0 +1,5 @@
+numpy>=1.26
+Pillow>=10.0
+torch>=2.5
+safetensors>=0.4.5
+huggingface_hub>=1.0

requirements.txt

@@ -0,0 +1,4 @@
+numpy>=1.26
+Pillow>=10.0
+onnxruntime>=1.20
+huggingface_hub>=1.0

thresholds.json

@@ -0,0 +1,6 @@
+{
+  "default": 0.7085427045822144,
+  "source": "200-point micro-F1 grid on 13,569 shared development tags; not final deployment calibration",
+  "score_calibration": false,
+  "training_diagnostic_threshold": 0.7927
+}

training_buckets.csv

@@ -0,0 +1,6 @@
+bucket,epoch7,epoch8,epoch9,epoch7_to_9_points
+300-499,37.72737681865692,37.53387928009033,37.342289090156555,-0.3850877285003662
+500-999,61.54675483703613,61.733537912368774,61.75047755241394,0.20372271537780762
+1000-4999,57.105010747909546,57.1607232093811,57.215815782547,0.11080503463745117
+5000-9999,52.05754041671753,52.18653082847595,52.18939185142517,0.1318514347076416
+10000+,57.25409388542175,57.3200523853302,57.348477840423584,0.09438395500183105

training_epochs.csv

@@ -0,0 +1,10 @@
+epoch,validation_time,exact_map,binned_map,macro_f1,micro_f1,val_loss,train_loss,samples,skipped
+1,2026-09-13T18:19:02.022533-04:00,55.92384934425354,55.73147535324097,43.22007596492767,64.86003398895264,0.001663603587076068,0.0016142689855769277,30000,0
+2,2026-09-14T04:16:38.378663-04:00,56.59424662590027,56.39519691467285,44.07239556312561,65.32300114631653,0.0016429137904196978,0.0015640060883015394,30000,0
+3,2026-09-14T13:52:19.915251-04:00,56.98235034942627,56.77585005760193,44.495999813079834,65.55598378181458,0.001634233514778316,0.0015406454913318157,30000,0
+4,2026-09-14T23:42:07.983341-04:00,57.2227418422699,57.01723098754883,44.791001081466675,65.70025682449341,0.0016256555682048202,0.0015247202245518565,30000,0
+5,2026-09-15T09:26:10.020244-04:00,57.34878182411194,57.14470148086548,44.962278008461,65.80770015716553,0.0016207325970754027,0.0015117722796276212,30000,0
+6,2026-09-15T19:23:48.602481-04:00,57.48898983001709,57.290494441986084,45.18154859542847,65.89643359184265,0.0016172741306945682,0.0015005783643573523,30000,0
+7,2026-09-16T05:10:51.323995-04:00,57.589566707611084,57.38818645477295,45.37670016288757,66.06296300888062,0.0016201940597966313,0.00149048981256783,30000,0
+8,2026-09-16T14:55:16.420660-04:00,57.68805742263794,57.48524069786072,45.54440975189209,66.12292528152466,0.0016188551671802998,0.0014811164001002908,30000,0
+9,2026-09-17T01:43:38.368238-04:00,57.71656036376953,57.506269216537476,45.742398500442505,66.21212959289551,0.0016158618964254856,0.0014653843827545643,30000,0

training_metadata.json

@@ -0,0 +1,273 @@
+{
+  "source_checkpoint_sha256": "d400c2e460e7287357f86083f84c187790a45b1d47985e908c52825373f888c3",
+  "parent_release": "Grio43/OppaiOracle-experimental-X1",
+  "parent_epoch": 21,
+  "parent_step": 127974,
+  "phase": 2,
+  "epoch": 9,
+  "step": 55719,
+  "completed_at": "2026-09-17T01:43:41.793971",
+  "model": {
+    "architecture_type": "vit",
+    "hidden_size": 896,
+    "num_hidden_layers": 18,
+    "num_attention_heads": 14,
+    "intermediate_size": 2320,
+    "image_size": 448,
+    "patch_size": 16,
+    "num_channels": 3,
+    "hidden_dropout_prob": 0.0,
+    "pos_dropout": 0.0,
+    "attention_dropout": 0.0,
+    "qk_norm": true,
+    "layer_scale_init": 0.1,
+    "drop_path_rate": 0.15,
+    "initializer_range": 0.02,
+    "layer_norm_eps": 1e-06,
+    "use_fp32_layernorm": false,
+    "use_flex_attention": true,
+    "attention_bias": true,
+    "flex_block_size": 128,
+    "token_ignore_threshold": 0.9,
+    "num_labels": 14996,
+    "num_groups": 20,
+    "tags_per_group": 10000,
+    "gradient_checkpointing": false,
+    "checkpoint_every_n_layers": 1
+  },
+  "training": {
+    "memory_format": "contiguous",
+    "num_epochs": 9,
+    "learning_rate": 1.5e-05,
+    "lr_scaling_mode": "sqrt",
+    "lr_base_batch_size": 256,
+    "weight_decay": 0.05,
+    "gradient_accumulation_steps": 21,
+    "world_size": 1,
+    "optimizer": "adamw8bit",
+    "adam_beta1": 0.9,
+    "adam_beta2": 0.999,
+    "adan_beta3": 0.99,
+    "adam_epsilon": 1e-07,
+    "scheduler": "wsd",
+    "warmup_epochs": 1,
+    "warmup_steps": null,
+    "wsd_cooldown_fraction": 0.15,
+    "wsd_plateau_enabled": false,
+    "fp32_head_optimizer": true,
+    "num_cycles": 1,
+    "cycle_decay": 0.9,
+    "lr_end": 0.0,
+    "use_amp": true,
+    "amp_dtype": "bfloat16",
+    "enable_anomaly_detection": false,
+    "max_grad_norm": 1.0,
+    "gradient_clipping": {
+      "enabled": true,
+      "max_norm": 1.0
+    },
+    "save_steps": 1500,
+    "save_total_limit": 3,
+    "save_best_only": false,
+    "resume_from": "latest",
+    "eval_steps": 0,
+    "logging_steps": 50,
+    "resume_reseed": true,
+    "allow_unverified_vocab_resume": false,
+    "on_lr_schedule_change": "error",
+    "tag_loss": {
+      "alpha": 1.0,
+      "gamma_neg": 3.0,
+      "gamma_neg_override": null,
+      "gamma_pos": 0.0,
+      "detach_focal_weight": true,
+      "label_smoothing": 0.0,
+      "clip": 0.05,
+      "class_weights": null,
+      "class_weight_strategy": null,
+      "class_weight_clip_min": 0.05,
+      "class_weight_clip_max": 5.0,
+      "class_weight_beta": 0.9999
+    },
+    "phase": 2,
+    "asl_schedule": {
+      "enabled": false,
+      "min_dwell_epochs": 3,
+      "phase1": {
+        "gamma_neg_min": 3.0,
+        "gamma_neg_max": 3.0,
+        "hold_epochs": 8
+      },
+      "phase2": {
+        "gamma_neg_min": 3.0,
+        "gamma_neg_max": 3.0,
+        "hold_epochs": 2
+      },
+      "phase3": {
+        "gamma_neg_min": 3.0,
+        "gamma_neg_max": 3.0,
+        "hold_epochs": 0
+      }
+    },
+    "asl_telemetry": {
+      "enabled": true,
+      "interval_updates": 100,
+      "log_every_updates": 500,
+      "ema_beta": 0.98,
+      "topk_hard": 10,
+      "num_deciles": 10,
+      "hist_min": 0.05,
+      "hist_max": 0.95,
+      "hist_bins": 18,
+      "watch_band_low": 0.2,
+      "watch_band_high": 0.5,
+      "epr_alarm_rel_drop": 0.05,
+      "epr_alarm_window_epochs": 2,
+      "exclude_rating_tags": true,
+      "sibling_groups_path": "./configs/confusable_groups.json",
+      "shadow_controller_enabled": true,
+      "shadow_lambda": 0.05,
+      "shadow_delta_p_target": 0.2
+    },
+    "device": "cuda",
+    "use_compile": true,
+    "compile_mode": "max-autotune-no-cudagraphs",
+    "compile_fullgraph": false,
+    "compile_dynamic": true,
+    "use_tensorboard": true,
+    "seed": 42,
+    "deterministic": false,
+    "benchmark": true,
+    "selection_metric": "val_mAP",
+    "early_stopping_patience": 4,
+    "early_stopping_threshold": 5e-07,
+    "early_stopping_burn_in_epochs": 2,
+    "early_stopping_burn_in_strategy": "median",
+    "early_stopping_mode": "advise",
+    "early_stopping_min_delta": 0.003,
+    "early_stopping_plateau_slope_threshold": 0.0001,
+    "early_stopping_plateau_noise_multiplier": 2.0,
+    "early_stopping_window": 5,
+    "early_stopping_confirm": 2,
+    "early_stopping_min_epochs": 10,
+    "early_stopping_regression_delta": 0.006,
+    "early_stopping_regression_confirm": 2,
+    "early_stopping_overfit_confirm": 3,
+    "early_stopping_overfit_window": 8,
+    "early_stopping_overfit_min_relative_change": 0.01,
+    "early_stopping_overfit_noise_multiplier": 2.0,
+    "early_stopping_overfit_ranking_tolerance": 0.0001,
+    "wsd_deterioration_actions": {
+      "peak_decline": "preserve",
+      "overfit_loss": "warn"
+    },
+    "early_stopping_phases": {
+      "phase1": {
+        "mode": null,
+        "window": 5,
+        "confirm": 2,
+        "min_epochs": 10,
+        "plateau_slope_threshold": null,
+        "plateau_noise_multiplier": null
+      },
+      "phase2": {
+        "mode": "advise",
+        "window": 5,
+        "confirm": 2,
+        "min_epochs": 10,
+        "plateau_slope_threshold": null,
+        "plateau_noise_multiplier": null
+      },
+      "phase3": {
+        "mode": "off",
+        "window": null,
+        "confirm": null,
+        "min_epochs": null,
+        "plateau_slope_threshold": null,
+        "plateau_noise_multiplier": null
+      }
+    },
+    "use_distillation": false,
+    "distillation_alpha": 0.7,
+    "distillation_temperature": 3.0
+  },
+  "data_recipe": {
+    "vocab_min_frequency": 550,
+    "vocab_min_frequency_by_category": {
+      "artist": 550,
+      "char": 550,
+      "copyright": 550,
+      "gen": 450,
+      "meta": 550
+    },
+    "random_flip_prob": 0.5,
+    "horizontal_flip_exclude_tags": [
+      "gen:left-handed",
+      "gen:left-to-right_manga",
+      "gen:right-over-left_kimono",
+      "gen:right-to-left_comic"
+    ],
+    "color_jitter_enabled": true,
+    "color_jitter_brightness": 0.22,
+    "color_jitter_brightness_p": 0.5,
+    "color_jitter_contrast": 0.15,
+    "color_jitter_contrast_p": 0.5,
+    "color_jitter_saturation": 0.06,
+    "color_jitter_saturation_p": 0.5,
+    "random_erasing_enabled": false,
+    "random_erasing_p": 0.35,
+    "random_erasing_scale_min": 0.007,
+    "random_erasing_scale_max": 0.03,
+    "random_erasing_ratio_min": 0.3,
+    "random_erasing_ratio_max": 2.0,
+    "random_rotation_enabled": true,
+    "random_rotation_p": 0.3,
+    "random_rotation_min_degrees": 2.0,
+    "random_rotation_max_degrees": 5.0,
+    "gaussian_blur_enabled": true,
+    "gaussian_blur_p": 0.15,
+    "gaussian_blur_kernel_size": 3,
+    "gaussian_blur_sigma_min": 0.1,
+    "gaussian_blur_sigma_max": 1.0
+  },
+  "train_images": 6239823,
+  "eval_images": 30000,
+  "dataset_maintenance": {
+    "danbooru_metadata_sync": "2026-08",
+    "deprecated_tags_removed": true,
+    "renamed_tags_updated_to_current_names": true,
+    "provenance": "model author"
+  },
+  "metrics": {
+    "train_loss": 0.0014653843827545643,
+    "val_loss": 0.0016158618964254856,
+    "val_f1_macro": 0.45742398500442505,
+    "val_f1_micro": 0.6621212959289551,
+    "val_mAP": 0.5750626921653748,
+    "selection_metric": "val_mAP"
+  },
+  "scheduler": {
+    "warmup_steps": 6191,
+    "total_steps": 55719,
+    "cooldown_fraction": 0.15,
+    "budget_cooldown_start": 48451,
+    "min_lr": 0.0,
+    "base_max_lr": 2.9764702249476646e-05,
+    "cooldown_start": 48451,
+    "cooldown_steps": 7268,
+    "cooldown_reason": "budget",
+    "cooldown_best_metric": 0.5750626921653748,
+    "base_lrs": [
+      2.9764702249476646e-05,
+      2.9764702249476646e-05
+    ],
+    "last_epoch": 55719,
+    "_step_count": 55720,
+    "_is_initial": false,
+    "_get_lr_called_within_step": false,
+    "_last_lr": [
+      0.0,
+      0.0
+    ]
+  }
+}

validation.json

@@ -0,0 +1,67 @@
+{
+  "model": "X1 448 final epoch 9",
+  "checkpoint": {
+    "checkpoint": "checkpoint_epoch_9_step_55719.pt",
+    "checkpoint_sha256": "d400c2e460e7287357f86083f84c187790a45b1d47985e908c52825373f888c3",
+    "epoch": 9,
+    "step": 55719,
+    "scheduler": {
+      "warmup_steps": 6191,
+      "total_steps": 55719,
+      "cooldown_fraction": 0.15,
+      "budget_cooldown_start": 48451,
+      "min_lr": 0.0,
+      "base_max_lr": 2.9764702249476646e-05,
+      "cooldown_start": 48451,
+      "cooldown_steps": 7268,
+      "cooldown_reason": "budget",
+      "cooldown_best_metric": 0.5750626921653748,
+      "base_lrs": [
+        2.9764702249476646e-05,
+        2.9764702249476646e-05
+      ],
+      "last_epoch": 55719,
+      "_step_count": 55720,
+      "_is_initial": false,
+      "_get_lr_called_within_step": false,
+      "_last_lr": [
+        0.0,
+        0.0
+      ]
+    },
+    "preprocessing": {
+      "normalize_mean": [
+        0.5,
+        0.5,
+        0.5
+      ],
+      "normalize_std": [
+        0.5,
+        0.5,
+        0.5
+      ],
+      "image_size": 448,
+      "patch_size": 16,
+      "color_order": "RGB"
+    }
+  },
+  "model_sha256": "c5fcab0fb53e07aee437f3ef29e19a1e60b3dc18380048fa6d99baf35fa6c8f3",
+  "native_image_size": 448,
+  "samples": 30000,
+  "sample_sha256": "20a60e7a747cdf9a698a464c9d324612aad2ca6b43a541db0813554eb9af8694",
+  "shared_tags": 13569,
+  "supported_shared_tags": 13394,
+  "binned_mAP": 0.5851927682386426,
+  "exact_mAP": 0.5875134140920024,
+  "best_grid_micro_f1": 0.6864101262842809,
+  "best_grid_threshold": 0.7085427045822144,
+  "precision_at_best_grid": 0.7065020665350971,
+  "recall_at_best_grid": 0.6674293596992867,
+  "all_supported_tags": 14767,
+  "all_tag_exact_mAP": 0.5774638913821051,
+  "all_tag_binned_mAP": 0.5750950549967043,
+  "inference_seconds": 347.91008079855237,
+  "total_seconds": 487.7315649986267,
+  "preprocessing": "FP32 RGB, native resolution, downscale-only centered LANCZOS letterbox, mean/std 0.5, pad 114, True=padding",
+  "verified_sidecars": 30000
+}