From 2c463ae3dc4c00c6d6af3bc0a5a355d37520a4a9 Mon Sep 17 00:00:00 2001 From: tazlin Date: Wed, 29 May 2024 07:00:54 -0400 Subject: [PATCH] fix: handle custom models in `keep_single_inference` This should prevent crashes when QR code/CN workflows are used with custom models --- horde_worker_regen/process_management/process_manager.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/horde_worker_regen/process_management/process_manager.py b/horde_worker_regen/process_management/process_manager.py index 51d2a0cf..c85d5bc0 100644 --- a/horde_worker_regen/process_management/process_manager.py +++ b/horde_worker_regen/process_management/process_manager.py @@ -519,7 +519,11 @@ def keep_single_inference( ) continue - model_info = stable_diffusion_model_reference.root[model] + model_info = stable_diffusion_model_reference.root.get(model) + if model_info is None: + logger.debug(f"Model {model} not found in stable diffusion model reference. Is it a custom model?") + continue + if model_info.baseline == STABLE_DIFFUSION_BASELINE_CATEGORY.stable_diffusion_xl and ( p.can_accept_job() or p.last_process_state == HordeProcessState.PRELOADING_MODEL