livlea blog

Live as if you were to die tomorrow. Learn as if you were to live forever. (Mohandas Karamchand Gandhi)

AMDの高速な画像生成 "Nitro-E" をRyzen AI Max+ 395で試す

AMDからNitro-Eという高速な画像生成手法が2025年10月にリリースされました。2025年11月8日時点では、ほとんどNitro-Eを試したという情報がなかったので試してみました。
最初に結論から述べると、画像生成することはできましたが、生成時間が遅いです。Ryzen AI Max+ 395を使うには、まだまだ未成熟ですね。

rocm.blogs.amd.com

Nitro-Eとは

Nitro-EとはAMDが開発した高速な画像生成手法です。非常に効率的なトレーニングに焦点を当てた、テキストから画像への拡散モデルで、学習にはAMD InstinctTM MI300X GPUを8基搭載する1つのノードでわずか1.5日しかかかりません。こちらは、ROCmのブログより転載したより速度と品質(GenEvalのスコア:指示された内容をどれだけ正確に忠実に画像として再現できるか)の比較グラフです。赤丸のE-MMDiT系がNitro-Eのことで、生成速度が早く品質も優れているとのこと。

ROCm Blogより転載

開発環境

PC:EVO-X2(128GBモデル。GPUに96GB割付)
OS:Ubuntu24.04 LTS

環境構築

公式の手順ではdocker imageを使うようになっていますが、ROCm 6.2.2版を指定しており古いのが気になったので、docker imageは使わずに独自で環境を準備しました。色々、手こずりましたが、画像生成するところまでは何とかこぎつけました。しかし、flash attentionのビルドがうまくいかずスキップしているので、この処理がボトルネックになってしまっている可能性があり、生成時間がとても長くなっています。解像度512x512の画像を1枚生成するのに、60秒ぐらいかかってしまいます。

github.com

細かい手順は記載せずに、コマンドと簡単な説明だけの記載に留めます。

git clone https://github.com/AMD-AGI/Nitro-E.git # 公式のGitHubからclone
cd Nitro-E # cloneしたNitro-Eに移動
uv init --python 3.12 # uvでプロジェクトを初期化
uv venv # 仮想環境の作成
source .venv/bin/activate # 仮想環境の有効化
uv pip install diffusers==0.32.2 transformers==4.49.0 accelerate==1.7.0 wandb torchmetrics pycocotools torchmetrics[image] mosaicml-streaming==0.11.0 beautifulsoup4 tabulate timm==0.9.1 pyarrow einops omegaconf sentencepiece==0.2.0 pandas==2.2.3 alive-progress ftfy peft

次にflash-attentionをインストールするのですが、ここがうまくいきません。現時点(2025年11月8日)でflash-attentionはRyzen AI Max+ 395に対応していないようです。コマンドだけ記載しますが、結局このやり方ではうまくいっていません。

cd .. # Nitro-Eの下にcloneしたくなかったので、Nitro-Eと同じ階層に移動して作業を進める。  
git clone https://github.com/ROCm/flash-attention.git # 公式レポジトリからclone
cd flash-attention
MAX_JOBS=`nproc` python setup.py install # Nitro-Eのためにuvで作成した仮想環境内で実行。nprocにて利用可能な全てのCPUコア数がわかるので、全てを使ってflash-attentionパッケージのコンパイルを行ってインストール。エラーで成功せず。

次にHugginfaceの設定をします。事前にHuggingfaceにログインしてアクセストークンを発行しておいてください。

uv pip install -U "huggingface_hub"
hf auth login # このコマンドを実行すると、アクセストークンの入力が求められるので入力

GitHubのReadmeに記載のあるFull-step Modelの推論のサンプルコードを適当な名前のファイルに記述して実行。私は inference.pyとしました。サンプルコードでは生成した画像を保存していないので、画像の保存も付け加えました。

# 画像生成のサンプルコード (inference.py)
import torch
from core.tools.inference_pipe import init_pipe
import datetime

