mirror of
https://github.com/invoke-ai/InvokeAI.git
synced 2026-04-23 03:00:31 -04: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:
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user