💥

Diffusers で使える便利なクラス・関数

2024/06/25に公開

Python

犬がDiffusersが便利であることを告げている画像

あまり知名度はありませんが、Diffusers には Diffusion Model を使うにあたって開発体験を向上させるための便利なクラスや関数がいくつか実装されています。

この記事では、それらの中でも特に便利だと感じるものを紹介します。以下の内容は Diffusers v0.29.1 時点のものです。

公式のドキュメントはこちらのサイトから見ることができます。

Utilities

複数画像を一つの画像にする - `make_image_grid()`

複数の画像に対して行数と列数を指定して一枚の画像にします。

Parameters

images: List[PIL.Image.Image]
rows: int
cols: int
resize: int = None - 全ての画像を一定のサイズ（正方形）にリサイズする。

Returns

PIL.Image.Image

Example

from diffusers.utils import load_image, make_image_grid  
from pathlib import Path  

paths = ["hoge.png", "fuga.jpg"]
images = [load_image(p) for p in paths]

image = make_iage_grid(images, rows=2, cols=2, resize=512)
image.save("buri.png")

make_image_grid() 関数を適用した時の例

画像を読み込む - `load_image()`

画像を読み込んで Pillow 形式にする関数です。

Parameters

image: Union[str, PIL.Image.Image] - Pillow画像フォーマットに変換する画像やファイルパス
convert_method: Callable[[PIL.Image.Image], PIL.Image.Image] = None - 画像を読み込んだ後に適用する変換方法。Noneに設定すると画像は "RGB "に変換されます。

Returns

PIL.Image.Image - Pillowの画像

Example

from diffusers.utils import load_image

# type: PIL.Image.Image
image = load_image('hoge.png')

`画像のリストを GIF にする - export_to_gif()`

画像のリストを GIF 形式にして返します。

Parameters

image: List
output_gif_path: str = None
fps: int = 10

Returns

str - 出力の GIF のファイルパス

Example

from diffusers.utils import load_image, export_to_gif  
from pathlib import Path  

paths = ["hoge.png", "fuga.jpg"]
images = [load_image(p) for p in paths]

gif_path = export_to_gif(images, "output.gif", fps=10)

`画像のリストを動画にする - export_to_video()`

画像のリストを動画形式にして返します。

Parameters

