InvokeAI icon indicating copy to clipboard operation
InvokeAI copied to clipboard

Add Embedding Parsing

Open blessedcoolant opened this issue 3 years ago • 2 comments

Embedding Parsing

This PR adds a tad bit more comprehensive parsing of the embeddings that are being loaded to InvokeAI.

Three new functions have been added.

  1. parse_embedding -- This is what you call.
  • The parent function for parsing. Checks if a path is a pt or a bin file and returns embedding_info respectively. Ends and warns a user if the supplied file is not a pt or a bin file.
  1. parse_embedding_pt
  • Takes an embedding location and loads it. Then it checks if the embedding contains two keys string_to_token and string_to_param which are signs to indicate that this is a valid embedding pt file.
  • If the above two keys are found, it will parse the rest of the embedding and return an embedding_info object that has three mandatory values: name, embedding and num_vectors_per_token that are assigned to this particular embedding.
  • Additionally this also checks if for trained_steps, trained_model_name, trained_model_checksum keys. These are generally found in embeddings trained with Auto1111's repo. They provide relevant info. This info is not utilized anywhere currently but can be used to effectively load or suggest the correct model to use with the embedding (as embeddings work best with models they were trained on).
  • num_vectors_per_token is calculated by measuring the tensor size. With this info, we no longer have to worry about manually editing the inference file to match different embeddings. They all just work now.
  1. parse_embeddings_bin
  • Checks for keys in the bin file and if none are found, returns warning the user that this is not a valid concepts file.
  • On parsing the bin file, it returns the above mentioned three mandatory keys that are required to load an embedding.
  • The num_vectors_per_token is defaulted to 1 because I believe all Hugging Face concepts are defaulted to that. I haven't found a single one to show otherwise.

Bug Fixes

  • Added manual checks when an invalid pt or bin file is tried to load. Right now the app was just crashing in this instance. As of this change, it'll just warn the user that the mentioned file is not a valid file or is corrupt and bypass it loading the rest of the embeds.
  • Added some failsafe checks for when scanning fails. Instead of crashing the app, it just bypasses instead.

Could not find any obvious bugs. And all embeddings I have tried -- those downloaded from the internet, those downloaded from HuggingFace Concepts and my own personally trained ones at different vector tokens of 1, 2, 4 and 8 work correctly.

Maybe I broke some stuff. So if you spot anything, let me know.

blessedcoolant avatar Dec 13 '22 18:12 blessedcoolant

Great work. This review will take a little longer than usual because I want to throw a few problematic .pt files at it and I've got a full work day today and this evening. I'll try to have it done by tomorrow at noon.

lstein avatar Dec 14 '22 19:12 lstein

Great work. This review will take a little longer than usual because I want to throw a few problematic .pt files at it and I've got a full work day today and this evening. I'll try to have it done by tomorrow at noon.

No rush. Whenever you are free. Let me know if it needs any work.

blessedcoolant avatar Dec 14 '22 19:12 blessedcoolant