add on-the-fly conversion of .ckpt to diffusers models

1. On-the-fly conversion code can be found in the file ldm/invoke/ckpt_to_diffusers.py.

2. A new !optimize command has been added to the CLI. Should be ported to Web GUI.

User experience on the CLI is this:

```
invoke> !optimize /home/lstein/invokeai/models/ldm/stable-diffusion-v1/sd-v1-4.ckpt
INFO: Converting legacy weights file /home/lstein/invokeai/models/ldm/stable-diffusion-v1/sd-v1-4.ckpt to optimized diffuser model.
      This operation will take 30-60s to complete.
Success. Optimized model is now located at /home/lstein/tmp/invokeai/models/optimized-ckpts/sd-v1-4
Writing new config file entry for sd-v1-4...

>> New configuration:
sd-v1-4:
  description: Optimized version of sd-v1-4
  format: diffusers
  path: /home/lstein/tmp/invokeai/models/optimized-ckpts/sd-v1-4

OK to import [n]? y
>> Verifying that new model loads...
>> Current VRAM usage:  2.60G
>> Offloading stable-diffusion-2.1 to CPU
>> Loading diffusers model from /home/lstein/tmp/invokeai/models/optimized-ckpts/sd-v1-4
  | Using faster float16 precision
You have disabled the safety checker for <class 'ldm.invoke.generator.diffusers_pipeline.StableDiffusionGeneratorPipeline'> by passing `safety_checker=None`. Ensure that you abide to the conditions of the Stable Diffusion \
license and do not expose unfiltered results in services or applications open to the public. Both the diffusers team and Hugging Face strongly recommend to keep the safety filter enabled in all public facing circumstances,\
 disabling it only for use-cases that involve analyzing network behavior or auditing its results. For more information, please have a look at huggingface/diffusers#254 .
  | training width x height = (512 x 512)
>> Model loaded in 3.48s
>> Max VRAM used to load the model: 2.17G
>> Current VRAM usage:2.17G
>> Textual inversions available:
>> Setting Sampler to k_lms (LMSDiscreteScheduler)
Keep model loaded? [y]
```

configs/stable-diffusion/v2-inference-v.yaml

@@ -0,0 +1,68 @@
+model:
+  base_learning_rate: 1.0e-4
+  target: ldm.models.diffusion.ddpm.LatentDiffusion
+  params:
+    parameterization: "v"
+    linear_start: 0.00085
+    linear_end: 0.0120
+    num_timesteps_cond: 1
+    log_every_t: 200
+    timesteps: 1000
+    first_stage_key: "jpg"
+    cond_stage_key: "txt"
+    image_size: 64
+    channels: 4
+    cond_stage_trainable: false
+    conditioning_key: crossattn
+    monitor: val/loss_simple_ema
+    scale_factor: 0.18215
+    use_ema: False # we set this to false because this is an inference only config
+
+    unet_config:
+      target: ldm.modules.diffusionmodules.openaimodel.UNetModel
+      params:
+        use_checkpoint: True
+        use_fp16: True
+        image_size: 32 # unused
+        in_channels: 4
+        out_channels: 4
+        model_channels: 320
+        attention_resolutions: [ 4, 2, 1 ]
+        num_res_blocks: 2
+        channel_mult: [ 1, 2, 4, 4 ]
+        num_head_channels: 64 # need to fix for flash-attn
+        use_spatial_transformer: True
+        use_linear_in_transformer: True
+        transformer_depth: 1
+        context_dim: 1024
+        legacy: False
+
+    first_stage_config:
+      target: ldm.models.autoencoder.AutoencoderKL
+      params:
+        embed_dim: 4
+        monitor: val/rec_loss
+        ddconfig:
+          #attn_type: "vanilla-xformers"
+          double_z: true
+          z_channels: 4
+          resolution: 256
+          in_channels: 3
+          out_ch: 3
+          ch: 128
+          ch_mult:
+          - 1
+          - 2
+          - 4
+          - 4
+          num_res_blocks: 2
+          attn_resolutions: []
+          dropout: 0.0
+        lossconfig:
+          target: torch.nn.Identity
+
+    cond_stage_config:
+      target: ldm.modules.encoders.modules.FrozenOpenCLIPEmbedder
+      params:
+        freeze: True
+        layer: "penultimate"

ldm/invoke/CLI.py

@@ -460,6 +460,17 @@ def do_command(command:str, gen, opt:Args, completer) -> tuple:
         completer.add_history(command)
         operation = None
 
