RFD4-Proteína

RFD4-Proteína는 Baker Lab과 NVIDIA의 생체 분자 설계를 위한 flow-matching 모델입니다. 결합 대상, 효모 모티프, 리간드, 핫스팟 또는 2차 구조를 유지하도록 분자 컨텍스트에 조건을 부여하여 단백질 서열과 전 원자 구조를 함께 생성합니다. 페어링된 transformer와 autoencoder는 백본을 위한 것과 로컬 잠재 변수를 위한 것, 총 두 개의 RFD4ProteinaScheduler로 통합됩니다.

pipeline은 준비된 feature tensor를 소비합니다. 주석이 달린 CIF 파일을 해당 tensor로 변환하고, 생성된 tensor를 다시 CIF 파일로 변환하는 작업은 rfd4-proteina-prep 패키지에 포함된 도메인 특화 전처리입니다.

사용 예시

… 준비 패키지를 고정된 리비전으로 설치합니다. 이 패키지는 모델이 인식하는 특징 텐서를 결정하므로, 고정되지 않은 설치의 경우 패키지 업데이트 시 생성되는 구조가 변경됩니다.

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

팁

모델 저장소에는 examples/ 아래에 각 작업별로 주석이 달린 입력 데이터와 생성된 디자인이 함께 제공됩니다: protein binder, antibody, VHH, enzyme, motif scaffolding, small-molecule binder, unconditional, inverse folding, partial diffusion.
partial_t_bb_ca와 partial_t_local_latents를 시작 시간이 [0, 1) 범위에 있도록 전달하여 노이즈가 있는 시작점에서 입력을 재설계하거나, 두 시작 시간을 모두 0.0으로 설정하고 inverse_folding=True를 사용하여 백본을 유지한 채 시퀀스를 설계할 수 있습니다. 두 방식 모두 입력을 인코딩하므로 깨끗한 coords_nm가 필요하며, 유한하지 않은 좌표는 거부합니다. 이는 일반적으로 미해결 잔기입니다.
단계 수는 recycle_depths 파이프라인 구성의 길이이며, 공개된 레시피의 경우 60(180 trunk passes)입니다. 파이프라인을 구성할 때 다른 recycle_depths 리스트를 전달하여 스케줄을 변경할 수 있습니다.

각 입력에 대해 여러 디자인을 생성하려면 num_designs_per_prompt를 설정하고 디자인별로 서로 다른 생성기를 하나씩 전달하세요. 출력에는 batch * num_designs_per_prompt개의 행이 있으며, r번째 행은 입력 행 r // num_designs_per_prompt의 디자인 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는 작성하는 행의 특징을 읽으므로, repeat 대신 repeat_interleave를 사용하여 동일한 방식으로 반복합니다:

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)

속도를 위해 torch.bfloat16로 로드합니다. 샘플러 상태와 특징 임베딩은 float32로 유지됩니다. 다른 dtype, 커널 또는 배치 크기를 사용하면 동일한 시드에서도 다른 설계가 나올 수 있습니다.
생성된 구조는 검증된 설계가 아닙니다. 의도된 작업을 위해 재접힘(refold) 및 점수화를 수행하세요.

LoRA

공개된 모델은 사전 학습된 베이스의 rank-32 LoRA 파인튜닝 버전으로, 어댑터가 가중치에 병합되어 있습니다. 베이스와 어댑터는 모델 저장소의 pretrain-base 리비전에 있으며, 파이프라인은 RFD4ProteinaLoraLoaderMixin을 통해 load_lora_weights로 어댑터를 로드합니다:

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

클래스 diffusers.RFD4ProteinaPipeline

< source >

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

매개변수

transformer (RFD4ProteinaTransformer) — 페어링된 백본과 로컬-잠재 속도 디노이저.
vae (RFD4ProteinaAutoencoder) — denoiser 체크포인트와 함께 사용되는 autoencoder입니다.
bb_ca_scheduler (RFD4ProteinaScheduler) — 백본 좌표를 통합합니다.
local_latents_scheduler (RFD4ProteinaScheduler) — 로컬 latents를 통합합니다.
recycle_depths (list[int], 선택 사항) — 각 디노이징 단계에 대한 트렁크 재활용 패스입니다. 이 리스트의 길이는 추론 단계의 기본 수입니다. 기본값은 60단계, 180패스로 구성된 공개된 스케줄입니다.

