Add CryoFM model weights and configurations

- Add CryoFM-S and CryoFM-L model variants
- Include model configs and safetensors checkpoints
- Add README with model description and usage examples

.gitattributes

@@ -33,3 +33,5 @@ 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/cryofm_archs.jpg filter=lfs diff=lfs merge=lfs -text
+assets/cryofm.gif filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,3 +1,165 @@
----
-license: apache-2.0
----
+---
+license: apache-2.0
+tags:
+- cryo-em
+- flow-matching
+- 3d-density-maps
+- foundation-model
+---
+
+# CryoFM: Flow-based Foundation Model for Cryo-EM Density Maps
+
+<div align="center">
+
+[![arXiv](https://img.shields.io/badge/arXiv-2410.08631-B31B1B?logo=arxiv&logoColor=white)](https://arxiv.org/abs/2410.08631)
+[![GitHub](https://img.shields.io/badge/GitHub-cryofm-181717?logo=github&logoColor=white)](https://github.com/ByteDance-Seed/cryofm)
+[![License](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)
+
+</div>
+
+<div align="center">
+  <img src="./assets/cryofm.gif" alt="CryoFM Demo" style="max-width: 100%; height: auto; width: 800px;"/>
+</div>
+
+## Model Description
+
+CryoFM1 is a flow-based foundation model for 3D cryo-electron microscopy (cryo-EM) density maps. The model employs a Hierarchical Diffusion Transformer (HDiT) architecture, specifically designed to learn deep priors of 3D cryo-EM densities. CryoFM1 supports various downstream tasks including density map denoising, anisotropy noise correction, missing wedge inpainting, and *ab initio* modeling.
+
+### Key Features
+
+- **Flow Matching Framework**: Uses flow matching for efficient and stable training
+- **HDiT Architecture**: Hierarchical Diffusion Transformer with local and global attention mechanisms
+- **Two Model Variants**: CryoFM-S (64³) and CryoFM-L (128³) for different resolution needs
+- **Downstream Task Support**: Denoising, anisotropy noise correction, missing wedge restoration, and more
+
+## Model Details
+
+CryoFM1 employs a Hierarchical Diffusion Transformer (HDiT) architecture that combines local neighborhood attention with global attention mechanisms. This design enables the model to effectively capture both fine-grained local structures and long-range dependencies in 3D cryo-EM density maps. The architecture processes 3D volumes through a hierarchical patch-based approach, progressively building representations at multiple scales.
+
+<div align="center">
+  <img src="./assets/cryofm_archs.jpg" alt="CryoFM Architecture" style="max-width: 100%; height: auto; width: 600px;"/>
+</div>
+
+The model is available in two variants optimized for different resolution requirements. The following table summarizes the key architectural and training parameters for each variant:
+
+| Parameter | CRYOFM-S | CRYOFM-L |
+|-----------|----------|----------|
+| **Parameters** | 335.18 M | 308.54 M |
+| **GFLOP/forward** | 395.87 | 427.26 |
+| **Training Steps** | 150k | 300k |
+| **Batch Size** | 128 | 128 |
+| **Precision** | bf16 | bf16 |
+| **Training Hardware** | 8×A100 | 8×A100 |
+| **Patchifying** | 4 | 4 |
+| **Levels (Local + Global Attention)** | 1 + 1 | 2 + 1 |
+| **Depth** | [4, 8] | [2, 2, 12] |
+| **Widths** | [768, 1536] | [320, 640, 1280] |
+| **Attention Heads (Width / Head Dim)** | [12, 24] | [5, 10, 20] |
+| **Attention Head Dim** | 64 | 64 |
+| **Neighborhood Kernel Size** | 7 | 7 |
+
+## Quick Start
+
+### Unconditional Generation
+
+CryoFM1 provides two model variants for different resolution needs:
+- **CryoFM-S**: Generates 64×64×64 voxel density maps at 1.5 Å/pixel resolution
+- **CryoFM-L**: Generates 128×128×128 voxel density maps at 3.0 Å/pixel resolution
+
+
+```python
+import torch
+from mmengine import Config
+from cryofm.core.utils.mrc_io import save_mrc
+from cryofm.projects.cryofm1.lit_modules import CryoFM1
+from cryofm.core.utils.sampling_fm import sample_from_fm
+
+# Choose model variant: "cryofm-s" or "cryofm-l"
+model_variant = "cryofm-s"  # or "cryofm-l"
+model_config = {
+    "cryofm-s": {
+        "config_path": "cryofm-v1/cryofm-s/config.yaml",
+        "model_path": "cryofm-v1/cryofm-s/model.safetensors",
+        "side_shape": 64,
+        "apix": 1.5
+    },
+    "cryofm-l": {
+        "config_path": "cryofm-v1/cryofm-l/config.yaml",
+        "model_path": "cryofm-v1/cryofm-l/model.safetensors",
+        "side_shape": 128,
+        "apix": 3.0
+    }
+}
+
+# Load configuration and model
+cfg = Config.fromfile(model_config[model_variant]["config_path"])
+lit_model = CryoFM1.load_from_safetensors(
+    model_config[model_variant]["model_path"], 
+    cfg=cfg
+)
+
+# Set up device and model
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+lit_model = lit_model.to(device)
+lit_model.eval()
+
+# Define vector field function for flow matching
+def v_xt_t(_xt, _t):
+    return lit_model(_xt, _t)
+
+# Generate samples
+# Note: Enable bfloat16 if your GPU supports it for better performance
+with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16):
+    out = sample_from_fm(
+        v_xt_t, 
+        lit_model.noise_scheduler, 
+        method="euler", 
+        num_steps=200, 
+        num_samples=3, 
+        device=device, 
+        side_shape=model_config[model_variant]["side_shape"]
+    )
+    # Apply z-scaling normalization if configured
+    if hasattr(lit_model.cfg, "z_scale") and lit_model.cfg.z_scale.mean is not None:
+        out = out * lit_model.cfg.z_scale.std + lit_model.cfg.z_scale.mean
+
+# Save generated density maps
+for i in range(3):
+    save_mrc(
+        out[i].float().cpu().numpy(), 
+        f"sample-{i}.mrc", 
+        apix=model_config[model_variant]["apix"]  # Angstroms per pixel
+    )
+```
+
+### Ethical Considerations
+
+This model is intended for scientific research and structural biology applications. Users should:
+- Ensure proper attribution when using generated structures
+- Validate generated structures through experimental verification
+- Be aware of potential biases in the training data
+
+## Citation
+
+If you use CryoFM1 in your research, please cite:
+
+```bibtex
+@inproceedings{
+  zhou2025cryofm,
+  title={Cryo{FM}: A Flow-based Foundation Model for Cryo-{EM} Densities},
+  author={Yi Zhou and Yilai Li and Jing Yuan and Quanquan Gu},
+  booktitle={The Thirteenth International Conference on Learning Representations},
+  year={2025},
+  url={https://openreview.net/forum?id=T4sMzjy7fO}
+}
+```
+
+## License
+
+This model is released under the Apache 2.0 License. See the [LICENSE](https://github.com/ByteDance-Seed/cryofm/blob/main/LICENSE) file for details.
+
+## Acknowledgments
+
+This work is developed by the ByteDance Seed Team. For more information, visit:
+- [Project Repository](https://github.com/ByteDance-Seed/cryofm)
+- [ByteDance Seed Team](https://seed.bytedance.com/)

cryofm-l/config.yaml

@@ -0,0 +1,46 @@
+ckpt_path: null
+ddpm:
+  prediction_type: v_prediction
+exp_name: 128-hdit_fm_scale_bf16
+hdit_model:
+  depths:
+  - 2
+  - 2
+  - 12
+  input_channels: 1
+  input_size:
+  - 128
+  - 128
+  - 128
+  patch_size:
+  - 4
+  - 4
+  - 4
+  self_attns:
+  - d_head: 64
+    kernel_size: 7
+    type: neighborhood
+  - d_head: 64
+    kernel_size: 7
+    type: neighborhood
+  - d_head: 64
+    type: global
+  type: image_transformer_v2
+  widths:
+  - 320
+  - 640
+  - 1280
+keep_last_k: null
+model_type: hdit
+num_val_samples: 3
+optimizer:
+  lr: 0.0001
+  warmup: 2000
+patch_size: 128
+process: fm
+seed: 42
+work_dir: work_dirs/128-hdit_fm_scale_bf16_00
+z_crop: null
+z_scale:
+  mean: 0.04
+  std: 0.09

cryofm-l/model.safetensors

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

cryofm-s/config.yaml

@@ -0,0 +1,42 @@
+ckpt_path: null
+ddpm:
+  prediction_type: v_prediction
+exp_name: 64-hdit_fm_scale_bf16
+hdit_model:
+  depths:
+  - 4
+  - 8
+  input_channels: 1
+  input_size:
+  - 64
+  - 64
+  - 64
+  patch_size:
+  - 4
+  - 4
+  - 4
+  self_attns:
+  - d_head: 64
+    kernel_size: 7
+    type: neighborhood
+  - d_head: 64
+    type: global
+  type: image_transformer_v2
+  widths:
+  - 768
+  - 1536
+keep_last_k: null
+mode: train
+model_type: hdit
+num_val_samples: 3
+optimizer:
+  lr: 0.0001
+  warmup: 2000
+patch_size: 64
+process: fm
+seed: 42
+work_dir: work_dirs/64-hdit_fm_scale_bf16_00
+z_crop: null
+z_scale:
+  mean: 0.04
+  std: 0.09

cryofm-s/model.safetensors

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