Model card

README.md

@@ -0,0 +1,98 @@
+---
+library_name: pytorch
+tags:
+- topology-optimization
+- structural-optimization
+- flow-matching
+- diffusion-transformer
+- pytorch_model_hub_mixin
+- model_hub_mixin
+datasets:
+- OpenTO/OpenTO
+pipeline_tag: image-to-image
+---
+
+# TopoDiT — Optimize Any Topology 2
+
+TopoDiT is the generative model of **Optimize Any Topology 2 (OAT2)**: a 688.6M-parameter
+conditional diffusion transformer, trained with flow matching, that generates
+minimum-compliance structural topologies in the latent space of the frozen
+[OpenTO/NFAE](https://huggingface.co/OpenTO/NFAE) neural-field autoencoder. It is the
+successor of the latent-diffusion U-Net of
+[Optimize Any Topology](https://arxiv.org/abs/2510.23667) (NeurIPS 2025) and takes the
+same inputs: domain shape, mesh cell size, target volume fraction, boundary conditions
+and loads — at any resolution and aspect ratio.
+
+**Code:** the `OptimizeAnyTopology2` repository (training, evaluation, GPU FEM).
+
+## Results (OpenTO test split, 5,000 problems, zero-shot, no CFG)
+
+| | OAT (2025) | **TopoDiT / OAT2** |
+|---|---|---|
+| failure rate, 1 sample | 40.0% | **21.8%** |
+| median compliance error, 1 sample | 3.94% | **0.46%** |
+| failure rate, best of 4 | 25.2% | **11.6%** |
+| median compliance error, best of 4 | 2.86% | **0.18%** |
+| failure rate, best of 4 + 10 PGD steps | 14.9% | **5.8%** |
+
+Identical FEM and statistics for both rows (CE = (C − C_gt)/C_gt; failure = CE ≥ 100%).
+
+## Architecture
+
+- Dense DiT over the 64×64×1 NFAE latent: patch size 4 → 256 tokens, 24 blocks,
+  width 1152, 16 heads; predicts the flow-matching velocity.
+- Global conditions (shape, cell size, volume fraction) → one vector concatenated with
+  the timestep embedding into adaLN-Zero modulation. Boundary conditions and loads →
+  attention-pooled spatial tokens (32-cell grid over the domain, every point kept),
+  read by cross-attention in every block. Continuous rectangle position embeddings,
+  QK-normalized attention.
+- Condition encoder: 4 layers, width 1152, 12 heads, token width 768.
+
+## Training
+
+OpenTO `labeled` + `NITO` (894k optimized structures), 50 epochs, 349,200 steps at
+effective batch 128, AdamW lr 1e-4 (cosine to 1e-6, 1k warm-up), weight decay 1e-4,
+logit-normal timesteps, classifier-free-guidance dropout per condition. Latent
+normalization statistics are stored on the model.
+
+## Usage
+
+```python
+import torch
+from datasets import load_dataset
+from OAT import NFAE, TopoDiT
+from OAT.DataUtils import OpenTO, DiffusionCollator, cached_full_grid_cell
+from OAT.Pipelines import FlowMatchPipeline
+
+device = 'cuda'
+model = TopoDiT.from_pretrained('OpenTO/TopoDiT').to(device).eval()
+nfae = NFAE.from_pretrained('OpenTO/NFAE').to(device).eval()
+
+data = load_dataset('OpenTO/OpenTO', split='test')
+ds = OpenTO(data, mode='diffusion', train=False)
+batch = DiffusionCollator()([ds[0]]).to(device)
+pipe = FlowMatchPipeline(shift=1.0)
+
+with torch.no_grad(), torch.autocast('cuda', torch.bfloat16):
+    z = pipe.inference(model, batch, num_sampling_steps=20, guidance_scale=1.0)
+    phi = nfae.decoder(model.denormalize(z.float()))
+    w, h = data[0]['topology'].size
+    coord, cell = cached_full_grid_cell(h, w)
+    density = nfae.renderer(phi, [coord[None].to(device)], [cell[None].to(device)])[0][0, 0]
+topology = density.float().cpu().numpy() > 0.5
+```
+
+Recommended sampling: 20 Euler steps, `guidance_scale=1.0` (CFG off); 5–10 steps are
+equally good zero-shot.
+
+## Citation
+
+```bibtex
+@inproceedings{nobari2025oat,
+  title     = {Optimize Any Topology: A Foundation Model for Shape- and Resolution-Free Structural Topology Optimization},
+  author    = {Heyrani Nobari, Amin and Regenwetter, Lyle and Picard, Cyril and Han, Ligong and Ahmed, Faez},
+  booktitle = {Advances in Neural Information Processing Systems (NeurIPS)},
+  year      = {2025},
+  note      = {arXiv:2510.23667}
+}
+```