RFD4-Proteína

RFD4-Proteína 是由 Baker Lab 和 NVIDIA 推出的用于生物分子设计的流匹配模型。它根据分子上下文条件，同时生成蛋白质序列和全原子结构，以保留结合靶标、酶基序、配体、热点或二级结构。一个成对的 Transformer 和自编码器通过两个 RFD4ProteinaScheduler 进行整合，一个用于主链，另一个用于局部潜在变量。

该 pipeline 会消费已准备好的特征张量。将带注释的 CIF 文件转换为这些张量，以及将生成的张量转换回 CIF 文件，属于领域特定的预处理，位于 rfd4-proteina-prep 包中。

使用示例

在固定版本上安装准备包。该包决定了模型所见的特征张量，因此如果未固定版本，当包更新时，生成的结构将会发生变化。

pip install diffusers git+https://huggingface.co/nvidia/RFD4-Proteina-prep@c9959efdebd75ac7fe9129fbbf0e291a6e0487d5

import torch
from huggingface_hub import hf_hub_download
from diffusers import RFD4ProteinaPipeline
from rfd4_proteina_prep import prepare_cif, write_cif

pipe = RFD4ProteinaPipeline.from_pretrained("nvidia/RFD4-Proteina-Diffusers", dtype=torch.bfloat16).to("cuda")

# An annotated CIF specifies what to hold fixed and what to design: here a PD-L1 target and a 50-120 residue binder.
input_cif = hf_hub_download("nvidia/RFD4-Proteina-Diffusers", "examples/protein_binder/input.cif")
features, atoms = prepare_cif(input_cif, seed=42, device=pipe.device)

# The masks and molecule tensors the pipeline always needs are arguments; the annotation tensors go in one mapping.
required = ["padding_mask", "is_in_context_conditioning", "supports_local_latent", "token_type", "atom37_mask"]
conditioning = {name: value for name, value in features.items() if name not in required}

output = pipe(
    padding_mask=features["padding_mask"],
    is_in_context_conditioning=features["is_in_context_conditioning"],
    supports_local_latent=features["supports_local_latent"],
    token_type=features["token_type"],
    atom37_mask=features["atom37_mask"],
    conditioning_features=conditioning,
    generator=torch.Generator("cuda").manual_seed(42),
)
write_cif(output, features, atoms, "design.cif")

提示

模型仓库在 examples/ 目录下为每个任务提供了带注释的输入和生成的设计：蛋白结合剂、抗体、VHH、酶、基序支架、小分子结合剂、无条件生成、逆折叠和局部扩散。
将 partial_t_bb_ca 和 partial_t_local_latents 的起始时间设为 [0, 1) 内的值，可从带噪声的起始状态重新设计输入；或者将 inverse_folding=True 与两个起始时间均设为 0.0，以保留骨架并设计序列。这两种方式都会对输入进行编码，因此需要干净的 coords_nm，并会拒绝非有限坐标，这些坐标通常对应于未解析的残基。
步数等于 recycle_depths 流水线配置的长度，已发布的配方为 60（180 次主干遍历）。在构建流水线时传入不同的 recycle_depths 列表即可更改该计划。

要为每个输入生成多个设计，请设置 num_designs_per_prompt，并为每个设计传入一个独立的 generator。输出共有 batch * num_designs_per_prompt 行，其中第 r 行是输入行 r // num_designs_per_prompt 的第 r % num_designs_per_prompt 个设计。

generators = [torch.Generator("cuda").manual_seed(seed) for seed in range(8)]
output = pipe(
    padding_mask=features["padding_mask"],
    is_in_context_conditioning=features["is_in_context_conditioning"],
    supports_local_latent=features["supports_local_latent"],
    token_type=features["token_type"],
    atom37_mask=features["atom37_mask"],
    conditioning_features=conditioning,
    num_designs_per_prompt=8,
    generator=generators,
)

write_cif 会读取所写入行的特征，因此请以相同方式重复它们，使用 repeat_interleave 而非 repeat：

