nodes/kijai-ComfyUI-WanVideoWrapper
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge pull request #952 from SD-inst/tdqm-fix

Browse Source 
Reduce tdqm spam
This commit is contained in:
Jukka Seppänen
2025-08-02 19:46:40 +03:00
committed by GitHub
parent e9b537f23c ed247d78b6
commit 158b62dafa
1 changed files with 11 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nodes_model_loading.py
+11 -2
View File
@@ -1041,6 +1041,7 @@ class WanVideoModelLoader:
log.info("Using accelerate to load and assign model weights to device...")
param_count = sum(1 for _ in transformer.named_parameters())
pbar = ProgressBar(param_count)
cnt = 0
for name, param in tqdm(transformer.named_parameters(),
desc=f"Loading transformer parameters to {transformer_load_device}",
total=param_count,
@@ -1052,10 +1053,13 @@ class WanVideoModelLoader:
if "patch_embedding" in name:
dtype_to_use = torch.float32
set_module_tensor_to_device(transformer, name, device=transformer_load_device, dtype=dtype_to_use, value=sd[name])
pbar.update(1)
cnt += 1
if cnt % 100 == 0:
pbar.update(100)
#for name, param in transformer.named_parameters():
# print(name, param.dtype, param.device, param.shape)
pbar.update_absolute(param_count)
comfy_model.diffusion_model = transformer
comfy_model.load_device = transformer_load_device
@@ -1134,6 +1138,7 @@ class WanVideoModelLoader:
patcher.model.diffusion_model = _replace_with_gguf_linear(patcher.model.diffusion_model, base_dtype, sd, patches=patcher.patches)
pbar = ProgressBar(param_count)
cnt = 0
for name, param in tqdm(patcher.model.diffusion_model.named_parameters(),
desc=f"Loading transformer parameters to {transformer_load_device}",
total=param_count,
@@ -1146,10 +1151,14 @@ class WanVideoModelLoader:
else:
dtype_to_use = base_dtype
set_module_tensor_to_device(patcher.model.diffusion_model, name, device=transformer_load_device, dtype=dtype_to_use, value=sd[name])
pbar.update(1)
cnt += 1
if cnt % 100 == 0:
pbar.update(100)
#for name, param in transformer.named_parameters():
# print(name, param.dtype, param.device, param.shape)
#patcher.load(device, full_load=True)
pbar.update_absolute(param_count)
patcher.model.is_patched = True