# RFD4-Proteína

RFD4-Proteína — это модель flow-matching от Baker Lab и NVIDIA для дизайна биомолекул. Она совместно генерирует последовательности белков и всеатомные структуры, используя молекулярный контекст для сохранения: мишени связывания, ферментного мотива, лиганда, горячих точек или вторичной структуры. Связанная трансформерная модель и автоэнкодер интегрируются с помощью двух [RFD4ProteinaScheduler](/docs/diffusers/main/ru/api/schedulers/rfd4_proteina#diffusers.RFD4ProteinaScheduler): одного для каркаса и одного для локальных латентов.

Пайплайн потребляет подготовленные тензоры признаков. Преобразование размеченного файла CIF в такие тензоры, а также обратное преобразование сгенерированных тензоров в файл CIF — это доменно-специфичная предобработка, которая находится в пакете [`rfd4-proteina-prep`](https://huggingface.co/nvidia/RFD4-Proteina-prep).

## Пример использования

Установите пакет подготовки с фиксированной ревизией. Он определяет тензоры признаков, которые видит модель, поэтому установка без фиксации версии приводит к изменению генерируемых структур при обновлении пакета.

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

```py
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`.

  ```py
  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`:

  ```py
  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`:

```py
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]]

#### diffusers.RFD4ProteinaPipeline[[diffusers.RFD4ProteinaPipeline]]

```python
diffusers.RFD4ProteinaPipeline(transformer: RFD4ProteinaTransformer, vae: RFD4ProteinaAutoencoder, bb_ca_scheduler: RFD4ProteinaScheduler, local_latents_scheduler: RFD4ProteinaScheduler, recycle_depths: list[int] | None = None)
```

[Источник](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/rfd4_proteina/pipeline_rfd4_proteina.py#L68)

**Параметры:**

transformer ([RFD4ProteinaTransformer](/docs/diffusers/main/ru/api/models/rfd4_proteina_transformer#diffusers.RFD4ProteinaTransformer)) : Спаренный бэкбон и локально-латентный денуизер скорости.

vae ([RFD4ProteinaAutoencoder](/docs/diffusers/main/ru/api/models/rfd4_proteina_autoencoder#diffusers.RFD4ProteinaAutoencoder)) : автоэнкодер, используемый вместе с чекпоинтом денойзера.

bb_ca_scheduler ([RFD4ProteinaScheduler](/docs/diffusers/main/ru/api/schedulers/rfd4_proteina#diffusers.RFD4ProteinaScheduler)) : Интегрирует координаты каркаса.

local_latents_scheduler ([RFD4ProteinaScheduler](/docs/diffusers/main/ru/api/schedulers/rfd4_proteina#diffusers.RFD4ProteinaScheduler)) : Интегрирует локальные латенты.

recycle_depths (`list[int]`, *optional*) : Пассы переработки ствола для каждого шага денойзинга. Его длина — это количество шагов инференса по умолчанию. По умолчанию используется опубликованный расписание из 60 шагов и 180 пассов.

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

Пайплайн принимает тензоры признаков, сгенерированные на этапе подготовки входных данных RFD4 (см. пакет `rfd4-proteina-prep` для аннотированных файлов CIF), и возвращает декодированные координаты atom37 и идентификаторы остатков. Координаты указаны в нанометрах.

#### __call__[[diffusers.RFD4ProteinaPipeline.__call__]]

```python
__call__(padding_mask: Tensor, is_in_context_conditioning: Tensor, supports_local_latent: Tensor, token_type: typing.Optional[torch.Tensor] = None, atom37_mask: typing.Optional[torch.Tensor] = None, coords_nm: typing.Optional[torch.Tensor] = None, conditioning_features: dict[str, typing.Any] | None = None, num_designs_per_prompt: int = 1, generator: typing.Union[torch.Generator, list[torch.Generator], NoneType] = None, partial_t_bb_ca: float | None = None, partial_t_local_latents: float | None = None, inverse_folding: bool = False, latents_bb_ca: typing.Optional[torch.Tensor] = None, latents_local_latents: typing.Optional[torch.Tensor] = None, attention_kwargs: dict[str, typing.Any] | None = None, callback_on_step_end: collections.abc.Callable[[diffusers.pipelines.pipeline_utils.DiffusionPipeline, int, torch.Tensor, dict], dict | None] | None = None, callback_on_step_end_tensor_inputs: list[str] = ['bb_ca', 'local_latents'], output_type: str = 'pt', return_dict: bool = True)
```

[Источник](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/rfd4_proteina/pipeline_rfd4_proteina.py#L287)

**Параметры:**

padding_mask (`torch.Tensor`) : Булево `(batch, tokens)`, true там, где токен является реальным, а не заполнителем.

is_in_context_conditioning (`torch.Tensor`) : Булево `(batch, tokens)`, true для токенов, которые заданы как контекст и остаются без изменений.

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`, *необязательный*) : Чистые координаты atom37 `(batch, tokens, 37, 3)` в нанометрах, необходимые для частичной диффузии и обратного сворачивания.

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`, *optional*) : Время начала в диапазоне [0, 1) для backbone, чтобы начать с зашумленной версии входа вместо чистого шума. Указывайте вместе с `partial_t_local_latents`; требуется `coords_nm` и `atom37_mask`.

partial_t_local_latents (`float`, *optional*) : Время начала в диапазоне [0, 1) для локальных латентов; см. `partial_t_bb_ca`.

inverse_folding (`bool`, по умолчанию `False`) : Сохранить входной backbone и спроектировать последовательность. Требует нулевые значения для обоих частичных начальных времен.

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`) : Возвращает типизированный вывод или кортеж его полей, отличных от None.