repeated = {k: v.repeat_interleave(8, dim=0) if torch.is_tensor(v) else v for k, v in features.items()}
for row in range(8):
    write_cif(output, repeated, atoms, f"design_{row}.cif", row=row)

以 torch.bfloat16 加载以提升速度。采样器状态和特征嵌入保持为 float32。使用不同的 dtype、kernel 或 batch size，即使 seed 相同，也可能导致不同的设计。
生成的结构并非经过验证的设计。请针对预期任务对其进行重新折叠和评分。

LoRA

发布的模型是对预训练基座模型进行的秩为 32 的 LoRA 微调，并已合并到权重中。基座模型和适配器位于模型仓库的 pretrain-base 修订版本中，流水线通过 RFD4ProteinaLoraLoaderMixin 使用 load_lora_weights 加载适配器：

pipe = RFD4ProteinaPipeline.from_pretrained(
    "nvidia/RFD4-Proteina-Diffusers", revision="pretrain-base", dtype=torch.bfloat16
).to("cuda")
pipe.load_lora_weights(
    "nvidia/RFD4-Proteina-Diffusers",
    revision="pretrain-base",
    subfolder="lora",
    weight_name="pytorch_lora_weights.safetensors",
)

RFD4ProteinaPipeline

类 diffusers.RFD4ProteinaPipeline

< source >

( transformer: RFD4ProteinaTransformervae: RFD4ProteinaAutoencoderbb_ca_scheduler: RFD4ProteinaSchedulerlocal_latents_scheduler: RFD4ProteinaSchedulerrecycle_depths: list[int] | None = None )

参数

transformer (RFD4ProteinaTransformer) — 成对的骨干网络与局部潜变量速度去噪器。
vae (RFD4ProteinaAutoencoder) — 与去噪器检查点配对的自编码器。
bb_ca_scheduler（RFD4ProteinaScheduler）— 整合骨架坐标。
local_latents_scheduler (RFD4ProteinaScheduler) — 整合局部潜变量。
recycle_depths (list[int], 可选) — 每个去噪步骤的 Trunk 回收次数。其长度即为默认的推理步数。默认值为发布的 60 步、180 次回收的调度方案。

从准备好的分子张量生成分子结构和序列。

该 pipeline 会消费由 RFD4 输入准备阶段生成的特征张量（有关带注释的 CIF 文件，请参阅 rfd4-proteina-prep 包），并返回解码后的 atom37 坐标和残基标识。坐标单位为纳米。

call

< source >

( padding_mask: Tensoris_in_context_conditioning: Tensorsupports_local_latent: Tensortoken_type: typing.Optional[torch.Tensor] = Noneatom37_mask: typing.Optional[torch.Tensor] = Nonecoords_nm: typing.Optional[torch.Tensor] = Noneconditioning_features: dict[str, typing.Any] | None = Nonenum_designs_per_prompt: int = 1generator: typing.Union[torch.Generator, list[torch.Generator], NoneType] = Nonepartial_t_bb_ca: float | None = Nonepartial_t_local_latents: float | None = Noneinverse_folding: bool = Falselatents_bb_ca: typing.Optional[torch.Tensor] = Nonelatents_local_latents: typing.Optional[torch.Tensor] = Noneattention_kwargs: dict[str, typing.Any] | None = Nonecallback_on_step_end: collections.abc.Callable[[diffusers.pipelines.pipeline_utils.DiffusionPipeline, int, torch.Tensor, dict], dict | None] | None = Nonecallback_on_step_end_tensor_inputs: list[str] = ['bb_ca', 'local_latents']output_type: str = 'pt'return_dict: bool = True ) → RFD4ProteinaPipelineOutput or tuple

参数

