RFD4-Proteína

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

パイプラインは、準備された特徴量テンソルを消費します。アノテーション付きの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 回の trunk パス）です。パイプラインを構築する際に別の recycle_depths リストを渡すことで、スケジュールを変更できます。

各入力に対して複数のデザインを生成するには、num_designs_per_promptを設定し、各デザインごとに異なるジェネレーターを1つずつ渡します。出力は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、カーネル、またはバッチサイズを使用すると、同じシードでも異なる設計になる可能性があります。
生成された構造は、検証済みの設計ではありません。意図されたタスクに合わせて、再フォールディングとスコアリングを実行してください。

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

class 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の入力準備によって生成された特徴量テンソルを消費し（注釈付き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[str] = ['bb_ca', 'local_latents']output_type: str = 'pt'return_dict: bool = True ) → RFD4ProteinaPipelineOutput or tuple

パラメータ

padding_mask (torch.Tensor) — (batch, tokens) の Boolean 型。トークンがパディングではなく実データである箇所が true になります。
is_in_context_conditioning (torch.Tensor) — (batch, tokens) のブール値。コンテキストとして与えられ、変更されないトークンに対して true です。
supports_local_latent (torch.Tensor) — (batch, tokens) のブール値。ローカル潜在変数を保持するトークンに対して true です。3 つのマスクはすべて同じ形状とデバイスを持ちます。
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], 省略可) — プリプロセッサが出力するキーの下にある、残りのトポロジーおよび条件付けテンソルです。シーケンス条件である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）、最終的なノイズ付き状態、最後のクリーンな推定値、および最後のシーケンスと挿入のロジット。

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

例:

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

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

RFD4ProteinaPipelineOutput

class 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) — ナノメートル単位での最終的なバックボーン状態 (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)。

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

Update on GitHub

Diffusers

RFD4-Proteína

使用例

ヒント

LoRA

RFD4ProteinaPipeline

class diffusers.RFD4ProteinaPipeline

__call__

check_inputs

noise_backbone

prepare_latents

RFD4ProteinaPipelineOutput

class diffusers.pipelines.rfd4_proteina.pipeline_output.RFD4ProteinaPipelineOutput

call