RFD4-Proteína

RFD4-Proteína adalah model flow-matching dari Baker Lab dan NVIDIA untuk desain biomolekuler. Model ini menghasilkan urutan protein dan struktur all-atom secara bersamaan, dikondisikan pada konteks molekul untuk mempertahankan: target pengikatan, motif enzim, ligan, hotspot, atau struktur sekunder. Transformer berpasangan dan autoencoder diintegrasikan oleh dua RFD4ProteinaScheduler, satu untuk backbone dan satu untuk latens lokal.

Pipeline mengonsumsi tensor fitur yang sudah disiapkan. Mengubah file CIF ber anotasi menjadi tensor tersebut, dan mengubah tensor yang dihasilkan kembali menjadi file CIF, merupakan pra-pemrosesan spesifik domain yang terdapat dalam paket rfd4-proteina-prep.

Contoh penggunaan

Install paket persiapan pada revisi yang tetap. Paket ini menentukan tensor fitur yang dilihat oleh model, sehingga instalasi tanpa pin akan mengubah struktur yang dihasilkan ketika paket diperbarui.

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

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")

Tips

Repositori model menyertakan input ber anotasi dan desain yang dihasilkan untuk setiap tugas di bawah examples/: protein binder, antibodi, VHH, enzim, motif scaffolding, small-molecule binder, unconditional, inverse folding, dan partial diffusion.
Lewatkan partial_t_bb_ca dan partial_t_local_latents dengan waktu mulai dalam [0, 1) untuk merancang ulang input dari awal yang berisik, atau inverse_folding=True dengan kedua waktu mulai pada 0.0 untuk mempertahankan backbone dan merancang sekuens. Keduanya meng-encode input, sehingga memerlukan coords_nm yang bersih dan menolak koordinat yang tidak finite, yang biasanya merupakan residu yang tidak terurai.
Jumlah langkah adalah panjang konfigurasi pipeline recycle_depths, yaitu 60 untuk resep yang dirilis (180 lintasan trunk). Ubah jadwal dengan meneruskan daftar recycle_depths yang berbeda saat membangun pipeline.

Untuk menghasilkan beberapa desain dari setiap input, atur num_designs_per_prompt dan berikan satu generator yang berbeda untuk setiap desain. Output memiliki batch * num_designs_per_prompt baris, dan baris r adalah desain r % num_designs_per_prompt dari baris input r // num_designs_per_prompt.

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 membaca fitur dari baris yang ditulisnya, jadi ulangi dengan cara yang sama, menggunakan repeat_interleave alih-alih repeat:

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)

Muat dalam torch.bfloat16 untuk kecepatan. State sampler dan embeddings fitur tetap dalam float32. dtype, kernel, atau ukuran batch yang berbeda dapat menghasilkan desain yang berbeda untuk seed yang sama.
Struktur yang dihasilkan bukanlah desain yang telah divalidasi. Lakukan refold dan skor untuk tugas yang dimaksud.

LoRA

Model yang dirilis adalah fine-tune LoRA rank-32 dari base pretrained, yang telah digabungkan ke dalam bobot. Base dan adapter berada pada revisi pretrain-base dari repositori model, dan pipeline memuat adapter dengan load_lora_weights melalui RFD4ProteinaLoraLoaderMixin:

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

kelas diffusers.RFD4ProteinaPipeline

< source >

( transformer: RFD4ProteinaTransformervae: RFD4ProteinaAutoencoderbb_ca_scheduler: RFD4ProteinaSchedulerlocal_latents_scheduler: RFD4ProteinaSchedulerrecycle_depths: list[int] | None = None )

Parameter

transformer (RFD4ProteinaTransformer) — Backbone berpasangan dan denoiser kecepatan laten lokal.
vae (RFD4ProteinaAutoencoder) — Autoencoder yang berpasangan dengan checkpoint denoiser.
bb_ca_scheduler (RFD4ProteinaScheduler) — Mengintegrasikan koordinat tulang punggung.
local_latents_scheduler (RFD4ProteinaScheduler) — Mengintegrasikan latents lokal.
recycle_depths (list[int], opsional) — Lolosan daur ulang trunk untuk setiap langkah denoising. Panjangnya adalah jumlah langkah inferensi default. Defaultnya adalah jadwal rilis 60 langkah dan 180 lolosan.

Generate struktur dan sekuens molekul dari tensor molekul yang telah disiapkan.

Pipeline mengonsumsi tensor fitur yang dihasilkan oleh persiapan input RFD4 (lihat paket rfd4-proteina-prep untuk file CIF ber anotasi) dan mengembalikan koordinat atom37 yang didekodekan serta identitas residu. Koordinat dalam satuan nanometer.

