RFD4-Proteína

RFD4-Proteína là một mô hình flow-matching từ Baker Lab và NVIDIA dành cho thiết kế sinh học phân tử. Nó tạo ra các trình tự protein và cấu trúc tất cả nguyên tử cùng lúc, có điều kiện theo ngữ cảnh phân tử để bảo toàn: mục tiêu liên kết, motif enzyme, phối tử, các điểm nóng hoặc cấu trúc bậc hai. Một transformer ghép cặp và một autoencoder được tích hợp bởi hai RFD4ProteinaScheduler, một cho khung xương và một cho các ẩn cục bộ.

Pipeline tiêu thụ các tensor đặc trưng đã được chuẩn bị. Việc chuyển đổi một tệp CIF có chú thích thành các tensor đó, và chuyển các tensor đã tạo thành tệp CIF, là tiền xử lý đặc thù theo lĩnh vực, nằm trong gói rfd4-proteina-prep.

Ví dụ sử dụng

Cài đặt gói chuẩn bị tại một phiên bản cố định. Gói này quyết định các tensor đặc trưng mà mô hình nhìn thấy, do đó việc cài đặt không cố định phiên bản sẽ làm thay đổi các cấu trúc được tạo ra khi gói được cập nhật.

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

Mẹo

Kho lưu trữ mô hình đi kèm một đầu vào có chú thích và một thiết kế được tạo cho từng tác vụ dưới examples/: protein binder, antibody, VHH, enzyme, motif scaffolding, small-molecule binder, unconditional, inverse folding và partial diffusion.
Truyền partial_t_bb_ca và partial_t_local_latents với thời điểm bắt đầu trong [0, 1) để thiết kế lại đầu vào từ một điểm bắt đầu có nhiễu, hoặc inverse_folding=True với cả hai thời điểm bắt đầu đều là 0.0 để giữ nguyên khung xương và thiết kế chuỗi. Cả hai đều mã hóa đầu vào, do đó cần coords_nm sạch và từ chối các tọa độ không hữu hạn, thường là các phần tử chưa được giải quyết.
Số lượng bước là độ dài của cấu hình pipeline recycle_depths, bằng 60 đối với công thức đã phát hành (180 lượt chạy trunk). Thay đổi lịch trình bằng cách truyền một danh sách recycle_depths khác khi xây dựng pipeline.

Để tạo ra nhiều thiết kế cho mỗi đầu vào, hãy đặt num_designs_per_prompt và truyền một generator riêng biệt cho mỗi thiết kế. Đầu ra có batch * num_designs_per_prompt hàng, và hàng r là thiết kế r % num_designs_per_prompt của hàng đầu vào 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 đọc các đặc trưng của hàng mà nó ghi, vì vậy hãy lặp lại chúng theo cách tương tự, sử dụng repeat_interleave thay vì 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)

Nạp vào torch.bfloat16 để tăng tốc. Trạng thái của sampler và các feature embeddings được giữ ở float32. Một dtype, kernel hoặc batch size khác có thể dẫn đến một thiết kế khác nhau cho cùng một seed.
Một cấu trúc được tạo ra không phải là một thiết kế đã được xác thực. Hãy gấp lại và chấm điểm nó cho tác vụ dự kiến.

LoRA

Mô hình được phát hành là bản fine-tune LoRA rank-32 của mô hình base đã được huấn luyện trước, được gộp vào trọng số. Mô hình base và adapter nằm trên revision pretrain-base của kho mô hình, và pipeline nạp adapter bằng load_lora_weights thông qua 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

lớp diffusers.RFD4ProteinaPipeline

< source >

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

Tham số

transformer (RFD4ProteinaTransformer) — Backbone ghép cặp và bộ khử nhiễu vận tốc latent cục bộ.
vae (RFD4ProteinaAutoencoder) — Autoencoder đi kèm với checkpoint của denoiser.
bb_ca_scheduler (RFD4ProteinaScheduler) — Tích hợp các tọa độ khung xương.
local_latents_scheduler (RFD4ProteinaScheduler) — Tích hợp các latents cục bộ.
recycle_depths (list[int], tùy chọn) — Số lượt tái sử dụng phần thân (trunk) cho mỗi bước khử nhiễu. Độ dài của danh sách này là số bước suy luận mặc định. Mặc định là lịch trình đã phát hành gồm 60 bước và 180 lượt.

Tạo các cấu trúc và trình tự phân tử từ các tensor phân tử đã được chuẩn bị.

Pipeline tiêu thụ các tensor đặc trưng do bước chuẩn bị đầu vào RFD4 tạo ra (xem gói rfd4-proteina-prep để biết các tệp CIF có chú thích) và trả về tọa độ atom37 đã giải mã cùng danh tính các residue. Tọa độ được tính bằng nanomet.

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

Tham số

padding_mask (torch.Tensor) — Boolean (batch, tokens), true tại các vị trí token thực sự thay vì padding.
is_in_context_conditioning (torch.Tensor) — Boolean (batch, tokens), true cho các token được cung cấp làm ngữ cảnh và giữ nguyên không thay đổi.
supports_local_latent (torch.Tensor) — Boolean (batch, tokens), true cho các token mang local latent. Cả ba mask đều chia sẻ cùng một shape và device.
token_type (torch.Tensor, tùy chọn) — (batch, tokens) mã loại token.
atom37_mask (torch.Tensor, tuỳ chọn) — (batch, tokens, 37), giá trị true tại các vị trí mà một ô atom được sử dụng.
coords_nm (torch.Tensor, tùy chọn) — Toạ độ atom37 sạch (batch, tokens, 37, 3) tính bằng nanomet, cần cho partial diffusion và inverse folding.
conditioning_features (dict[str, Any], tùy chọn) — Các tensor topology và conditioning còn lại dưới các key mà preprocessor phát ra. Nó phải chứa điều kiện chuỗi, condition_sequence_1_residue, và molecule_iid. Việc đọc các phần còn lại được cố định bởi cấu hình đặc trưng của checkpoint. Đây là một mapping vì preprocessor phát ra khoảng 30 tensor tùy chọn dưới các tên cố định, và một chữ ký có nhiều tham số như vậy sẽ ràng buộc pipeline với cách đặt tên của preprocessor.
num_designs_per_prompt (int, mặc định là 1) — Số lượng thiết kế cần tạo cho mỗi hàng đầu vào. Các đầu vào được lặp lại dọc theo trục batch, do đó đầu ra có batch * num_designs_per_prompt hàng và hàng r là thiết kế r % num_designs_per_prompt của hàng đầu vào r // num_designs_per_prompt.
generator (torch.Generator hoặc list[torch.Generator], tuỳ chọn) — Một generator riêng biệt cho mỗi hàng đầu ra, để tạo ra các thiết kế có thể tái lập. Các generator sẽ tiến lên và luồng ngẫu nhiên toàn cục của PyTorch được giữ nguyên. Bỏ qua tham số này để lấy mẫu từ luồng toàn cục.
partial_t_bb_ca (float, tùy chọn) — Thời điểm bắt đầu trong [0, 1) cho phần backbone, để bắt đầu từ một phiên bản có nhiễu của đầu vào thay vì nhiễu thuần túy. Cung cấp cùng với partial_t_local_latents; yêu cầu coords_nm và atom37_mask.
partial_t_local_latents (float, tùy chọn) — Thời điểm bắt đầu trong [0, 1) cho các latents cục bộ; xem partial_t_bb_ca.
inverse_folding (bool, mặc định là False) — Giữ nguyên khung xương (backbone) đầu vào và thiết kế trình tự. Yêu cầu cả hai thời điểm bắt đầu một phần đều bằng 0.
latents_bb_ca (torch.Tensor, tuỳ chọn) — Trạng thái backbone khởi đầu đã tạo trước (batch * num_designs_per_prompt, tokens, 3), dùng để thay thế cho trạng thái mà cuộc gọi này sẽ lấy mẫu nếu không có. Lịch trình không thay đổi. Cung cấp cùng với latents_local_latents.
latents_local_latents (torch.Tensor, tùy chọn) — Các local latents khởi đầu đã được tạo trước (batch * num_designs_per_prompt, tokens, latent_dim).
attention_kwargs (dict, tùy chọn) — Được truyền cho transformer; scale đặt tỷ lệ LoRA.
callback_on_step_end (Callable, tùy chọn) — Được gọi dưới dạng callback_on_step_end(pipeline, step, timestep, callback_kwargs) sau mỗi bước. callback_kwargs chứa các tensor được đặt tên trong callback_on_step_end_tensor_inputs; một dict được trả về có thể thay thế "bb_ca" và "local_latents".
output_type (str, mặc định là "pt") — "pt" trả về tensor torch và "np" trả về mảng numpy.
callback_on_step_end_tensor_inputs (list[str], mặc định là ["bb_ca", "local_latents"]) — Các tensor được truyền vào callback: "bb_ca", "local_latents", "last_clean_bb_ca", "last_clean_local_latents", "sequence_logits", "insertion_logits".
return_dict (bool, mặc định là True) — Trả về một đầu ra có kiểu, hoặc một tuple chứa các trường không phải None của nó.

