Update CheckpointLoaderSimpleText.py

2023-04-06 21:29:47 +09:00
parent 059d59455c
commit 41be495672
1 changed files with 4 additions and 1 deletions
@@ -43,7 +43,10 @@ class CheckpointLoaderSimpleText:
        print(f"ckpt_path",ckpt_path)
        if ckpt_path is None:
            print(f"{ckpt_name} is none")
-            (name,fullpath)=filenameget(os.path.join(models_dir, "checkpoints")+"/**/"+ckpt_name+"*.safetensors")
+            if ckpt_name.endswith(".safetensors") or ckpt_name.endswith(".ckpt"):
+                (name,fullpath)=filenameget(os.path.join(models_dir, "checkpoints")+"/**/"+ckpt_name)
+            else:
+                (name,fullpath)=filenameget(os.path.join(models_dir, "checkpoints")+"/**/"+ckpt_name+"*.safetensors")
            if fullpath is None:
                print(f"{name} is none")
                return