mirrors/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-09-12 12:37:01 +00:00
Code Issues Packages Projects Releases Wiki Activity

Remove pytorch_lightning dependency.

Browse Source
This commit is contained in:
comfyanonymous
2023-06-13 10:11:33 -04:00
parent cb180b9998
commit 735ac4cf81
3 changed files with 15 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
comfy/utils.py
View File

@@ -1,6 +1,7 @@
import torch
import math
import struct
import comfy.checkpoint_pickle
def load_torch_file(ckpt, safe_load=False):
if ckpt.lower().endswith(".safetensors"):
@@ -14,7 +15,7 @@ def load_torch_file(ckpt, safe_load=False):
if safe_load:
pl_sd = torch.load(ckpt, map_location="cpu", weights_only=True)
else:
pl_sd = torch.load(ckpt, map_location="cpu")
pl_sd = torch.load(ckpt, map_location="cpu", pickle_module=comfy.checkpoint_pickle)
if "global_step" in pl_sd:
print(f"Global Step: {pl_sd['global_step']}")
if "state_dict" in pl_sd: