ByteDanceの画像生成AIのSDXL-LightningのGoogle Colabでの使い方

ByteDanceといえば、tiktokの会社ですね。そこが、画像生成AIで、軽くて早く実行できるというSDXL-Lightningを出しましたので、Google Colabでの使い方を紹介します。

ByteDance/SDXL-Lightning · Hugging Face

今回は、Google ColabでA100で上記のコードを少し修正して、画像を30枚作成できるコードとなります。

!pip install diffusers
import torch
from diffusers import StableDiffusionXLPipeline, UNet2DConditionModel, EulerDiscreteScheduler
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file

base = "stabilityai/stable-diffusion-xl-base-1.0"
repo = "ByteDance/SDXL-Lightning"
ckpt = "sdxl_lightning_4step_unet.safetensors" # Use the correct ckpt for your step setting!

# Load model.
unet = UNet2DConditionModel.from_config(base, subfolder="unet").to("cuda", torch.float16)
unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device="cuda"))
pipe = StableDiffusionXLPipeline.from_pretrained(base, unet=unet, torch_dtype=torch.float16, variant="fp16").to("cuda")

# Ensure sampler uses "trailing" timesteps.
pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")

# Ensure using the same inference steps as the loaded model and CFG set to 0.
for i in range(30) :
  pipe("A girl smiling", num_inference_steps=4, guidance_scale=0).images[0].save("output"+str(i)+".png")

上記を実行しますと、30枚、Google Colabの/content配下に作成されます。

作成された画像は下記となります。3枚だけ表示しています。

所感としては、作成は軽くて早く出来た印象を受けます。一つのプロンプトしか試していませんが、生成される画像の質も高いようです。Tiktok運営しているByteDanceだから、画像生成AIのための学習データはたくさんあるように見受けられます。