VAEs as CKPT?

[mark-wd]

Some models like Counterfeit have their VAEs as CKPT files. The advice is to put them alongside the model in the "Stable-diffusion" folder with the same name and that should work. However, these show up like normal Checkpoints in the list causing some confusion.

Q: If I put those in the VAE folder, are they auto-loaded by the UI in case there is a name match? Would the correct setting for that for VAEs be "Auto"?

[vladmandic]

That's the idea. Does it work? Don't know.
But if it doesn't, open an issue as I'd like that to work.

[mark-wd]

In my environment, if I put a CKPT or PT with a matching name in the models/VAE folder, it does not get loaded automatically upon start or when switching models when Auto is set in the VAE dropdown. Can you test in your environment @vladmandic?

[vladmandic]

ok, i've just added that.

[Tillerz]

I just checked the code:

    vae_paths = [
        os.path.join(sd_models.model_path, '**/*.vae.ckpt'),
        os.path.join(sd_models.model_path, '**/*.vae.pt'),
        os.path.join(sd_models.model_path, '**/*.vae.safetensors'),
        os.path.join(shared.opts.vae_dir, '**/*.ckpt'),
        os.path.join(shared.opts.vae_dir, '**/*.pt'),
        os.path.join(shared.opts.vae_dir, '**/*.safetensors'),

It looks like the VAE must be named <something>.vae.ckpt to get found at all, if it's not in the vae/ folder.

And I just found this:

    if is_automatic:
        for named_vae_location in [
            os.path.join(vae_path, os.path.splitext(os.path.basename(checkpoint_file))[0] + ".vae.pt"), 
            os.path.join(vae_path, os.path.splitext(os.path.basename(checkpoint_file))[0] + ".vae.ckpt"),
            os.path.join(vae_path, os.path.splitext(os.path.basename(checkpoint_file))[0] + ".vae.safetensors")
        ]:
            if os.path.isfile(named_vae_location):
                return named_vae_location, 'in VAE dir'

So the VAE must be named <modelname>.vae.<extension> to get found automatically.

[mark-wd]

I've tested now and with vlad's fix, VAEs load properly from the directory when there is a match, at least in the case of Counterfeit and its CKPT.

Now ideally, on model change it would also unload the VAE, as it's unlikely that the next model will need that specific VAE.

[mark-wd]

One final thing that would really optimize this for good would be the option to always load the 560000 VAE (or whatever the user wants) unless there is a matching VAE for that model.

[Tillerz]

I think that would be easy to do?

 if is_automatic:
        for named_vae_location in [
            os.path.join(vae_path, os.path.splitext(os.path.basename(checkpoint_file))[0] + ".vae.pt"), 
            os.path.join(vae_path, os.path.splitext(os.path.basename(checkpoint_file))[0] + ".vae.ckpt"),
            os.path.join(vae_path, os.path.splitext(os.path.basename(checkpoint_file))[0] + ".vae.safetensors")
        ]:
            if os.path.isfile(named_vae_location):
                return named_vae_location, 'in VAE dir'

That code should just return the 560000 one if it doesn't find anything. Then you would just set VAE to default and it would work. Probably needs a line of code or two, maybe a dozen if the "fallback VAE" shall be somewhere in the Settings. :D

[vladmandic]

if vae is not found for a model and you had vae loaded before, you should already see Restoring base VAE.
implementing fallbacks is doable, but then we're getting into what should be fallback, does fallback exist, etc...i'm not sure if its worth it?

[Tillerz]

I remember there was a toggle to "load selected VAE only when there is no embedded one", that could have gotten expanded to "and if there's no matching VAE in the list of search folders", but that toggle was refactored out, I believe. :D

[mark-wd]

Yeah, I wonder why that toggle ceased to exist. It wasn't a bad idea to begin with.