call

< source >

( padding_mask: Tensoris_in_context_conditioning: Tensorsupports_local_latent: Tensortoken_type: typing.Optional[torch.Tensor] = Noneatom37_mask: typing.Optional[torch.Tensor] = Nonecoords_nm: typing.Optional[torch.Tensor] = Noneconditioning_features: dict[str, typing.Any] | None = Nonenum_designs_per_prompt: int = 1generator: typing.Union[torch.Generator, list[torch.Generator], NoneType] = Nonepartial_t_bb_ca: float | None = Nonepartial_t_local_latents: float | None = Noneinverse_folding: bool = Falselatents_bb_ca: typing.Optional[torch.Tensor] = Nonelatents_local_latents: typing.Optional[torch.Tensor] = Noneattention_kwargs: dict[str, typing.Any] | None = Nonecallback_on_step_end: collections.abc.Callable[[diffusers.pipelines.pipeline_utils.DiffusionPipeline, int, torch.Tensor, dict], dict | None] | None = Nonecallback_on_step_end_tensor_inputs: list[str] = ['bb_ca', 'local_latents']output_type: str = 'pt'return_dict: bool = True ) → RFD4ProteinaPipelineOutput or tuple

Parameter

padding_mask (torch.Tensor) — Boolean (batch, tokens), bernilai true di mana token tersebut asli, bukan padding.
is_in_context_conditioning (torch.Tensor) — Boolean (batch, tokens), bernilai true untuk token yang diberikan sebagai konteks dan tidak diubah.
supports_local_latent (torch.Tensor) — Boolean (batch, tokens), bernilai true untuk token yang membawa latent lokal. Ketiga mask berbagi satu bentuk dan perangkat yang sama.
token_type (torch.Tensor, opsional) — (batch, tokens) kode tipe token.
atom37_mask (torch.Tensor, opsional) — (batch, tokens, 37), bernilai true di mana slot atom digunakan.
coords_nm (torch.Tensor, opsional) — Koordinat atom37 bersih (batch, tokens, 37, 3) dalam nanometer, yang diperlukan oleh partial diffusion dan inverse folding.
conditioning_features (dict[str, Any], opsional) — Tensor topologi dan kondisioning yang tersisa di bawah kunci yang dipancarkan oleh preprocessor. Harus memuat kondisi sekuens, condition_sequence_1_residue, dan molecule_iid. Sisa tensor mana yang dibaca ditentukan oleh konfigurasi fitur checkpoint. Ini berupa satu pemetaan karena preprocessor memancarkan sekitar 30 tensor opsional di bawah nama-nama tetap, dan tanda tangan dengan begitu banyak argumen akan mengikat pipeline ke penamaan preprocessor.
num_designs_per_prompt (int, defaultnya 1) — Jumlah desain yang akan dihasilkan untuk setiap baris input. Input diulang sepanjang sumbu batch, sehingga output memiliki batch * num_designs_per_prompt baris dan baris r adalah desain r % num_designs_per_prompt dari baris input r // num_designs_per_prompt.
generator (torch.Generator atau list[torch.Generator], opsional) — Satu generator berbeda untuk setiap baris output, untuk desain yang dapat direproduksi. Generator-generator tersebut maju dan aliran acak global PyTorch dibiarkan tidak tersentuh. Lewatkan parameter ini untuk mengambil dari aliran global.
partial_t_bb_ca (float, opsional) — Waktu mulai dalam [0, 1) untuk backbone, untuk memulai dari versi ber-noise dari input alih-alih noise murni. Berikan bersama dengan partial_t_local_latents; membutuhkan coords_nm dan atom37_mask.
partial_t_local_latents (float, opsional) — Waktu mulai dalam [0, 1) untuk latens lokal; lihat partial_t_bb_ca.
inverse_folding (bool, defaultnya False) — Pertahankan backbone input dan rancang sekuensnya. Memerlukan nol untuk kedua waktu mulai parsial.
latents_bb_ca (torch.Tensor, opsional) — Keadaan awal backbone yang telah dibuat sebelumnya (batch * num_designs_per_prompt, tokens, 3), untuk menggantikan keadaan yang seharusnya disampel oleh panggilan ini. Jadwal tidak berubah. Berikan bersama dengan latents_local_latents.
latents_local_latents (torch.Tensor, opsional) — Latens lokal awal yang telah dibuat sebelumnya (batch * num_designs_per_prompt, tokens, latent_dim).
attention_kwargs (dict, opsional) — Diteruskan ke transformer; scale mengatur skala LoRA.
callback_on_step_end (Callable, opsional) — Dipanggil sebagai callback_on_step_end(pipeline, step, timestep, callback_kwargs) setelah setiap langkah. callback_kwargs menyimpan tensor yang namanya tercantum dalam callback_on_step_end_tensor_inputs; dict yang dikembalikan dapat menggantikan "bb_ca" dan "local_latents".
output_type (str, defaultnya "pt") — "pt" mengembalikan tensor torch dan "np" mengembalikan array numpy.
callback_on_step_end_tensor_inputs (list[str], defaultnya ["bb_ca", "local_latents"]) — Tensor yang diteruskan ke callback: "bb_ca", "local_latents", "last_clean_bb_ca", "last_clean_local_latents", "sequence_logits", "insertion_logits".
return_dict (bool, defaultnya True) — Mengembalikan output bertipe, atau tuple dari field-field non-None-nya.

