Hugging Face's logo

Diffusers documentation

RFD4ProteinaAutoencoder

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

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→