revise sd_model_checkpoint

AIRI-AUBE · Dec 7, 2022 · 1eb35d5 · 1eb35d5
1 parent 366b4c1
commit 1eb35d5
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 15 deletions.
diff --git a/modules/sd_models.py b/modules/sd_models.py
@@ -74,17 +74,20 @@ def __init__(self, sd_model_name, sd_model_hash, sd_model_checkpoint, sd_checkpo
                     shared.opts.data['sd_model_checkpoint'] = title
 
                 checkpoints_list[title] = CheckpointInfo(filename, title, h, short_model_name, config)
-            
+
             sd_model_checkpoint = shared.opts.data['sd_model_checkpoint']
-            sd_checkpoint_info = checkpoints_list[sd_model_checkpoint]
-            sd_model_name = checkpoints_list[sd_model_checkpoint].model_name
-            sd_model_hash = checkpoints_list[sd_model_checkpoint].hash
-            shared.sd_model = SDModel(
-                sd_model_name,
-                sd_model_hash,
-                sd_model_checkpoint,
-                sd_checkpoint_info
-            )
+            if sd_model_checkpoint:
+                sd_checkpoint_info = checkpoints_list[sd_model_checkpoint]
+                sd_model_name = checkpoints_list[sd_model_checkpoint].model_name
+                sd_model_hash = checkpoints_list[sd_model_checkpoint].hash
+                shared.sd_model = SDModel(
+                    sd_model_name,
+                    sd_model_hash,
+                    sd_model_checkpoint,
+                    sd_checkpoint_info
+                )
+            else:
+                shared.sd_model = None
     else:
         model_list = modelloader.load_models(model_path=model_path, command_path=shared.cmd_opts.ckpt_dir, ext_filter=[".ckpt"])
 

diff --git a/modules/ui.py b/modules/ui.py
@@ -2172,14 +2172,15 @@ def user_signup(signup_username, signup_password, signup_email):
                 return {
                     signup_output: gr.update(value='Signup failed, please check and retry again')
                 }                    
-                
+
         def user_signout():
             shared.username=''
             opts.data = shared.default_options
-            for key in sd_models.checkpoints_list:
-                if sd_models.checkpoints_list[key].title == opts.data['sd_model_checkpoint']:
-                    shared.sd_model.sd_model_name = sd_models.checkpoints_list[key].model_name
-                    break            
+            if 'sd_model_checkpoint' in opts.data:
+                for key in sd_models.checkpoints_list:
+                    if sd_models.checkpoints_list[key].title == opts.data['sd_model_checkpoint']:
+                        shared.sd_model.sd_model_name = sd_models.checkpoints_list[key].model_name
+                        break            
 
             response = {
                 user_login_row : gr.update(visible=False),