PixArt-α LoRA Training¶

You can also check configs/pixart_alpha_lora/README.md file.

Configs¶

All configuration files are placed under the configs/pixart_alpha_lora folder.

Following is the example config fixed from the pixart_alpha_1024_lora_pokemon_blip config file in configs/pixart_alpha_lora/pixart_alpha_1024_lora_pokemon_blip.py:

from mmengine.config import read_base

from diffengine.engine.hooks import PeftSaveHook, VisualizationHook

with read_base():
    from .._base_.datasets.pokemon_blip_pixart import *
    from .._base_.default_runtime import *
    from .._base_.models.pixart_alpha_1024_lora import *
    from .._base_.schedules.stable_diffusion_50e import *

optim_wrapper.update(
    dtype="bfloat16")

custom_hooks = [
    dict(
        type=VisualizationHook,
        prompt=["yoda pokemon"] * 4,
        height=1024,
        width=1024),
    dict(type=PeftSaveHook),
]

Run LoRA training¶

Run LoRA training:

# single gpu
$ diffengine train ${CONFIG_FILE}
# multi gpus
$ NPROC_PER_NODE=${GPU_NUM} diffengine train ${CONFIG_FILE}

# Example.
$ diffengine train pixart_alpha_1024_lora_pokemon_blip

Inference with diffusers¶

Once you have trained a model, specify the path to the saved model and utilize it for inference using the diffusers.pipeline module.

from pathlib import Path

import torch
from diffusers import PixArtAlphaPipeline, AutoencoderKL
from peft import PeftModel

checkpoint = Path('work_dirs/pixart_alpha_1024_lora_pokemon_blip/step20850')
prompt = 'yoda pokemon'

vae = AutoencoderKL.from_pretrained(
    'stabilityai/sd-vae-ft-ema',
)
pipe = PixArtAlphaPipeline.from_pretrained(
    "PixArt-alpha/PixArt-XL-2-1024-MS",
    vae=vae,
    torch_dtype=torch.float32,
).to("cuda")
pipe.transformer = PeftModel.from_pretrained(pipe.transformer, checkpoint / "transformer", adapter_name="default")

img = pipe(
    prompt,
    width=1024,
    height=1024,
    num_inference_steps=50,
).images[0]
img.save("demo.png")

Results Example¶

pixart_alpha_1024_lora_pokemon_blip¶

example1

You can check configs/pixart_alpha_lora/README.md for more details.