improve importation and conversion of legacy checkpoint files

[lstein]

A long-standing issue with importing legacy checkpoints (both ckpt and safetensors) is that the user has to identify the correct config file, either by providing its path or by selecting which type of model the checkpoint is (e.g. "v1 inpainting"). In addition, some users wish to provide custom VAEs for use with the model. Currently this is done in the WebUI by importing the model, editing it, and then typing in the path to the VAE.

Model configuration file selection

To improve the user experience, the model manager's heuristic_import() method has been enhanced as follows:

1. When initially called, the caller can pass a config file path, in which case it will be used.

2. If no config file provided, the method looks for a .yaml file in the same directory as the model which bears the same basename. e.g.

   my-new-model.safetensors
   my-new-model.yaml

The yaml file is then used as the configuration file for importation and conversion.

3. If no such file is found, then the method opens up the checkpoint and probes it to determine whether it is V1, V1-inpaint or V2. If it is a V1 format, then the appropriate v1-inference.yaml config file is used. Unfortunately there are two V2 variants that cannot be distinguished by introspection.

4. If the probe algorithm is unable to determine the model type, then its last-ditch effort is to execute an optional callback function that can be provided by the caller. This callback, named config_file_callback receives the path to the legacy checkpoint and returns the path to the config file to use. The CLI uses to put up a multiple choice prompt to the user. The WebUI could use this to prompt the user to choose from a radio-button selection.

5. If the config file cannot be determined, then the import is abandoned.

Custom VAE Selection

The user can attach a custom VAE to the imported and converted model by copying the desired VAE into the same directory as the file to be imported, and giving it the same basename. E.g.:

    my-new-model.safetensors
    my-new-model.vae.pt

For this to work, the VAE must end with ".vae.pt", ".vae.ckpt", or ".vae.safetensors". The indicated VAE will be converted into diffusers format and stored with the converted models file, so the ".pt" file can be deleted after conversion.

No facility is currently provided to swap a diffusers VAE at import time, but this can be done after the fact using the WebUI and CLI's model editing functions.

Note that this is the same fix that was applied to the 2.3 branch in #3043 . This applies to main.

[damian0815]

lgtm (but i didn't test it)

.. one thing that might be worth considering/adding is the ability to save the converted diffusers model as half-precision, perhaps based on the precision of the input ckpt, or on user selection - because at the moment a 2.5GB v2.1 ckpt file will be converted to a 5GB diffusers folder. (to do this you call pipeline.to(torch.float16) before calling save_pretrained())