パイプライン、モデル、スケジューラーの理解

🧨 Diffusers は、ユースケースに合わせた拡散システムを構築するための、ユーザーフレンドリーかつ柔軟なツールボックスとして設計されています。このツールボックスの中核をなすのは、モデルとスケジューラーです。DiffusionPipeline はこれらのコンポーネントを便利にバンドルしていますが、パイプラインを分解して、モデルとスケジューラーを個別に使用し、新しい拡散システムを作成することもできます。

このチュートリアルでは、基本的なパイプラインから始め、Stable Diffusion パイプラインへと進みながら、モデルとスケジューラを使用して推論用の拡散システムを構築する方法を学びます。

基本的なパイプラインの分解

パイプラインは、推論のためにモデルを実行する簡単で迅速な方法であり、画像を生成するには4行のコードで十分です。

>>> from diffusers import DDPMPipeline

>>> ddpm = DDPMPipeline.from_pretrained("google/ddpm-cat-256", use_safetensors=True).to("cuda")  # or "mps", "xpu", "cpu"
>>> image = ddpm(num_inference_steps=25).images[0]
>>> image

とても簡単でしたね。では、pipelineはどのようにしてこれを成し遂げたのでしょうか？pipelineを分解し、内部で何が起きているのか見ていきましょう。

上記の例では、パイプラインにはUNet2DModelモデルとDDPMSchedulerが含まれています。パイプラインは、望ましい出力サイズと同じサイズのランダムノイズを取得し、モデルを複数回通すことで画像のノイズ除去を行います。各タイムステップにおいて、モデルはノイズ残差を予測し、スケジューラはそれを用いてノイズの少ない画像を予測します。パイプラインは、指定された推論ステップ数の終了に達するまで、このプロセスを繰り返します。

モデルとスケジューラを個別に指定してパイプラインを再構築するには、独自のデノイジングプロセスを実装します。

モデルとスケジューラーを読み込みます:

>>> from diffusers import DDPMScheduler, UNet2DModel

>>> scheduler = DDPMScheduler.from_pretrained("google/ddpm-cat-256")
>>> model = UNet2DModel.from_pretrained("google/ddpm-cat-256", use_safetensors=True).to("cuda")  # or "mps", "xpu", "cpu"

デノイジングプロセスを実行するタイムステップ数を設定します:

>>> scheduler.set_timesteps(50)

スケジューラのタイムステップを設定すると、この例では50個の要素が等間隔に配置されたテンソルが作成されます。各要素は、モデルが画像のノイズ除去を行うタイムステップに対応します。後でデノイジングループを作成する際、このテンソルを反復して画像のノイズ除去を行います。

>>> scheduler.timesteps
tensor([980, 960, 940, 920, 900, 880, 860, 840, 820, 800, 780, 760, 740, 720,
    700, 680, 660, 640, 620, 600, 580, 560, 540, 520, 500, 480, 460, 440,
    420, 400, 380, 360, 340, 320, 300, 280, 260, 240, 220, 200, 180, 160,
    140, 120, 100,  80,  60,  40,  20,   0])

望ましい出力と同じ形状のランダムなノイズを作成します:

>>> import torch

>>> sample_size = model.config.sample_size
>>> noise = torch.randn((1, 3, sample_size, sample_size), device="cuda")  # or "mps", "xpu", "cpu"

次に、タイムステップを反復するためのループを書きます。各タイムステップにおいて、モデルはUNet2DModel.forward()パスを実行し、ノイズ付き残差を返します。スケジューラのstep()メソッドは、ノイズ付き残差、タイムステップ、および入力を受け取り、1つ前のタイムステップの画像を予測します。この出力は、デノイジングループにおけるモデルの次の入力となり、timesteps配列の末尾に到達するまで繰り返されます。

>>> input = noise

>>> for t in scheduler.timesteps:
...     with torch.no_grad():
...         noisy_residual = model(input, t).sample
...     previous_noisy_sample = scheduler.step(noisy_residual, t, input).prev_sample
...     input = previous_noisy_sample

これがデノイジングの全プロセスであり、この同じパターンを用いて任意の拡散システムを記述できます。

最後のステップは、ノイズ除去された出力を画像に変換することです:

>>> from PIL import Image
>>> import numpy as np

