jinx/ComfyUI
1
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-08-02 23:14:49 +08:00
Code Issues Activity

Replace prints with logging and add --verbose argument.

Browse Source
This commit is contained in:
comfyanonymous
2024-03-10 11:37:08 -04:00
parent 4656273e72
commit 65397ce601
12 changed files with 90 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
comfy/utils.py
View File

@@ -5,6 +5,7 @@ import comfy.checkpoint_pickle
import safetensors.torch
import numpy as np
from PIL import Image
import logging
def load_torch_file(ckpt, safe_load=False, device=None):
if device is None:
@@ -14,14 +15,14 @@ def load_torch_file(ckpt, safe_load=False, device=None):
else:
if safe_load:
if not 'weights_only' in torch.load.__code__.co_varnames:
print("Warning torch.load doesn't support weights_only on this pytorch version, loading unsafely.")
logging.warning("Warning torch.load doesn't support weights_only on this pytorch version, loading unsafely.")
safe_load = False
if safe_load:
pl_sd = torch.load(ckpt, map_location=device, weights_only=True)
else:
pl_sd = torch.load(ckpt, map_location=device, pickle_module=comfy.checkpoint_pickle)
if "global_step" in pl_sd:
print(f"Global Step: {pl_sd['global_step']}")
logging.info(f"Global Step: {pl_sd['global_step']}")
if "state_dict" in pl_sd:
sd = pl_sd["state_dict"]
else: