consistency_models

一貫性モデル

このリポジトリには、ImageNet-64、LSUN Bedroom-256、および LSUN Cat-256 で大規模な実験を行うために PyTorch を使用して実装された整合性モデルのコードベースが含まれています。私たちのリポジトリは、当初 MIT ライセンスに基づいてリリースされた openai/guided-diffusion に基づいています。私たちの修正により、一貫性蒸留、一貫性トレーニング、および論文で説明されているいくつかのサンプリングおよび編集アルゴリズムのサポートが可能になりました。

CIFAR-10 実験のリポジトリは JAX にあり、openai/consistency_models_cifar10 にあります。

事前トレーニングされたモデル

紙面では主要モデルのチェックポイントを公開しました。これらのモデルを使用する前に、対応するモデル カードを確認して、これらのモデルの使用目的と制限事項を理解してください。

各モデル チェックポイントのダウンロード リンクは次のとおりです。

• ImageNet-64 の EDM: edm_imagenet64_ema.pt

• l2 メトリックを使用した ImageNet-64 の CD: cd_imagenet64_l2.pt

• LPIPS メトリックを使用した ImageNet-64 の CD: cd_imagenet64_lpips.pt

• ImageNet-64 の CT: ct_imagenet64.pt

• LSUN ベッドルーム 256 の EDM: edm_bedroom256_ema.pt

• l2 メトリックを使用した LSUN Bedroom-256 の CD: cd_bedroom256_l2.pt

• LPIPS メトリックを使用した LSUN Bedroom-256 の CD: cd_bedroom256_lpips.pt

• LSUN ベッドルーム 256 の CT: ct_bedroom256.pt

• LSUN Cat-256 の EDM: edm_cat256_ema.pt

• l2 メトリックを使用した LSUN Cat-256 の CD: cd_cat256_l2.pt

• LPIPS メトリックを使用した LSUN Cat-256 の CD: cd_cat256_lpips.pt

• LSUN Cat-256 の CT: ct_cat256.pt

依存関係

このコードベース内のすべてのパッケージをその依存関係とともにインストールするには、次を実行します。

 pip install -e 。

Docker を使用してインストールするには、次のコマンドを実行します。

 cd docker && make build && make run

モデルのトレーニングとサンプリング

EDM トレーニング、整合性蒸留、整合性トレーニング、シングルステップ生成、マルチステップ生成の例を scripts/launch.sh で提供します。

評価

さまざまな生成モデルを比較するには、FID、精度、再現率、および開始スコアを使用します。これらのメトリクスはすべて、 .npz (numpy) ファイルに保存されたサンプルのバッチを使用して計算できます。 openai/guided-diffusion で説明されているのと同じ方法で、cm/evaluations/evaluator.py を使用してサンプルを評価できます。参照データセット バッチも提供されています。

で使用しますか?ディフューザー

整合性モデルは ? でサポートされています。 ConsistencyModelPipelineクラスを介したディフューザー。以下に例を示します。

 import torchfrom diffusers import ConsistencyModelPipelinedevice = "cuda"# cd_imagenet64_l2 チェックポイントをロードします。model_id_or_path = "openai/diffusers-cd_imagenet64_l2"pipe = ConsistencyModelPipeline.from_pretrained(model_id_or_path, torch_dtype=torch.float16)pipe.to(device)# ワンステップ サンプリングimage = Pipe(num_inference_steps=1).images[0]image.save("consistency_model_onestep_sample.png")# ワンステップ サンプリング、クラス条件付き画像生成# ImageNet-64クラスラベル 145 はオウサマペンギンに対応しますclass_id = 145class_id = torch.tensor(class_id, dtype=torch.long)image = Pipe(num_inference_steps=1, class_labels=class_id).images[0]image.save("consistency_model_onestep_sample_penguin.png")# マルチステップ サンプリング、クラス条件付き画像生成# タイムステップ明示的に指定できます。以下の特定のタイムステップは、元の Github リポジトリからのものです。# https://github.com/openai/consistency_models/blob/main/scripts/launch.sh#L77image = Pipe(timesteps=[22, 0], class_labels=class_id) .images[0]image.save("consistency_model_multistep_sample_penguin.png")

pipe.unetでtorch.compile()使用すると、推論プロセスをさらに高速化できます (PyTorch 2.0 からのみサポート)。詳細については、公式ドキュメントを確認してください。この支援は?に寄付されました。 dg845とayushtuesのディフューザー。

引用

このメソッドやコードが役立つと思われる場合は、引用することを検討してください。

 @article{song2023consistency, title={一貫性モデル}、author={Song、Yang と Dhariwal、Prafulla と Chen、Mark と Sutskever、Ilya}、journal={arXiv preprint arXiv:2303.01469}、year={2023}、
}