Mengembalikan

RFD4ProteinaPipelineOutput atau tuple

Tensor atom37 yang telah didekodekan (coords_nm dalam nanometer, residue_type, emit_mask, atom_mask), keadaan berisik akhir, estimasi bersih terakhir, serta logits urutan dan sisipan terakhir.

Generate struktur dan sekuens dari tensor molekuler yang telah disiapkan.

Contoh:

>>> 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

< source >

( padding_maskis_in_context_conditioningsupports_local_latentatom37_maskcoords_nmconditioning_featuresnum_designs_per_promptgeneratorpartial_tinverse_foldinglatentscallback_on_step_end_tensor_inputsoutput_type )

Validasi masker molekul, generator per-desain, difusi parsial, dan callback.

noise_backbone

< source >

( target_backbonebb_maskconditioninggeneratorsindex )

Tambahkan noise pada backbone yang dipertahankan hingga waktu yang dijadwalkan pada langkah index, untuk inverse folding.

prepare_latents

< source >

( bb_mask: Tensorlatent_mask: Tensorconditioning: dict[str, torch.Tensor]generators: list[torch.Generator] | Noneclean_bb: typing.Optional[torch.Tensor] = Noneclean_latents: typing.Optional[torch.Tensor] = Nonestart_times: tuple[float, float] | None = None )

Gambar keadaan awal: noise murni, atau input bersih yang diberi noise hingga start_times (backbone, latents). clean_latents dapat diabaikan jika waktu awal latents adalah nol.

Undian mengikuti urutan sampler asli pada setiap generator baris: untuk generasi penuh, noise backbone dan offset stokastiknya muncul sebelum noise laten, sedangkan difusi parsial mengundi kedua noise tersebut sebelum offset.

RFD4ProteinaPipelineOutput

kelas diffusers.pipelines.rfd4_proteina.pipeline_output.RFD4ProteinaPipelineOutput

< source >

( coords_nm: Tensorresidue_type: Tensoremit_mask: Tensoratom_mask: Tensorbb_ca: Tensorlocal_latents: Tensorclean_bb_ca: Tensorclean_local_latents: Tensorsequence_logits: Tensorinsertion_logits: Tensor )

Parameter

coords_nm (torch.Tensor) — Koordinat atom37 yang didekodekan (B, N, 37, 3).
residue_type (torch.Tensor) — Indeks kosakata residu (B, N); -1 di mana tidak ada residu yang dihasilkan.
emit_mask (torch.Tensor) — Mask boolean (B, N) dari token-token yang dirancang.
atom_mask (torch.Tensor) — Mask boolean (B, N, 37) dari atom yang valid.
bb_ca (torch.Tensor) — Keadaan akhir backbone (B, N, 3) dalam nanometer.
local_latents (torch.Tensor) — Keadaan lokal-latent akhir (B, N, D).
clean_bb_ca (torch.Tensor) — Prediksi backbone dari keadaan bersih pada langkah denoising terakhir.
clean_local_latents (torch.Tensor) — Prediksi latensi lokal dari keadaan bersih pada langkah denoising terakhir.
sequence_logits (torch.Tensor) — Logit residu per token akhir (B, N, 147).
insertion_logits (torch.Tensor) — Prediksi penyisipan akhir (B, N, N) untuk motif tanpa indeks.

Hasil inferensi molekuler. Koordinat dalam nanometer.

Update on GitHub

Diffusers

RFD4-Proteína

Contoh penggunaan

Tips

LoRA

RFD4ProteinaPipeline

kelas diffusers.RFD4ProteinaPipeline

__call__

check_inputs

noise_backbone

prepare_latents

RFD4ProteinaPipelineOutput

kelas diffusers.pipelines.rfd4_proteina.pipeline_output.RFD4ProteinaPipelineOutput

call