Repository: AUTOMATIC1111/stable-diffusion-webui

Issue:
[Bug]: Cannot switch checkpoints
### Is there an existing issue for this?

- [X] I have searched the existing issues and checked the recent builds/commits

### What happened?

I cannot change the checkpoint in the WebUI anymore since updating today. This is the error message I get:

LatentDiffusion: Running in eps-prediction mode
Traceback (most recent call last):
  File "E:\sd\stable-diffusion-webui\venv\lib\site-packages\gradio\routes.py", line 284, in run_predict
    output = await app.blocks.process_api(
  File "E:\sd\stable-diffusion-webui\venv\lib\site-packages\gradio\blocks.py", line 982, in process_api
    result = await self.call_function(fn_index, inputs, iterator)
  File "E:\sd\stable-diffusion-webui\venv\lib\site-packages\gradio\blocks.py", line 824, in call_function
    prediction = await anyio.to_thread.run_sync(
  File "E:\sd\stable-diffusion-webui\venv\lib\site-packages\anyio\to_thread.py", line 31, in run_sync
    return await get_asynclib().run_sync_in_worker_thread(
  File "E:\sd\stable-diffusion-webui\venv\lib\site-packages\anyio\_backends\_asyncio.py", line 937, in run_sync_in_worker_thread
    return await future
  File "E:\sd\stable-diffusion-webui\venv\lib\site-packages\anyio\_backends\_asyncio.py", line 867, in run
    result = context.run(func, *args)
  File "E:\sd\stable-diffusion-webui\modules\ui.py", line 443, in update_token_counter
    tokens, token_count, max_length = max([model_hijack.tokenize(prompt) for prompt in prompts], key=lambda args: args[1])
  File "E:\sd\stable-diffusion-webui\modules\ui.py", line 443, in <listcomp>
    tokens, token_count, max_length = max([model_hijack.tokenize(prompt) for prompt in prompts], key=lambda args: args[1])
  File "E:\sd\stable-diffusion-webui\modules\sd_hijack.py", line 116, in tokenize
    _, remade_batch_tokens, _, _, _, token_count = self.clip.process_text([text])
AttributeError: 'NoneType' object has no attribute 'process_text'
DiffusionWrapper has 859.52 M params.
making attention of type 'vanilla' with 512 in_channels
Working with z of shape (1, 4, 32, 32) = 4096 dimensions.
making attention of type 'vanilla' with 512 in_channels
Traceback (most recent call last):
  File "E:\sd\stable-diffusion-webui\venv\lib\site-packages\gradio\routes.py", line 284, in run_predict
    output = await app.blocks.process_api(
  File "E:\sd\stable-diffusion-webui\venv\lib\site-packages\gradio\blocks.py", line 982, in process_api
    result = await self.call_function(fn_index, inputs, iterator)
  File "E:\sd\stable-diffusion-webui\venv\lib\site-packages\gradio\blocks.py", line 824, in call_function
    prediction = await anyio.to_thread.run_sync(
  File "E:\sd\stable-diffusion-webui\venv\lib\site-packages\anyio\to_thread.py", line 31, in run_sync
    return await get_asynclib().run_sync_in_worker_thread(
  File "E:\sd\stable-diffusion-webui\venv\lib\site-packages\anyio\_backends\_asyncio.py", line 937, in run_sync_in_worker_thread
    return await future
  File "E:\sd\stable-diffusion-webui\venv\lib\site-packages\anyio\_backends\_asyncio.py", line 867, in run
    result = context.run(func, *args)
  File "E:\sd\stable-diffusion-webui\modules\ui.py", line 1662, in <lambda>
    fn=lambda value, k=k: run_settings_single(value, key=k),
  File "E:\sd\stable-diffusion-webui\modules\ui.py", line 1504, in run_settings_single
    opts.data_labels[key].onchange()
  File "E:\sd\stable-diffusion-webui\webui.py", line 41, in f
    res = func(*args, **kwargs)
  File "E:\sd\stable-diffusion-webui\webui.py", line 83, in <lambda>
    shared.opts.onchange("sd_model_checkpoint", wrap_queued_call(lambda: modules.sd_models.reload_model_weights()))
  File "E:\sd\stable-diffusion-webui\modules\sd_models.py", line 285, in reload_model_weights
    load_model(checkpoint_info)
  File "E:\sd\stable-diffusion-webui\modules\sd_models.py", line 254, in load_model
    load_model_weights(sd_model, checkpoint_info)
  File "E:\sd\stable-diffusion-webui\modules\sd_models.py", line 169, in load_model_weights
    sd_vae.restore_base_vae(model)
  File "E:\sd\stable-diffusion-webui\modules\sd_vae.py", line 54, in restore_base_vae
    if base_vae is not None and checkpoint_info == model.sd_checkpoint_info:
  File "E:\sd\stable-diffusion-webui\venv\lib\site-packages\torch\nn\modules\module.py", line 1207, in __getattr__
    raise AttributeError("'{}' object has no attribute '{}'".format(
AttributeError: 'LatentDiffusion' object has no attribute 'sd_checkpoint_info'

Then, I cannot generate any images. I get this error message:

Error completing request
Arguments: ('ppp', '', 'None', 'None', 36, 2, False, False, 1, 6, 7, -1.0, -1.0, 0, 0, 0, False, 512, 512, False, 0.7, 0, 0, 0, 0.9, 5, '0.0001', False, 'None', '', 0.1, False, '<div id="dynamic-prompting">\n    <h3><strong>Combinations</strong></h3>\n\n    Choose a number of terms from a list, in this case we choose two artists: \n    <code class="codeblock">{2$$artist1|artist2|artist3}</code><br/>\n\n    If $$ is not provided, then 1$$ is assumed.<br/><br/>\n\n    If the chosen number of terms is greater than the available terms, then some terms will be duplicated, otherwise chosen terms will be unique. This is useful in the case of wildcards, e.g.\n    <code class="codeblock">{2$$__artist__}</code> is equivalent to <code class="codeblock">{2$$__artist__|__artist__}</code><br/><br/>\n\n    A range can be provided:\n    <code class="codeblock">{1-3$$artist1|artist2|artist3}</code><br/>\n    In this case, a random number of artists between 1 and 3 is chosen.<br/><br/>\n\n    Wildcards can be used and the joiner can also be specified:\n    <code class="codeblock">{{1-$$and$$__adjective__}}</code><br/>\n\n    Here, a random number between 1 and 3 words from adjective.txt will be chosen and joined together with the word \'and\' instead of the default comma.\n\n    <br/><br/>\n\n    <h3><strong>Wildcards</strong></h3>\n    \n\n    <br/>\n    If the groups wont drop down click <strong onclick="check_collapsibles()" style="cursor: pointer">here</strong> to fix the issue.\n\n    <br/><br/>\n\n    <code class="codeblock">WILDCARD_DIR: E:\\sd\\stable-diffusion-webui\\extensions\\sd-dynamic-prompts\\wildcards</code><br/>\n    <small onload="check_collapsibles()">You can add more wildcards by creating a text file with one term per line and name is mywildcards.txt. Place it in E:\\sd\\stable-diffusion-webui\\extensions\\sd-dynamic-prompts\\wildcards. <code class="codeblock">__&#60;folder&#62;/mywildcards__</code> will then become available.</small>\n</div>\n\n', False, 1, False, 100, 0.7, False, False, 'Not set', 'Not set', 'Not set', 'Not set', 'No focus', 'None', False, False, False, '', 1, '', 0, '', True, False, False, '', 'None', 30, 4, 0, 0, False, 'None', '<br>', 'None', 30, 4, 0, 0, 4, 0.4, True, 32, 1.0, 2.0, 'a painting in', 'style', 'picture frame, portrait photo', None) {}
Traceback (most recent call last):
  File "E:\sd\stable-diffusion-webui\venv\lib\site-packages\gradio\routes.py", line 284, in run_predict
    output = await app.blocks.process_api(
  File "E:\sd\stable-diffusion-webui\venv\lib\site-packages\gradio\blocks.py", line 982, in process_api
    result = await self.call_function(fn_index, inputs, iterator)
  File "E:\sd\stable-diffusion-webui\venv\lib\site-packages\gradio\blocks.py", line 824, in call_function
    prediction = await anyio.to_thread.run_sync(
  File "E:\sd\stable-diffusion-webui\venv\lib\site-packages\anyio\to_thread.py", line 31, in run_sync
    return await get_asynclib().run_sync_in_worker_thread(
  File "E:\sd\stable-diffusion-webui\venv\lib\site-packages\anyio\_backends\_asyncio.py", line 937, in run_sync_in_worker_thread
    return await future
  File "E:\sd\stable-diffusion-webui\venv\lib\site-packages\anyio\_backends\_asyncio.py", line 867, in run
    result = context.run(func, *args)
  File "E:\sd\stable-diffusion-webui\modules\ui.py", line 443, in update_token_counter
    tokens, token_count, max_length = max([model_hijack.tokenize(prompt) for prompt in prompts], key=lambda args: args[1])
  File "E:\sd\stable-diffusion-webui\modules\ui.py", line 443, in <listcomp>
    tokens, token_count, max_length = max([model_hijack.tokenize(prompt) for prompt in prompts], key=lambda args: args[1])
  File "E:\sd\stable-diffusion-webui\modules\sd_hijack.py", line 116, in tokenize
    _, remade_batch_tokens, _, _, _, token_count = self.clip.process_text([text])
AttributeError: 'NoneType' object has no attribute 'process_text'
Traceback (most recent call last):
  File "E:\sd\stable-diffusion-webui\modules\ui.py", line 185, in f
    res = list(func(*args, **kwargs))
  File "E:\sd\stable-diffusion-webui\webui.py", line 54, in f
    res = func(*args, **kwargs)
  File "E:\sd\stable-diffusion-webui\modules\txt2img.py", line 48, in txt2img
    processed = process_images(p)
  File "E:\sd\stable-diffusion-webui\modules\processing.py", line 423, in process_images
    res = process_images_inner(p)
  File "E:\sd\stable-diffusion-webui\modules\processing.py", line 441, in process_images_inner
    processed = Processed(p, [], p.seed, "")
  File "E:\sd\stable-diffusion-webui\modules\processing.py", line 220, in __init__
    self.sd_model_hash = shared.sd_model.sd_model_hash
AttributeError: 'NoneType' object has no attribute 'sd_model_hash'



### Steps to reproduce the problem

1. Change the checkpoint in the dropdown at the top of the WebUI.

I had the inpainting model loaded last, and now I cannot switch to any other. Then, I cannot generate.

### What should have happened?

It should have normally switched it to a checkpoint.

### Commit where the problem happens

98947d173e3f1667eba29c904f681047dea9de90

### What platforms do you use to access UI ?

Windows

### What browsers do you use to access the UI ?

Microsoft Edge

### Command Line Arguments

```Shell
--precision full --medvram --no-half --ckpt-dir "C:\SD\models" --listen --enable-insecure-extension-access --xformers --vae-path "C:\SD\moremodels\v1-5-pruned-emaonly.vae.pt" --api --cors-allow-origins=*
```


### Additional information, context and logs

It seems to be related to this PR: #4514


Discussion:
trying to reproduce the problem, what is your checkpoint cache setting set to ?

@R-N  do you think we can just remove the restore_base_vae() as you mentioned ?
> trying to reproduce the problem, what is your checkpoint cache setting set to ?
> 
> @R-N do you think we can just remove the restore_base_vae() as you mentioned ?

I had the cache set to 2. I set it to 0 and that fixed the problem
i think adding back in the _and hasattr(model, "sd_checkpoint_info")_ check would fix the problem.
if the model has no sd_checkpoint_info yet, do not try to restore.

honestly i am not sure if the restore is still needed now after we changed the caching a bit.

line 168
```
    if cache_enabled and hasattr(model, "sd_checkpoint_info"):
        sd_vae.restore_base_vae(model)
```
could you test that ?
> @R-N do you think we can just remove the restore_base_vae() as you mentioned ?

Yeah that's what I suggested in #4514 . If you have a checkpoint with config file to trigger it, please try it.

> But that restore_base_vae call was from when the caching was done at the start of load_model_weights. It was called so that the caching won't cache the separately loaded VAE. Now that the caching is done right after the checkpoint weight was loaded, that restore_base_vae call can just be removed.

Resolving patch:
diff --git a/modules/sd_models.py b/modules/sd_models.py
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -165,16 +165,9 @@ def load_model_weights(model, checkpoint_info, vae_file="auto"):
 
     cache_enabled = shared.opts.sd_checkpoint_cache > 0
 
-    if cache_enabled:
-        sd_vae.restore_base_vae(model)
-
-    vae_file = sd_vae.resolve_vae(checkpoint_file, vae_file=vae_file)
-
     if cache_enabled and checkpoint_info in checkpoints_loaded:
         # use checkpoint cache
-        vae_name = sd_vae.get_filename(vae_file) if vae_file else None
-        vae_message = f" with {vae_name} VAE" if vae_name else ""
-        print(f"Loading weights [{sd_model_hash}]{vae_message} from cache")
+        print(f"Loading weights [{sd_model_hash}] from cache")
         model.load_state_dict(checkpoints_loaded[checkpoint_info])
     else:
         # load from file
@@ -220,6 +213,7 @@ def load_model_weights(model, checkpoint_info, vae_file="auto"):
     model.sd_model_checkpoint = checkpoint_file
     model.sd_checkpoint_info = checkpoint_info
 
+    vae_file = sd_vae.resolve_vae(checkpoint_file, vae_file=vae_file)
     sd_vae.load_vae(model, vae_file)
 
 
diff --git a/modules/sd_vae.py b/modules/sd_vae.py
--- a/modules/sd_vae.py
+++ b/modules/sd_vae.py
@@ -91,7 +91,7 @@ def get_vae_from_settings(vae_file="auto"):
         # if VAE selected but not found, fallback to auto
         if vae_file not in default_vae_values and not os.path.isfile(vae_file):
             vae_file = "auto"
-            print("Selected VAE doesn't exist")
+            print(f"Selected VAE doesn't exist: {vae_file}")
     return vae_file
 
 
@@ -101,15 +101,15 @@ def resolve_vae(checkpoint_file=None, vae_file="auto"):
     # if vae_file argument is provided, it takes priority, but not saved
     if vae_file and vae_file not in default_vae_list:
         if not os.path.isfile(vae_file):
+            print(f"VAE provided as function argument doesn't exist: {vae_file}")
             vae_file = "auto"
-            print("VAE provided as function argument doesn't exist")
     # for the first load, if vae-path is provided, it takes priority, saved, and failure is reported
     if first_load and shared.cmd_opts.vae_path is not None:
         if os.path.isfile(shared.cmd_opts.vae_path):
             vae_file = shared.cmd_opts.vae_path
             shared.opts.data['sd_vae'] = get_filename(vae_file)
         else:
-            print("VAE provided as command line argument doesn't exist")
+            print(f"VAE provided as command line argument doesn't exist: {vae_file}")
     # fallback to selector in settings, if vae selector not set to act as default fallback
     if not shared.opts.sd_vae_as_default:
         vae_file = get_vae_from_settings(vae_file)
@@ -117,20 +117,20 @@ def resolve_vae(checkpoint_file=None, vae_file="auto"):
     if vae_file == "auto" and shared.cmd_opts.vae_path is not None:
         if os.path.isfile(shared.cmd_opts.vae_path):
             vae_file = shared.cmd_opts.vae_path
-            print("Using VAE provided as command line argument")
+            print(f"Using VAE provided as command line argument: {vae_file}")
     # if still not found, try look for ".vae.pt" beside model
     model_path = os.path.splitext(checkpoint_file)[0]
     if vae_file == "auto":
         vae_file_try = model_path + ".vae.pt"
         if os.path.isfile(vae_file_try):
             vae_file = vae_file_try
-            print("Using VAE found beside selected model")
+            print(f"Using VAE found similar to selected model: {vae_file}")
     # if still not found, try look for ".vae.ckpt" beside model
     if vae_file == "auto":
         vae_file_try = model_path + ".vae.ckpt"
         if os.path.isfile(vae_file_try):
             vae_file = vae_file_try
-            print("Using VAE found beside selected model")
+            print(f"Using VAE found similar to selected model: {vae_file}")
     # No more fallbacks for auto
     if vae_file == "auto":
         vae_file = None
@@ -146,6 +146,7 @@ def load_vae(model, vae_file=None):
     # save_settings = False
 
     if vae_file:
+        assert os.path.isfile(vae_file), f"VAE file doesn't exist: {vae_file}"
         print(f"Loading VAE weights from: {vae_file}")
         vae_ckpt = torch.load(vae_file, map_location=shared.weight_load_location)
         vae_dict_1 = {k: v for k, v in vae_ckpt["state_dict"].items() if k[0:4] != "loss" and k not in vae_ignore_keys}
diff --git a/modules/shared.py b/modules/shared.py
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -334,7 +334,7 @@ def options_section(section_identifier, options_dict):
 options_templates.update(options_section(('sd', "Stable Diffusion"), {
     "sd_model_checkpoint": OptionInfo(None, "Stable Diffusion checkpoint", gr.Dropdown, lambda: {"choices": modules.sd_models.checkpoint_tiles()}, refresh=sd_models.list_models),
     "sd_checkpoint_cache": OptionInfo(0, "Checkpoints to cache in RAM", gr.Slider, {"minimum": 0, "maximum": 10, "step": 1}),
-    "sd_vae": OptionInfo("auto", "SD VAE", gr.Dropdown, lambda: {"choices": list(sd_vae.vae_list)}, refresh=sd_vae.refresh_vae_list),
+    "sd_vae": OptionInfo("auto", "SD VAE", gr.Dropdown, lambda: {"choices": sd_vae.vae_list}, refresh=sd_vae.refresh_vae_list),
     "sd_vae_as_default": OptionInfo(False, "Ignore selected VAE for stable diffusion checkpoints that have their own .vae.pt next to them"),
     "sd_hypernetwork": OptionInfo("None", "Hypernetwork", gr.Dropdown, lambda: {"choices": ["None"] + [x for x in hypernetworks.keys()]}, refresh=reload_hypernetworks),
     "sd_hypernetwork_strength": OptionInfo(1.0, "Hypernetwork strength", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.001}),


Test patch: