Crate pyke_diffusers

pyke-diffusers is a modular library for pretrained diffusion model inference using ONNX Runtime, inspired by HuggingFace diffusers.

ONNX Runtime provides optimized inference for both CPUs and GPUs, including both NVIDIA & AMD GPUs via DirectML.

pyke-diffusers is focused on ease of use, with an API closely modeled after HuggingFace diffusers:

use std::sync::Arc;

use pyke_diffusers::{
	EulerDiscreteScheduler, OrtEnvironment, SchedulerOptimizedDefaults, StableDiffusionOptions, StableDiffusionPipeline,
	StableDiffusionTxt2ImgOptions
};

let environment = Arc::new(OrtEnvironment::builder().build()?);
let mut scheduler = EulerDiscreteScheduler::stable_diffusion_v1_optimized_default()?;
let pipeline =
	StableDiffusionPipeline::new(&environment, "./stable-diffusion-v1-5/", StableDiffusionOptions::default())?;

let imgs = pipeline.txt2img("photo of a red fox", &mut scheduler, StableDiffusionTxt2ImgOptions::default())?;

See StableDiffusionPipeline for more info on the Stable Diffusion pipeline.

Re-exports

pub use self::pipelines::*;

pub use self::schedulers::*;

Modules

pipelines

Diffusion pipelines.

schedulers

The schedule functions, denoted Schedulers in the library, take in the output of a trained model, a sample which the diffusion process is iterating on, and a timestep, returning a denoised sample.

Structs

CUDADeviceOptions

Device options for the CUDA execution provider.

DiffusionDeviceControl

Select which device each model should be placed on.

OrtEnvironment

An Environment is the main entry point of the ONNX Runtime.

Enums

ArenaExtendStrategy

The strategy to use for extending the device memory arena.

CuDNNConvolutionAlgorithmSearch

The type of search done for cuDNN convolution algorithms.

DiffusionDevice

A device on which to place a diffusion model on.