device = torch.device('cuda:0')
dtype = torch.bfloat16
repo_name = "amd/Nitro-E"
resolution = 512
ckpt_name = 'Nitro-E-512px.safetensors'
grpo_name = 'ckpt_grpo_512px' #for grpo post training model
use_grpo = True

# for 1024px model
# resolution = 1024
# ckpt_name = 'Nitro-E-1024px.safetensors'
# grpo_name = 'ckpt_grpo_1024px' 

if use_grpo:
    pipe = init_pipe(device, dtype, resolution, repo_name=repo_name, ckpt_name=ckpt_name, ckpt_path_grpo=grpo_name)
else:
    pipe = init_pipe(device, dtype, resolution, repo_name=repo_name, ckpt_name=ckpt_name)
prompt = 'A hot air balloon in the shape of a heart grand canyon'
images = pipe(prompt=prompt, width=resolution, height=resolution, num_inference_steps=20, guidance_scale=4.5).images

if images:
    print(f"Number of images: {len(images)} ")
    first_image = images[0]

    now = datetime.datetime.now()
    timestamp = now.strftime("%Y%m%d_%H%M%S")
    output_filename = f"img_{timestamp}.png"
    first_image.save(output_filename)
    
    print(f"image saved to {output_filename}")
else:
    print("No image generated.")

生成テスト

このコードを実行します。

python inference.py # 画像生成の推論コードを実行
# 私は上記を実行した際に tranformers と hugginface_hub のバージョンのコンフリクトのエラーがでたので、以下でtransformersのアップグレードしました。
# uv pip install transformers -U 

一応、画像生成自体はできました。でも解像度512x512の画像生成に1分ぐらいかかります。(初回は、モデルのダウンロード等があるのでもっと時間はかかりました)
これでは、全然実用的ではないですね。Distilled Modelでも画像生成できましたが、こちらも生成時間はほぼ同じで1分ぐらいかかりました。flash attentionが使えるようにできれば、これが劇的に早くなるのか。。。私のやり方がどこか間違っているのか。。。AMD公式のROCmのブログでは Strix Halo (Ryzen AI Max+ 395) で0.16秒で生成できたとの記載があるのですが、本当なのか。。。その実行方法と環境構築手順も含めてGitHubに書いてもらいたいですね。

Nitro-EのFull-step Modelでの生成画像

Nitro-EのDistilled Modelでの生成画像

ちなみに、docker imageを使ってdocker上でも実行を試しましたが、同じく1分ぐらいかかってしまいました。メモとして環境構築の手順だけ記載しておきます。

dockerを使う場合(rocm7.1版を使っています)

sudo docker pull rocm/pytorch:rocm7.1_ubuntu24.04_py3.12_pytorch_release_2.8.0

sudo docker run -it --network=host --device=/dev/kfd --device=/dev/dri --group-add=video --ipc=host --cap-add=SYS_PTRACE --security-opt seccomp=unconfined --shm-size 8G -v $HOME/Nitro-E:/Nitro-E -w /Nitro-E rocm/pytorch /bin/bash

# 以下、dockerでつくった環境にて
pip install diffusers==0.32.2 transformers==4.49.0 accelerate==1.7.0 wandb torchmetrics pycocotools torchmetrics[image] mosaicml-streaming==0.11.0 beautifulsoup4 tabulate timm==0.9.1 pyarrow einops omegaconf sentencepiece==0.2.0 pandas==2.2.3 alive-progress ftfy peft

cd /home/ubuntu
git clone https://github.com/ROCm/flash-attention.git
cd flash-attention
MAX_JOBS=`nproc` python setup.py install # ここでもエラーでインストールできませんでした

まとめ

色々試しましたが、結局、Ryzen AI Max+ 395のUbuntuマシンで、Nitro-Eを使った高速な画像生成はできませんでした。画像生成自体は成功しましたが、1分ぐらいかかっています。もうしばらく、AMDのAPU向けの最適化が進むのを待つ必要がありそうです。