Skip to content

Commit

Permalink
add more reliable fallback method for determining BnbQuantizedLlmInt8b
Browse files Browse the repository at this point in the history
  • Loading branch information
Lincoln Stein committed Sep 1, 2024
1 parent 70e7c41 commit d4eda7e
Showing 1 changed file with 12 additions and 2 deletions.
14 changes: 12 additions & 2 deletions invokeai/backend/model_manager/probe.py
Original file line number Diff line number Diff line change
Expand Up @@ -763,10 +763,20 @@ def get_format(self) -> ModelFormat:
path = self.model_path / "text_encoder_2"
if (path / "model.safetensors.index.json").exists():
return ModelFormat.T5Encoder
files = path.glob("*.safetensors")
files = list(path.glob("*.safetensors"))
if len(files) == 0:
raise InvalidModelConfigException(f"{self.model_path.as_posix()}: no .safetensors files found")

# shortcut: look for the quantization in the name
if any(x for x in files if "llm_int8" in x.as_posix()):
return ModelFormat.BnbQuantizedLlmInt8b
raise f"{self.model_path.as_posix()}: unknown model format"

# more reliable path: probe contents for a 'SCB' key
ckpt = read_checkpoint_meta(files[0], scan=True)
if any("SCB" in x for x in ckpt.keys()):
return ModelFormat.BnbQuantizedLlmInt8b

raise InvalidModelConfigException(f"{self.model_path.as_posix()}: unknown model format")


class ONNXFolderProbe(PipelineFolderProbe):
Expand Down

0 comments on commit d4eda7e

Please sign in to comment.