Jetson Orin NanoでStable Diffusionを動かす！画像生成入門【TensorRT最適化】
Jetson Orin NanoでStable Diffusionを動かしたいけれど、メモリ不足やビルドエラーで挫折した経験はないだろうか。この記事では実機を使った手順と、TensorRTで高速化する方法をまとめて解説する。
この記事でわかること
—
動作環境と必要なもの
今回の検証環境は以下の通りだ。
| 項目 | 内容 |
|——|——|
| ボード | Jetson Orin Nano Super 8GB |
| JetPack | 6.2 (L4T 36.4.3) |
| ストレージ | NVMe SSD 512GB |
| モデル | Stable Diffusion v1.5 |
Jetson Orin Nano Superは67 TOPS (INT8)の高性能で8GB LPDDR5メモリを搭載。4GBモデルは避け、8GBモデルを強く推奨する。ユニファイドメモリ構造のため、SD v1.5でも安定動作が可能だ。
ストレージはモデルファイルが数GBあるため、NVMe SSD Western Digital SN770 500GB NVMe（Amazon） を使うことを前提にしている。microSDカード Samsung microSDXC 128GB (Endurance)（Amazon） でも動作はするが、モデルのロード時間が大幅に伸びるため実用的ではない。
—
スワップ領域の設定（必須）
Jetson Orin Nanoは物理メモリとGPUメモリが共有されているため、Stable Diffusionの動作中にメモリが枯渇しやすい。まず最初にスワップを増やしておく。
デフォルトのスワップサイズを確認する。
free -h
スワップが2GB程度しかない場合は以下の手順で8GBに拡張する。
# 既存のスワップを無効化
sudo swapoff -a
# 8GBのスワップファイルを作成
sudo fallocate -l 8G /mnt/swapfile
sudo chmod 600 /mnt/swapfile
sudo mkswap /mnt/swapfile
sudo swapon /mnt/swapfile
# 再起動後も有効にする
echo '/mnt/swapfile none swap sw 0 0' | sudo tee -a /etc/fstab
スワップはNVMe SSD上に置くことで読み書き速度を確保できる。
—
環境構築：Pythonと依存ライブラリのインストール
仮想環境の準備
システムのPythonを汚さないために仮想環境を使う。
sudo apt update
sudo apt install -y python3-pip python3-venv
python3 -m venv ~/sd_env
source ~/sd_env/bin/activate
PyTorchのインストール
Jetson向けのPyTorchはNVIDIAが公式ビルドを提供している。通常のpipではARM+CUDA対応のバイナリが入らないため、必ずNVIDIAのリポジトリから取得する。
pip install --upgrade pip
# JetPack 6.2 / PyTorch 2.4対応版
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121
インストール後、CUDAが認識されているかを確認する。
python3 -c "import torch; print(torch.cuda.is_available()); print(torch.version.cuda)"
True と表示されれば成功だ。
diffusersと関連ライブラリのインストール
pip install diffusers transformers accelerate safetensors
pip install opencv-python Pillow
—
Stable Diffusionの動作確認（シンプル版）
まずTensorRT最適化なしの素の状態で動作確認する。
# generate_simple.py
from diffusers import StableDiffusionPipeline
import torch
model_id = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionPipeline.from_pretrained(
model_id,
torch_dtype=torch.float16,
safety_checker=None,
)
pipe = pipe.to("cuda")
# メモリ最適化
pipe.enable_attention_slicing()
prompt = "a photo of a cat on a sunny beach, high quality"
image = pipe(prompt, num_inference_steps=20).images[0]
image.save("output.png")
print("生成完了")
python3 generate_simple.py
初回実行時はHugging Faceからモデルがダウンロードされるためネット接続が必要だ。モデルのサイズは約4GBある。
この状態での生成速度は筆者の環境で512×512・20ステップあたり約70秒だった。実用には厳しい速度なので、次のステップでTensorRT最適化を行う。
—
TensorRTによる高速化
TensorRTのインストール確認
JetPack 6.2にはTensorRT 10.xが含まれているため、追加インストールは不要なことが多い。バージョンを確認する。
dpkg -l | grep tensorrt
python3 -c "import tensorrt; print(tensorrt.__version__)"
TensorRTビルド時にメモリ不足やlayernorm警告が出る場合は、ONNX opset 17以上でモデルを再エクスポートすると安定する。
torch2trtのインストール
UNetをTensorRTエンジンに変換するために torch2trt を使う。
cd ~
git clone https://github.com/NVIDIA-AI-IOT/torch2trt
cd torch2trt
pip install .
ビルドには10〜15分かかる。エラーが出る場合は cuda-toolkit の開発ヘッダが不足している場合が多い。
sudo apt install -y cuda-toolkit-12-4
TensorRTエンジンの変換スクリプト
# convert_trt.py
import torch
from diffusers import StableDiffusionPipeline
from torch2trt import torch2trt
model_id = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionPipeline.from_pretrained(
model_id,
torch_dtype=torch.float16,
safety_checker=None,
)
pipe = pipe.to("cuda")
# UNetのTensorRT変換
unet = pipe.unet
unet.eval()
# ダミー入力の作成（512x512、バッチサイズ1）
sample = torch.randn(2, 4, 64, 64).half().cuda()
timestep = torch.tensor([1]).cuda()
encoder_hidden = torch.randn(2, 77, 768).half().cuda()
print("TensorRT変換を開始します（10〜20分かかります）")
unet_trt = torch2trt(
unet,
[sample, timestep, encoder_hidden],
fp16_mode=True,
max_batch_size=2,
)
torch.save(unet_trt.state_dict(), "unet_trt.pth")
print("変換完了: unet_trt.pth を保存しました")
python3 convert_trt.py
変換中はGPU使用率が100%に張り付く。冷却ファン Noctua NF-A4x10 5V PWMファン（Amazon） をしっかり取り付けていないとサーマルスロットリングが発生するため注意が必要だ。
TensorRTエンジンを使った推論
# generate_trt.py
from diffusers import StableDiffusionPipeline
from torch2trt import TRTModule
import torch
model_id = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionPipeline.from_pretrained(
model_id,
torch_dtype=torch.float16,
safety_checker=None,
)
pipe = pipe.to("cuda")
# TRTエンジンに差し替え
unet_trt = TRTModule()
unet_trt.load_state_dict(torch.load("unet_trt.pth"))
pipe.unet = unet_trt
prompt = "a photo of a cat on a sunny beach, high quality"
image = pipe(prompt, num_inference_steps=20).images[0]
image.save("output_trt.png")
print("TRT推論完了")
python3 generate_trt.py
—
実測パフォーマンス比較
筆者の環境で計測した結果を以下に示す。
| 条件 | 生成時間（512×512, 20step） |
|——|————————–|
| 通常（float16） | 約70秒 |
| TensorRT最適化後 | 約28秒 |
TensorRT化によって約2.5倍の高速化を確認できた。初回のエンジン変換に時間がかかるが、一度変換してしまえば同じエンジンを使い回せる。SuperモデルではDLAコア併用でさらに高速化可能。
—
よくあるエラーと対処法
「CUDA out of memory」が出る
スワップの設定を再確認する。また enable_attention_slicing() に加えて以下も試す。
pipe.enable_sequential_cpu_offload()
torch2trtのビルドが失敗する
CUDAのバージョンとPyTorchのバージョンが一致していない場合に起きやすい。nvcc --version と torch.version.cuda が同じバージョンを示しているか確認する。TensorRTメモリ不足時はFP32レイヤー強制やopset17使用を検討。
変換後のエンジンで型エラーが出る
fp16_mode=True にしている場合、推論時の入力テンソルも .half() にする必要がある。
—
まとめ
Jetson Orin NanoでのStable Diffusion動作をまとめると次の通りだ。
Jetson Orin NanoはRaspberry Pi 5などと比べてGPUを持っているのが最大の強みだ。画像生成AIのような重い処理でその差が明確に出る。エッジでの画像生成を試してみたい方はぜひ挑戦してほしい。
—
次はこちら
TensorRT最適化をさらに深く理解したい方、あるいはほかのAIモデルをJetson上で動かしたい方は以下の記事も参考にしてほしい。
コメント