1. 概要
こちらのブログに記載した内容の続きになります。
上記のブログでは、2025年12月21日(日)に開催予定の第10回 岐阜AI勉強会のため、Black Forest Labs の FLUX.1 [schnell]をデバッグ実行しました。今回は、ノイズ除去処理 (denoising) を適用した画像を 1 ステップずつ出力するスクリプトを用意して動かしました。
調査前はこちらのブログに記載した方法と同様、git clone したスクリプトを修正する方法での実装を考えていました。diffusers には git clone したスクリプトを修正しなくても FluxPipeline に callback 関数を渡して生成途中の潜在空間の多次元配列 (latent) を参照するインターフェースが用意されています。そのため、callback 関数を使用して生成途中の画像を出力するようにしました。
Google Colab PRO で A100 GPU を選択して実行しました。
2. 生成途中の画像の出力例
下の動画は FLUX.1 [schnell] で 4 step のノイズ除去処理 (denoising) を適用していく途中の画像を並べた例になります。画像生成に使用したプロンプトは下記のプロンプトになります。4 枚の画像が切り替わっていく動画になります。(下記のスクリプトで得られた約 3MB の GIF ファイルを約 100KB の MP4 ファイルに変換したものです。)
The phrase “Gifu AI Study Group” formed from small colorful geometric shapes.
3. callback 関数を渡して生成途中の画像を出力するスクリプト
こちらのリンク先に用意した Google Colab のページのコードセルを順に実行すると生成途中の画像を含めた 4 枚の画像が切り替わっていく動画が得られます。
git clone したコードのデバッグ実行を試したこちらのブログの Google Colab のページを流用したため git clone したスクリプトを実行しています。callback 関数を使用しているので pip で diffusers をインストールした場合でも動作するスクリプトになっています。
下記のコードセルが callback 関数を使用して生成途中の画像を含めた 4 枚の画像が切り替わっていく GIF 画像を生成するコードセルになります。
from PIL import Image
import os
prompt = 'The phrase "Gifu AI Study Group" formed from small colorful geometric shapes.'
height = 1024
width = 1024
out_dir = "steps"
os.makedirs(out_dir, exist_ok=True)
# 生成途中の画像を追加する GIF ファイル用の空のリスト
gif_frames = []
def callback(pipe, step_index, timestep, callback_kwargs):
latents = callback_kwargs["latents"]
with torch.no_grad():
latents = pipe._unpack_latents(latents, height, width, pipe.vae_scale_factor)
latents = (latents / pipe.vae.config.scaling_factor) + pipe.vae.config.shift_factor
images = pipe.vae.decode(latents, return_dict=False)[0]
images = pipe.image_processor.postprocess(images, output_type="pil")
images[0].save(f"steps/step_{step_index:03d}.png")
gif_frames.append(images[0])
return callback_kwargs
pipe(
prompt,
height=height,
width=width,
guidance_scale=0.0,
num_inference_steps=4,
max_sequence_length=256,
callback_on_step_end=callback
)
# ===== GIF を作成 =====
gif_path = "Gifu-AI-Study-Group.gif"
gif_frames[0].save(
gif_path,
save_all=True,
append_images=gif_frames[1:],
duration=500, # ms / frame
loop=0
)
print(f"GIF saved to {gif_path}")
上記のコードのハイライトした行は下記のスクリプトの
src/diffusers/pipelines/flux/pipeline_flux.py
下記のハイライトした行に対応する変換をしています。callback 関数内では output_type=”pil” を指定し、潜在空間の多次元配列 (tensor) のリスト latents を PIL（Python Imaging Library）形式の画像のインスタンスのリストに変換しています。
def __call__(
self,
...
):
...
if output_type == "latent":
image = latents
else:
latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
image = self.vae.decode(latents, return_dict=False)[0]
image = self.image_processor.postprocess(image, output_type=output_type)
...
4. FluxPipeline の callback 関数について
diffusers の Pipeline の callback 関数については Hugging Face のこちらのページに解説があります。
下記のスクリプトの FluxPipeline の __call__ メソッドのコメントには下記のように記載されています。
src/diffusers/pipelines/flux/pipeline_flux.py
- callback_on_step_end (`Callable`, *optional*):
- A function that calls at the end of each denoising steps during the inference. The function is called with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`.
- callback_on_step_end_tensor_inputs (`List`, *optional*):
- The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class.
FluxPipeline の __call__ メソッドでは下記のように参照されています。
def __call__(
self,
...
callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
callback_on_step_end_tensor_inputs: List[str] = ["latents"],
...
):
...
with self.progress_bar(total=num_inference_steps) as progress_bar:
for i, t in enumerate(timesteps):
...
if callback_on_step_end is not None:
callback_kwargs = {}
for k in callback_on_step_end_tensor_inputs:
callback_kwargs[k] = locals()[k]
callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
latents = callback_outputs.pop("latents", latents)
prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
上記のハイライトした行で callback_on_step_end(self, i, t, callback_kwargs) に渡される引数は下記のようになっています。
- self: FluxPipeline のインスタンス
- i: 現在のステップ番号
- t: 現在のステップのタイムステップ値
- callback_kwargs: 該当ステップのテンソルなどをローカル変数名を key, インスタンスを value として格納した dict 型のインスタンス。__call__ メソッドに callback_on_step_end_tensor_inputs を渡さない場合は、latents のインスタンスだけがセットされています。
上記 3. の callback 関数内ではこれらのインスタンスを参照し、ノイズ除去の各ステップが完了した時点の画像を生成しています。
今回試したスクリプトでは git clone したスクリプトは修正しませんでしたが、その場合でも diffusers のスクリプトをローカルの PC で見られるようにしておいたほうが実装は進め易い気がしました。