From 118695d68cd7cc37df1de66ffe66c0f83dc6c819 Mon Sep 17 00:00:00 2001 From: Sean Sube Date: Sat, 18 Feb 2023 11:57:18 -0600 Subject: [PATCH] fix(api): add error handling for optimizations --- api/onnx_web/diffusion/load.py | 27 ++++++++++++++++++++++----- 1 file changed, 22 insertions(+), 5 deletions(-) diff --git a/api/onnx_web/diffusion/load.py b/api/onnx_web/diffusion/load.py index 53cd6af0..38afb0ab 100644 --- a/api/onnx_web/diffusion/load.py +++ b/api/onnx_web/diffusion/load.py @@ -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(