Diffusers documentation
RFD4ProteinaTransformer
RFD4ProteinaTransformer
RFD4-Proteína 的去噪器。它根据准备好的分子张量，预测采样器所集成的两种模态的速度——骨架 CA 坐标以及逐 token 的局部潜变量。输入为命名字张量以及一个条件映射，其内容由检查点配置决定；结构解析不在此类中处理。
from diffusers import RFD4ProteinaTransformer
transformer = RFD4ProteinaTransformer.from_pretrained(
"nvidia/RFD4-Proteina-Diffusers", subfolder="transformer"
)RFD4ProteinaTransformer
class diffusers.RFD4ProteinaTransformer
< source >( nlayers: int = 14token_dim: int = 768nheads: int = 12dim_cond: int = 256pair_repr_dim: int = 384latent_dim: int = 32token_type_emb_dim: int = 32cath_emb_dim: int = 32t_emb_dim: int = 256seq_sep_dim: int = 64xt_pair_dist_dim: int = 30xt_pair_dist_min: float = 0.1xt_pair_dist_max: float = 3.0feats_seq: tuple[str, ...] = ('xt_bb_ca', 'xt_local_latents', 'token_type', 'sequence_probabilities', 'cond_a37coors_nm', 'cond_a37coors_internal_nm', 'rasa_binned', 'secondary_structure', 'hbond_modernized', 'pistack', 'pication', 'saltbridge', 'hydrophobic', 'metal_coord', 'na_ss_state', 'hotspot', 'atom_charge', 'atomic_element', 'coord_group_seq', 'chirality', 'chain_centroid_xyz')feats_cond_seq: tuple[str, ...] = ('time_emb_bb_ca', 'time_emb_local_latents', 'cath_code', 'recycle_remaining')feats_pair_repr: tuple[str, ...] = ('xt_bb_ca_pair_dists', 'bond_mask', 'bond_order', 'is_aromatic', 'index_group_pair', 'molecule_idx_pair', 'molecule_entity_pair', 'token_bond_path_length', 'canonical_polymer_ordering', 'na_ss_pair')feats_pair_cond: tuple[str, ...] = ('time_emb_bb_ca', 'time_emb_local_latents')update_pair_repr_every_n: int = 2tri_mult_c: int = 256num_sequence_classes: int = 147num_token_types: int = 7recycle_conditioning_dim: int = 1 )
参数
- nlayers (
int，默认为14) — 主干层数。 - token_dim (
int，默认为768) — Token 表示的宽度。 - nheads（
int，默认为12）— 每层的注意力头数。 - dim_cond (
int，默认为256) — 条件向量（时间、回收索引、CATH 标签）的宽度。 - pair_repr_dim (
int，默认为384) — 配对表示的宽度。 - latent_dim (
int，默认为32) — 每个 token 的局部隐向量的宽度。 - token_type_emb_dim (
int，默认为32) — Token 类型嵌入的宽度。 - cath_emb_dim（
int，默认为32）— 每个 CATH 类别/架构/拓扑嵌入的宽度。 - t_emb_dim（
int，默认为256）— 正弦时间嵌入的宽度。 - seq_sep_dim（
int，默认为64）— 键路径长度分箱的数量。 - xt_pair_dist_dim（
int，默认为30）— 主干网络成对距离分箱的数量。 - xt_pair_dist_min (
float，默认为0.1)，xt_pair_dist_max (float，默认为3.0) — 纳米为单位的成对距离范围。 - feats_seq, feats_cond_seq, feats_pair_repr, feats_pair_cond (
list[str]) — 投影到 token、conditioning、pair 和 pair-conditioning 输入中的特征名称。 - update_pair_repr_every_n (
int，默认为2) — 每隔多少个层，使用层间 token 更新一次 pair 表示。 - tri_mult_c（
int，默认值为256）— 在这些更新中，三角形乘法的隐藏宽度。 - num_token_types (
int，默认为7) — 标记类型词汇表的大小：蛋白质、RNA、DNA、原子化、未知、聚糖和非标准氨基酸。 - recycle_conditioning_dim（
int，默认为1）— 当feats_cond_seq中包含recycle_remaining特征时，该 one-hot 的宽度。 - num_sequence_classes (
int，默认为147) — 序列读出所预测的残基词汇表大小。
RFD4 速度去噪器，采用配对的骨干网络和本地潜在模态。
模型在其自身的 dtype 下运行：加载时使用 dtype=torch.bfloat16 以进行低精度推理。传递给 forward 的采样状态保持为 float32；特征嵌入在 float32 下计算，并在输入投影处进行类型转换。
forward
< source >( bb_ca: Tensorlocal_latents: Tensorbb_ca_time: Tensorlocal_latents_time: Tensorpadding_mask: Tensortoken_type: typing.Optional[torch.Tensor] = Noneatom37_mask: typing.Optional[torch.Tensor] = Noneconditioning: dict[str, typing.Any] | None = Nonerecycle_depth: int = 0attention_kwargs: dict[str, typing.Any] | None = Nonereturn_dict: bool = True ) → RFD4ProteinaTransformerOutput or tuple
参数
- bb_ca (
torch.Tensor) — 带噪声的主干坐标(batch, tokens, 3)，单位为纳米。 - local_latents (
torch.Tensor) — 带噪声的局部潜变量(batch, tokens, latent_dim)。 - bb_ca_time (
torch.Tensor) —(batch,)主干网络的时间，取值范围为 [0, 1]。 - local_latents_time (
torch.Tensor) —(batch,)局部潜在时间，取值范围为 [0, 1]。 - padding_mask (
torch.Tensor) —(batch, tokens)，其中真实 token 的位置为 true，填充位置为 false。 - token_type (
torch.Tensor，可选) —(batch, tokens)的 token 类型编码。 - atom37_mask (
torch.Tensor, 可选) —(batch, tokens, 37)，在原子槽位被使用时为 true。 - conditioning (
dict[str, Any]，可选) — 其余的拓扑和条件张量， 位于预处理器输出的键下。哪些键会被读取由检查点的feats_seq、feats_cond_seq、feats_pair_repr和feats_pair_cond决定；如果某个已配置的特征请求的键不存在， 该特征将保持未加条件状态，而不会引发错误。 - recycle_depth (
int，默认为0) — 主干网络的遍历次数；设为 0 时仍会执行一次遍历。 - attention_kwargs (
dict，可选) — 仅读取scale，作为本次调用的 LoRA 缩放系数。 - return_dict (
bool，默认为True) — 返回带类型的输出，或一个由非 None 字段组成的元组。
返回值
RFD4ProteinaTransformerOutput 或 tuple
Backbone/局部潜在速度以及可选的序列/插入 logits，在调度器集成之前。
从准备好的分子张量中预测速度，而不推进采样器。
RFD4ProteinaTransformerOutput
class diffusers.models.transformers.transformer_rfd4_proteina.RFD4ProteinaTransformerOutput
< source >( backbone: Tensorlocal_latents: typing.Optional[torch.Tensor] = Nonesequence_logits: typing.Optional[torch.Tensor] = Noneinsertion_logits: typing.Optional[torch.Tensor] = None )
参数
在调度器集成之前，去噪器的预测结果。