padding_mask (torch.Tensor) — 布尔型 (batch, tokens)，在 token 为实际内容而非填充的位置为 true。
is_in_context_conditioning (torch.Tensor) — 布尔型 (batch, tokens)，对于作为上下文提供且保持不变的 token 为 true。
supports_local_latent (torch.Tensor) — 布尔型 (batch, tokens)，对于携带局部潜在变量的 token 为 true。这三个掩码共享相同的形状和设备。
token_type (torch.Tensor，可选) — (batch, tokens) 的 token 类型编码。
atom37_mask (torch.Tensor, 可选) — (batch, tokens, 37)，在原子槽位被使用时为 true。
coords_nm (torch.Tensor, 可选) — 以纳米为单位的干净 (batch, tokens, 37, 3) atom37 坐标，供部分扩散和反向折叠使用。
conditioning_features (dict[str, Any], 可选) — 预处理器输出的键所对应的其余拓扑和条件张量。它必须包含序列条件 condition_sequence_1_residue 以及 molecule_iid。其余哪些字段会被读取，由检查点的特征配置决定。之所以采用单一映射结构，是因为预处理器会以固定的名称输出约 30 个可选张量；如果签名中包含如此多的参数，会将流水线与预处理器的命名约定耦合在一起。
num_designs_per_prompt (int，默认为 1) — 为每个输入行生成的设计数量。输入会沿批次轴重复，因此输出共有 batch * num_designs_per_prompt 行，其中第 r 行是输入行 r // num_designs_per_prompt 的第 r % num_designs_per_prompt 个设计。
generator (torch.Generator 或 list[torch.Generator]，可选) — 每个输出行对应一个独立的生成器，以实现可复现的设计。这些生成器会向前推进，而全局 PyTorch 随机流保持不变。如果省略此参数，则从全局随机流中采样。
partial_t_bb_ca (float, 可选) — 主干网络的起始时间，取值范围为 [0, 1)，用于从输入的加噪版本而非纯噪声开始生成。需与 partial_t_local_latents 一起使用；要求提供 coords_nm 和 atom37_mask。
partial_t_local_latents (float, 可选) — 局部潜变量在 [0, 1) 范围内的起始时间；参见 partial_t_bb_ca。
inverse_folding (bool，默认为 False) — 保留输入的骨架并设计序列。要求两个部分起始时间均为零。
latents_bb_ca (torch.Tensor, 可选) — 预生成的起始主干状态 (batch * num_designs_per_prompt, tokens, 3)，用于替代本次调用中原本会采样的状态。调度保持不变。请将其与 latents_local_latents 一起传入。
latents_local_latents (torch.Tensor, 可选) — 预生成的起始局部潜变量 (batch * num_designs_per_prompt, tokens, latent_dim)。
attention_kwargs (dict, 可选) — 传递给 transformer；scale 用于设置 LoRA 的缩放系数。
callback_on_step_end (Callable，可选) — 在每一步之后，以 callback_on_step_end(pipeline, step, timestep, callback_kwargs) 的形式调用。 callback_kwargs 包含 callback_on_step_end_tensor_inputs 中指定的张量；返回的字典可以替换 "bb_ca" 和 "local_latents"。
output_type (str，默认为 "pt") — "pt" 返回 torch 张量，"np" 返回 numpy 数组。
callback_on_step_end_tensor_inputs（list[str]，默认为 ["bb_ca", "local_latents"]）— 传递给回调的张量："bb_ca"、"local_latents"、"last_clean_bb_ca"、 "last_clean_local_latents"、"sequence_logits"、"insertion_logits"。
return_dict (bool，默认为 True) — 返回带类型的输出，或返回其非 None 字段的元组。

返回值

RFD4ProteinaPipelineOutput 或 tuple

解码后的 atom37 张量（coords_nm 以纳米为单位，residue_type、emit_mask、atom_mask）、最终的含噪状态、最后一次干净估计，以及最后的序列和插入 logits。

从准备好的分子张量生成结构和序列。

示例：

>>> import torch
>>> from diffusers import RFD4ProteinaPipeline
>>> from rfd4_proteina_prep import prepare_cif, write_cif