+    elif command.startswith('!optimize'):
+        path = shlex.split(command)
+        if len(path) < 2:
+            print('** please provide a path to a .ckpt file')
+        elif not os.path.exists(path[1]):
+            print(f'** {path[1]}: file not found')
+        else:
+            optimize_model(path[1], gen, opt, completer)
+        completer.add_history(command)
+        operation = None
+
     elif command.startswith('!edit'):
         path = shlex.split(command)
         if len(path) < 2:
@@ -532,6 +543,7 @@ def add_weights_to_config(model_path:str, gen, opt, completer):
 
     new_config = {}
     new_config['weights'] = model_path
+    new_config['format'] = 'ckpt'
 
     done = False
     while not done:
@@ -579,6 +591,36 @@ def add_weights_to_config(model_path:str, gen, opt, completer):
     if write_config_file(opt.conf, gen, model_name, new_config, make_default=make_default):
         completer.add_model(model_name)
 
+def optimize_model(model_path:str, gen, opt, completer):
+    from ldm.invoke.ckpt_to_diffuser import convert_ckpt_to_diffuser
+    import transformers
+    basename = os.path.basename(os.path.splitext(model_path)[0])
+    dump_path = os.path.join(Globals.root, 'models','optimized-ckpts',basename)
+    if os.path.exists(dump_path):
+        print(f'ERROR: The path {dump_path} already exists. Please move or remove it and try again.')
+        return
+
+    print(f'INFO: Converting legacy weights file {model_path} to optimized diffuser model.')
+    print(f'      This operation will take 30-60s to complete.')
+    try:
+        verbosity =transformers.logging.get_verbosity()
+        transformers.logging.set_verbosity_error()
+        convert_ckpt_to_diffuser(model_path, dump_path)
+        transformers.logging.set_verbosity(verbosity)
+        print(f'Success. Optimized model is now located at {dump_path}')
+        print(f'Writing new config file entry for {basename}...')
+        model_name = basename
+        new_config = dict(
+            path=dump_path,
+            description=f'Optimized version of {basename}',
+            format='diffusers',
+        )
+        if write_config_file(opt.conf, gen, model_name, new_config):
+            completer.add_model(model_name)
+    except Exception as e:
+        print(f'** Conversion failed: {str(e)}')
+        traceback.print_exc()
+
 def del_config(model_name:str, gen, opt, completer):
     current_model = gen.model_name
     if model_name == current_model:
@@ -601,7 +643,7 @@ def edit_config(model_name:str, gen, opt, completer):
     conf = config[model_name]
     new_config = {}
     completer.complete_extensions(('.yaml','.yml','.ckpt','.vae.pt'))
-    for field in ('description', 'weights', 'vae', 'config', 'width','height'):
+    for field in ('description', 'weights', 'vae', 'config', 'width', 'height', 'format'):
         completer.linebuffer = str(conf[field]) if field in conf else ''
         new_value = input(f'{field}: ')
         new_config[field] = int(new_value) if field in ('width','height') else new_value
@@ -625,8 +667,12 @@ def write_config_file(conf_path, gen, model_name, new_config, clobber=False, mak
         gen.model_cache.add_model(model_name, new_config, clobber)
         assert gen.set_model(model_name) is not None, 'model failed to load'
     except AssertionError as e:
+        traceback.print_exc()
         print(f'** aborting **')
-        gen.model_cache.del_model(model_name)
+        try:
+            gen.model_cache.del_model(model_name)
+        except Exception:
+            pass
         return False
 
     if make_default:

ldm/invoke/args.py

@@ -724,11 +724,12 @@ def _create_dream_cmd_parser(self):
             !NN retrieves the NNth command from the history
 
             *Model manipulation*
-            !models                                 -- list models in configs/models.yaml
-            !switch <model_name>                    -- switch to model named <model_name>
-            !import_model path/to/weights/file.ckpt -- adds a model to your config
-            !edit_model <model_name>                -- edit a model's description
-            !del_model <model_name>                 -- delete a model
+            !models                                   -- list models in configs/models.yaml
+            !switch <model_name>                      -- switch to model named <model_name>
+            !import_model path/to/weights/file.ckpt   -- adds a .ckpt model to your config
+            !optimize_model path/to/weights/file.ckpt -- converts a .ckpt file model a diffusers model
+            !edit_model <model_name>                  -- edit a model's description
+            !del_model <model_name>                   -- delete a model
             """
         )
         render_group     = parser.add_argument_group('General rendering')