feat(api): return more data for embeddings

This commit is contained in:
Philpax
2023-01-02 12:21:22 +11:00
parent b5819d9bf1
commit c65909ad16
3 changed files with 28 additions and 8 deletions

View File

@@ -330,9 +330,22 @@ class Api:
def get_embeddings(self):
db = sd_hijack.model_hijack.embedding_db
def convert_embedding(embedding):
return {
"step": embedding.step,
"sd_checkpoint": embedding.sd_checkpoint,
"sd_checkpoint_name": embedding.sd_checkpoint_name,
"shape": embedding.shape,
"vectors": embedding.vectors,
}
def convert_embeddings(embeddings):
return {embedding.name: convert_embedding(embedding) for embedding in embeddings.values()}
return {
"loaded": sorted(db.word_embeddings.keys()),
"skipped": sorted(db.skipped_embeddings),
"loaded": convert_embeddings(db.word_embeddings),
"skipped": convert_embeddings(db.skipped_embeddings),
}
def refresh_checkpoints(self):