From b3b36e5a69e2409a5d6f327b66e023114b2a25f3 Mon Sep 17 00:00:00 2001 From: Kohaku-Blueleaf <59680068+KohakuBlueleaf@users.noreply.github.com> Date: Sat, 7 Jun 2025 05:21:37 +0800 Subject: [PATCH] Use tqdm for training loop --- comfy_extras/nodes_train.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/comfy_extras/nodes_train.py b/comfy_extras/nodes_train.py index 308db2a51..6e5cac3cc 100644 --- a/comfy_extras/nodes_train.py +++ b/comfy_extras/nodes_train.py @@ -9,6 +9,7 @@ import torch from PIL import Image, ImageDraw, ImageFont from PIL.PngImagePlugin import PngInfo import torch.utils.checkpoint +import tqdm import comfy.samplers import comfy.sd @@ -48,7 +49,6 @@ class TrainSampler(comfy.samplers.Sampler): if "does not require grad and does not have a grad_fn" in str(e): logging.info("WARNING: This is likely due to the model is loaded in inference mode.") loss.backward() - logging.info(f"Current Training Loss: {loss.item():.6f}") if self.loss_callback: self.loss_callback(loss.item()) @@ -478,9 +478,16 @@ class TrainLoraNode: elif loss_function == "SmoothL1": criterion = torch.nn.SmoothL1Loss() + # setup models + for m in find_all_highest_child_module_with_forward(mp.model.diffusion_model): + patch(m) + comfy.model_management.load_models_gpu([mp], memory_required=1e20, force_full_load=True) + # Setup sampler and guider like in test script loss_map = {"loss": []} - loss_callback = lambda loss: loss_map["loss"].append(loss) + def loss_callback(loss): + loss_map["loss"].append(loss) + pbar.set_postfix({"loss": f"{loss:.4f}"}) train_sampler = TrainSampler( criterion, optimizer, loss_callback=loss_callback ) @@ -490,15 +497,10 @@ class TrainLoraNode: # yoland: this currently resize to the first image in the dataset - # setup before training - for m in find_all_highest_child_module_with_forward(mp.model.diffusion_model): - patch(m) - comfy.model_management.load_models_gpu([mp], memory_required=1e20, force_full_load=True) - # Training loop torch.cuda.empty_cache() try: - for step in range(steps): + for step in (pbar:=tqdm.trange(steps, desc="Training LoRA", smoothing=0.01)): # Generate random sigma sigma = mp.model.model_sampling.percent_to_sigma( torch.rand((1,)).item()