mirror of
https://github.com/ostris/ai-toolkit.git
synced 2026-01-26 16:39:47 +00:00
21 lines
520 B
Python
21 lines
520 B
Python
from diffusers import AutoencoderKL
|
|
|
|
|
|
def load_vae(vae_path, dtype):
|
|
try:
|
|
vae = AutoencoderKL.from_pretrained(
|
|
vae_path,
|
|
torch_dtype=dtype,
|
|
)
|
|
except Exception as e:
|
|
try:
|
|
vae = AutoencoderKL.from_pretrained(
|
|
vae_path,
|
|
subfolder="vae",
|
|
torch_dtype=dtype,
|
|
)
|
|
except Exception as e:
|
|
raise ValueError(f"Failed to load VAE from {vae_path}: {e}")
|
|
vae.to(dtype)
|
|
return vae
|