# RFD4-Proteína

RFD4-Proteínaは、Baker LabとNVIDIAによる生体分子設計のためのフローマッチングモデルです。結合ターゲット、酵素モチーフ、リガンド、ホットスポット、または二次構造といった分子コンテキストを条件として与えることで、タンパク質配列と全原子構造を同時に生成します。バックボーン用とローカル潜在変数用の2つの[RFD4ProteinaScheduler](/docs/diffusers/main/ja/api/schedulers/rfd4_proteina#diffusers.RFD4ProteinaScheduler)によって、ペアリングされたTransformerとオートエンコーダーが統合されています。

パイプラインは、準備された特徴量テンソルを消費します。アノテーション付きの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 回の trunk パス）です。パイプラインを構築する際に別の `recycle_depths` リストを渡すことで、スケジュールを変更できます。
* 各入力に対して複数のデザインを生成するには、`num_designs_per_prompt`を設定し、各デザインごとに異なるgeneratorを1つずつ渡します。出力は`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` ではなく `repeat_interleave` を使用して、同じ方法で繰り返しを行ってください:

  ```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、カーネル、バッチサイズが異なる場合、同じシードでも異なる設計になる可能性があります。
* 生成された構造は、検証済みの設計ではありません。意図されたタスクに合わせて、再フォールディングとスコアリングを実行してください。

## LoRA

公開されているモデルは、事前学習済みベースモデルの rank-32 LoRA 微調整版であり、その重みにマージされています。ベースモデルとアダプターは、モデルリポジトリの `pretrain-base` リビジョンにあり、パイプラインは `RFD4ProteinaLoraLoaderMixin` を通じて `load_lora_weights` でアダプターを読み込みます:

```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/ja/api/models/rfd4_proteina_transformer#diffusers.RFD4ProteinaTransformer)) : ペアリングされたバックボーンとローカル潜在変数ベロシティデノイザー。

vae ([RFD4ProteinaAutoencoder](/docs/diffusers/main/ja/api/models/rfd4_proteina_autoencoder#diffusers.RFD4ProteinaAutoencoder)) : デノイザーのチェックポイントと組み合わせて使用されるオートエンコーダー。

bb_ca_scheduler ([RFD4ProteinaScheduler](/docs/diffusers/main/ja/api/schedulers/rfd4_proteina#diffusers.RFD4ProteinaScheduler)) : バックボーン座標を統合します。

local_latents_scheduler ([RFD4ProteinaScheduler](/docs/diffusers/main/ja/api/schedulers/rfd4_proteina#diffusers.RFD4ProteinaScheduler)) : ローカル潜在変数を統合します。

recycle_depths (`list[int]`, *optional*) : 各デノイジングステップにおける Trunk リサイクルのパス数。その長さは、推論ステップのデフォルト数です。デフォルトでは、リリース済みのスケジュールである 60 ステップ、180 パスを使用します。

準備済みの分子テンソルから分子構造と配列を生成します。

パイプラインは、RFD4 の入力準備によって生成された特徴量テンソルを消費し（注釈付き CIF ファイルについては `rfd4-proteina-prep` パッケージを参照）、デコードされた atom37 座標と残基 ID を返します。座標はナノメートル単位です。

#### __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`) : Boolean `(batch, tokens)`。トークンがパディングではなく実データであることを示す。

is_in_context_conditioning (`torch.Tensor`) : Boolean `(batch, tokens)`。コンテキストとして与えられ、変更されないトークンに対して true になります。

supports_local_latent (`torch.Tensor`) : Boolean `(batch, tokens)`。ローカル潜在変数を保持するトークンに対して true になります。3 つのマスクはいずれも同じ形状とデバイスを持ちます。

token_type (`torch.Tensor`, *任意*) : `(batch, tokens)` の token-type コード。

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個の任意テンソルを出力するため、それだけの数の引数を持つシグネチャにすると、パイプラインがプリプロセッサの命名規則に縛られてしまうため、1つのマッピングとして定義されています。

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`, *任意*) : ローカル潜在変数の開始時刻（[0, 1) の範囲）。`partial_t_bb_ca` を参照してください。

inverse_folding (`bool`、デフォルトは`False`) : 入力バックボーンを保持し、シーケンスを設計します。部分開始時刻の両方がゼロである必要があります。

latents_bb_ca (`torch.Tensor`, *任意*) : この呼び出しでサンプリングされるものを置き換えるために、事前生成された開始バックボーン状態 `(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`, *任意*) : transformer に渡されます。`scale` は LoRA のスケールを設定します。

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`）、最終的なノイズ付き状態、最後のクリーンな推定値、および最後のシーケンスと挿入のロジット。

準備済みの分子テンソルから構造とシーケンスを生成します。

例:
```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`までノイズが加えられたクリーンな入力（backbone、latents）です。latent の開始時刻がゼロの場合、`clean_latents`は省略できます。

draws は、すべての行ジェネレーターにおいてネイティブサンプラーの順序に従います。完全な生成では、バックボーンノイズとその確率的オフセットが潜在ノイズの前に配置され、部分拡散では両方のノイズがオフセットの前に配置されます。

## 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`) : 有効な原子の Boolean マスク `(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`) : 最終的なトークンごとの残差 logits `(B, N, 147)`。

insertion_logits (`torch.Tensor`) : インデックス未設定のモチーフに対する最終的な挿入予測 `(B, N, N)`。

分子推論の結果。座標はナノメートル単位です。