wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-02-07 03:52:32 +08:00
Code Issues Actions 15 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 safetensors
import torch import torch
import torch.utils.checkpoint import torch.utils.checkpoint
import tqdm from tqdm.auto import trange
from PIL import Image, ImageDraw, ImageFont from PIL import Image, ImageDraw, ImageFont
from typing_extensions import override from typing_extensions import override
@ -18,6 +18,7 @@ import folder_paths
import node_helpers import node_helpers
from comfy.weight_adapter import adapters, adapter_maps from comfy.weight_adapter import adapters, adapter_maps
from comfy_api.latest import ComfyExtension, io, ui from comfy_api.latest import ComfyExtension, io, ui
from comfy.utils import ProgressBar
def make_batch_extra_option_dict(d, indicies, full_size=None): 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"] cond = model_wrap.conds["positive"]
dataset_size = sigmas.size(0) dataset_size = sigmas.size(0)
torch.cuda.empty_cache() torch.cuda.empty_cache()
ui_pbar = ProgressBar(self.total_steps)
for i in ( for i in (
pbar := tqdm.trange( pbar := trange(
self.total_steps, self.total_steps,
desc="Training LoRA", desc="Training LoRA",
smoothing=0.01, smoothing=0.01,
@ -203,6 +205,7 @@ class TrainSampler(comfy.samplers.Sampler):
if (i + 1) % self.grad_acc == 0: if (i + 1) % self.grad_acc == 0:
self.optimizer.step() self.optimizer.step()
self.optimizer.zero_grad() self.optimizer.zero_grad()
ui_pbar.update(1)
torch.cuda.empty_cache() torch.cuda.empty_cache()
return torch.zeros_like(latent_image) return torch.zeros_like(latent_image)