# RFD4-Proteína

RFD4-Proteína 是由 Baker Lab 和 NVIDIA 推出的用于生物分子设计的流匹配模型。它在分子上下文的条件下，同时生成蛋白质序列和全原子结构，以保留结合靶标、酶基序、配体、热点或二级结构。该模型通过两个 [RFD4ProteinaScheduler](/docs/diffusers/v0.43.0/zh/api/schedulers/rfd4_proteina#diffusers.RFD4ProteinaScheduler) 整合了配对的 Transformer 和自编码器，一个用于骨架，另一个用于局部潜变量。

该 pipeline 会消费已准备好的特征张量。将带注释的 CIF 文件转换为这些张量，以及将生成的张量转换回 CIF 文件，属于领域特定的预处理，位于 [`rfd4-proteina-prep`](https://huggingface.co/nvidia/RFD4-Proteina-prep) 包中。

## 使用示例

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

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

```py
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`，并为每个设计传入一个不同的生成器。输出共有 `batch * num_designs_per_prompt` 行，其中第 `r` 行是输入行 `r // num_designs_per_prompt` 的第 `r % num_designs_per_prompt` 个设计。

  ```py
  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`：

  ```py
  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` 加载适配器：

```py
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]]

#### diffusers.RFD4ProteinaPipeline[[diffusers.RFD4ProteinaPipeline]]

```python
diffusers.RFD4ProteinaPipeline(transformer: RFD4ProteinaTransformer, vae: RFD4ProteinaAutoencoder, bb_ca_scheduler: RFD4ProteinaScheduler, local_latents_scheduler: RFD4ProteinaScheduler, recycle_depths: list[int] | None = None)
```

[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/pipelines/rfd4_proteina/pipeline_rfd4_proteina.py#L68)

**参数：**

transformer ([RFD4ProteinaTransformer](/docs/diffusers/v0.43.0/zh/api/models/rfd4_proteina_transformer#diffusers.RFD4ProteinaTransformer))：成对的骨干网络与局部潜变量速度去噪器。

vae ([RFD4ProteinaAutoencoder](/docs/diffusers/v0.43.0/zh/api/models/rfd4_proteina_autoencoder#diffusers.RFD4ProteinaAutoencoder))：与去噪器检查点配对的自编码器。

bb_ca_scheduler ([RFD4ProteinaScheduler](/docs/diffusers/v0.43.0/zh/api/schedulers/rfd4_proteina#diffusers.RFD4ProteinaScheduler))：整合骨架坐标。

local_latents_scheduler ([RFD4ProteinaScheduler](/docs/diffusers/v0.43.0/zh/api/schedulers/rfd4_proteina#diffusers.RFD4ProteinaScheduler))：整合局部潜变量。

recycle_depths (`list[int]`, *可选*) : 每个去噪步骤的 Trunk 回收次数。其长度即为默认的推理步数。默认值为发布的 60 步、180 次回收的计划。

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

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

#### __调用__[[diffusers.RFD4ProteinaPipeline.__call__]]

```python
__call__(padding_mask: Tensor, is_in_context_conditioning: Tensor, supports_local_latent: Tensor, token_type: typing.Optional[torch.Tensor] = None, atom37_mask: typing.Optional[torch.Tensor] = None, coords_nm: typing.Optional[torch.Tensor] = None, conditioning_features: dict[str, typing.Any] | None = None, num_designs_per_prompt: int = 1, generator: typing.Union[torch.Generator, list[torch.Generator], NoneType] = None, partial_t_bb_ca: float | None = None, partial_t_local_latents: float | None = None, inverse_folding: bool = False, latents_bb_ca: typing.Optional[torch.Tensor] = None, latents_local_latents: typing.Optional[torch.Tensor] = None, attention_kwargs: dict[str, typing.Any] | None = None, callback_on_step_end: collections.abc.Callable[[diffusers.pipelines.pipeline_utils.DiffusionPipeline, int, torch.Tensor, dict], dict | None] | None = None, callback_on_step_end_tensor_inputs: list[str] = ['bb_ca', 'local_latents'], output_type: str = 'pt', return_dict: bool = True)
```

[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/pipelines/rfd4_proteina/pipeline_rfd4_proteina.py#L287)

**参数：**

padding_mask (`torch.Tensor`) : 布尔型 `(batch, tokens)`，标记哪些 token 是真实 token 而非填充。

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 坐标，供 partial diffusion 和 inverse folding 使用。

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。

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

示例：
```py
>>> 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[[diffusers.RFD4ProteinaPipeline.check_inputs]]

```python
check_inputs(padding_mask, is_in_context_conditioning, supports_local_latent, atom37_mask, coords_nm, conditioning_features, num_designs_per_prompt, generator, partial_t, inverse_folding, latents, callback_on_step_end_tensor_inputs, output_type)
```

[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/pipelines/rfd4_proteina/pipeline_rfd4_proteina.py#L130)

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

#### noise_backbone[[diffusers.RFD4ProteinaPipeline.noise_backbone]]

```python
noise_backbone(target_backbone, bb_mask, conditioning, generators, index)
```

[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/pipelines/rfd4_proteina/pipeline_rfd4_proteina.py#L278)

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

#### prepare_latents[[diffusers.RFD4ProteinaPipeline.prepare_latents]]

```python
prepare_latents(bb_mask: Tensor, latent_mask: Tensor, conditioning: dict[str, torch.Tensor], generators: list[torch.Generator] | None, clean_bb: typing.Optional[torch.Tensor] = None, clean_latents: typing.Optional[torch.Tensor] = None, start_times: tuple[float, float] | None = None)
```

[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/pipelines/rfd4_proteina/pipeline_rfd4_proteina.py#L243)

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

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

## RFD4ProteinaPipelineOutput[[diffusers.pipelines.rfd4_proteina.pipeline_output.RFD4ProteinaPipelineOutput]]

#### diffusers.pipelines.rfd4_proteina.pipeline_output.RFD4ProteinaPipelineOutput[[diffusers.pipelines.rfd4_proteina.pipeline_output.RFD4ProteinaPipelineOutput]]

```python
diffusers.pipelines.rfd4_proteina.pipeline_output.RFD4ProteinaPipelineOutput(coords_nm: Tensor, residue_type: Tensor, emit_mask: Tensor, atom_mask: Tensor, bb_ca: Tensor, local_latents: Tensor, clean_bb_ca: Tensor, clean_local_latents: Tensor, sequence_logits: Tensor, insertion_logits: Tensor)
```

[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/pipelines/rfd4_proteina/pipeline_output.py#L12)

**参数：**

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)`。

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