>>> pipe = RFD4ProteinaPipeline.from_pretrained("nvidia/RFD4-Proteina-Diffusers", dtype=torch.bfloat16)
>>> pipe.to("cuda")

>>> features, atoms = prepare_cif("input.cif", seed=42, device=pipe.device)
>>> required = ["padding_mask", "is_in_context_conditioning", "supports_local_latent", "token_type"]
>>> required += ["atom37_mask"]
>>> output = pipe(
...     padding_mask=features["padding_mask"],
...     is_in_context_conditioning=features["is_in_context_conditioning"],
...     supports_local_latent=features["supports_local_latent"],
...     token_type=features["token_type"],
...     atom37_mask=features["atom37_mask"],
...     conditioning_features={name: value for name, value in features.items() if name not in required},
...     generator=torch.Generator("cuda").manual_seed(42),
... )
>>> write_cif(output, features, atoms, "design.cif")

check_inputs

< source >

( padding_maskis_in_context_conditioningsupports_local_latentatom37_maskcoords_nmconditioning_featuresnum_designs_per_promptgeneratorpartial_tinverse_foldinglatentscallback_on_step_end_tensor_inputsoutput_type )

验证分子掩码、按设计划分的生成器、部分扩散以及回调。

noise_backbone

< source >

( target_backbonebb_maskconditioninggeneratorsindex )

将保留的主干加噪至 index 步的预定时间，用于反向折叠。

prepare_latents

< source >

( bb_mask: Tensorlatent_mask: Tensorconditioning: dict[str, torch.Tensor]generators: list[torch.Generator] | Noneclean_bb: typing.Optional[torch.Tensor] = Noneclean_latents: typing.Optional[torch.Tensor] = Nonestart_times: tuple[float, float] | None = None )

绘制初始状态：纯噪声，或加噪至 start_times 的干净输入（backbone、latents）。当 latent 起始时间为零时，可以省略 clean_latents。

每次采样均遵循原生采样器在每个行生成器上的顺序：对于完整生成，骨干噪声及其随机偏移量先于潜在噪声出现；而对于部分扩散，则先采样两种噪声，再采样偏移量。

RFD4ProteinaPipelineOutput

class diffusers.pipelines.rfd4_proteina.pipeline_output.RFD4ProteinaPipelineOutput

< source >

( coords_nm: Tensorresidue_type: Tensoremit_mask: Tensoratom_mask: Tensorbb_ca: Tensorlocal_latents: Tensorclean_bb_ca: Tensorclean_local_latents: Tensorsequence_logits: Tensorinsertion_logits: Tensor )

参数

coords_nm (torch.Tensor) — 解码后的 atom37 坐标 (B, N, 37, 3)。
residue_type (torch.Tensor) — 残基词汇表索引 (B, N)；未输出残基的位置为 -1。
emit_mask (torch.Tensor) — 布尔掩码 (B, N)，指示哪些 token 已被设计。
atom_mask (torch.Tensor) — 有效原子的布尔掩码，形状为 (B, N, 37)。
bb_ca (torch.Tensor) — 以纳米为单位的最终主干状态 (B, N, 3)。
local_latents (torch.Tensor) — 最终的局部潜在状态 (B, N, D)。
clean_bb_ca (torch.Tensor) — 最后一步去噪时干净状态的骨干网络预测值。
clean_local_latents (torch.Tensor) — 最后一步去噪时干净状态的局部潜变量预测。
sequence_logits (torch.Tensor) — 最终的逐 token 残差 logits (B, N, 147)。
insertion_logits (torch.Tensor) — 针对未索引基序的最终插入预测 (B, N, N)。

分子推理结果。坐标单位为纳米。

Update on GitHub

Diffusers

RFD4-Proteína

使用示例

提示

LoRA

RFD4ProteinaPipeline

类 diffusers.RFD4ProteinaPipeline

__call__

check_inputs

noise_backbone

prepare_latents

RFD4ProteinaPipelineOutput

class diffusers.pipelines.rfd4_proteina.pipeline_output.RFD4ProteinaPipelineOutput

call