wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-02-03 02:00:29 +08:00
Code Issues Actions 9 Packages Projects Releases Wiki Activity

add ui pbar

Browse Source
This commit is contained in:
Kohaku-Blueleaf 2025-11-08 15:42:18 +08:00
parent b3784a7da1
commit b85ed81792
1 changed files with 5 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
comfy_extras/nodes_train.py
View File

@ -5,7 +5,7 @@ import numpy as np
import safetensors
import torch
import torch.utils.checkpoint
import tqdm
from tqdm.auto import trange
from PIL import Image, ImageDraw, ImageFont
from typing_extensions import override
@ -18,6 +18,7 @@ import folder_paths
import node_helpers
from comfy.weight_adapter import adapters, adapter_maps
from comfy_api.latest import ComfyExtension, io, ui
from comfy.utils import ProgressBar
def make_batch_extra_option_dict(d, indicies, full_size=None):
@ -129,8 +130,9 @@ class TrainSampler(comfy.samplers.Sampler):
cond = model_wrap.conds["positive"]
dataset_size = sigmas.size(0)
torch.cuda.empty_cache()
ui_pbar = ProgressBar(self.total_steps)
for i in (
pbar := tqdm.trange(
pbar := trange(
self.total_steps,
desc="Training LoRA",
smoothing=0.01,
@ -203,6 +205,7 @@ class TrainSampler(comfy.samplers.Sampler):
if (i + 1) % self.grad_acc == 0:
self.optimizer.step()
self.optimizer.zero_grad()
ui_pbar.update(1)
torch.cuda.empty_cache()
return torch.zeros_like(latent_image)