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

概述

CryoFM2 是一个基于流模型的冷冻电镜密度图生成式基础模型。它在经过筛选的 EMDB 半图数据集上进行预训练，以学习高质量冷冻电镜密度图的通用先验知识，并可针对下游任务进行微调。

该模型学习从简单的高斯分布到复杂的冷冻电镜密度图分布的连续映射，实现了稳定的生成能力和灵活的适应能力。CryoFM2 还可作为贝叶斯先验，与任务特定的似然函数自然融合，支持各向异性感知优化、非均匀重建和可控密度修饰等应用。

模型详情

CryoFM2 在经过筛选的 EMDB 半图数据集上进行预训练，以学习高质量冷冻电镜密度图的通用先验知识。该模型可针对多种下游任务进行微调，如密度图增强和后处理。

预训练架构：

微调架构（用于 EMhancer/EMReady 风格的后处理）：

架构

架构类型：3D UNet
输入尺寸：64×64×64 体素
输入通道数：预训练模型为 2，微调模型为 3
输出通道数：1
下采样模块：DownBlock3D、DownBlock3D、AttnDownBlock3D、AttnDownBlock3D
上采样模块：AttnUpBlock3D、AttnUpBlock3D、UpBlock3D、UpBlock3D
模块输出通道数：(64, 128, 256, 512)
每个模块的层数：2
注意力头维度：8
归一化：GroupNorm（32 组）
激活函数：SiLU
时间嵌入：位置编码

模型变体

cryofm2-pretrain：用于通用密度图生成的无条件预训练模型
cryofm2-emhancer：用于密度图增强的微调模型（EMhancer 风格）
cryofm2-emready：用于密度图增强的微调模型（EMReady 风格）

体验 CryoFM2

安装

使用 CryoFM2 之前，您需要配置环境并安装软件包。按照以下步骤开始：

# 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)

微调模型（EMhancer/EMReady）：

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)

密度图修改

CryoFM2 支持使用预训练模型作为贝叶斯先验进行多种密度图修改操作。支持的操作包括：

去噪（denoise）：去除密度图中的噪声
修复（inpaint）：填充缺失区域（例如，缺失楔）
去噪修复（denoise inpaint）：结合去噪和修复功能
非均匀权重（non-uniform weight）：在重建过程中应用非均匀加权

基本用法：

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

对于修复任务，您需要提供一个 RELION starfile 路径：

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

密度图后处理

CryoFM2 提供了针对不同风格密度图增强的微调模型，类似于 EMhancer 和 EMReady。

EMhancer 风格增强

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

EMReady 样式增强

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

参数：

-i：输入密度图文件（MRC 格式）
-o：输出目录
--model-dir：包含 config.yaml 和 model.safetensors 的模型目录路径
--output-tag：风格标签（1 表示 EMhancer，0 表示 EMReady）
--cfg-weight：无分类器引导权重（可选，默认值因模型而异）

性能提示

多 GPU 推理：使用 accelerate launch 在多 GPU 上实现更快推理：

NCCL_DEBUG=ERROR accelerate launch --num_processes=${NUM_GPUS} --main_process_port=8881 \
    python -m cryofm.projects.cryofm2.cond_sampling ...

混合精度：在支持时使用 --bf16 标志以减少内存占用并加快推理速度。
批量处理：根据您的 GPU 内存容量调整批处理大小。

局限性

输入尺寸固定为 64×64×64 体素
模型性能可能因输入密度图质量而异
微调模型针对特定增强风格进行了优化

伦理考量

本模型旨在用于科学研究和结构生物学应用。用户应：

使用生成的结构时确保适当的引用
通过实验验证来确认生成结构的有效性
意识到训练数据中可能存在的偏差
负责任地使用模型，并遵循科学最佳实践

引用

如果您发现 CryoFM2 有用，请引用：

@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}
}

许可协议

本模型基于 Apache 2.0 许可协议发布。详情请参见 LICENSE 文件。

致谢

本研究成果由字节跳动 Seed 团队开发。如需更多信息，请访问：