video_frames: Union[List[np.ndarray]
output_video_path: str = None
fps: int = 10

Returns

str

Example

from diffusers.utils import load_image, export_to_video 
from pathlib import Path  

paths = ["hoge.png", "fuga.jpg"]
images = [load_image(p) for p in paths]

video_path = export_to_video(images, "output.mp4", fps=10)

VAE Image Processor

PIL 画像と VAE を繋ぐプロセスを簡単にしてくれます。
以下公式の記述を翻訳したもの:

VaeImageProcessorは、StableDiffusionPipelines向けに統一されたAPIを提供し、VAE の Encode のための画像入力の準備や Decode 後の出力の後処理を行います。これには、リサイズや正規化、PIL 画像、PyTorch、およびNumPy 配列間の変換が含まれます。

VaeImageProcessorを備えたすべてのパイプラインは、PIL画像、PyTorch テンソル、またはNumPy 配列を画像入力として受け取り、ユーザーが指定するoutput_type 引数に基づいて出力を返します。Encode された latent を直接パイプラインに渡し、output_type 引数（例: output_type="latent"）を使用して latent として出力を返すことができます。これにより、生成された latent を別のパイプラインに入力として渡し、latent 空間を離れることなく使用することが可能になります。また、異なるパイプライン間でPyTorchテンソルを直接渡すことで、複数のパイプラインを一緒に使用するのが非常に簡単になります。

`VAEImageProcessor`

VAE Image Processor のクラス。こいつをインスタンスにして諸々の前処理・後処理関数を使っていきます。

Parameters

do_resize: bool = True - 画像のサイズを vae_scale_factor の倍数にダウンサンプルするかどうか。
vae_scale_factor: int = 8
vae_latent_channels: int = 4
resample: str = "lanczos"
do_normalize: bool = True - 画像を -1 ~ 1 に正規化するかどうか。
do_binarize: bool = False - 画像を 0, 1 に二値化するかどうか。
do_convert_rgb: bool = False - 画像を RGB フォーマットに統一するかどうか。
do_convert_grayscale: bool = False

画像の前処理を行う - `preprocess()`

preprocess 関数は PIL 画像から VAE への入力を作る部分の前処理を担います。

Parameters

image: Union[PIL.Image.Image, np.ndarray, torch.Tensor, List] - 入力画像。PIL 画像、NumPy 配列、PyTorch テンソル、またはそのリストを受け取ります。
height: int = None - 前処理された画像の高さ (デフォルト: None)
width: int = None - 前処理された画像の幅 (デフォルト: None)
resize_mode: str = 'default' - リサイズモード (デフォルト: 'default')
crops_coords: Optional[List[Tuple[int, int, int, int]]] = None - バッチ内の各画像のクロップ座標 (デフォルト: None)

画像の後処理を行う - `postprocess()`

postprocess 関数はテンソルから指定された形式に画像を後処理します。

Parameters

image: torch.Tensor - 入力画像 (テンソル)
output_type: str = 'pil' - 出力形式 (デフォルト: 'pil')
do_denormalize: Optional[List[bool]] = None - 画像を [0,1] の範囲にデノーマライズするかどうか (デフォルト: None)

Returns

Union[PIL.Image.Image, np.ndarray, torch.Tensor] - 後処理された画像

Example

from diffusers.image_processor import VaeImageProcessor  
from diffusers.utils import load_image  
   
# PIL画像の読み込み  
image = load_image('hoge.jpg')  

# VAEImageProcessorのインスタンス化 (引数がいろいろあるので適宜変更してください)  
processor = VaeImageProcessor()  

# 画像の前処理 PIL -> Tensor    
pre_processed_image = processor.preprocess(image, width=512, height=512)  
print(pre_processed_image.shape)  # torch.Size([1, 3, 512, 512])  
print(type(pre_processed_image))  # <class 'torch.Tensor'>  

# 画像の後処理 Tensor -> List[PIL]    
post_processed_image = processor.postprocess(pre_processed_image)  
print(len(post_processed_image))  # 1
print(post_processed_image[0].size)  # (512, 512)  
print(type(post_processed_image[0])) # <class 'PIL.Image.Image'>

Overlay を適用する - `apply_overlay()`

apply_overlay 関数はインペイントされたイメージを元のイメージに重ね合わせるための関数です。

Parameters

mask: Image - 適用するマスク画像
init_image: Image - 初期の画像
image: Image - インペイント後の画像
crop_coords: Optional[Tuple[int, int, int, int]] = None - クロップする座標(オプション)

Returns

PIL.Image.Image - Overlay が適用された画像

Example

from diffusers.image_processor import VaeImageProcessor  
from diffusers.utils import load_image  
  
# 初期画像とマスク画像の読み込み  
init_image = load_image('images/input.png')  
mask_image = load_image('images/mask.png')  
inp_image = load_image('images/inpaint.png')  

# VAEImageProcessorのインスタンス化  
processor = VaeImageProcessor()  

# Overlay の適用  
overlayed_image = processor.apply_overlay(mask=mask_image, init_image=init_image, image=inp_image)  
overlayed_image.save("images/overlay.png")

apply_overlay()関数を適用した結果

クロップ領域を取得する - `get_crop_region()`

get_crop_region 関数はマスクされたエリアを含む矩形領域を見つけ、その領域を画像の元のアスペクト比に合わせて展開します。

Parameters

mask_image: PIL.Image.Image - マスク画像
width: int - 処理する画像の幅
height: int - 処理する画像の高さ
pad: int = 0 - クロップ領域に追加するパディング (デフォルト: 0)

Returns

tuple - (x1, y1, x2, y2) の形式で、矩形領域が返されます。

Example

from diffusers.image_processor import VaeImageProcessor  
from diffusers.utils import load_image  
  
# 画像とマスク画像の読み込み  
image = load_image('images/input.png')  
mask_image = load_image('images/mask.png')  
  
# VAEImageProcessorのインスタンス化  
processor = VaeImageProcessor()  
  
# 画像のサイズを取得  
width, height = image.size  
  
# マスクされたエリアのクロップ領域を取得  
crop_coords = processor.get_crop_region(mask_image=mask_image, width=width, height=height, pad=10)  
print(crop_coords)  # (x1, y1, x2, y2) = (134, 121, 376, 363)

デフォルトの高さと幅を取得する - `get_default_height_width()`

get_default_height_width 関数は、入力画像を下記の係数に合わせてダウンサンプルした次の整数倍に高さと幅を調整します。VAE のスケールファクターが 8 だが、画像が 8の倍数では無い時などに使えます。

Parameters

image: Union[Image, np.ndarray, torch.Tensor] - 入力画像。PIL画像、NumPy配列、またはPyTorchテンソルのいずれか。
height: int = None - 前処理された画像の高さ (オプション)
width: int = None - 前処理された画像の幅 (オプション)

Returns

(height: int, width: int) - 計算された高さと幅

Example

from diffusers.image_processor import VaeImageProcessor  
from diffusers.utils import load_image  
import numpy as np  
  
# 画像の読み込み  
image = load_image('images/hoge.jpg')  
  
# VAEImageProcessorのインスタンス化  
processor = VaeImageProcessor(vae_scale_factor=8)  
  
# デフォルトの高さと幅を取得  
default_height, default_width = processor.get_default_height_width(image=image)  
  
print(f'Original: {image.size}')  # (1366, 853)  
print(f'Default: ({default_width}, {default_height})') # (1360, 848)

Video Processor

Video Processorは、ビデオパイプライン用の統一されたAPIを提供し、VAE の Encode のための入力準備や Decode 後の出力の後処理を行います。このクラスは VaeImageProcessor を継承しており、リサイズ、正規化、および PIL画像、PyTorch、NumPy 配列間の変換等を含みます。

`VideoProcessor`

動画処理を行うクラス。VaeImageProcessor を継承しています。

Parameters

do_resize: bool = True - 画像のサイズを vae_scale_factor の倍数にダウンサンプルするかどうか。
vae_scale_factor: int = 8
vae_latent_channels: int = 4
resample: str = "lanczos"
do_normalize: bool = True - 画像を -1 ~ 1 に正規化するかどうか。
do_binarize: bool = False - 画像を 0, 1 に二値化するかどうか。
do_convert_rgb: bool = False - 画像を RGB フォーマットに統一するかどうか。
do_convert_grayscale: bool = False

動画の前処理を行う - `preprocess_video()`

preprocess_video関数は、入力動画を前処理します。

Parameters

video: Union[List[PIL.Image], List[List[PIL.Image]], torch.Tensor, np.ndarray, List[torch.Tensor], List[np.ndarray]] - 入力動画。次のいずれかになります:
- PIL 画像のリスト
- PIL 画像のリストのリスト
- 4D Torch テンソル (各テンソルの期待される形状: (num_frames, num_channels, height, width))
- 4D NumPy 配列 (各配列の期待される形状: (num_frames, height, width, num_channels))
- 4D Torch テンソルのリスト (各テンソルの期待される形状: (num_frames, num_channels, height, width))
- 4D NumPy 配列のリスト (各配列の期待される形状: (num_frames, height, width, num_channels))
- 5D NumPy 配列 (各配列の期待される形状: (batch_size, num_frames, height, width, num_channels))
- 5D Torch テンソル (各テンソルの期待される形状: (batch_size, num_frames, num_channels, height, width))
height: int = None - 前処理済みフレームの高さ (オプション)
width: int = None - 前処理済みフレームの幅 (オプション)

動画の後処理を行う - `postprocess_video()`

postprocess_video関数は、指定された形式に動画のテンソルを後処理します。

Parameters

video: torch.Tensor - 入力動画
output_type: str = 'np' - 出力形式 (デフォルト: 'np')

Returns

Union[List[np.ndarray], List[PIL.Image]] - 後処理された動画フレームのリスト

Example

from diffusers.utils import load_image  
from diffusers.video_processor import VideoProcessor  
  
# PIL画像リストとして動画のフレームを読み込み  
video_frames = [load_image(f'images/1.png') for i in range(10)]  
  
# VideoProcessorのインスタンス化  
video_processor = VideoProcessor()  
  
# 動画の前処理 PIL -> Tensor
pre_processed_video = video_processor.preprocess_video(video_frames, height=512, width=512)  
print(pre_processed_video.shape)  # torch.Size([1, 3, 10, 512, 512])  
print(type(pre_processed_video))  # <class 'torch.Tensor'>

# 動画の後処理 Tensor -> list: (batch_size, n_frames, n_images)
post_processed_video = video_processor.postprocess_video(pre_processed_video, output_type='pil')  
print(len(post_processed_video))  # 1  
print(len(post_processed_video[0]))  # 10  
print(type(post_processed_video[0][0]))  # <class 'PIL.Image.Image'>

Attention Processor

Stable Diffusion における U-Net や SD3 における Transformer の Attention の処理を上書きすることができます。詳しくは以下の記事をどうぞ。

Utilities

Discussion