mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
Update model_probe to work with diffuser-format SD TI embeddings. (#5301)
## What type of PR is this? (check all applicable) - [ ] Refactor - [ ] Feature - [x] Bug Fix - [ ] Optimization - [ ] Documentation Update - [ ] Community Node Submission ## Have you updated all relevant documentation? - [x] Yes (N/A) - [ ] No ## Description This change enables the model probe to work with TI embeddings that have the follow state_dict structure: ```python { "<any_key>": torch.Tensor(...), # where the tensor has shape (N, embedding_dim) } ``` ## QA Instructions, Screenshots, Recordings I can't imagine an embedding format that would previously have passed the model probe, and would now fail after this change. That being said, I'll exercise a bunch of existing TIs before merging. - [x] Exercise existing TI formats ## Added/updated tests? - [ ] Yes - [x] No : _We could really benefit from tests for all of the supported TI formats... but I'm not taking on that project right now._
This commit is contained in:
commit
e469e24a58
@ -389,7 +389,7 @@ class TextualInversionCheckpointProbe(CheckpointProbeBase):
|
||||
elif "clip_g" in checkpoint:
|
||||
token_dim = checkpoint["clip_g"].shape[-1]
|
||||
else:
|
||||
token_dim = list(checkpoint.values())[0].shape[0]
|
||||
token_dim = list(checkpoint.values())[0].shape[-1]
|
||||
if token_dim == 768:
|
||||
return BaseModelType.StableDiffusion1
|
||||
elif token_dim == 1024:
|
||||
|
Loading…
Reference in New Issue
Block a user