Release Anima Style Embedding v1.0

.gitattributes

@@ -33,3 +33,7 @@ saved_model/**/* 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
+assets/generalization_and_gallery_size.png filter=lfs diff=lfs merge=lfs -text
+assets/learned_vs_frozen_baseline.png filter=lfs diff=lfs merge=lfs -text
+assets/reference_count_curve.png filter=lfs diff=lfs merge=lfs -text
+assets/training_validation_curves.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,172 @@
+---
+library_name: pytorch
+pipeline_tag: feature-extraction
+tags:
+  - anime
+  - style-embedding
+  - image-retrieval
+  - reference-image
+  - siglip2
+  - anima
+  - pytorch
+---
+
+# Anima Style Embedding v1.0
+
+**A multi-reference style encoder for person-centric anime and illustration images.**
+
+Anima Style Embedding converts one to eight reference images into a normalized 512-dimensional style vector. It combines multi-level SigLIP2 features from the full image and face with internal Anima features, then grounds the result across the complete reference set. The embedding is designed for open-set style retrieval and reference-driven style transfer systems rather than a fixed bank of known artists.
+
+| Model size | References | Embedding | Modalities |
+|---:|---:|---:|---|
+| 17.49M parameters | 1–8 images | 512-D, L2-normalized | Full image + face + Anima internals |
+
+![Learned encoder versus frozen backbone](assets/learned_vs_frozen_baseline.png)
+
+## Highlights
+
+- **Open-set style representation:** unseen artists can be embedded directly from reference images.
+- **Native multi-reference grounding:** Reference Consensus Re-reading lets each image be interpreted in the context of the complete 1–8 image set.
+- **Full-image and face awareness:** the encoder combines global composition and rendering cues with face-specific style information.
+- **Anima-aware features:** three internal Anima blocks complement SigLIP2 visual features with representations aligned to the target generator.
+- **Strong retrieval gains:** unseen synthetic-style MRR improves from 32.0% with the frozen SigLIP2 Full+Face baseline to 60.8% with the learned encoder on the 4,500-way validation gallery.
+
+## Performance
+
+Retrieval uses cosine similarity between normalized query embeddings and five-image style prototypes.
+
+### 4,500-way validation — unseen identities
+
+The unseen subset contains 500 identities retrieved against a gallery of 4,500 styles.
+
+| Source | Top-1 | Top-5 | MRR |
+|---|---:|---:|---:|
+| Synthetic | 51.8% | 71.7% | 60.8% |
+| Human | 56.4% | 75.3% | 65.0% |
+
+### 500-way sealed test — unseen identities
+
+| Source | Top-1 | Top-5 | MRR |
+|---|---:|---:|---:|
+| Synthetic | 73.8% | 90.3% | 81.1% |
+| Human | 76.1% | 92.0% | 83.0% |
+
+The sealed test uses a smaller gallery than validation, so its absolute scores should be read independently.
+
+![Generalization and gallery-size sensitivity](assets/generalization_and_gallery_size.png)
+
+### Reference count
+
+More references consistently improve human-style grounding. Synthetic K=5 uses the locked controlled/matched/cross-content protocol and is shown separately from the generic K diagnostics.
+
+![Reference count evaluation](assets/reference_count_curve.png)
+
+## Architecture
+
+The encoder projects full-image, face, and Anima tokens into a shared 512-wide space and reads them with four eight-head mixer blocks. Functional Style Factorization introduces four auxiliary style queries during training. Reference Consensus Re-reading computes a consensus across the reference set, re-reads every image under that context, and forms an agreement-weighted prototype.
+
+| Input | Shape per batch |
+|---|---|
+| Full-image SigLIP2 features | `[B, R, 30, 1152]` |
+| Face SigLIP2 features | `[B, R, 30, 1152]` |
+| Face validity mask | `[B, R]` |
+| Anima internal features | `[B, R, 3, 4096]` |
+| Reference mask | `[B, R]` |
+
+SigLIP2 features use [`timm/vit_so400m_patch16_siglip_gap_512.v2_webli`](https://huggingface.co/timm/vit_so400m_patch16_siglip_gap_512.v2_webli) at 512×512 and three backbone depths. Anima features use 768×768 inputs and blocks 8, 18, and 26 at sigma 0.1.
+
+## Quick start
+
+```python
+import json
+
+import torch
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
+
+from modeling_anima_style_embedding import UnifiedStyleEncoder
+
+repo_id = "Baragi-AI/Anima-Style-Embedding"
+config = json.load(open(hf_hub_download(repo_id, "config.json"), encoding="utf-8"))
+cfg = config["architecture"]
+
+encoder = UnifiedStyleEncoder(
+    width=cfg["width"],
+    blocks=cfg["blocks"],
+    heads=cfg["heads"],
+    embedding_dim=cfg["embedding_dim"],
+    dropout=cfg["dropout"],
+    use_anima=True,
+    reference_grounding=cfg["reference_grounding"],
+    functional_factors=cfg["functional_factors"],
+)
+encoder.load_state_dict(
+    load_file(hf_hub_download(repo_id, "model.safetensors")),
+    strict=True,
+)
+encoder.eval()
+```
+
+```python
+with torch.inference_mode():
+    style_embedding, per_image_embeddings = encoder(
+        full_features,
+        reference_mask,
+        face_features=face_features,
+        face_mask=face_mask,
+        anima_features=anima_features,
+    )
+
+assert style_embedding.shape == (full_features.shape[0], 512)
+```
+
+`modeling_anima_style_embedding.py` contains the matching encoder definition. The checkpoint expects precomputed SigLIP2 and Anima features; RGB and Anima feature extraction are supplied by downstream integrations.
+
+## Training
+
+The model was trained on 320,000 person-centric style records: 160,000 Anima-generated images and 160,000 human-created illustrations across 8,000 training identities. Synthetic and human examples are mixed within each batch. Training uses four RTX 5090 GPUs, BF16 DistributedDataParallel, 16-way episodes, two queries per prototype, variable 1–8 image support, global negatives, and weak hard-negative sampling. The 4,500-step checkpoint produced the best full-validation score.
+
+![Training dynamics](assets/training_validation_curves.png)
+
+The training record is available in [Weights & Biases](https://wandb.ai/1wndrla17-kyung-hee-university/anima-style-embedding/runs/5c9ae843a14f46529fe9aaa604a78d3b). Machine-readable evaluation results are included under `metrics/`.
+
+## Applications
+
+- reference-driven style transfer for Anima;
+- artist-style retrieval and clustering;
+- aggregation of multiple references into a single style prototype;
+- open-set style recognition for previously unseen artists;
+- similarity scoring and reference-set curation.
+
+## Files
+
+| File | Contents |
+|---|---|
+| `model.safetensors` | Encoder weights |
+| `config.json` | Architecture and feature contract |
+| `modeling_anima_style_embedding.py` | PyTorch model definition |
+| `metrics/` | Training curves and retrieval results |
+| `SHA256SUMS.json` | Release checksums |
+
+## Limitations
+
+- The model is optimized for person-centric anime and illustration images; other domains are not evaluated.
+- Retrieval metrics measure representation quality, not generated-image fidelity.
+- Style and artist-associated content preferences are not completely disentangled.
+- Anima features must follow the documented model, block, resolution, and sigma configuration.
+
+## License
+
+No model license has been selected for v1.0. The training images are not distributed with the checkpoint.
+
+## Citation
+
+```bibtex
+@misc{baragi2026animastyleembedding,
+  title        = {Anima Style Embedding v1.0},
+  author       = {{Baragi AI}},
+  year         = {2026},
+  howpublished = {Hugging Face model repository},
+  url          = {https://huggingface.co/Baragi-AI/Anima-Style-Embedding}
+}
+```

SHA256SUMS.json

@@ -0,0 +1,4 @@
+{
+  "model.safetensors": "6cdf126152927f6e31dc856a8d10b8decee3bb20ce4c0ca449206db65d4633ec",
+  "config.json": "7323cbf6061896453a88cd0f6302831b96278a52eb8b0b3c1500d84552a1e68f"
+}

config.json

@@ -0,0 +1,75 @@
+{
+  "model_type": "anima_style_embedding",
+  "version": "1.0",
+  "architecture": {
+    "width": 512,
+    "blocks": 4,
+    "heads": 8,
+    "embedding_dim": 512,
+    "dropout": 0.1,
+    "reference_grounding": true,
+    "functional_factors": true,
+    "reference_count": [
+      1,
+      8
+    ],
+    "trainable_parameters": 17488353
+  },
+  "inputs": {
+    "siglip2_full": {
+      "shape": [
+        30,
+        1152
+      ],
+      "dtype": "bfloat16"
+    },
+    "siglip2_face": {
+      "shape": [
+        30,
+        1152
+      ],
+      "dtype": "bfloat16",
+      "optional": true
+    },
+    "anima": {
+      "shape": [
+        3,
+        4096
+      ],
+      "dtype": "bfloat16"
+    }
+  },
+  "vision_backbone": {
+    "repo_id": "timm/vit_so400m_patch16_siglip_gap_512.v2_webli",
+    "image_size": 512,
+    "layer_fractions": [
+      0.25,
+      0.55,
+      1.0
+    ],
+    "pool_sizes": [
+      2,
+      4,
+      2
+    ]
+  },
+  "anima_features": {
+    "image_size": 768,
+    "blocks": [
+      8,
+      18,
+      26
+    ],
+    "sigma": 0.1,
+    "noise_seed": 20260715
+  },
+  "training": {
+    "checkpoint_step": 4500,
+    "face_dropout": 0.3,
+    "query_count": 2,
+    "optimizer": "AdamW(fused)",
+    "scheduler": "cosine_with_linear_warmup",
+    "learning_rate": 0.0001,
+    "packed_manifest_sha256": "c1267b9fd45c6e0d222b91d6ae2291b36b763277cc39cbc32484874ffeae9b32"
+  }
+}

metrics/model_card_metrics.json

@@ -0,0 +1,169 @@
+{
+  "selected_checkpoint": {
+    "checkpoint": "/workspace/runs/style_embedding_combined_lr1e4_clip10_5k/best.pt",
+    "checkpoint_sha256": "16345685b72c9bb018833ea9d42d18ebdaafa585b2a2eb136f267006c1de5663",
+    "step": 4500,
+    "validation_checkpoint_score": 0.6301507651805878
+  },
+  "validation_k5": {
+    "synthetic": {
+      "styles": 4500.0,
+      "queries": 22500.0,
+      "top1": 0.5625777840614319,
+      "top5": 0.7585333585739136,
+      "mrr": 0.6518997550010681,
+      "margin": -0.006614952813833952,
+      "positive_cosine": 0.7365614771842957,
+      "hardest_negative_cosine": 0.7431764006614685,
+      "verification_auroc": 0.5534452199935913,
+      "embedding_std": 0.04349411278963089,
+      "embedding_effective_rank": 24.163583755493164,
+      "seen_queries": 20000.0,
+      "seen_top1": 0.5681999921798706,
+      "seen_top5": 0.7637499570846558,
+      "seen_mrr": 0.6573799848556519,
+      "seen_margin": -0.004584707319736481,
+      "unseen_queries": 2500.0,
+      "unseen_top1": 0.5175999999046326,
+      "unseen_top5": 0.7167999744415283,
+      "unseen_mrr": 0.6080579161643982,
+      "unseen_margin": -0.02285691164433956,
+      "matched_queries": 4500.0,
+      "matched_top1": 0.48955556750297546,
+      "matched_top5": 0.6902222037315369,
+      "matched_mrr": 0.5840909481048584,
+      "matched_margin": -0.030851511284708977,
+      "cross_queries": 18000.0,
+      "cross_top1": 0.5808333158493042,
+      "cross_top5": 0.7756111025810242,
+      "cross_mrr": 0.6688519716262817,
+      "cross_margin": -0.0005558129050768912
+    },
+    "human": {
+      "styles": 4500.0,
+      "queries": 22500.0,
+      "top1": 0.6493777632713318,
+      "top5": 0.8219555616378784,
+      "mrr": 0.7280733585357666,
+      "margin": 0.022814957424998283,
+      "positive_cosine": 0.7811680436134338,
+      "hardest_negative_cosine": 0.7583531141281128,
+      "verification_auroc": 0.6369090676307678,
+      "embedding_std": 0.043804656714200974,
+      "embedding_effective_rank": 27.383806228637695,
+      "seen_queries": 20000.0,
+      "seen_top1": 0.660099983215332,
+      "seen_top5": 0.830549955368042,
+      "seen_mrr": 0.7378283143043518,
+      "seen_margin": 0.026178810745477676,
+      "unseen_queries": 2500.0,
+      "unseen_top1": 0.5636000037193298,
+      "unseen_top5": 0.7531999945640564,
+      "unseen_mrr": 0.6500336527824402,
+      "unseen_margin": -0.004095893818885088
+    }
+  },
+  "sealed_test_k5": {
+    "synthetic": {
+      "styles": 500.0,
+      "queries": 2500.0,
+      "top1": 0.7375999689102173,
+      "top5": 0.9031999707221985,
+      "mrr": 0.8114251494407654,
+      "margin": 0.06940444558858871,
+      "positive_cosine": 0.7258960008621216,
+      "hardest_negative_cosine": 0.6564915776252747,
+      "verification_auroc": 0.7113653421401978,
+      "embedding_std": 0.04344247281551361,
+      "embedding_effective_rank": 23.896039962768555,
+      "unseen_queries": 2500.0,
+      "unseen_top1": 0.7375999689102173,
+      "unseen_top5": 0.9031999707221985,
+      "unseen_mrr": 0.8114251494407654,
+      "unseen_margin": 0.06940444558858871,
+      "matched_queries": 500.0,
+      "matched_top1": 0.6600000262260437,
+      "matched_top5": 0.8520000576972961,
+      "matched_mrr": 0.748103678226471,
+      "matched_margin": 0.045967720448970795,
+      "cross_queries": 2000.0,
+      "cross_top1": 0.7570000290870667,
+      "cross_top5": 0.9160000681877136,
+      "cross_mrr": 0.8272556662559509,
+      "cross_margin": 0.0752636268734932
+    },
+    "human": {
+      "styles": 500.0,
+      "queries": 2500.0,
+      "top1": 0.761199951171875,
+      "top5": 0.9199999570846558,
+      "mrr": 0.830402135848999,
+      "margin": 0.08207859843969345,
+      "positive_cosine": 0.7568042874336243,
+      "hardest_negative_cosine": 0.6747256517410278,
+      "verification_auroc": 0.7453367114067078,
+      "embedding_std": 0.04376835376024246,
+      "embedding_effective_rank": 26.91031265258789,
+      "unseen_queries": 2500.0,
+      "unseen_top1": 0.761199951171875,
+      "unseen_top5": 0.9199999570846558,
+      "unseen_mrr": 0.830402135848999,
+      "unseen_margin": 0.08207859843969345
+    }
+  },
+  "frozen_full_face_validation_k5": {
+    "synthetic": {
+      "styles": 4500.0,
+      "queries": 22500.0,
+      "top1": 0.2617333233356476,
+      "top5": 0.3952000141143799,
+      "mrr": 0.3287946879863739,
+      "margin": -0.058847684413194656,
+      "positive_cosine": 0.3081073760986328,
+      "hardest_negative_cosine": 0.36695507168769836,
+      "verification_auroc": 0.3488680124282837,
+      "seen_queries": 20000.0,
+      "seen_top1": 0.2628999948501587,
+      "seen_top5": 0.3959999978542328,
+      "seen_mrr": 0.3298960030078888,
+      "seen_margin": -0.0586077980697155,
+      "unseen_queries": 2500.0,
+      "unseen_top1": 0.2523999810218811,
+      "unseen_top5": 0.3887999951839447,
+      "unseen_mrr": 0.31998398900032043,
+      "unseen_margin": -0.06076674535870552,
+      "matched_queries": 4500.0,
+      "matched_top1": 0.10533333569765091,
+      "matched_top5": 0.1851111203432083,
+      "matched_mrr": 0.14898858964443207,
+      "matched_margin": -0.11573364585638046,
+      "cross_queries": 18000.0,
+      "cross_top1": 0.3008333444595337,
+      "cross_top5": 0.4477222263813019,
+      "cross_mrr": 0.37374621629714966,
+      "cross_margin": -0.04462619498372078
+    },
+    "human": {
+      "styles": 4500.0,
+      "queries": 22500.0,
+      "top1": 0.43657779693603516,
+      "top5": 0.5786222219467163,
+      "mrr": 0.5056873559951782,
+      "margin": -0.029125619679689407,
+      "positive_cosine": 0.4710022509098053,
+      "hardest_negative_cosine": 0.5001279711723328,
+      "verification_auroc": 0.4406580626964569,
+      "seen_queries": 20000.0,
+      "seen_top1": 0.4320499897003174,
+      "seen_top5": 0.5762999653816223,
+      "seen_mrr": 0.5021449327468872,
+      "seen_margin": -0.030033092945814133,
+      "unseen_queries": 2500.0,
+      "unseen_top1": 0.47279998660087585,
+      "unseen_top5": 0.5971999764442444,
+      "unseen_mrr": 0.5340264439582825,
+      "unseen_margin": -0.021865803748369217
+    }
+  },
+  "protocol_note": "Synthetic K=5 uses the locked five-controlled support, one matched and four cross-content query protocol. Other K values use deterministic generic record ordering."
+}

metrics/retrieval_metrics.csv

@@ -0,0 +1,21 @@
+split,source,k,styles,queries,top1,top5,mrr,margin,seen_top1,seen_top5,seen_mrr,unseen_top1,unseen_top5,unseen_mrr,matched_mrr,cross_mrr,verification_auroc,embedding_effective_rank
+validation,synthetic,1,4500.0,40500.0,0.25738272070884705,0.4447901248931885,0.3484785854816437,-0.13694550096988678,0.26072221994400024,0.448972225189209,0.35212650895118713,0.23066666722297668,0.4113333523273468,0.3192950487136841,,,0.2625024914741516,24.163583755493164
+validation,synthetic,2,4500.0,36000.0,0.38269445300102234,0.5837222337722778,0.479046106338501,-0.0735325887799263,0.3878125250339508,0.5885000228881836,0.48388049006462097,0.34175002574920654,0.5455000400543213,0.4403711259365082,,,0.3855718970298767,24.163583755493164
+validation,synthetic,4,4500.0,27000.0,0.4913703501224518,0.6868888735771179,0.5827372670173645,-0.032036565244197845,0.4964166581630707,0.6924583315849304,0.5878510475158691,0.45100000500679016,0.6423333287239075,0.5418276190757751,,,0.4890711009502411,24.163583755493164
+validation,synthetic,5,4500.0,22500.0,0.5625777840614319,0.7585333585739136,0.6518997550010681,-0.006614952813833952,0.5681999921798706,0.7637499570846558,0.6573799848556519,0.5175999999046326,0.7167999744415283,0.6080579161643982,0.5840909481048584,0.6688519716262817,0.5534452199935913,24.163583755493164
+validation,synthetic,8,4500.0,9000.0,0.5631111264228821,0.7501111030578613,0.6487501263618469,-0.006703122518956661,0.5692500472068787,0.7573750615119934,0.6548616290092468,0.5139999985694885,0.6920000314712524,0.5998580455780029,,,0.557915985584259,24.163583755493164
+validation,human,1,4500.0,40500.0,0.37967902421951294,0.5846666693687439,0.47633036971092224,-0.07382111996412277,0.3872777819633484,0.5945833325386047,0.48492416739463806,0.31888890266418457,0.5053333640098572,0.40757980942726135,,,0.38891345262527466,27.383806228637695
+validation,human,2,4500.0,36000.0,0.5184444785118103,0.7187777757644653,0.6118195652961731,-0.018171189352869987,0.5287187695503235,0.729593813419342,0.6219512224197388,0.4362500309944153,0.6322500109672546,0.5307666063308716,,,0.516968846321106,27.383806228637695
+validation,human,4,4500.0,27000.0,0.6229629516601562,0.8089629411697388,0.7071906924247742,0.01628497801721096,0.6336249709129333,0.8175833225250244,0.7167206406593323,0.5376666784286499,0.7400000095367432,0.6309514045715332,,,0.6160270571708679,27.383806228637695
+validation,human,5,4500.0,22500.0,0.6493777632713318,0.8219555616378784,0.7280733585357666,0.022814957424998283,0.660099983215332,0.830549955368042,0.7378283143043518,0.5636000037193298,0.7531999945640564,0.6500336527824402,,,0.6369090676307678,27.383806228637695
+validation,human,8,4500.0,9000.0,0.6846666932106018,0.8516666889190674,0.7596060037612915,0.03352390229701996,0.6943750381469727,0.85875004529953,0.7680479884147644,0.6070000529289246,0.7950000166893005,0.6920705437660217,,,0.6696112155914307,27.383806228637695
+test,synthetic,1,500.0,4500.0,0.41688889265060425,0.6495555639266968,0.5274384617805481,-0.07223178446292877,,,,0.41688889265060425,0.6495555639266968,0.5274384617805481,,,0.41443032026290894,23.896039962768555
+test,synthetic,2,500.0,4000.0,0.5692500472068787,0.7830000519752502,0.6672611236572266,0.00026985956355929375,,,,0.5692500472068787,0.7830000519752502,0.6672611236572266,,,0.557541012763977,23.896039962768555
+test,synthetic,4,500.0,3000.0,0.6786666512489319,0.8600000143051147,0.7601169943809509,0.04286628216505051,,,,0.6786666512489319,0.8600000143051147,0.7601169943809509,,,0.6542216539382935,23.896039962768555
+test,synthetic,5,500.0,2500.0,0.7375999689102173,0.9031999707221985,0.8114251494407654,0.06940444558858871,,,,0.7375999689102173,0.9031999707221985,0.8114251494407654,0.748103678226471,0.8272556662559509,0.7113653421401978,23.896039962768555
+test,synthetic,8,500.0,1000.0,0.7070000171661377,0.8790000677108765,0.7862902879714966,0.06344308704137802,,,,0.7070000171661377,0.8790000677108765,0.7862902879714966,,,0.7011929750442505,23.896039962768555
+test,human,1,500.0,4500.0,0.5155555605888367,0.7431111335754395,0.6204044222831726,-0.0215375367552042,,,,0.5155555605888367,0.7431111335754395,0.6204044222831726,,,0.5082485675811768,26.91031265258789
+test,human,2,500.0,4000.0,0.6565000414848328,0.8540000319480896,0.7454702258110046,0.03674145042896271,,,,0.6565000414848328,0.8540000319480896,0.7454702258110046,,,0.6430025100708008,26.91031265258789
+test,human,4,500.0,3000.0,0.737666666507721,0.9136666655540466,0.8140823245048523,0.07469645887613297,,,,0.737666666507721,0.9136666655540466,0.8140823245048523,,,0.726514458656311,26.91031265258789
+test,human,5,500.0,2500.0,0.761199951171875,0.9199999570846558,0.830402135848999,0.08207859843969345,,,,0.761199951171875,0.9199999570846558,0.830402135848999,,,0.7453367114067078,26.91031265258789
+test,human,8,500.0,1000.0,0.7820000648498535,0.9250000715255737,0.8470030426979065,0.08853543549776077,,,,0.7820000648498535,0.9250000715255737,0.8470030426979065,,,0.7604439854621887,26.91031265258789

metrics/test_k_curve.json

@@ -0,0 +1,204 @@
+{
+  "synthetic": {
+    "k1": {
+      "styles": 500.0,
+      "queries": 4500.0,
+      "top1": 0.41688889265060425,
+      "top5": 0.6495555639266968,
+      "mrr": 0.5274384617805481,
+      "margin": -0.07223178446292877,
+      "positive_cosine": 0.5964593291282654,
+      "hardest_negative_cosine": 0.668691098690033,
+      "verification_auroc": 0.41443032026290894,
+      "embedding_std": 0.04344247281551361,
+      "embedding_effective_rank": 23.896039962768555,
+      "unseen_queries": 4500.0,
+      "unseen_top1": 0.41688889265060425,
+      "unseen_top5": 0.6495555639266968,
+      "unseen_mrr": 0.5274384617805481,
+      "unseen_margin": -0.07223178446292877
+    },
+    "k2": {
+      "styles": 500.0,
+      "queries": 4000.0,
+      "top1": 0.5692500472068787,
+      "top5": 0.7830000519752502,
+      "mrr": 0.6672611236572266,
+      "margin": 0.00026985956355929375,
+      "positive_cosine": 0.6695296764373779,
+      "hardest_negative_cosine": 0.6692599058151245,
+      "verification_auroc": 0.557541012763977,
+      "embedding_std": 0.04344247281551361,
+      "embedding_effective_rank": 23.896039962768555,
+      "unseen_queries": 4000.0,
+      "unseen_top1": 0.5692500472068787,
+      "unseen_top5": 0.7830000519752502,
+      "unseen_mrr": 0.6672611236572266,
+      "unseen_margin": 0.00026985956355929375
+    },
+    "k4": {
+      "styles": 500.0,
+      "queries": 3000.0,
+      "top1": 0.6786666512489319,
+      "top5": 0.8600000143051147,
+      "mrr": 0.7601169943809509,
+      "margin": 0.04286628216505051,
+      "positive_cosine": 0.7160320281982422,
+      "hardest_negative_cosine": 0.6731657981872559,
+      "verification_auroc": 0.6542216539382935,
+      "embedding_std": 0.04344247281551361,
+      "embedding_effective_rank": 23.896039962768555,
+      "unseen_queries": 3000.0,
+      "unseen_top1": 0.6786666512489319,
+      "unseen_top5": 0.8600000143051147,
+      "unseen_mrr": 0.7601169943809509,
+      "unseen_margin": 0.04286628216505051
+    },
+    "k5": {
+      "styles": 500.0,
+      "queries": 2500.0,
+      "top1": 0.7375999689102173,
+      "top5": 0.9031999707221985,
+      "mrr": 0.8114251494407654,
+      "margin": 0.06940444558858871,
+      "positive_cosine": 0.7258960008621216,
+      "hardest_negative_cosine": 0.6564915776252747,
+      "verification_auroc": 0.7113653421401978,
+      "embedding_std": 0.04344247281551361,
+      "embedding_effective_rank": 23.896039962768555,
+      "unseen_queries": 2500.0,
+      "unseen_top1": 0.7375999689102173,
+      "unseen_top5": 0.9031999707221985,
+      "unseen_mrr": 0.8114251494407654,
+      "unseen_margin": 0.06940444558858871,
+      "matched_queries": 500.0,
+      "matched_top1": 0.6600000262260437,
+      "matched_top5": 0.8520000576972961,
+      "matched_mrr": 0.748103678226471,
+      "matched_margin": 0.045967720448970795,
+      "cross_queries": 2000.0,
+      "cross_top1": 0.7570000290870667,
+      "cross_top5": 0.9160000681877136,
+      "cross_mrr": 0.8272556662559509,
+      "cross_margin": 0.0752636268734932
+    },
+    "k8": {
+      "styles": 500.0,
+      "queries": 1000.0,
+      "top1": 0.7070000171661377,
+      "top5": 0.8790000677108765,
+      "mrr": 0.7862902879714966,
+      "margin": 0.06344308704137802,
+      "positive_cosine": 0.739560604095459,
+      "hardest_negative_cosine": 0.676117479801178,
+      "verification_auroc": 0.7011929750442505,
+      "embedding_std": 0.04344247281551361,
+      "embedding_effective_rank": 23.896039962768555,
+      "unseen_queries": 1000.0,
+      "unseen_top1": 0.7070000171661377,
+      "unseen_top5": 0.8790000677108765,
+      "unseen_mrr": 0.7862902879714966,
+      "unseen_margin": 0.06344308704137802
+    }
+  },
+  "human": {
+    "k1": {
+      "styles": 500.0,
+      "queries": 4500.0,
+      "top1": 0.5155555605888367,
+      "top5": 0.7431111335754395,
+      "mrr": 0.6204044222831726,
+      "margin": -0.0215375367552042,
+      "positive_cosine": 0.6388872861862183,
+      "hardest_negative_cosine": 0.6604248285293579,
+      "verification_auroc": 0.5082485675811768,
+      "embedding_std": 0.04376835376024246,
+      "embedding_effective_rank": 26.91031265258789,
+      "unseen_queries": 4500.0,
+      "unseen_top1": 0.5155555605888367,
+      "unseen_top5": 0.7431111335754395,
+      "unseen_mrr": 0.6204044222831726,
+      "unseen_margin": -0.0215375367552042
+    },
+    "k2": {
+      "styles": 500.0,
+      "queries": 4000.0,
+      "top1": 0.6565000414848328,
+      "top5": 0.8540000319480896,
+      "mrr": 0.7454702258110046,
+      "margin": 0.03674145042896271,
+      "positive_cosine": 0.7047365307807922,
+      "hardest_negative_cosine": 0.6679950952529907,
+      "verification_auroc": 0.6430025100708008,
+      "embedding_std": 0.04376835376024246,
+      "embedding_effective_rank": 26.91031265258789,
+      "unseen_queries": 4000.0,
+      "unseen_top1": 0.6565000414848328,
+      "unseen_top5": 0.8540000319480896,
+      "unseen_mrr": 0.7454702258110046,
+      "unseen_margin": 0.03674145042896271
+    },
+    "k4": {
+      "styles": 500.0,
+      "queries": 3000.0,
+      "top1": 0.737666666507721,
+      "top5": 0.9136666655540466,
+      "mrr": 0.8140823245048523,
+      "margin": 0.07469645887613297,
+      "positive_cosine": 0.7484809756278992,
+      "hardest_negative_cosine": 0.6737845540046692,
+      "verification_auroc": 0.726514458656311,
+      "embedding_std": 0.04376835376024246,
+      "embedding_effective_rank": 26.91031265258789,
+      "unseen_queries": 3000.0,
+      "unseen_top1": 0.737666666507721,
+      "unseen_top5": 0.9136666655540466,
+      "unseen_mrr": 0.8140823245048523,
+      "unseen_margin": 0.07469645887613297
+    },
+    "k5": {
+      "styles": 500.0,
+      "queries": 2500.0,
+      "top1": 0.761199951171875,
+      "top5": 0.9199999570846558,
+      "mrr": 0.830402135848999,
+      "margin": 0.08207859843969345,
+      "positive_cosine": 0.7568042874336243,
+      "hardest_negative_cosine": 0.6747256517410278,
+      "verification_auroc": 0.7453367114067078,
+      "embedding_std": 0.04376835376024246,
+      "embedding_effective_rank": 26.91031265258789,
+      "unseen_queries": 2500.0,
+      "unseen_top1": 0.761199951171875,
+      "unseen_top5": 0.9199999570846558,
+      "unseen_mrr": 0.830402135848999,
+      "unseen_margin": 0.08207859843969345
+    },
+    "k8": {
+      "styles": 500.0,
+      "queries": 1000.0,
+      "top1": 0.7820000648498535,
+      "top5": 0.9250000715255737,
+      "mrr": 0.8470030426979065,
+      "margin": 0.08853543549776077,
+      "positive_cosine": 0.7681199312210083,
+      "hardest_negative_cosine": 0.6795844435691833,
+      "verification_auroc": 0.7604439854621887,
+      "embedding_std": 0.04376835376024246,
+      "embedding_effective_rank": 26.91031265258789,
+      "unseen_queries": 1000.0,
+      "unseen_top1": 0.7820000648498535,
+      "unseen_top5": 0.9250000715255737,
+      "unseen_mrr": 0.8470030426979065,
+      "unseen_margin": 0.08853543549776077
+    }
+  },
+  "checkpoint_score": 0.8209136426448822,
+  "checkpoint_config": {
+    "packed_manifest_sha256": "c1267b9fd45c6e0d222b91d6ae2291b36b763277cc39cbc32484874ffeae9b32",
+    "use_anima": true,
+    "face_dropout": 0.3,
+    "episode_pairs": 8,
+    "seed": 20260719
+  }
+}

metrics/validation_k_curve.json

@@ -0,0 +1,254 @@
+{
+  "synthetic": {
+    "k1": {
+      "styles": 4500.0,
+      "queries": 40500.0,
+      "top1": 0.25738272070884705,
+      "top5": 0.4447901248931885,
+      "mrr": 0.3484785854816437,
+      "margin": -0.13694550096988678,
+      "positive_cosine": 0.6166356205940247,
+      "hardest_negative_cosine": 0.7535811066627502,
+      "verification_auroc": 0.2625024914741516,
+      "embedding_std": 0.04349411278963089,
+      "embedding_effective_rank": 24.163583755493164,
+      "seen_queries": 36000.0,
+      "seen_top1": 0.26072221994400024,
+      "seen_top5": 0.448972225189209,
+      "seen_mrr": 0.35212650895118713,
+      "seen_margin": -0.13477584719657898,
+      "unseen_queries": 4500.0,
+      "unseen_top1": 0.23066666722297668,
+      "unseen_top5": 0.4113333523273468,
+      "unseen_mrr": 0.3192950487136841,
+      "unseen_margin": -0.15430252254009247
+    },
+    "k2": {
+      "styles": 4500.0,
+      "queries": 36000.0,
+      "top1": 0.38269445300102234,
+      "top5": 0.5837222337722778,
+      "mrr": 0.479046106338501,
+      "margin": -0.0735325887799263,
+      "positive_cosine": 0.6810839176177979,
+      "hardest_negative_cosine": 0.7546165585517883,
+      "verification_auroc": 0.3855718970298767,
+      "embedding_std": 0.04349411278963089,
+      "embedding_effective_rank": 24.163583755493164,
+      "seen_queries": 32000.0,
+      "seen_top1": 0.3878125250339508,
+      "seen_top5": 0.5885000228881836,
+      "seen_mrr": 0.48388049006462097,
+      "seen_margin": -0.07155531644821167,
+      "unseen_queries": 4000.0,
+      "unseen_top1": 0.34175002574920654,
+      "unseen_top5": 0.5455000400543213,
+      "unseen_mrr": 0.4403711259365082,
+      "unseen_margin": -0.08935073763132095
+    },
+    "k4": {
+      "styles": 4500.0,
+      "queries": 27000.0,
+      "top1": 0.4913703501224518,
+      "top5": 0.6868888735771179,
+      "mrr": 0.5827372670173645,
+      "margin": -0.032036565244197845,
+      "positive_cosine": 0.7242447733879089,
+      "hardest_negative_cosine": 0.756281316280365,
+      "verification_auroc": 0.4890711009502411,
+      "embedding_std": 0.04349411278963089,
+      "embedding_effective_rank": 24.163583755493164,
+      "seen_queries": 24000.0,
+      "seen_top1": 0.4964166581630707,
+      "seen_top5": 0.6924583315849304,
+      "seen_mrr": 0.5878510475158691,
+      "seen_margin": -0.030110228806734085,
+      "unseen_queries": 3000.0,
+      "unseen_top1": 0.45100000500679016,
+      "unseen_top5": 0.6423333287239075,
+      "unseen_mrr": 0.5418276190757751,
+      "unseen_margin": -0.047447264194488525
+    },
+    "k5": {
+      "styles": 4500.0,
+      "queries": 22500.0,
+      "top1": 0.5625777840614319,
+      "top5": 0.7585333585739136,
+      "mrr": 0.6518997550010681,
+      "margin": -0.006614952813833952,
+      "positive_cosine": 0.7365614771842957,
+      "hardest_negative_cosine": 0.7431764006614685,
+      "verification_auroc": 0.5534452199935913,
+      "embedding_std": 0.04349411278963089,
+      "embedding_effective_rank": 24.163583755493164,
+      "seen_queries": 20000.0,
+      "seen_top1": 0.5681999921798706,
+      "seen_top5": 0.7637499570846558,
+      "seen_mrr": 0.6573799848556519,
+      "seen_margin": -0.004584707319736481,
+      "unseen_queries": 2500.0,
+      "unseen_top1": 0.5175999999046326,
+      "unseen_top5": 0.7167999744415283,
+      "unseen_mrr": 0.6080579161643982,
+      "unseen_margin": -0.02285691164433956,
+      "matched_queries": 4500.0,
+      "matched_top1": 0.48955556750297546,
+      "matched_top5": 0.6902222037315369,
+      "matched_mrr": 0.5840909481048584,
+      "matched_margin": -0.030851511284708977,
+      "cross_queries": 18000.0,
+      "cross_top1": 0.5808333158493042,
+      "cross_top5": 0.7756111025810242,
+      "cross_mrr": 0.6688519716262817,
+      "cross_margin": -0.0005558129050768912
+    },
+    "k8": {
+      "styles": 4500.0,
+      "queries": 9000.0,
+      "top1": 0.5631111264228821,
+      "top5": 0.7501111030578613,
+      "mrr": 0.6487501263618469,
+      "margin": -0.006703122518956661,
+      "positive_cosine": 0.7508071660995483,
+      "hardest_negative_cosine": 0.7575103044509888,
+      "verification_auroc": 0.557915985584259,
+      "embedding_std": 0.04349411278963089,
+      "embedding_effective_rank": 24.163583755493164,
+      "seen_queries": 8000.0,
+      "seen_top1": 0.5692500472068787,
+      "seen_top5": 0.7573750615119934,
+      "seen_mrr": 0.6548616290092468,
+      "seen_margin": -0.004510539583861828,
+      "unseen_queries": 1000.0,
+      "unseen_top1": 0.5139999985694885,
+      "unseen_top5": 0.6920000314712524,
+      "unseen_mrr": 0.5998580455780029,
+      "unseen_margin": -0.024243783205747604
+    }
+  },
+  "human": {
+    "k1": {
+      "styles": 4500.0,
+      "queries": 40500.0,
+      "top1": 0.37967902421951294,
+      "top5": 0.5846666693687439,
+      "mrr": 0.47633036971092224,
+      "margin": -0.07382111996412277,
+      "positive_cosine": 0.6724859476089478,
+      "hardest_negative_cosine": 0.7463070750236511,
+      "verification_auroc": 0.38891345262527466,
+      "embedding_std": 0.043804656714200974,
+      "embedding_effective_rank": 27.383806228637695,
+      "seen_queries": 36000.0,
+      "seen_top1": 0.3872777819633484,
+      "seen_top5": 0.5945833325386047,
+      "seen_mrr": 0.48492416739463806,
+      "seen_margin": -0.06964229047298431,
+      "unseen_queries": 4500.0,
+      "unseen_top1": 0.31888890266418457,
+      "unseen_top5": 0.5053333640098572,
+      "unseen_mrr": 0.40757980942726135,
+      "unseen_margin": -0.10725177824497223
+    },
+    "k2": {
+      "styles": 4500.0,
+      "queries": 36000.0,
+      "top1": 0.5184444785118103,
+      "top5": 0.7187777757644653,
+      "mrr": 0.6118195652961731,
+      "margin": -0.018171189352869987,
+      "positive_cosine": 0.7339839339256287,
+      "hardest_negative_cosine": 0.7521551847457886,
+      "verification_auroc": 0.516968846321106,
+      "embedding_std": 0.043804656714200974,
+      "embedding_effective_rank": 27.383806228637695,
+      "seen_queries": 32000.0,
+      "seen_top1": 0.5287187695503235,
+      "seen_top5": 0.729593813419342,
+      "seen_mrr": 0.6219512224197388,
+      "seen_margin": -0.014534689486026764,
+      "unseen_queries": 4000.0,
+      "unseen_top1": 0.4362500309944153,
+      "unseen_top5": 0.6322500109672546,
+      "unseen_mrr": 0.5307666063308716,
+      "unseen_margin": -0.04726320877671242
+    },
+    "k4": {
+      "styles": 4500.0,
+      "queries": 27000.0,
+      "top1": 0.6229629516601562,
+      "top5": 0.8089629411697388,
+      "mrr": 0.7071906924247742,
+      "margin": 0.01628497801721096,
+      "positive_cosine": 0.7733137607574463,
+      "hardest_negative_cosine": 0.7570288181304932,
+      "verification_auroc": 0.6160270571708679,
+      "embedding_std": 0.043804656714200974,
+      "embedding_effective_rank": 27.383806228637695,
+      "seen_queries": 24000.0,
+      "seen_top1": 0.6336249709129333,
+      "seen_top5": 0.8175833225250244,
+      "seen_mrr": 0.7167206406593323,
+      "seen_margin": 0.019602350890636444,
+      "unseen_queries": 3000.0,
+      "unseen_top1": 0.5376666784286499,
+      "unseen_top5": 0.7400000095367432,
+      "unseen_mrr": 0.6309514045715332,
+      "unseen_margin": -0.01025400310754776
+    },
+    "k5": {
+      "styles": 4500.0,
+      "queries": 22500.0,
+      "top1": 0.6493777632713318,
+      "top5": 0.8219555616378784,
+      "mrr": 0.7280733585357666,
+      "margin": 0.022814957424998283,
+      "positive_cosine": 0.7811680436134338,
+      "hardest_negative_cosine": 0.7583531141281128,
+      "verification_auroc": 0.6369090676307678,
+      "embedding_std": 0.043804656714200974,
+      "embedding_effective_rank": 27.383806228637695,
+      "seen_queries": 20000.0,
+      "seen_top1": 0.660099983215332,
+      "seen_top5": 0.830549955368042,
+      "seen_mrr": 0.7378283143043518,
+      "seen_margin": 0.026178810745477676,
+      "unseen_queries": 2500.0,
+      "unseen_top1": 0.5636000037193298,
+      "unseen_top5": 0.7531999945640564,
+      "unseen_mrr": 0.6500336527824402,
+      "unseen_margin": -0.004095893818885088
+    },
+    "k8": {
+      "styles": 4500.0,
+      "queries": 9000.0,
+      "top1": 0.6846666932106018,
+      "top5": 0.8516666889190674,
+      "mrr": 0.7596060037612915,
+      "margin": 0.03352390229701996,
+      "positive_cosine": 0.793936014175415,
+      "hardest_negative_cosine": 0.7604120373725891,
+      "verification_auroc": 0.6696112155914307,
+      "embedding_std": 0.043804656714200974,
+      "embedding_effective_rank": 27.383806228637695,
+      "seen_queries": 8000.0,
+      "seen_top1": 0.6943750381469727,
+      "seen_top5": 0.85875004529953,
+      "seen_mrr": 0.7680479884147644,
+      "seen_margin": 0.03656964749097824,
+      "unseen_queries": 1000.0,
+      "unseen_top1": 0.6070000529289246,
+      "unseen_top5": 0.7950000166893005,
+      "unseen_mrr": 0.6920705437660217,
+      "unseen_margin": 0.009157972410321236
+    }
+  },
+  "checkpoint_score": 0.6290457844734192,
+  "checkpoint_config": {
+    "packed_manifest_sha256": "c1267b9fd45c6e0d222b91d6ae2291b36b763277cc39cbc32484874ffeae9b32",
+    "use_anima": true,
+    "face_dropout": 0.3,
+    "episode_pairs": 8,
+    "seed": 20260719
+  }
+}

model.safetensors

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

modeling_anima_style_embedding.py

@@ -0,0 +1,341 @@
+from __future__ import annotations
+
+import math
+
+import torch
+from torch import Tensor, nn
+from torch.nn import functional as F
+
+
+def masked_mean_prototype(per_image: Tensor, mask: Tensor) -> Tensor:
+    if per_image.ndim != 3 or mask.shape != per_image.shape[:2]:
+        raise ValueError("expected per_image [batch,references,width] and matching mask")
+    mask = mask.to(device=per_image.device, dtype=torch.bool)
+    counts = mask.sum(dim=1)
+    if torch.any(counts < 1) or torch.any(counts > 8):
+        raise ValueError("each reference set must contain 1-8 images")
+    weights = mask.unsqueeze(-1).to(per_image.dtype)
+    mean = (per_image * weights).sum(dim=1) / counts.unsqueeze(-1)
+    return F.normalize(mean, dim=-1)
+
+
+class StyleMixerBlock(nn.Module):
+    def __init__(self, width: int, heads: int, dropout: float = 0.0) -> None:
+        super().__init__()
+        self.query_norm = nn.LayerNorm(width)
+        self.token_norm = nn.LayerNorm(width)
+        self.attention = nn.MultiheadAttention(width, heads, batch_first=True)
+        self.attention_dropout = nn.Dropout(dropout)
+        self.output_norm = nn.LayerNorm(width)
+        self.mlp = nn.Sequential(
+            nn.Linear(width, 4 * width),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(4 * width, width),
+            nn.Dropout(dropout),
+        )
+
+    def forward(self, query: Tensor, tokens: Tensor, token_mask: Tensor) -> Tensor:
+        value, _ = self.attention(
+            self.query_norm(query),
+            self.token_norm(tokens),
+            self.token_norm(tokens),
+            key_padding_mask=~token_mask,
+            need_weights=False,
+        )
+        query = query + self.attention_dropout(value)
+        return query + self.mlp(self.output_norm(query))
+
+
+class UnifiedStyleEncoder(nn.Module):
+    """Encode cached full/face vision tokens and optional Anima descriptors."""
+
+    def __init__(
+        self,
+        *,
+        external_dim: int = 1152,
+        external_tokens: int = 30,
+        anima_dim: int = 4096,
+        anima_tokens: int = 3,
+        width: int = 384,
+        blocks: int = 2,
+        heads: int = 6,
+        embedding_dim: int = 512,
+        use_anima: bool = True,
+        dropout: float = 0.0,
+        reference_grounding: bool = False,
+        functional_factors: bool = False,
+    ) -> None:
+        super().__init__()
+        if (
+            blocks < 1
+            or width % heads
+            or (reference_grounding and embedding_dim % heads)
+            or not 0.0 <= dropout < 1.0
+        ):
+            raise ValueError("blocks must be positive and width must be divisible by heads")
+        if functional_factors and embedding_dim != 512:
+            raise ValueError("functional factorization currently requires a 512-D embedding")
+        self.external_tokens = external_tokens
+        self.anima_tokens = anima_tokens if use_anima else 0
+        self.external_projector = nn.Sequential(
+            nn.LayerNorm(external_dim),
+            nn.Linear(external_dim, width),
+        )
+        self.anima_projector = (
+            nn.Sequential(nn.LayerNorm(anima_dim), nn.Linear(anima_dim, width))
+            if use_anima
+            else None
+        )
+        self.full_positions = nn.Parameter(torch.randn(external_tokens, width) * 0.02)
+        self.face_positions = nn.Parameter(torch.randn(external_tokens, width) * 0.02)
+        self.anima_positions = (
+            nn.Parameter(torch.randn(anima_tokens, width) * 0.02) if use_anima else None
+        )
+        self.reference_grounding = reference_grounding
+        self.functional_factors = functional_factors
+        query_count = 5 if functional_factors else 1
+        self.style_query = nn.Parameter(torch.randn(1, query_count, width) * 0.02)
+        self.blocks = nn.ModuleList(
+            [StyleMixerBlock(width, heads, dropout) for _ in range(blocks)]
+        )
+        if functional_factors:
+            self.factor_output = nn.Sequential(nn.LayerNorm(width), nn.Linear(width, 96))
+            self.shared_output = nn.Sequential(nn.LayerNorm(width), nn.Linear(width, 128))
+            self.output = nn.Sequential(nn.LayerNorm(512), nn.Linear(512, embedding_dim))
+            self.factor_decoders = nn.ModuleList([nn.Linear(96, external_dim) for _ in range(4)])
+        else:
+            self.output = nn.Sequential(nn.LayerNorm(width), nn.Linear(width, embedding_dim))
+        if reference_grounding:
+            self.consensus_query = nn.Parameter(torch.randn(1, 1, embedding_dim) * 0.02)
+            self.consensus_attention = nn.MultiheadAttention(
+                embedding_dim, heads, batch_first=True, dropout=dropout
+            )
+            self.consensus_condition = nn.Linear(embedding_dim, width)
+            self.consensus_gamma = nn.Parameter(torch.tensor(-2.9444))
+
+    def _read_tokens(
+        self, tokens: Tensor, token_mask: Tensor, conditioning: Tensor | None = None
+    ) -> tuple[Tensor, Tensor | None]:
+        query = self.style_query.expand(tokens.shape[0], -1, -1)
+        if conditioning is not None:
+            query = query + self.consensus_condition(conditioning)[:, None]
+        for block in self.blocks:
+            query = block(query, tokens, token_mask)
+        if self.functional_factors:
+            factors = F.normalize(self.factor_output(query[:, :4]), dim=-1)
+            shared = self.shared_output(query[:, 4])
+            embedding = self.output(torch.cat((factors.flatten(1), shared), dim=-1))
+            return F.normalize(embedding, dim=-1), factors
+        return F.normalize(self.output(query[:, 0]), dim=-1), None
+
+    def encode_images(
+        self,
+        full_features: Tensor,
+        *,
+        face_features: Tensor | None = None,
+        face_mask: Tensor | None = None,
+        anima_features: Tensor | None = None,
+        conditioning: Tensor | None = None,
+    ) -> Tensor:
+        """Encode independent images without padding them into reference sets."""
+
+        if full_features.ndim != 3 or full_features.shape[1] != self.external_tokens:
+            raise ValueError("full_features must have shape [images,external_tokens,width]")
+        images = full_features.shape[0]
+        full = self.external_projector(full_features) + self.full_positions
+        tokens = [full]
+        masks = [torch.ones(full.shape[:2], dtype=torch.bool, device=full.device)]
+
+        if face_features is not None:
+            if face_features.shape != full_features.shape:
+                raise ValueError("face_features must match full_features")
+            if face_mask is None or face_mask.shape != (images,):
+                raise ValueError("face_mask is required and must contain one value per image")
+            face = self.external_projector(face_features)
+            tokens.append(face + self.face_positions)
+            masks.append(face_mask[:, None].to(torch.bool).expand(-1, self.external_tokens))
+
+        if self.anima_projector is not None:
+            expected = (images, self.anima_tokens)
+            if anima_features is None or anima_features.shape[:2] != expected:
+                raise ValueError("anima_features are required by this encoder configuration")
+            anima = self.anima_projector(anima_features)
+            tokens.append(anima + self.anima_positions)
+            masks.append(torch.ones(anima.shape[:2], dtype=torch.bool, device=anima.device))
+
+        token_tensor = torch.cat(tokens, dim=1)
+        token_mask = torch.cat(masks, dim=1)
+        return self._read_tokens(token_tensor, token_mask, conditioning)[0]
+
+    def encode_images_with_factors(
+        self,
+        full_features: Tensor,
+        *,
+        face_features: Tensor | None = None,
+        face_mask: Tensor | None = None,
+        anima_features: Tensor | None = None,
+        conditioning: Tensor | None = None,
+    ) -> tuple[Tensor, Tensor | None]:
+        if full_features.ndim != 3 or full_features.shape[1] != self.external_tokens:
+            raise ValueError("full_features must have shape [images,external_tokens,width]")
+        images = full_features.shape[0]
+        full = self.external_projector(full_features) + self.full_positions
+        tokens = [full]
+        masks = [torch.ones(full.shape[:2], dtype=torch.bool, device=full.device)]
+        if face_features is not None:
+            if face_features.shape != full_features.shape:
+                raise ValueError("face_features must match full_features")
+            if face_mask is None or face_mask.shape != (images,):
+                raise ValueError("face_mask is required and must contain one value per image")
+            tokens.append(self.external_projector(face_features) + self.face_positions)
+            masks.append(face_mask[:, None].to(torch.bool).expand(-1, self.external_tokens))
+        if self.anima_projector is not None:
+            if anima_features is None or anima_features.shape[:2] != (images, self.anima_tokens):
+                raise ValueError("anima_features are required by this encoder configuration")
+            tokens.append(self.anima_projector(anima_features) + self.anima_positions)
+            masks.append(torch.ones(anima_features.shape[:2], dtype=torch.bool, device=full.device))
+        return self._read_tokens(torch.cat(tokens, dim=1), torch.cat(masks, dim=1), conditioning)
+
+    def encode_reference_sets(
+        self,
+        full_features: Tensor,
+        reference_mask: Tensor,
+        *,
+        face_features: Tensor | None = None,
+        face_mask: Tensor | None = None,
+        anima_features: Tensor | None = None,
+    ) -> tuple[Tensor, Tensor, Tensor | None]:
+        batch, references = full_features.shape[:2]
+        flat_face = None if face_features is None else face_features.flatten(0, 1)
+        flat_face_mask = None if face_mask is None else face_mask.flatten()
+        flat_anima = None if anima_features is None else anima_features.flatten(0, 1)
+        per_image, factors = self.encode_images_with_factors(
+            full_features.flatten(0, 1),
+            face_features=flat_face,
+            face_mask=flat_face_mask,
+            anima_features=flat_anima,
+        )
+        per_image = per_image.reshape(batch, references, -1)
+        factors = None if factors is None else factors.reshape(batch, references, 4, -1)
+        base = masked_mean_prototype(per_image, reference_mask)
+        if not self.reference_grounding:
+            return base, per_image, factors
+        consensus, _ = self.consensus_attention(
+            self.consensus_query.expand(batch, -1, -1),
+            per_image,
+            per_image,
+            key_padding_mask=~reference_mask,
+            need_weights=False,
+        )
+        consensus = F.normalize(consensus[:, 0], dim=-1)
+        reread, factors = self.encode_images_with_factors(
+            full_features.flatten(0, 1),
+            face_features=flat_face,
+            face_mask=flat_face_mask,
+            anima_features=flat_anima,
+            conditioning=consensus[:, None].expand(-1, references, -1).reshape(batch * references, -1),
+        )
+        reread = reread.reshape(batch, references, -1)
+        factors = None if factors is None else factors.reshape(batch, references, 4, -1)
+        agreement = F.cosine_similarity(reread, consensus[:, None], dim=-1)
+        centered = agreement - (
+            (agreement * reference_mask).sum(1, keepdim=True)
+            / reference_mask.sum(1, keepdim=True).clamp_min(1)
+        )
+        weights = (1.0 + 0.25 * torch.tanh(4.0 * centered)) * reference_mask
+        grounded = F.normalize((reread * weights[..., None]).sum(1) / weights.sum(1, keepdim=True), dim=-1)
+        gamma = torch.sigmoid(self.consensus_gamma)
+        return F.normalize(base + gamma * (grounded - base), dim=-1), reread, factors
+
+    def functional_response_loss(self, factor_codes: Tensor, full_features: Tensor) -> Tensor:
+        """Reconstruct artist response after removing a shared prompt/seed cell mean."""
+        if not self.functional_factors:
+            return full_features.new_zeros(())
+        groups = ((2, 3, 4, 5), (0, 6, 24), tuple(range(8, 24)), (1, 7, 25))
+        losses = []
+        for index, token_indices in enumerate(groups):
+            teacher = full_features[..., list(token_indices), :].float().mean(dim=-2)
+            teacher = F.normalize(teacher - teacher.mean(dim=1, keepdim=True), dim=-1)
+            prediction = F.normalize(self.factor_decoders[index](factor_codes[..., index, :]).float(), dim=-1)
+            losses.append((1.0 - F.cosine_similarity(prediction, teacher, dim=-1)).mean())
+        return torch.stack(losses).mean()
+
+    def forward(
+        self,
+        full_features: Tensor,
+        reference_mask: Tensor,
+        *,
+        face_features: Tensor | None = None,
+        face_mask: Tensor | None = None,
+        anima_features: Tensor | None = None,
+    ) -> tuple[Tensor, Tensor]:
+        if full_features.ndim != 4 or full_features.shape[2] != self.external_tokens:
+            raise ValueError("full_features must have shape [batch,references,external_tokens,width]")
+        batch, references = full_features.shape[:2]
+        if reference_mask.shape != (batch, references):
+            raise ValueError("reference_mask shape does not match full_features")
+
+        prototype, per_image, _ = self.encode_reference_sets(
+            full_features,
+            reference_mask,
+            face_features=face_features,
+            face_mask=face_mask,
+            anima_features=anima_features,
+        )
+        return prototype, per_image
+
+
+class AngularPrototypeLoss(nn.Module):
+    """Episode prototype softmax with an optional ArcFace-style positive margin."""
+
+    def __init__(
+        self, initial_scale: float = 10.0, initial_bias: float | None = None
+    ) -> None:
+        super().__init__()
+        if initial_scale <= 0:
+            raise ValueError("initial_scale must be positive")
+        self.raw_scale = nn.Parameter(torch.tensor(math.log(math.expm1(initial_scale))))
+        # initial_bias remains accepted for feasibility-script compatibility. A
+        # common softmax bias is a no-op and is not a trainable production parameter.
+
+    def forward(
+        self,
+        queries: Tensor,
+        prototypes: Tensor,
+        targets: Tensor,
+        *,
+        margin: float = 0.0,
+    ) -> tuple[Tensor, Tensor]:
+        if queries.ndim < 2 or prototypes.ndim != queries.ndim:
+            raise ValueError("queries and prototypes must have matching [episode,...,width] ranks")
+        if queries.shape[:-2] != prototypes.shape[:-2] or queries.shape[-1] != prototypes.shape[-1]:
+            raise ValueError("queries and prototypes must have matching episode axes and widths")
+        if targets.shape != queries.shape[:-1]:
+            raise ValueError("targets must contain one class index per query")
+        # Autocast can round a near-perfect BF16 cosine to exactly one. ArcFace's
+        # sqrt(1-cos²) derivative is then singular, so margin math and CE stay FP32.
+        cosine = (
+            F.normalize(queries, dim=-1)
+            @ F.normalize(prototypes, dim=-1).transpose(-1, -2)
+        ).float()
+        loss_cosine = cosine
+        if margin:
+            if not 0.0 <= margin < math.pi / 2:
+                raise ValueError("margin must be in [0, pi/2)")
+            positive = cosine.gather(-1, targets.unsqueeze(-1)).clamp(-1 + 1e-6, 1 - 1e-6)
+            positive_margin = positive * math.cos(margin) - torch.sqrt(1 - positive.square()) * math.sin(margin)
+            loss_cosine = cosine.scatter(-1, targets.unsqueeze(-1), positive_margin)
+        scale = F.softplus(self.raw_scale).clamp_max(100.0)
+        loss_logits = scale * loss_cosine
+        metric_logits = scale * cosine
+        return (
+            F.cross_entropy(loss_logits.flatten(0, -2), targets.flatten()),
+            metric_logits,
+        )
+
+
+def subset_consistency_loss(first: Tensor, second: Tensor) -> Tensor:
+    if first.shape != second.shape or first.ndim != 2:
+        raise ValueError("subset prototypes must have equal [batch,width] shapes")
+    return (1.0 - F.cosine_similarity(first, second, dim=-1)).mean()

requirements.txt

@@ -0,0 +1,3 @@
+torch>=2.4
+safetensors>=0.4
+huggingface_hub>=0.28