준비된 분자 텐서에서 분자 구조와 시퀀스를 생성합니다.

파이프라인은 RFD4 입력 준비 단계에서 생성된 특징 텐서를 소비합니다(주석이 달린 CIF 파일은 rfd4-proteina-prep 패키지를 참조하세요). 그리고 디코딩된 atom37 좌표와 잔기 신원을 반환합니다. 좌표는 나노미터 단위입니다.

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 = ['bb_ca', 'local_latents']output_type: str = 'pt'return_dict: bool = True ) → RFD4ProteinaPipelineOutput or tuple

매개변수

padding_mask (torch.Tensor) — Boolean (batch, tokens), 토큰이 패딩이 아닌 실제 토큰인 위치에서 true입니다.
is_in_context_conditioning (torch.Tensor) — Boolean (batch, tokens). 컨텍스트로 제공되어 변경되지 않는 토큰에 대해 true입니다.
supports_local_latent (torch.Tensor) — 로컬 latent을 가진 토큰에 대해 true인 Boolean (batch, tokens)입니다. 세 가지 mask 모두 동일한 shape와 device를 공유합니다.
token_type (torch.Tensor, 선택 사항) — (batch, tokens) 토큰 유형 코드.
atom37_mask (torch.Tensor, 선택 사항) — (batch, tokens, 37), 원자 슬롯이 사용된 위치에서 true입니다.
coords_nm (torch.Tensor, 선택 사항) — 부분 확산(partial diffusion)과 역 폴딩(inverse folding)에 필요한 나노미터 단위 클린 (batch, tokens, 37, 3) atom37 좌표입니다.
conditioning_features (dict[str, Any], 선택) — preprocessor가 출력하는 키 아래에 있는 나머지 토폴로지 및 조건화 텐서입니다. 시퀀스 조건, condition_sequence_1_residue, molecule_iid를 반드시 포함해야 합니다. 나머지 중 어떤 항목이 읽히는지 고정된 이름으로 약 30개의 선택적 텐서를 출력하는 preprocessor의 특성상, 해당 개수만큼의 인수를 가진 시그니처는 파이프라인을 preprocessor의 명명 규칙에 묶어 버리기 때문에, 체크포인트의 기능 구성에 의해 결정됩니다.
num_designs_per_prompt (int, 기본값 1) — 각 입력 행에 대해 생성할 디자인의 수입니다. 입력은 배치 축을 따라 반복되므로, 출력은 batch * num_designs_per_prompt개의 행을 가지며, r번째 행은 입력 행 r // num_designs_per_prompt의 r % num_designs_per_prompt번째 디자인입니다.
generator (torch.Generator 또는 list[torch.Generator], 선택 사항) — 재현 가능한 설계를 위해 출력 행마다 서로 다른 생성기를 하나씩 사용합니다. 생성기가 진행되고 전역 PyTorch 랜덤 스트림은 그대로 유지됩니다. 전역 스트림에서 값을 추출하려면 이 인수를 생략하십시오.
partial_t_bb_ca (float, 선택 사항) — 순수 노이즈 대신 입력의 노이즈 버전에서 시작하도록 백본의 시작 시간을 [0, 1) 범위로 지정합니다. partial_t_local_latents와 함께 제공해야 하며, coords_nm과 atom37_mask가 필요합니다.
partial_t_local_latents (float, 선택 사항) — 로컬 latents의 시작 시간. [0, 1) 범위이며, partial_t_bb_ca를 참조하세요.
inverse_folding (bool, 기본값 False) — 입력 백본을 유지하고 시퀀스를 설계합니다. 부분 시작 시간이 모두 0이어야 합니다.
latents_bb_ca (torch.Tensor, 선택 사항) — 이 호출에서 샘플링될 값을 대체하기 위해 사전 생성된 시작 백본 상태 (batch * num_designs_per_prompt, tokens, 3)입니다. 스케줄은 변경되지 않습니다. latents_local_latents와 함께 제공하세요.
latents_local_latents (torch.Tensor, 선택 사항) — 미리 생성된 시작 로컬 latents (batch * num_designs_per_prompt, tokens, latent_dim).
attention_kwargs (dict, 선택 사항) — transformer에 전달되며, scale은 LoRA scale을 설정합니다.
callback_on_step_end (Callable, 선택 사항) — 각 단계 이후에 callback_on_step_end(pipeline, step, timestep, callback_kwargs)로 호출됩니다. callback_kwargs에는 callback_on_step_end_tensor_inputs에 명시된 이름의 텐서가 포함되며, 반환된 dict는 "bb_ca"와 "local_latents"를 대체할 수 있습니다.
output_type (str, 기본값 "pt") — "pt"는 torch 텐서를, "np"는 numpy 배열을 반환합니다.
callback_on_step_end_tensor_inputs (list[str], 기본값 ["bb_ca", "local_latents"]) — 콜백에 전달되는 텐서: "bb_ca", "local_latents", "last_clean_bb_ca", "last_clean_local_latents", "sequence_logits", "insertion_logits".
return_dict (bool, 기본값 True) — 타입이 지정된 출력을 반환하거나, None이 아닌 필드의 튜플을 반환합니다.

