From 3829afec365b748e330da33b00a0e363f8c8ab71 Mon Sep 17 00:00:00 2001
From: linkoid <36754150+linkoid@users.noreply.github.com>
Date: Fri, 26 May 2023 15:08:53 -0400
Subject: [PATCH 1/2] Remove exit() from select_checkpoint()

Raising a FileNotFoundError instead.
---
 modules/sd_models.py | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/modules/sd_models.py b/modules/sd_models.py
index 91b3eb11..1871cc97 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -164,6 +164,7 @@ def model_hash(filename):
 
 
 def select_checkpoint():
+    """Raises `FileNotFoundError` if no checkpoints are found."""
     model_checkpoint = shared.opts.sd_model_checkpoint
 
     checkpoint_info = checkpoint_alisases.get(model_checkpoint, None)
@@ -171,14 +172,14 @@ def select_checkpoint():
         return checkpoint_info
 
     if len(checkpoints_list) == 0:
-        print("No checkpoints found. When searching for checkpoints, looked at:", file=sys.stderr)
+        error_message = "No checkpoints found. When searching for checkpoints, looked at:"
         if shared.cmd_opts.ckpt is not None:
-            print(f" - file {os.path.abspath(shared.cmd_opts.ckpt)}", file=sys.stderr)
-        print(f" - directory {model_path}", file=sys.stderr)
+            error_message += f"\n - file {os.path.abspath(shared.cmd_opts.ckpt)}"
+        error_message += f"\n - directory {model_path}"
         if shared.cmd_opts.ckpt_dir is not None:
-            print(f" - directory {os.path.abspath(shared.cmd_opts.ckpt_dir)}", file=sys.stderr)
-        print("Can't run without a checkpoint. Find and place a .ckpt or .safetensors file into any of those locations. The program will exit.", file=sys.stderr)
-        exit(1)
+            error_message += f"\n - directory {os.path.abspath(shared.cmd_opts.ckpt_dir)}"
+        error_message += "Can't run without a checkpoint. Find and place a .ckpt or .safetensors file into any of those locations."
+        raise FileNotFoundError(error_message)
 
     checkpoint_info = next(iter(checkpoints_list.values()))
     if model_checkpoint is not None:

From 1f0fdede176989f151da6b97bd9a140b7f0af6e5 Mon Sep 17 00:00:00 2001
From: linkoid <36754150+linkoid@users.noreply.github.com>
Date: Fri, 26 May 2023 15:15:59 -0400
Subject: [PATCH 2/2] Show full traceback in get_sd_model()

to reveal if an error is caused by an extension
---
 modules/errors.py    | 8 ++++++--
 modules/sd_models.py | 2 +-
 2 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/modules/errors.py b/modules/errors.py
index f6b80dbb..da4694f8 100644
--- a/modules/errors.py
+++ b/modules/errors.py
@@ -12,9 +12,13 @@ def print_error_explanation(message):
     print('=' * max_len, file=sys.stderr)
 
 
-def display(e: Exception, task):
+def display(e: Exception, task, *, full_traceback=False):
     print(f"{task or 'error'}: {type(e).__name__}", file=sys.stderr)
-    print(traceback.format_exc(), file=sys.stderr)
+    te = traceback.TracebackException.from_exception(e)
+    if full_traceback:
+        # include frames leading up to the try-catch block
+        te.stack = traceback.StackSummary(traceback.extract_stack()[:-2] + te.stack)
+    print(*te.format(), sep="", file=sys.stderr)
 
     message = str(e)
     if "copying a param with shape torch.Size([640, 1024]) from checkpoint, the shape in current model is torch.Size([640, 768])" in message:
diff --git a/modules/sd_models.py b/modules/sd_models.py
index 1871cc97..3e7fc7e3 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -424,7 +424,7 @@ class SdModelData:
                 try:
                     load_model()
                 except Exception as e:
-                    errors.display(e, "loading stable diffusion model")
+                    errors.display(e, "loading stable diffusion model", full_traceback=True)
                     print("", file=sys.stderr)
                     print("Stable diffusion model failed to load", file=sys.stderr)
                     self.sd_model = None