apextoaster
/
onnx-web
1
0
Fork 0
Code Actions

fix(api): add error handling for optimizations

This commit is contained in:
Sean Sube 2023-02-18 11:57:18 -06:00
parent ab6462d095
commit 118695d68c
Signed by: ssube
GPG Key ID: 3EED7B957D362AF1
1 changed files with 22 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

27
api/onnx_web/diffusion/load.py
View File

@ -94,24 +94,41 @@ def optimize_pipeline(
) -> None:
if "attention-slicing" in server.optimizations:
logger.debug("enabling attention slicing on SD pipeline")
pipe.enable_attention_slicing()
try:
pipe.enable_attention_slicing()
except Exception as e:
logger.warning("error enabling attention slicing: %s", e)
if "vae-slicing" in server.optimizations:
logger.debug("enabling VAE slicing on SD pipeline")
pipe.enable_vae_slicing()
try:
pipe.enable_vae_slicing()
except Exception as e:
logger.warning("error enabling VAE slicing: %s", e)
if "sequential-cpu-offload" in server.optimizations:
logger.debug("enabling sequential CPU offload on SD pipeline")
pipe.enable_sequential_cpu_offload()
try:
pipe.enable_sequential_cpu_offload()
except Exception as e:
logger.warning("error enabling sequential CPU offload: %s", e)
elif "model-cpu-offload" in server.optimizations:
# TODO: check for accelerate
logger.debug("enabling model CPU offload on SD pipeline")
pipe.enable_model_cpu_offload()
try:
pipe.enable_model_cpu_offload()
except Exception as e:
logger.warning("error enabling model CPU offload: %s", e)
if "memory-efficient-attention" in server.optimizations:
# TODO: check for xformers
logger.debug("enabling memory efficient attention for SD pipeline")
pipe.enable_xformers_memory_efficient_attention()
try:
pipe.enable_xformers_memory_efficient_attention()
except Exception as e:
logger.warning("error enabling memory efficient attention: %s", e)
def load_pipeline(