Bugfixes for multi-vector token handling

1 files changed, 9 insertions, 4 deletions

diff --git a/infer.py b/infer.py
index 4bcaff5..f88245a 100644
--- a/infer.py
+++ b/infer.py
@@ -7,6 +7,8 @@ import cmd
 from pathlib import Path
 import torch
 import json
+import traceback
+
 from PIL import Image
 from slugify import slugify
 from diffusers import (
@@ -165,8 +167,8 @@ def run_parser(parser, defaults, input=None):
     conf_args = argparse.Namespace()
 
     if args.config is not None:
-        args = load_config(args.config)
-        args = parser.parse_args(namespace=argparse.Namespace(**args))
+        conf_args = load_config(args.config)
+        conf_args = parser.parse_known_args(namespace=argparse.Namespace(**conf_args))[0]
 
     res = defaults.copy()
     for dict in [vars(conf_args), vars(args)]:
@@ -295,6 +297,7 @@ class CmdParse(cmd.Cmd):
             elements = shlex.split(line)
         except ValueError as e:
             print(str(e))
+            return
 
         if elements[0] == 'q':
             return True
@@ -306,9 +309,11 @@ class CmdParse(cmd.Cmd):
                 print('Try again with a prompt!')
                 return
         except SystemExit:
+            traceback.print_exc()
             self.parser.print_help()
+            return
         except Exception as e:
-            print(e)
+            traceback.print_exc()
             return
 
         try:
@@ -316,7 +321,7 @@ class CmdParse(cmd.Cmd):
         except KeyboardInterrupt:
             print('Generation cancelled.')
         except Exception as e:
-            print(e)
+            traceback.print_exc()
             return
 
     def do_exit(self, line):