Stable Diffusion Accelerated (SDA)
Note: Most of the code is borrowed from nvidia's repo
This is a simplified library to use stable diffusion with TensorRT
Usage
Import
from sda.diffusion_mode_pipeline import DiffusionModePipelineInitiate Pipeline
pipe = DiffusionModePipeline(
max_batch_size=16,
device="cuda",
verbose=False,
min_image_shape=256,
max_image_shape=1024,
text_max_len=77
)Remember that the models must match the configuration
Load Models
pipe.loadEngine(
path="/shared/models/tensor/unet.plan",
alias="SD2"
)
pipe.loadEngine(
path="/shared/models/tensor/clip.plan",
alias="SD2"
)
pipe.loadEngine(
path="/shared/models/tensor/vae.plan",
alias="SD2"
)You can load multiple models on a single GPU
Prepare and Generate
img = pipe.infer(
prompt=["A car"],
negative_prompt=[""],
steps=50,
models={"unet": "SD2", "textencoder": "SD2", "vae": "SD2"},
scheduler="DDIM",
cfg=7.5,
image_height=512,
image_width=512,
seed=54321
)
img.save("output.png")Configuration File
The compiler isn't here yet, but you can just run nvidia's repo for now, and then add a json file right next to the models.
For example, for anime_unet.plan, anime_unet.json should have:
{
"alias": "STABLE_DIFFUSION-2_1",
"type": "unet",
"version": 2,
"inpaint": false,
"max_batch_size": 4,
"fp16": true
}Adjust it to the configuration used during compilation