RFD4-Proteína

RFD4-Proteína — это модель flow-matching от Baker Lab и NVIDIA для дизайна биомолекул. Она совместно генерирует последовательности белков и всеатомные структуры, используя молекулярный контекст для сохранения: мишени связывания, ферментного мотива, лиганда, горячих точек или вторичной структуры. Связанная трансформерная модель и автоэнкодер интегрируются с помощью двух RFD4ProteinaScheduler: одного для каркаса и одного для локальных латентов.

Пайплайн потребляет подготовленные тензоры признаков. Преобразование размеченного файла CIF в такие тензоры, а также обратное преобразование сгенерированных тензоров в файл 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/: белковый биндер, антитело, VHH, фермент, мотивное каркасирование, маломолекулярный биндер, безусловная генерация, обратное сворачивание и частичная диффузия.
Передайте partial_t_bb_ca и partial_t_local_latents со стартовыми временами в диапазоне [0, 1), чтобы перепроектировать вход, начиная с зашумленного состояния, либо укажите inverse_folding=True и установите оба стартовых времени в 0.0, чтобы сохранить каркас и спроектировать последовательность. Оба параметра кодируют вход, поэтому требуют чистых coords_nm и отклоняют координаты, которые не являются конечными, что обычно соответствует неразрешённым остаткам.
Количество шагов равно длине конфигурации конвейера recycle_depths, для опубликованного рецепта это 60 (180 проходов по стволу). Чтобы изменить расписание, передайте другой список 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_interleave, а не 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)
```
Загружайте в torch.bfloat16 для скорости. Состояние сэмплера и эмбеддинги признаков остаются в float32. Другой dtype, ядро или размер батча могут привести к другому результату при том же seed.
Сгенерированная структура не является валидированным дизайном. Выполните перегиб и оценку для целевой задачи.

LoRA

Опубликованная модель представляет собой LoRA-дообучение ранга 32 предобученной базовой модели, объединённое с её весами. Базовая модель и адаптер находятся на ревизии pretrain-base репозитория модели, а конвейер загружает адаптер с помощью load_lora_weights через 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

класс diffusers.RFD4ProteinaPipeline

< source >

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

Параметры

transformer (RFD4ProteinaTransformer) — Спаренный бэкбон и локально-латентный денойзер скорости.
vae (RFD4ProteinaAutoencoder) — автоэнкодер, соответствующий контрольной точке денойзера.
bb_ca_scheduler (RFD4ProteinaScheduler) — интегрирует координаты каркаса.
local_latents_scheduler (RFD4ProteinaScheduler) — Интегрирует локальные латенты.
recycle_depths (list[int], необязательный) — Количество проходов переработки ствола для каждого шага денойзинга. Длина списка определяет количество шагов вывода по умолчанию. По умолчанию используется опубликованный график из 60 шагов и 180 проходов.

Генерация молекулярных структур и последовательностей из подготовленных молекулярных тензоров.

Пайплайн принимает тензоры признаков, сгенерированные при подготовке входных данных RFD4 (см. пакет rfd4-proteina-prep для аннотированных файлов CIF), и возвращает декодированные координаты 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[str] = ['bb_ca', 'local_latents']output_type: str = 'pt'return_dict: bool = True ) → RFD4ProteinaPipelineOutput or tuple

Параметры

padding_mask (torch.Tensor) — Булево (batch, tokens), true там, где токен является реальным, а не заполнителем.
is_in_context_conditioning (torch.Tensor) — Булево значение (batch, tokens), истинное для токенов, которые заданы как контекст и остаются неизменными.
supports_local_latent (torch.Tensor) — Булево значение (batch, tokens), истинное для токенов, несущих локальный латент. Все три маски имеют одинаковую форму и устройство.
token_type (torch.Tensor, необязательный) — (batch, tokens) коды типов токенов.
atom37_mask (torch.Tensor, необязательный) — (batch, tokens, 37), true там, где используется слот атома.
coords_nm (torch.Tensor, необязательный) — Очищенные координаты (batch, tokens, 37, 3) для atom37 в нанометрах, необходимые для частичной диффузии и обратного сворачивания.
conditioning_features (dict[str, Any], необязательный) — Остальные тензоры топологии и условий под ключами, которые выдаёт препроцессор. В нём должны содержаться условие последовательности, condition_sequence_1_residue и molecule_iid. Какие из остальных тензоров будут прочитаны, определяется конфигурацией признаков в контрольной точке. Это одно отображение, потому что препроцессор выдаёт около 30 необязательных тензоров с фиксированными именами, а сигнатура с таким количеством аргументов привязала бы конвейер к именованию препроцессора.
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) для backbone, чтобы начать с зашумленной версии входа вместо чистого шума. Указывайте вместе с partial_t_local_latents; требуется coords_nm и atom37_mask.
partial_t_local_latents (float, необязательный) — Начальное время в диапазоне [0, 1) для локальных латентов; см. partial_t_bb_ca.
inverse_folding (bool, по умолчанию False) — Сохранить входной каркас и спроектировать последовательность. Требует нулевые значения для обоих частичных начальных времен.
latents_bb_ca (torch.Tensor, необязательный) — Предварительно сгенерированное начальное состояние backbone (batch * num_designs_per_prompt, tokens, 3), заменяющее то, которое было бы сэмплировано в этом вызове. Расписание не изменяется. Передавайте вместе с latents_local_latents.
latents_local_latents (torch.Tensor, необязательный) — Предварительно сгенерированные начальные локальные латенты (batch * num_designs_per_prompt, tokens, latent_dim).
attention_kwargs (dict, необязательный) — Передаётся в трансформер; scale задаёт масштаб LoRA.
callback_on_step_end (Callable, необязательный) — Вызывается как callback_on_step_end(pipeline, step, timestep, callback_kwargs) после каждого шага. callback_kwargs содержит тензоры, указанные в callback_on_step_end_tensor_inputs; возвращаемый словарь может заменить "bb_ca" и "local_latents".
output_type (str, по умолчанию "pt") — "pt" возвращает тензоры torch, а "np" — массивы numpy.
callback_on_step_end_tensor_inputs (list[str], по умолчанию ["bb_ca", "local_latents"]) — Тензоры, передаваемые в callback: "bb_ca", "local_latents", "last_clean_bb_ca", "last_clean_local_latents", "sequence_logits", "insertion_logits".
return_dict (bool, по умолчанию True) — возвращает типизированный объект вывода или кортеж его непустых полей.

Возвращает

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 )

Проверка молекулярных масок, генераторов для каждого дизайна, частичной диффузии и обратных вызовов.

noise_backbone

< source >

( target_backbonebb_maskconditioninggeneratorsindex )

Добавьте шум к сохранённому стволу до запланированного времени шага index для обратного сворачивания.

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 )

Начерните начальное состояние: чистый шум или чистый вход, зашумлённый до start_times (бэкбон, латенты). 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) — Индексы словаря остатков (B, N); -1 там, где остаток не выводится.
emit_mask (torch.Tensor) — Булева маска (B, N) токенов, которые были спроектированы.
atom_mask (torch.Tensor) — Булева маска (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