반환값

RFD4ProteinaPipelineOutput 또는 tuple

디코딩된 atom37 텐서(coords_nm는 나노미터 단위, residue_type, emit_mask, atom_mask), 최종 노이즈 상태, 마지막 클린 추정치, 그리고 마지막 시퀀스 및 삽입 로짓을 반환합니다.

…

예:

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

Validate molecular masks, per-design generators, partial diffusion and callbacks.

noise_backbone

< source >

( target_backbonebb_maskconditioninggeneratorsindex )

index 단계의 예약된 시점까지 보존된 백본에 노이즈를 추가하여 역 폴딩을 수행합니다.

prepare_latents

< source >

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

시작 상태를 그립니다: 순수 노이즈, 또는 start_times까지 노이즈가 추가된 깨끗한 입력(백본, 잠재 변수). 잠재 변수 시작 시간이 0인 경우 clean_latents는 생략할 수 있습니다.

모든 행 생성기에서 샘플링은 네이티브 샘플러의 순서를 따릅니다. 전체 생성의 경우, 백본 노이즈와 그 확률적 오프셋이 잠재 노이즈보다 먼저 추출되며, 부분 확산의 경우 두 노이즈가 모두 오프셋보다 먼저 추출됩니다.

RFD4ProteinaPipelineOutput

클래스 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 )

매개변수

coords_nm (torch.Tensor) — 디코딩된 atom37 좌표 (B, N, 37, 3).
residue_type (torch.Tensor) — Residue 어휘 인덱스 (B, N). residue가 출력되지 않는 위치는 -1입니다.
emit_mask (torch.Tensor) — 디자인된 토큰의 Boolean mask (B, N).
atom_mask (torch.Tensor) — 유효한 원자에 대한 Boolean mask (B, N, 37).
bb_ca (torch.Tensor) — 나노미터 단위 최종 backbone 상태 (B, N, 3).
local_latents (torch.Tensor) — 최종 로컬-라텐트 상태 (B, N, D).
clean_bb_ca (torch.Tensor) — 마지막 디노이징 단계에서 클린 상태의 백본 예측값.
clean_local_latents (torch.Tensor) — 마지막 디노이징 단계에서 클린 상태의 로컬-라텐트 예측값.
sequence_logits (torch.Tensor) — 최종 토큰별 잔차 로짓 (B, N, 147).
insertion_logits (torch.Tensor) — 인덱스가 없는 모티프에 대한 최종 삽입 예측 (B, N, N).

분자 추론 결과. 좌표는 나노미터 단위입니다.

Update on GitHub

Diffusers

RFD4-Proteína

사용 예시

팁

LoRA

RFD4ProteinaPipeline

클래스 diffusers.RFD4ProteinaPipeline

__call__

check_inputs

noise_backbone

prepare_latents

RFD4ProteinaPipelineOutput

클래스 diffusers.pipelines.rfd4_proteina.pipeline_output.RFD4ProteinaPipelineOutput

call