From e9e9d75f9df67df25be80f77f95bcbefa9575014 Mon Sep 17 00:00:00 2001 From: Sean Sube Date: Thu, 20 Apr 2023 18:00:35 -0500 Subject: [PATCH] fix(api): correct return value after fully converting checkpoints --- api/onnx_web/convert/diffusion/original.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/api/onnx_web/convert/diffusion/original.py b/api/onnx_web/convert/diffusion/original.py index 03599eb1..80f996d4 100644 --- a/api/onnx_web/convert/diffusion/original.py +++ b/api/onnx_web/convert/diffusion/original.py @@ -1645,7 +1645,8 @@ def extract_checkpoint( if not os.path.exists(full_path): result_status = f"Missing model directory, removing model: {full_path}" shutil.rmtree(db_config.model_dir, ignore_errors=False, onerror=None) - break + return False + remove_dirs = ["logging", "samples"] for rd in remove_dirs: rem_dir = os.path.join(db_config.model_dir, rd) @@ -1655,6 +1656,7 @@ def extract_checkpoint( os.makedirs(rem_dir) logger.info(result_status) + return True @torch.no_grad()