From fc2098834d51f673eabdc58e311cff4e6e6bfe3e Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Tue, 17 Jan 2023 08:11:19 -0500 Subject: [PATCH] support direct loading of .safetensors models - Small fix to allow ckpt files with the .safetensors suffix to be directly loaded, rather than undergo a conversion step first. --- ldm/invoke/model_manager.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ldm/invoke/model_manager.py b/ldm/invoke/model_manager.py index 880d75476f..eeb2208c91 100644 --- a/ldm/invoke/model_manager.py +++ b/ldm/invoke/model_manager.py @@ -142,7 +142,7 @@ class ModelManager(object): Return true if this is a legacy (.ckpt) model ''' info = self.model_info(model_name) - if 'weights' in info and info['weights'].endswith('.ckpt'): + if 'weights' in info and info['weights'].endswith(('.ckpt','.safetensors')): return True return False