**Возвращает:** `RFD4ProteinaPipelineOutput` или `tuple`

Декодированные тензоры atom37 (`coords_nm` в нанометрах, `residue_type`, `emit_mask`, `atom_mask`), финальное зашумленное состояние, последняя чистая оценка, а также последние логиты последовательности и вставки.

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

Примеры:
```py
>>> 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[[diffusers.RFD4ProteinaPipeline.check_inputs]]

```python
check_inputs(padding_mask, is_in_context_conditioning, supports_local_latent, atom37_mask, coords_nm, conditioning_features, num_designs_per_prompt, generator, partial_t, inverse_folding, latents, callback_on_step_end_tensor_inputs, output_type)
```

[Источник](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/rfd4_proteina/pipeline_rfd4_proteina.py#L130)

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

#### noise_backbone[[diffusers.RFD4ProteinaPipeline.noise_backbone]]

```python
noise_backbone(target_backbone, bb_mask, conditioning, generators, index)
```

[Источник](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/rfd4_proteina/pipeline_rfd4_proteina.py#L278)

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

#### prepare_latents[[diffusers.RFD4ProteinaPipeline.prepare_latents]]

```python
prepare_latents(bb_mask: Tensor, latent_mask: Tensor, conditioning: dict[str, torch.Tensor], generators: list[torch.Generator] | None, clean_bb: typing.Optional[torch.Tensor] = None, clean_latents: typing.Optional[torch.Tensor] = None, start_times: tuple[float, float] | None = None)
```

[Источник](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/rfd4_proteina/pipeline_rfd4_proteina.py#L243)

Начерните начальное состояние: чистый шум или чистый вход, зашумлённый до `start_times` (бэкбон, латенты). `clean_latents` можно опустить, если начальное время латентов равно нулю.

Выборки следуют порядку нативного сэмплера в каждом генераторе строк: при полной генерации шум бэкбона и его стохастические смещения приходят раньше шума латентов, тогда как при частичной диффузии оба шума выбираются до смещений.

## RFD4ProteinaPipelineOutput[[diffusers.pipelines.rfd4_proteina.pipeline_output.RFD4ProteinaPipelineOutput]]

#### diffusers.pipelines.rfd4_proteina.pipeline_output.RFD4ProteinaPipelineOutput[[diffusers.pipelines.rfd4_proteina.pipeline_output.RFD4ProteinaPipelineOutput]]

```python
diffusers.pipelines.rfd4_proteina.pipeline_output.RFD4ProteinaPipelineOutput(coords_nm: Tensor, residue_type: Tensor, emit_mask: Tensor, atom_mask: Tensor, bb_ca: Tensor, local_latents: Tensor, clean_bb_ca: Tensor, clean_local_latents: Tensor, sequence_logits: Tensor, insertion_logits: Tensor)
```

[Источник](https://github.com/huggingface/diffusers/blob/main/src/diffusers/pipelines/rfd4_proteina/pipeline_output.py#L12)

**Параметры:**

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`) : Итоговое состояние каркаса `(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)` для мотивов без индексов.

Результаты молекулярного вывода. Координаты указаны в нанометрах.