Hugging Face's logo

Diffusers documentation

RFD4ProteinaTransformer

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v0.41.0).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

RFD4ProteinaTransformer

The denoiser of RFD4-Proteína. It predicts velocities for the two modalities the sampler integrates — backbone CA positions and per-token local latents — from prepared molecular tensors. Inputs are named tensors plus a conditioning mapping whose contents the checkpoint configuration selects; structure parsing is outside this class.

from diffusers import RFD4ProteinaTransformer

transformer = RFD4ProteinaTransformer.from_pretrained(
    "nvidia/RFD4-Proteina-Diffusers", subfolder="transformer"
)

RFD4ProteinaTransformer

class diffusers.RFD4ProteinaTransformer

< >

( 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 = ('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 = ('time_emb_bb_ca', 'time_emb_local_latents', 'cath_code', 'recycle_remaining')feats_pair_repr: tuple = ('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 = ('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 )

Parameters

  • nlayers (int, defaults to 14) — Number of trunk layers.
  • token_dim (int, defaults to 768) — Token representation width.
  • nheads (int, defaults to 12) — Attention heads per layer.
  • dim_cond (int, defaults to 256) — Width of the conditioning vector (time, recycle index, CATH label).
  • pair_repr_dim (int, defaults to 384) — Pair representation width.
  • latent_dim (int, defaults to 32) — Width of the per-token local latent.
  • token_type_emb_dim (int, defaults to 32) — Token-type embedding width.
  • cath_emb_dim (int, defaults to 32) — Width of each CATH class/architecture/topology embedding.
  • t_emb_dim (int, defaults to 256) — Sinusoidal time embedding width.
  • seq_sep_dim (int, defaults to 64) — Number of bond-path-length bins.
  • xt_pair_dist_dim (int, defaults to 30) — Number of backbone pair-distance bins.
  • xt_pair_dist_min (float, defaults to 0.1), xt_pair_dist_max (float, defaults to 3.0) — Pair-distance range in nanometers.
  • feats_seq, feats_cond_seq, feats_pair_repr, feats_pair_cond (list[str]) — Names of the features projected into the token, conditioning, pair and pair-conditioning inputs.
  • update_pair_repr_every_n (int, defaults to 2) — How often the pair representation is updated from the tokens between layers.
  • tri_mult_c (int, defaults to 256) — Hidden width of the triangle multiplication in those updates.
  • num_token_types (int, defaults to 7) — Size of the token-type vocabulary: protein, RNA, DNA, atomized, unknown, glycan and non-canonical amino acid.
  • recycle_conditioning_dim (int, defaults to 1) — Width of the recycle_remaining one-hot, when that feature is in feats_cond_seq.
  • num_sequence_classes (int, defaults to 147) — Size of the residue vocabulary the sequence readout predicts.

RFD4 velocity denoiser with paired backbone and local-latent modalities.

The model runs in its own dtype: load it with dtype=torch.bfloat16 for reduced-precision inference. The sampling state passed to forward stays in float32; feature embeddings are computed in float32 and cast at the input projection.

forward

< >

( 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

Parameters

  • bb_ca (torch.Tensor) — Noisy backbone coordinates (batch, tokens, 3) in nanometers.
  • local_latents (torch.Tensor) — Noisy local latents (batch, tokens, latent_dim).
  • bb_ca_time (torch.Tensor) — (batch,) backbone time in [0, 1].
  • local_latents_time (torch.Tensor) — (batch,) local-latent time in [0, 1].
  • padding_mask (torch.Tensor) — (batch, tokens), true where a token is real rather than padding.
  • token_type (torch.Tensor, optional) — (batch, tokens) token-type codes.
  • atom37_mask (torch.Tensor, optional) — (batch, tokens, 37), true where an atom slot is used.
  • conditioning (dict[str, Any], optional) — The remaining topology and conditioning tensors, under the keys the preprocessor emits. Which of them are read is fixed by the checkpoint’s feats_seq, feats_cond_seq, feats_pair_repr and feats_pair_cond; a key a configured feature asks for but that is absent leaves that feature unconditioned rather than raising.
  • recycle_depth (int, defaults to 0) — Number of trunk passes; zero still executes one pass.
  • attention_kwargs (dict, optional) — Only scale is read, as the LoRA scale for this call.
  • return_dict (bool, defaults to True) — Return a typed output, or a tuple of non-None fields.

Returns

RFD4ProteinaTransformerOutput or tuple

Backbone/local-latent velocities and optional sequence/insertion logits, before scheduler integration.

Predict velocities from prepared molecular tensors without advancing the sampler.

RFD4ProteinaTransformerOutput

class diffusers.models.transformers.transformer_rfd4_proteina.RFD4ProteinaTransformerOutput

< >

( backbone: Tensorlocal_latents: typing.Optional[torch.Tensor] = Nonesequence_logits: typing.Optional[torch.Tensor] = Noneinsertion_logits: typing.Optional[torch.Tensor] = None )

Parameters

  • backbone (torch.Tensor) — Backbone velocity shaped (batch, tokens, 3), in nanometers per unit time.
  • local_latents (torch.Tensor) — Latent velocity shaped (batch, tokens, latent_dim).
  • sequence_logits (torch.Tensor, optional) — Residue logits shaped (batch, tokens, 147).
  • insertion_logits (torch.Tensor, optional) — Insertion assignment logits shaped (batch, tokens, tokens).

Denoiser predictions before scheduler integration.

Update on GitHub

←RFD4ProteinaAutoencoder SanaTransformer2DModel→