Trả về

RFD4ProteinaPipelineOutput hoặc tuple

Các tensor atom37 đã giải mã (coords_nm tính bằng nanomet, residue_type, emit_mask, atom_mask), trạng thái nhiễu cuối cùng, ước lượng sạch cuối cùng, cùng với logits chèn và logits chuỗi cuối cùng.

Tạo các cấu trúc và chuỗi từ các tensor phân tử đã được chuẩn bị.

Ví dụ:

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

Xác thực các mặt nạ phân tử, trình tạo theo từng thiết kế, khuếch tán một phần và các callback.

noise_backbone

< source >

( target_backbonebb_maskconditioninggeneratorsindex )

Thêm nhiễu vào phần khung xương được bảo toàn tại thời điểm đã lên lịch của bước index, cho quá trình gấp ngược.

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 )

Vẽ trạng thái ban đầu: nhiễu thuần túy, hoặc đầu vào sạch đã được thêm nhiễu đến start_times (backbone, latents). Có thể bỏ qua clean_latents khi thời gian bắt đầu của latent là zero.

Các lần rút mẫu tuân theo thứ tự của bộ tạo mẫu gốc trên mỗi bộ tạo hàng: đối với một lần sinh đầy đủ, nhiễu nền tảng và các độ dịch ngẫu nhiên của nó xuất hiện trước nhiễu tiềm ẩn, trong khi khuếch tán một phần rút cả hai loại nhiễu trước các độ dịch.

RFD4ProteinaPipelineOutput

lớp 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 )

Tham số

coords_nm (torch.Tensor) — Tọa độ atom37 đã giải mã (B, N, 37, 3).
residue_type (torch.Tensor) — Chỉ số từ vựng residue (B, N); -1 tại các vị trí không phát ra residue.
emit_mask (torch.Tensor) — Mặt nạ Boolean (B, N) của các token đã được thiết kế.
atom_mask (torch.Tensor) — Mặt nạ Boolean (B, N, 37) của các nguyên tử hợp lệ.
bb_ca (torch.Tensor) — Trạng thái backbone cuối cùng (B, N, 3) tính bằng nanomet.
local_latents (torch.Tensor) — Trạng thái local-latent cuối cùng (B, N, D).
clean_bb_ca (torch.Tensor) — Dự đoán của backbone cho trạng thái sạch ở bước khử nhiễu cuối cùng.
clean_local_latents (torch.Tensor) — Dự đoán local-latent của trạng thái sạch ở bước khử nhiễu cuối cùng.
sequence_logits (torch.Tensor) — Logits residue cho từng token cuối cùng (B, N, 147).
insertion_logits (torch.Tensor) — Dự đoán chèn cuối cùng (B, N, N) cho các motif không có chỉ số.

Kết quả suy luận phân tử. Tọa độ được tính bằng nanomet.

Update on GitHub

Diffusers

RFD4-Proteína

Ví dụ sử dụng

Mẹo

LoRA

RFD4ProteinaPipeline

lớp diffusers.RFD4ProteinaPipeline

__call__

check_inputs

noise_backbone

prepare_latents

RFD4ProteinaPipelineOutput

lớp diffusers.pipelines.rfd4_proteina.pipeline_output.RFD4ProteinaPipelineOutput

call