apextoaster
/
onnx-web
1
0
Fork 0
Code Actions

fix(api): load pretrained VAE from original path

This commit is contained in:
Sean Sube 2023-12-08 23:19:52 -06:00
parent 4fd50984f0
commit e7da2cf8a6
Signed by: ssube
GPG Key ID: 3EED7B957D362AF1
2 changed files with 4 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
api/onnx_web/convert/diffusion/diffusion.py
View File

@ -381,11 +381,11 @@ def convert_diffusion_diffusers(
if replace_vae is not None:
vae_path = path.join(conversion.model_path, replace_vae)
vae_file = check_ext(replace_vae, RESOLVE_FORMATS)
vae_file = check_ext(vae_path, RESOLVE_FORMATS)
if vae_file[0]:
pipeline.vae = AutoencoderKL.from_single_file(vae_path)
else:
pipeline.vae = AutoencoderKL.from_pretrained(vae_path)
pipeline.vae = AutoencoderKL.from_pretrained(replace_vae)
if is_torch_2_0:
pipeline.unet.set_attn_processor(AttnProcessor())

4
api/onnx_web/convert/diffusion/diffusion_xl.py
View File

@ -69,8 +69,8 @@ def convert_diffusion_diffusers_xl(
logger.debug("loading VAE from single tensor file: %s", vae_path)
pipeline.vae = AutoencoderKL.from_single_file(vae_path)
else:
logger.debug("loading pretrained VAE from path: %s", vae_path)
pipeline.vae = AutoencoderKL.from_pretrained(vae_path)
logger.debug("loading pretrained VAE from path: %s", replace_vae)
pipeline.vae = AutoencoderKL.from_pretrained(replace_vae)
if path.exists(temp_path):
logger.debug("torch model already exists for %s: %s", source, temp_path)