Replace prints with logging and add --verbose argument. · code/app-python-comfyui@65397ce · GitHub
Skip to content

Commit 65397ce

Browse files
Replace prints with logging and add --verbose argument.
1 parent 4656273 commit 65397ce

12 files changed

Lines changed: 90 additions & 65 deletions

‎comfy/cli_args.py‎

Lines changed: 10 additions & 0 deletions

‎comfy/clip_vision.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
import os
33
import torch
44
import json
5+
import logging
56

67
import comfy.ops
78
import comfy.model_patcher
@@ -99,7 +100,7 @@ def load_clipvision_from_sd(sd, prefix="", convert_keys=False):
99100
clip = ClipVisionModel(json_config)
100101
m, u = clip.load_sd(sd)
101102
if len(m) > 0:
102-
print("missing clip vision:", m)
103+
logging.warning("missing clip vision: {}".format(m))
103104
u = set(u)
104105
keys = list(sd.keys())
105106
for k in keys:

‎comfy/controlnet.py‎

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import torch
22
import math
33
import os
4+
import logging
45
import comfy.utils
56
import comfy.model_management
67
import comfy.model_detection
@@ -367,7 +368,7 @@ def load_controlnet(ckpt_path, model=None):
367368

368369
leftover_keys = controlnet_data.keys()
369370
if len(leftover_keys) > 0:
370-
print("leftover keys:", leftover_keys)
371+
logging.warning("leftover keys: {}".format(leftover_keys))
371372
controlnet_data = new_sd
372373

373374
pth_key = 'control_model.zero_convs.0.0.weight'
@@ -382,7 +383,7 @@ def load_controlnet(ckpt_path, model=None):
382383
else:
383384
net = load_t2i_adapter(controlnet_data)
384385
if net is None:
385-
print("error checkpoint does not contain controlnet or t2i adapter data", ckpt_path)
386+
logging.error("error checkpoint does not contain controlnet or t2i adapter data {}".format(ckpt_path))
386387
return net
387388

388389
if controlnet_config is None:
@@ -417,7 +418,7 @@ def load_controlnet(ckpt_path, model=None):
417418
cd = controlnet_data[x]
418419
cd += model_sd[sd_key].type(cd.dtype).to(cd.device)
419420
else:
420-
print("WARNING: Loaded a diff controlnet without a model. It will very likely not work.")
421+
logging.warning("WARNING: Loaded a diff controlnet without a model. It will very likely not work.")
421422

422423
class WeightsLoader(torch.nn.Module):
423424
pass
@@ -426,7 +427,12 @@ class WeightsLoader(torch.nn.Module):
426427
missing, unexpected = w.load_state_dict(controlnet_data, strict=False)
427428
else:
428429
missing, unexpected = control_model.load_state_dict(controlnet_data, strict=False)
429-
print(missing, unexpected)
430+
431+
if len(missing) > 0:
432+
logging.warning("missing controlnet keys: {}".format(missing))
433+
434+
if len(unexpected) > 0:
435+
logging.info("unexpected controlnet keys: {}".format(unexpected))
430436

431437
global_average_pooling = False
432438
filename = os.path.splitext(ckpt_path)[0]
@@ -536,9 +542,9 @@ def load_t2i_adapter(t2i_data):
536542

537543
missing, unexpected = model_ad.load_state_dict(t2i_data)
538544
if len(missing) > 0:
539-
print("t2i missing", missing)
545+
logging.warning("t2i missing {}".format(missing))
540546

541547
if len(unexpected) > 0:
542-
print("t2i unexpected", unexpected)
548+
logging.info("t2i unexpected {}".format(unexpected))
543549

544550
return T2IAdapter(model_ad, model_ad.input_channels, compression_ratio, upscale_algorithm)

‎comfy/diffusers_convert.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import re
22
import torch
3+
import logging
34

45
# conversion code from https://github.com/huggingface/diffusers/blob/main/scripts/convert_diffusers_to_original_stable_diffusion.py
56

@@ -177,7 +178,7 @@ def convert_vae_state_dict(vae_state_dict):
177178
for k, v in new_state_dict.items():
178179
for weight_name in weights_to_convert:
179180
if f"mid.attn_1.{weight_name}.weight" in k:
180-
print(f"Reshaping {k} for SD format")
181+
logging.info(f"Reshaping {k} for SD format")
181182
new_state_dict[k] = reshape_weight_for_sd(v)
182183
return new_state_dict
183184

‎comfy/lora.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import comfy.utils
2+
import logging
23

