Skip to content

Commit 1db1264

Browse files
fix(mm): tis use existing weight_files method
1 parent 7ab6042 commit 1db1264

File tree

1 file changed

+1
-2
lines changed

1 file changed

+1
-2
lines changed

invokeai/backend/model_manager/config.py

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -607,7 +607,6 @@ class ControlNetCheckpointConfig(CheckpointConfigBase, ControlAdapterConfigBase,
607607
class TextualInversionConfigBase(ABC, BaseModel):
608608
type: Literal[ModelType.TextualInversion] = ModelType.TextualInversion
609609

610-
KNOWN_SUFFIXES: ClassVar = {"bin", "safetensors", "pt", "ckpt"}
611610
KNOWN_KEYS: ClassVar = {"string_to_param", "emb_params", "clip_g"}
612611

613612
@classmethod
@@ -621,7 +620,7 @@ def file_looks_like_embedding(cls, mod: ModelOnDisk, path: Path | None = None) -
621620
if p.is_dir():
622621
return False
623622

624-
if p.name in [f"learned_embeds.{s}" for s in cls.KNOWN_SUFFIXES]:
623+
if p.name in [f"learned_embeds.{s}" for s in mod.weight_files()]:
625624
return True
626625

627626
state_dict = mod.load_state_dict(p)

0 commit comments

Comments
 (0)