>>> image = (input / 2 + 0.5).clamp(0, 1).squeeze()
>>> image = (image.permute(1, 2, 0) * 255).round().to(torch.uint8).cpu().numpy()
>>> image = Image.fromarray(image)
>>> image

次のセクションでは、より複雑な Stable Diffusion パイプラインを分解しながら、あなたのスキルを試すことになります。手順はおおむね同じです。必要なコンポーネントを初期化し、timestep 配列を作成するためのタイムステップ数を設定します。timestep 配列はデノイジングループで使用され、この配列の各要素に対して、モデルはノイズの少ない画像を予測します。デノイジングループは timestep を反復し、各タイムステップでノイズ付きの残差を出力し、スケジューラがそれを前回のタイムステップでのノイズの少ない画像の予測に使用します。このプロセスは、timestep 配列の末尾に到達するまで繰り返されます。

さっそく試してみましょう!

Stable Diffusion パイプラインの分解

Stable Diffusion は、テキストから画像を生成する潜在拡散モデルです。実際のピクセル空間ではなく、画像の低次元表現を扱うため、潜在拡散モデルと呼ばれ、メモリ効率が高くなります。エンコーダーが画像をより小さい表現に圧縮し、デコーダーがその圧縮された表現を画像に戻します。テキストから画像を生成するモデルでは、テキスト埋め込みを生成するためにトークナイザーとエンコーダーが必要です。前回の例から、UNet モデルとスケジューラーが必要であることはすでにわかっています。

ご覧のとおり、これはUNetモデルのみを含むDDPMパイプラインよりもすでに複雑です。Stable Diffusionモデルは、3つの独立した事前学習済みモデルで構成されています。

💡 VAE、UNet、テキストエンコーダーモデルの動作の詳細については、How does Stable Diffusion work?ブログをお読みください。

Stable Diffusion パイプラインに必要なものが分かったところで、from_pretrained() メソッドを使ってこれらのコンポーネントをすべて読み込みます。これらは事前学習済みのstable-diffusion-v1-5/stable-diffusion-v1-5チェックポイントにあり、各コンポーネントは個別のサブフォルダに保存されています:

>>> from PIL import Image
>>> import torch
>>> from transformers import CLIPTextModel, CLIPTokenizer
>>> from diffusers import AutoencoderKL, UNet2DConditionModel, PNDMScheduler

>>> vae = AutoencoderKL.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="vae", use_safetensors=True)
>>> tokenizer = CLIPTokenizer.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="tokenizer")
>>> text_encoder = CLIPTextModel.from_pretrained(
...     "CompVis/stable-diffusion-v1-4", subfolder="text_encoder", use_safetensors=True
... )
>>> unet = UNet2DConditionModel.from_pretrained(
...     "CompVis/stable-diffusion-v1-4", subfolder="unet", use_safetensors=True
... )

デフォルトのPNDMSchedulerの代わりに、UniPCMultistepSchedulerに置き換えて、別のスケジューラを差し込むのがいかに簡単かを確認します:

>>> from diffusers import UniPCMultistepScheduler

>>> scheduler = UniPCMultistepScheduler.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="scheduler")

推論を高速化するために、スケジューラとは異なり学習可能な重みを持つモデルを GPU に移動します:

>>> torch_device = "cuda"  # or "mps", "xpu", "cpu"
>>> vae.to(torch_device)
>>> text_encoder.to(torch_device)
>>> unet.to(torch_device)

テキスト埋め込みの作成

次のステップは、テキストをトークン化して埋め込みを生成することです。このテキストは、UNetモデルを条件付け、拡散プロセスをインプットプロンプトに似た方向へ導くために使用されます。

💡 guidance_scale パラメータは、画像生成時にプロンプトにどの程度の重みを付与するかを決定します。

別のものを生成したい場合は、お好みのプロンプトを自由に選択してください!

>>> prompt = ["a photograph of an astronaut riding a horse"]
>>> height = 512  # default height of Stable Diffusion
>>> width = 512  # default width of Stable Diffusion
>>> num_inference_steps = 25  # Number of denoising steps
>>> guidance_scale = 7.5  # Scale for classifier-free guidance
>>> generator = torch.manual_seed(0)  # Seed generator to create the initial latent noise
>>> batch_size = len(prompt)

テキストをトークン化し、プロンプトから埋め込みを生成します:

>>> text_input = tokenizer(
...     prompt, padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt"
... )

>>> with torch.no_grad():
...     text_embeddings = text_encoder(text_input.input_ids.to(torch_device))[0]

また、パディングトークンの埋め込みである無条件テキスト埋め込みを生成する必要があります。これらは、条件付きのtext_embeddingsと同じ形状(batch_sizeとseq_length)である必要があります:

>>> max_length = text_input.input_ids.shape[-1]
>>> uncond_input = tokenizer([""] * batch_size, padding="max_length", max_length=max_length, return_tensors="pt")
>>> uncond_embeddings = text_encoder(uncond_input.input_ids.to(torch_device))[0]

2回のフォワードパスを避けるために、条件付きと条件なしの埋め込みをバッチに連結しましょう:

>>> text_embeddings = torch.cat([uncond_embeddings, text_embeddings])

ランダムノイズを作成

次に、拡散プロセスの起点となる初期のランダムノイズを生成します。これは画像の潜在表現であり、徐々にノイズ除去されます。この時点では、latent画像は最終的な画像サイズよりも小さいですが、モデルが後で最終的な512x512の画像サイズに変換するため、問題ありません。

💡 vae モデルには3つのダウンサンプリング層があるため、高さ（height）と幅（width）は8で割られます。次のコードを実行して確認できます。
2 ** (len(vae.config.block_out_channels) - 1) == 8

>>> latents = torch.randn(
...     (batch_size, unet.config.in_channels, height // 8, width // 8),
...     generator=generator,
...     device=torch_device,
... )

画像のノイズを除去します

初期ノイズ分布であるノイズスケール値 sigma で入力をスケーリングすることから始めます。UniPCMultistepScheduler などの改善されたスケジューラーには、この値が必要です。

>>> latents = latents * scheduler.init_noise_sigma

最後のステップは、latents内の純粋なノイズを、プロンプトで記述された画像へと段階的に変換するデノイジングループを作成することです。デノイジングループには、3つの処理が必要です。

デノイジング中に使用するスケジューラのタイムステップを設定します。
タイムステップを反復します。
各タイムステップで、UNetモデルを呼び出してノイズ残差を予測し、それをスケジューラに渡して前のノイズ付きサンプルを計算します。

>>> from tqdm.auto import tqdm

>>> scheduler.set_timesteps(num_inference_steps)

>>> for t in tqdm(scheduler.timesteps):
...     # expand the latents if we are doing classifier-free guidance to avoid doing two forward passes.
...     latent_model_input = torch.cat([latents] * 2)

...     latent_model_input = scheduler.scale_model_input(latent_model_input, timestep=t)

...     # predict the noise residual
...     with torch.no_grad():
...         noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample

...     # perform guidance
...     noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
...     noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

...     # compute the previous noisy sample x_t -> x_t-1
...     latents = scheduler.step(noise_pred, t, latents).prev_sample

画像をデコードします

最後のステップは、vaeを使用して潜在表現を画像にデコードし、sampleでデコードされた出力を取得することです。

# scale and decode the image latents with vae
latents = 1 / 0.18215 * latents
with torch.no_grad():
    image = vae.decode(latents).sample

最後に、生成された画像を確認するために、画像をPIL.Imageに変換します。

>>> image = (image / 2 + 0.5).clamp(0, 1).squeeze()
>>> image = (image.permute(1, 2, 0) * 255).to(torch.uint8).cpu().numpy()
>>> image = Image.fromarray(image)
>>> image

次のステップ

基本的なパイプラインから複雑なパイプラインまで、ご確認いただけたように、独自の拡散システムを構築するために本当に必要なのは、ノイズ除去ループだけです。このループでは、スケジューラーのタイムステップを設定し、それらを反復しながら、UNet モデルを呼び出してノイズ残差を予測し、それをスケジューラーに渡して前のノイズ付きサンプルを計算するという処理を交互に行います。

まさに、🧨 Diffusers は、モデルとスケジューラーを用いて独自の拡散システムを直感的かつ簡単に構築できるように設計されています。

次のステップとして、以下をお試しください:

🧨 Diffusers にパイプラインを構築して貢献する方法を学びましょう。皆さんがどのようなアイデアを生み出すのか、今からとても楽しみです!
ライブラリ内の既存のパイプラインを探索し、モデルとスケジューラを個別に使用してパイプラインを分解し、ゼロから構築できるか試してみてください。

Update on GitHub