Add & use modules.errors.print_error where currently printing exception info by hand

1 files changed, 14 insertions, 12 deletions

diff --git a/modules/safe.py b/modules/safe.py
index e8f50774..b596f565 100644
--- a/modules/safe.py
+++ b/modules/safe.py
@@ -2,8 +2,6 @@
 

 import pickle

 import collections

-import sys

-import traceback

 

 import torch

 import numpy

@@ -11,6 +9,8 @@ import _codecs
 import zipfile

 import re

 

+from modules.errors import print_error

+

 # PyTorch 1.13 and later have _TypedStorage renamed to TypedStorage

 TypedStorage = torch.storage.TypedStorage if hasattr(torch.storage, 'TypedStorage') else torch.storage._TypedStorage

 

@@ -136,17 +136,20 @@ def load_with_extra(filename, extra_handler=None, *args, **kwargs):
             check_pt(filename, extra_handler)

 

     except pickle.UnpicklingError:

-        print(f"Error verifying pickled file from {filename}:", file=sys.stderr)

-        print(traceback.format_exc(), file=sys.stderr)

-        print("-----> !!!! The file is most likely corrupted !!!! <-----", file=sys.stderr)

-        print("You can skip this check with --disable-safe-unpickle commandline argument, but that is not going to help you.\n\n", file=sys.stderr)

+        print_error(

+            f"Error verifying pickled file from {filename}\n"

+            "-----> !!!! The file is most likely corrupted !!!! <-----\n"

+            "You can skip this check with --disable-safe-unpickle commandline argument, but that is not going to help you.\n\n",

+            exc_info=True,

+        )

         return None

-

     except Exception:

-        print(f"Error verifying pickled file from {filename}:", file=sys.stderr)

-        print(traceback.format_exc(), file=sys.stderr)

-        print("\nThe file may be malicious, so the program is not going to read it.", file=sys.stderr)

-        print("You can skip this check with --disable-safe-unpickle commandline argument.\n\n", file=sys.stderr)

+        print_error(

+            f"Error verifying pickled file from {filename}\n"

+            f"The file may be malicious, so the program is not going to read it.\n"

+            f"You can skip this check with --disable-safe-unpickle commandline argument.\n\n",

+            exc_info=True,

+        )

         return None

 

     return unsafe_torch_load(filename, *args, **kwargs)

@@ -190,4 +193,3 @@ with safe.Extra(handler):
 unsafe_torch_load = torch.load

 torch.load = load

 global_extra_handler = None

-