Hugging Face's logo

Diffusers documentation

RFD4ProteinaAutoencoder

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

RFD4ProteinaAutoencoder

The autoencoder of RFD4-Proteína. encode maps atom37 coordinates to per-token local latents, and decode turns latents and backbone positions back into atom37 structure and residue identities. It is registered in the pipeline as vae.

from diffusers import RFD4ProteinaAutoencoder

vae = RFD4ProteinaAutoencoder.from_pretrained("nvidia/RFD4-Proteina-Diffusers", subfolder="vae")

RFD4ProteinaAutoencoder

class diffusers.RFD4ProteinaAutoencoder

< >

( atom37_mask: Sequenceallowed_residue_types: Sequencelatent_dim: int = 32n_residue_types: int = 147num_token_types: int = 7nlayers: int = 12token_dim: int = 768nheads: int = 12dim_cond: int = 128pair_repr_dim: int = 256token_type_emb_dim: int = 32encoder_feats_seq: tuple = ('token_type', 'x1_a37coors_nm', 'x1_a37coors_nm_rel', 'chain_idx_seq', 'incident_bond_type_counts', 'atom37_element_charge_embedding')decoder_feats_seq: tuple = ('ca_coors_nm', 'z_latent_seq', 'token_type', 'cond_residue_type') )

Parameters

  • latent_dim (int, defaults to 32) — Width of the per-token local latent.
  • n_residue_types (int, defaults to 147) — Residue vocabulary size.
  • num_token_types (int, defaults to 7) — Size of the token-type vocabulary: protein, RNA, DNA, atomized, unknown, glycan and non-canonical amino acid.
  • allowed_residue_types (list[list[bool]]) — (num_token_types, n_residue_types) table of the residue classes each token type may decode to.
  • nlayers (int, defaults to 12) — Trunk layers in each of the encoder and decoder.
  • token_dim (int, defaults to 768) — Token representation width.
  • nheads (int, defaults to 12) — Attention heads per layer.
  • dim_cond (int, defaults to 128) — Width of the (unused, zero) conditioning vector.
  • pair_repr_dim (int, defaults to 256) — Pair width the trunk layers are constructed with. The released autoencoder uses no pair bias, so this holds no parameters.
  • token_type_emb_dim (int, defaults to 32) — Token-type embedding width.
  • encoder_feats_seq, decoder_feats_seq (list[str]) — Features projected into the encoder and decoder inputs.
  • atom37_mask (list[list[bool]]) — Valid atom37 slots for each residue type, shaped (n_residue_types, 37).

Paired molecular encoder and atom37 decoder.

Coordinates supplied to and returned by this model are in nanometers.

encode

< >

( coords_nm: Tensorpadding_mask: Tensortoken_type: typing.Optional[torch.Tensor] = Noneatom37_mask: typing.Optional[torch.Tensor] = Noneconditioning: dict[str, typing.Any] | None = Nonegenerator: typing.Union[torch.Generator, list[torch.Generator], NoneType] = Nonereturn_dict: bool = True ) → RFD4ProteinaEncoderOutput or tuple

Parameters

  • coords_nm (torch.Tensor) — (batch, tokens, 37, 3) atom37 coordinates in nanometers.
  • 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.
  • generator (torch.Generator or list[torch.Generator], optional) — One generator per batch row for independent stochastic latent samples. Without one, sampling draws from the global stream.
  • return_dict (bool, defaults to True) — Return a typed output, or a tuple of its fields.

Returns

RFD4ProteinaEncoderOutput or tuple

Latent, mean and log-scale tensors shaped (batch, tokens, latent_dim).

Encode prepared coordinates and conditioning into sampled per-token local latents.

decode

< >

( z_latent: Tensorca_coors_nm: Tensorpadding_mask: Tensortoken_type: typing.Optional[torch.Tensor] = Noneatom37_mask: typing.Optional[torch.Tensor] = Noneconditioning: dict[str, typing.Any] | None = Nonereturn_dict: bool = True ) → RFD4ProteinaDecoderOutput or tuple

Parameters

  • z_latent (torch.Tensor) — (batch, tokens, latent_dim) per-token local latents.
  • ca_coors_nm (torch.Tensor) — (batch, tokens, 3) backbone CA positions in nanometers.
  • 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. The sequence condition is read from here.
  • return_dict (bool, defaults to True) — Return a typed output, or a tuple of its fields.

Returns

RFD4ProteinaDecoderOutput or tuple

Coordinates (batch, tokens, 37, 3) in nanometers, residue types (batch, tokens), emission mask (batch, tokens) and atom mask (batch, tokens, 37). Non-emitted atom coordinates are zero.

Decode per-token latents and backbone positions into atom37 structure.

forward

< >

( z_latent: Tensorca_coors_nm: Tensorpadding_mask: Tensortoken_type: typing.Optional[torch.Tensor] = Noneatom37_mask: typing.Optional[torch.Tensor] = Noneconditioning: dict[str, typing.Any] | None = Nonereturn_dict: bool = True ) → RFD4ProteinaDecoderOutput or tuple

Parameters

  • z_latent (torch.Tensor) — (batch, tokens, latent_dim) per-token local latents.
  • ca_coors_nm (torch.Tensor) — (batch, tokens, 3) backbone CA positions in nanometers.
  • 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.
  • return_dict (bool, defaults to True) — Return a typed output, or a tuple of its fields.

Returns

RFD4ProteinaDecoderOutput or tuple

Decoded coordinates, residue types and atom/emit masks.

Decode per-token latents and backbone positions; equivalent to calling decode().

RFD4ProteinaEncoderOutput

class diffusers.models.autoencoders.autoencoder_rfd4_proteina.RFD4ProteinaEncoderOutput

< >

( latent: Tensormean: typing.Optional[torch.Tensor] = Nonelog_scale: typing.Optional[torch.Tensor] = None )

Parameters

  • latent (torch.Tensor) — Sample mean + noise * exp(log_scale), zero outside emitted latent tokens.
  • mean (torch.Tensor) — Gaussian mean, zero outside emitted latent tokens.
  • log_scale (torch.Tensor) — Log standard deviation, zero outside emitted latent tokens.

Sampled latent and Gaussian parameters, each shaped (batch, tokens, latent_dim).

RFD4ProteinaDecoderOutput

class diffusers.models.autoencoders.autoencoder_rfd4_proteina.RFD4ProteinaDecoderOutput

< >

( coordinates: Tensorresidue_type: typing.Optional[torch.Tensor] = Noneemit_mask: typing.Optional[torch.Tensor] = Noneatom_mask: typing.Optional[torch.Tensor] = None )

Parameters

  • coordinates (torch.Tensor) — Coordinates (batch, tokens, 37, 3) in nanometers; non-emitted atoms are zero.
  • residue_type (torch.Tensor) — Residue vocabulary indices shaped (batch, tokens).
  • emit_mask (torch.Tensor) — Boolean emitted-token mask shaped (batch, tokens).
  • atom_mask (torch.Tensor) — Boolean valid-atom mask shaped (batch, tokens, 37).

Decoded molecular coordinates and discrete atom37 topology.

Update on GitHub

←QwenImageTransformer2DModel RFD4ProteinaTransformer→