wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-01-26 22:30:19 +08:00
Code Issues Actions 9 Packages Projects Releases Wiki Activity

fix pylint notifs

Browse Source
This commit is contained in:
doctorpangloss 2025-06-17 14:49:06 -07:00
parent 52c522d019
commit 42dd7a59e3
1 changed files with 4 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
comfy_extras/nodes/nodes_train.py
View File

@ -19,6 +19,7 @@ from comfy.cmd import folder_paths
from comfy import node_helpers from comfy import node_helpers
from comfy.cli_args import args from comfy.cli_args import args
from comfy.comfy_types.node_typing import IO from comfy.comfy_types.node_typing import IO
from comfy.execution_context import current_execution_context
from comfy.weight_adapter import adapters from comfy.weight_adapter import adapters
from . import nodes_custom_sampler from . import nodes_custom_sampler
@ -470,6 +471,8 @@ class TrainLoraNode:
criterion = torch.nn.HuberLoss() criterion = torch.nn.HuberLoss()
elif loss_function == "SmoothL1": elif loss_function == "SmoothL1":
criterion = torch.nn.SmoothL1Loss() criterion = torch.nn.SmoothL1Loss()
else:
criterion = None
# setup models # setup models
for m in find_all_highest_child_module_with_forward(mp.model.diffusion_model): for m in find_all_highest_child_module_with_forward(mp.model.diffusion_model):
@ -493,7 +496,7 @@ class TrainLoraNode:
# Training loop # Training loop
torch.cuda.empty_cache() torch.cuda.empty_cache()
try: try:
for step in (pbar:=tqdm.trange(steps, desc="Training LoRA", smoothing=0.01, disable=not comfy.utils.PROGRESS_BAR_ENABLED)): for step in (pbar:=tqdm.trange(steps, desc="Training LoRA", smoothing=0.01, disable=not current_execution_context().server.receive_all_progress_notifications)):
# Generate random sigma # Generate random sigma
sigma = mp.model.model_sampling.percent_to_sigma( sigma = mp.model.model_sampling.percent_to_sigma(
torch.rand((1,)).item() torch.rand((1,)).item()