diff --git a/CHANGES.md b/CHANGES.md index ac31a612..2151b8e8 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -20,6 +20,7 @@ - A `What's New?` tab in the UI ### Detailed changelog +* 2.4.7 - 18 Nov 2022 - Don't crash if a VAE file fails to load * 2.4.7 - 17 Nov 2022 - Fix a bug where Face Correction (GFPGAN) would fail on cuda:N (i.e. GPUs other than cuda:0), as well as fail on CPU if the system had an incompatible GPU. * 2.4.6 - 16 Nov 2022 - Fix a regression in VRAM usage during startup, which caused 'Out of Memory' errors when starting on GPUs with 4gb (or less) VRAM * 2.4.5 - 16 Nov 2022 - Add checkbox for "Open browser on startup". diff --git a/ui/sd_internal/runtime.py b/ui/sd_internal/runtime.py index 2308298e..486d2179 100644 --- a/ui/sd_internal/runtime.py +++ b/ui/sd_internal/runtime.py @@ -140,15 +140,24 @@ def load_model_ckpt(): _, _ = modelFS.load_state_dict(sd, strict=False) if thread_data.vae_file is not None: - for model_extension in ['.ckpt', '.vae.pt']: - if os.path.exists(thread_data.vae_file + model_extension): - print(f"Loading VAE weights from: {thread_data.vae_file}{model_extension}") - vae_ckpt = torch.load(thread_data.vae_file + model_extension, map_location="cpu") - vae_dict = {k: v for k, v in vae_ckpt["state_dict"].items() if k[0:4] != "loss"} - modelFS.first_stage_model.load_state_dict(vae_dict, strict=False) - break - else: - print(f'Cannot find VAE file: {thread_data.vae_file}{model_extension}') + try: + loaded = False + for model_extension in ['.ckpt', '.vae.pt']: + if os.path.exists(thread_data.vae_file + model_extension): + print(f"Loading VAE weights from: {thread_data.vae_file}{model_extension}") + vae_ckpt = torch.load(thread_data.vae_file + model_extension, map_location="cpu") + vae_dict = {k: v for k, v in vae_ckpt["state_dict"].items() if k[0:4] != "loss"} + modelFS.first_stage_model.load_state_dict(vae_dict, strict=False) + loaded = True + break + + if not loaded: + print(f'Cannot find VAE: {thread_data.vae_file}') + thread_data.vae_file = None + except: + print(traceback.format_exc()) + print(f'Could not load VAE: {thread_data.vae_file}') + thread_data.vae_file = None modelFS.eval() # if thread_data.device != 'cpu': diff --git a/ui/sd_internal/task_manager.py b/ui/sd_internal/task_manager.py index 45f90f9f..4b96438c 100644 --- a/ui/sd_internal/task_manager.py +++ b/ui/sd_internal/task_manager.py @@ -436,7 +436,7 @@ def stop_render_thread(device): try: device_manager.validate_device_id(device, log_prefix='stop_render_thread') except: - print(traceback.format_exec()) + print(traceback.format_exc()) return False if not manager_lock.acquire(blocking=True, timeout=LOCK_TIMEOUT): raise Exception('stop_render_thread' + ERR_LOCK_FAILED)