CryoFM2：冷冻电镜密度的生成式基础模型

cryofm-v2/
├── inference.py            # 服务化推理脚本（HTTP 服务 + 自测模式）
├── make_screenshots.py     # 截图/水印生成脚本（maggie_Ha 水印）
├── requirements.txt        # 运行环境依赖（基于本机实测版本）
├── workflow.txt            # 适配流程文本（用于生成适配过程截图）
├── result.txt              # 完整中文测试日志
├── cryofm2_npu/            # 自包含推理包（昇腾适配）
│   ├── __init__.py         #   对外导出 API
│   ├── unet3d.py           #   3D UNet 模型（复刻官方架构，权重严格加载）
│   ├── scheduler.py        #   流匹配调度器 FMScheduler
│   ├── sampler.py          #   Euler / midpoint 采样
│   ├── model.py            #   无条件/条件模型前向封装
│   └── mrc_io.py           #   MRC 文件读写（自包含，MRC2014）
├── cryofm2-pretrain/       # 无条件预训练权重（config.yaml + model.safetensors）
├── cryofm2-emhancer/       # EMhancer 风格微调权重（in_channels=3, class_embedding）
├── cryofm2-emready/        # EMReady 风格微调权重（in_channels=3, class_embedding）
├── outputs/                # 推理输出
│   ├── cryofm2-pretrain_selftest.mrc  # 自测生成的 64³ 密度图
│   ├── selftest.log        # 完整中文自测日志
│   └── server.log          # HTTP 服务日志
└── assets/                 # 部署截图（全部带 maggie_Ha 水印）
    ├── npu_device_call.png #   NPU 设备调用截图
    ├── model_result.png    #   测试结果截图
    ├── agent_workflow.png  #   适配过程截图
    ├── cryofm2_overview.jpg
    ├── cryofm2_arch-pretrain.jpg
    └── cryofm2_arch-finetune.jpg

source /usr/local/Ascend/ascend-toolkit/set_env.sh
pip install -r requirements.txt

# 默认变体 cryofm2-pretrain，端口 8090
python inference.py --variant cryofm2-pretrain --port 8090

# 条件模型（EMhancer / EMReady）
python inference.py --variant cryofm2-emhancer --output-tag 1 --port 8091
python inference.py --variant cryofm2-emready  --output-tag 0 --port 8092

# 健康检查
curl http://127.0.0.1:8090/health

# 生成 1 个 64³ 密度图（50 步 Euler）
curl -X POST http://127.0.0.1:8090/generate \
     -H "Content-Type: application/json" \
     -d '{"num_samples":1,"num_steps":50,"method":"euler","side_shape":64}'

python inference.py --variant cryofm2-pretrain --selftest --num-steps 50

python make_screenshots.py \
    --npu-info npu_info.txt \
    --result result.txt \
    --workflow workflow.txt

# Clone the repository
git clone https://github.com/ByteDance-Seed/cryofm.git
cd cryofm

# Create a new conda environment for CryoFM (recommended)
conda create -n cryofm python=3.10 -y
conda activate cryofm

# Install CryoFM
pip install .

import torch
from mmengine import Config

from cryofm.core.utils.mrc_io import save_mrc
from cryofm.core.utils.sampling_fm import sample_from_fm
from cryofm.projects.cryofm2.lit_modules import CryoFM2Uncond

# Update the path to your model directory
model_dir = "path/to/cryofm-v2/cryofm2-pretrain"
cfg = Config.fromfile(f"{model_dir}/config.yaml")
lit_model = CryoFM2Uncond.load_from_safetensors(f"{model_dir}/model.safetensors", cfg=cfg)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

lit_model = lit_model.to(device)
lit_model.eval()
def v_xt_t(_xt, _t):
    return lit_model(_xt, _t)

# Enable bfloat16 for faster inference if your GPU supports it
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=lit_model.device, 
        side_shape=64
    )
    # Apply 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 samples
for i in range(3):
    save_mrc(out[i].float().cpu().numpy(), f"sample-{i}.mrc", voxel_size=1.5)

import torch
from mmengine import Config

from cryofm.core.utils.mrc_io import save_mrc
from cryofm.core.utils.sampling_fm import sample_from_fm
from cryofm.projects.cryofm2.lit_modules import CryoFM2Cond

# Choose style: "emhancer" or "emready"
style = "emhancer"
model_dir = f"path/to/cryofm-v2/cryofm2-{style}"
cfg = Config.fromfile(f"{model_dir}/config.yaml")
lit_model = CryoFM2Cond.load_from_safetensors(f"{model_dir}/model.safetensors", cfg=cfg)
output_tag = 1 if style == "emhancer" else 0

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

lit_model = lit_model.to(device)
lit_model.eval()
def v_xt_t(_xt, _t):
    bs = _xt.shape[0]
    unconditional_generation_conds = {
        "input_cond": None,
        "output_cond": torch.tensor([output_tag] * bs).to(device),
        "vol_cond": None,  # dimension should be [bs, d, h, w]
    }
    return lit_model(_xt, _t, generation_conds=unconditional_generation_conds)

# Enable bfloat16 for faster inference if your GPU supports it
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=lit_model.device, 
        side_shape=64
    )
    # Apply 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 samples
for i in range(3):
    save_mrc(out[i].float().cpu().numpy(), f"{style}-sample-{i}.mrc", voxel_size=1.5)

python -m cryofm.projects.cryofm2.uncond_sampling \
    -i1 half_map_1.mrc \
    -i2 half_map_2.mrc \
    -o ./output \
    --model-dir path/to/cryofm-v2/cryofm2-pretrain \
    --op denoise \
    --norm-grad \
    --use-lamb-w

python -m cryofm.projects.cryofm2.uncond_sampling \
    -i1 half_map_1.mrc \
    -i2 half_map_2.mrc \
    -o ./output \
    --model-dir path/to/cryofm-v2/cryofm2-pretrain \
    --op inpaint \
    --data-starfile-path path/to/relion_data.star \
    --norm-grad \
    --use-lamb-w

python -m cryofm.projects.cryofm2.cond_sampling \
    -i input_map.mrc \
    -o ./output_emhancer \
    --model-dir path/to/cryofm-v2/cryofm2-emhancer \
    --output-tag 1

python -m cryofm.projects.cryofm2.cond_sampling \
    -i input_map.mrc \
    -o ./output_emready \
    --model-dir path/to/cryofm-v2/cryofm2-emready \
    --output-tag 0 \
    --cfg-weight 0.5

@article{
Li2025.12.29.696802,
author={Li, Yilai and Yuan, Jing and Zhou, Yi and Wang, Zhenghua and Chen, Suyi and Yang, Fengyu and Ling, Haibin and Kovalsky, Shahar Z and Zheng, Xiaoqing and Gu, Quanquan},
title={A Generative Foundation Model for Cryo-EM Densities},
elocation-id={2025.12.29.696802},
year={2025},
doi={10.64898/2025.12.29.696802},
publisher={Cold Spring Harbor Laboratory},
URL={https://www.biorxiv.org/content/early/2025/12/29/2025.12.29.696802},
eprint={https://www.biorxiv.org/content/early/2025/12/29/2025.12.29.696802.full.pdf},
journal={bioRxiv}
}