[SDXL Turbo] 爆速画像生成AIで柴犬ジェネレーターを作ってみました

テキストプロンプトのみで、思い通りの画像を生成するのは、実は、結構難しい作業だったりします。この解決策の模索として、とにかく多数の画像を生成して、その中から、イメージにマッチするものを選ぶという作業を試してみました。

#Stable Diffusion

#生成AI

#Jetson

SIN

2024.10.05

1 はじめに

製造ビジネステクノロジー部の平内（SIN）です。

テキストから画像が生成できるAIは、イメージを言葉で伝えるだけで、高品質な画像が生成できることから、様々なシーンでの活用が検討されていると思います。

しかし、テキストプロンプトのみで、思い通りの画像を生成するのは、実は、結構難しい作業だったりします。「なかなか、思った通りの絵にならない」というところです。

そして、その解決策の１つが、とにかく多数の画像を生成して、その中から、イメージにマッチするものを選ぶというものです。

今回は、Stability AIの SDXL Turboを使用して、この作業を試してみました。

001

SDXL Turbo は、敵対的拡散蒸留 Adversarial Diffusion Distillation：ADD と呼ばれる新しい蒸留技術に基づき、たったの1ステップで、超高速に画像を出力できます。

最初に、試してみた「柴犬ジェネレーター」を見てやってください。

プロンプトには、デフォルトで「柴犬」が設定されており、入力したテキスト及び、「ANIME」「PHOTO」の選択が、プロンプトに付加され、モデルに送られています。

GENERATEボタンを押す度に、1枚ずつ画像が生成されます。そして、Auto を ON にすることで、勝手にGENERATEボタンが次々押されるようになります。

注意：SDXL Turboの商用利用は、https://stability.ai/license を参照ください。

2 構成

試してみた、SDXL Turbo　の利用形態は、おおむね以下のとおりです。

画像生成は、Jetson AGX Orinで行っています
SDXL Turboは、HuggingFace Diffusersから利用しています
FastAPIでPythonによるRestAPIを構成し、画像生成等を行っています
ブラウザのUIは、ローカルでホストしているReactで動作しています
プロンプト入力は、日本語になっていますが、APIに送る前に、Amazon Translateで英語に翻訳しています

002

3 HuggingFace Diffusers

モデルの実行には、HuggingFace の Diffusers を使用させていただきました。

Diffusersは、トレーニング済みの拡散モデルを扱うライブラリで、推論のためのパイプラインを数行のコードで構成できます。

AutoPipelineForText2Imageで、Hugging Face Hub の任意のチェックポイントを使用できるので、ここで SDXL Turboを指定することでモデルをダウンロードしています。

また、高速化のため、精度を落としましたが（torch_dtype=torch.float16）、生成される画像は、殆どかわりませんでした。

SDXL Turboは、通常の StableDeffusion のモデルとは、パラメータが違い、ネガティブプロンプトなどが無いことに注意が必要です。

from diffusers import AutoPipelineForText2Image

pipe = AutoPipelineForText2Image.from_pretrained(
    "stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16"
)
pipe.to("cuda")

seed = random.randint(1, 10000)
generator = torch.Generator("cuda").manual_seed(seed)

image = pipe(
    prompt="prompt",
    num_inference_steps=1,
    generator=generator,
    guidance_scale=0.0,
).images[0]
image.save(output_filename)

4 FastAPI

APIを構成するために、FastAPIを使用しています。

/generateにPOSTを受けると、指定されたパラメータに基づいて、Diffusersのパイプラインで画像を生成しています。

class Params(BaseModel):
    prompt: str
    steps: int
    seed: int

# 画像生成のエンドポイント
@app.post("/generate")
async def generate(params: Params):
    prompt = params.prompt
    timestamp = int(time.time() * 1000)
    output_filename = f"{BASE_PATH}/{OUTPUT_PATH}/{timestamp}.png"
    seed = random.randint(1, 10000) if params.seed == -1 else params.seed
    steps = params.steps

    # 画像を生成
    generate_data(
        pipe,
        prompt,
        seed,
        steps,
        output_filename,
    )

    # 生成された画像のURLを返す
    url_path = f"/{OUTPUT_PATH}/{timestamp}.png"
    return {"url": url_path}

生成された画像は、ファイル名を指定して、/output からGETで取得するようになっています。

# 指定されたファイル名の出力ファイルを取得するエンドポイント
@app.get("/output/{filename}")
async def get_output_file(filename):
    file_path = f"{BASE_PATH}/{OUTPUT_PATH}/{filename}"
    if os.path.exists(file_path):
        return FileResponse(file_path, media_type="image/png")
    else:
        return {"error": "File not found"}

全てのコードです。
app.py

5 Docker

Jetson AGX Orin上では、Dockerで環境を構築しています。

JetPack6.1用のベースイメージとして、dustynv/jetson-inference:r36.3.0を使用することで、Pytorch等の環境構築が完了できます。

環境変数「XDG_CACHE_HOME」は、Diffusersがモデルをダウンロードするディレクトリをマウントしたディスクにするために設定しています。

FROM dustynv/jetson-inference:r36.3.0
RUN pip3 install diffusers==0.30.3 transformers==4.45.1 ftfy accelerate fastapi uvicorn
COPY api.py start.sh /
RUN chmod 755 /start.sh
# diffusersでダウンロードされるモデルをマウントしたディスクにするためキャッシュ指定設定する
ENV XDG_CACHE_HOME="/data"
CMD ["/bin/bash","/start.sh"]

Dockerfile
start.sh
README.md

x001

# python3 --version
Python 3.10.12

# python3 -c "import torch;print(torch.__version__)"
2.2.0

# python3 -c "import torch;print(torch.cuda.is_available())"
True