34
LORA_CLIP_MAP = {
45
"mlp.fc1": "mlp_fc1",
@@ -156,7 +157,7 @@ def load_lora(lora, to_load):
156157

157158
for x in lora.keys():
158159
if x not in loaded_keys:
159-
print("lora key not loaded", x)
160+
logging.warning("lora key not loaded: {}".format(x))
160161
return patch_dict
161162

162163
def model_lora_keys_clip(model, key_map={}):

‎comfy/model_base.py‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import torch
2+
import logging
23
from comfy.ldm.modules.diffusionmodules.openaimodel import UNetModel, Timestep
34
from comfy.ldm.cascade.stage_c import StageC
45
from comfy.ldm.cascade.stage_b import StageB
@@ -66,8 +67,8 @@ def __init__(self, model_config, model_type=ModelType.EPS, device=None, unet_mod
6667
if self.adm_channels is None:
6768
self.adm_channels = 0
6869
self.inpaint_model = False
69-
print("model_type", model_type.name)
70-
print("adm", self.adm_channels)
70+
logging.warning("model_type {}".format(model_type.name))
71+
logging.info("adm {}".format(self.adm_channels))
7172

7273
def apply_model(self, x, t, c_concat=None, c_crossattn=None, control=None, transformer_options={}, **kwargs):
7374
sigma = t
@@ -183,10 +184,10 @@ def load_model_weights(self, sd, unet_prefix=""):
183184
to_load = self.model_config.process_unet_state_dict(to_load)
184185
m, u = self.diffusion_model.load_state_dict(to_load, strict=False)
185186
if len(m) > 0:
186-
print("unet missing:", m)
187+
logging.warning("unet missing: {}".format(m))
187188

188189
if len(u) > 0:
189-
print("unet unexpected:", u)
190+
logging.warning("unet unexpected: {}".format(u))
190191
del to_load
191192
return self
192193

‎comfy/model_detection.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import comfy.supported_models
22
import comfy.supported_models_base
3+
import logging
34

45
def count_blocks(state_dict_keys, prefix_string):
56
count = 0
@@ -186,7 +187,7 @@ def model_config_from_unet_config(unet_config):
186187
if model_config.matches(unet_config):
187188
return model_config(unet_config)
188189

189-
print("no match", unet_config)
190+
logging.error("no match {}".format(unet_config))
190191
return None
191192

192193
def model_config_from_unet(state_dict, unet_key_prefix, use_base_if_no_match=False):

‎comfy/model_management.py‎

Lines changed: 20 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import psutil
2+
import logging
23
from enum import Enum
34
from comfy.cli_args import args
45
import comfy.utils
@@ -29,7 +30,7 @@ class CPUState(Enum):
2930
xpu_available = False
3031

3132
if args.deterministic:
32-
print("Using deterministic algorithms for pytorch")
33+
logging.warning("Using deterministic algorithms for pytorch")
3334
torch.use_deterministic_algorithms(True, warn_only=True)
3435

3536
directml_enabled = False
@@ -41,7 +42,7 @@ class CPUState(Enum):
4142
directml_device = torch_directml.device()
4243
else:
4344
directml_device = torch_directml.device(device_index)
44-
print("Using directml with device:", torch_directml.device_name(device_index))
45+
logging.warning("Using directml with device: {}".format(torch_directml.device_name(device_index)))
4546
# torch_directml.disable_tiled_resources(True)
4647
lowvram_available = False #TODO: need to find a way to get free memory in directml before this can be enabled by default.
4748

@@ -117,10 +118,10 @@ def get_total_memory(dev=None, torch_total_too=False):
117118

118119
total_vram = get_total_memory(get_torch_device()) / (1024 * 1024)
119120
total_ram = psutil.virtual_memory().total / (1024 * 1024)
120-
print("Total VRAM {:0.0f} MB, total RAM {:0.0f} MB".format(total_vram, total_ram))
121+
logging.warning("Total VRAM {:0.0f} MB, total RAM {:0.0f} MB".format(total_vram, total_ram))
121122
if not args.normalvram and not args.cpu:
122123
if lowvram_available and total_vram <= 4096:
123-
print("Trying to enable lowvram mode because your GPU seems to have 4GB or less. If you don't want this use: --normalvram")
124+
logging.warning("Trying to enable lowvram mode because your GPU seems to have 4GB or less. If you don't want this use: --normalvram")
124125
set_vram_to = VRAMState.LOW_VRAM
125126

126127
try:
@@ -143,12 +144,10 @@ def get_total_memory(dev=None, torch_total_too=False):
143144
pass
144145
try:
145146
XFORMERS_VERSION = xformers.version.__version__
146-
print("xformers version:", XFORMERS_VERSION)
147+
logging.warning("xformers version: {}".format(XFORMERS_VERSION))
147148
if XFORMERS_VERSION.startswith("0.0.18"):
148-
print()
149-
print("WARNING: This version of xformers has a major bug where you will get black images when generating high resolution images.")
150-
print("Please downgrade or upgrade xformers to a different version.")
151-
print()
149+
logging.warning("\nWARNING: This version of xformers has a major bug where you will get black images when generating high resolution images.")
150+
logging.warning("Please downgrade or upgrade xformers to a different version.\n")
152151
XFORMERS_ENABLED_VAE = False
153152
except:
154153
pass
@@ -213,11 +212,11 @@ def is_nvidia():
213212
FORCE_FP32 = False
214213
FORCE_FP16 = False
215214
if args.force_fp32:
216-
print("Forcing FP32, if this improves things please report it.")
215+
logging.warning("Forcing FP32, if this improves things please report it.")
217216
FORCE_FP32 = True
218217

219218
if args.force_fp16:
220-
print("Forcing FP16.")
219+
logging.warning("Forcing FP16.")
221220
FORCE_FP16 = True
222221

223222
if lowvram_available:
@@ -231,12 +230,12 @@ def is_nvidia():
231230
if cpu_state == CPUState.MPS:
232231
vram_state = VRAMState.SHARED
233232

234-
print(f"Set vram state to: {vram_state.name}")
233+
logging.warning(f"Set vram state to: {vram_state.name}")
235234

236235
DISABLE_SMART_MEMORY = args.disable_smart_memory
237236

238237
if DISABLE_SMART_MEMORY:
239-
print("Disabling smart memory management")
238+
logging.warning("Disabling smart memory management")
240239

241240
def get_torch_device_name(device):
242241
if hasattr(device, 'type'):
@@ -254,11 +253,11 @@ def get_torch_device_name(device):
254253
return "CUDA {}: {}".format(device, torch.cuda.get_device_name(device))
255254

256255
try:
257-
print("Device:", get_torch_device_name(get_torch_device()))
256+
logging.warning("Device: {}".format(get_torch_device_name(get_torch_device())))
258257
except:
259-
print("Could not pick default device.")
258+
logging.warning("Could not pick default device.")
260259

261-
print("VAE dtype:", VAE_DTYPE)
260+
logging.warning("VAE dtype: {}".format(VAE_DTYPE))
262261

263262
current_loaded_models = []
264263

@@ -301,7 +300,7 @@ def model_load(self, lowvram_model_memory=0):
301300
raise e
302301

303302
if lowvram_model_memory > 0:
304-
print("loading in lowvram mode", lowvram_model_memory/(1024 * 1024))
303+
logging.warning("loading in lowvram mode {}".format(lowvram_model_memory/(1024 * 1024)))
305304
mem_counter = 0
306305
for m in self.real_model.modules():
307306
if hasattr(m, "comfy_cast_weights"):
@@ -314,7 +313,7 @@ def model_load(self, lowvram_model_memory=0):
314313
elif hasattr(m, "weight"): #only modules with comfy_cast_weights can be set to lowvram mode
315314
m.to(self.device)
316315
mem_counter += module_size(m)
317-
print("lowvram: loaded module regularly", m)
316+
logging.warning("lowvram: loaded module regularly {}".format(m))
318317

319318
self.model_accelerated = True
320319

@@ -348,7 +347,7 @@ def unload_model_clones(model):
348347
to_unload = [i] + to_unload
349348

350349
for i in to_unload:
351-
print("unload clone", i)
350+
logging.warning("unload clone {}".format(i))
352351
current_loaded_models.pop(i).model_unload()
353352

354353
def free_memory(memory_required, device, keep_loaded=[]):
@@ -390,7 +389,7 @@ def load_models_gpu(models, memory_required=0):
390389
models_already_loaded.append(loaded_model)
391390
else:
392391
if hasattr(x, "model"):
393-
print(f"Requested to load {x.model.__class__.__name__}")
392+
logging.warning(f"Requested to load {x.model.__class__.__name__}")
394393
models_to_load.append(loaded_model)
395394

396395
if len(models_to_load) == 0:
@@ -400,7 +399,7 @@ def load_models_gpu(models, memory_required=0):
400399
free_memory(extra_mem, d, models_already_loaded)
401400
return
402401

403-
print(f"Loading {len(models_to_load)} new model{'s' if len(models_to_load) > 1 else ''}")
402+
logging.warning(f"Loading {len(models_to_load)} new model{'s' if len(models_to_load) > 1 else ''}")
404403

405404
total_memory_required = {}
406405
for loaded_model in models_to_load:

‎comfy/model_patcher.py‎

Lines changed: 11 additions & 7 deletions

0 commit comments

Comments
 (0)