pbar fixes

nodes.py

@@ -715,7 +715,8 @@ class FluxTrainLoop:
             initial_global_step = network_trainer.global_step
 
             target_global_step = network_trainer.global_step + steps
-            #pbar = comfy.utils.ProgressBar(steps)
+            comfy_pbar = comfy.utils.ProgressBar(steps)
+            network_trainer.comfy_pbar = comfy_pbar
             while network_trainer.global_step < target_global_step:
                 steps_done = training_loop(
                     break_at_steps = target_global_step,
@@ -726,7 +727,7 @@ class FluxTrainLoop:
                 # Also break if the global steps have reached the max train steps
                 if network_trainer.global_step >= network_trainer.args.max_train_steps:
                     break
-
+            
             trainer = {
                 "network_trainer": network_trainer,
                 "training_loop": training_loop,

train_network.py

@@ -1074,12 +1074,13 @@ class NetworkTrainer:
         self.lr_scheduler = lr_scheduler
         self.save_model = save_model
         self.remove_model = remove_model
+        self.comfy_pbar = None
         
         progress_bar = tqdm(range(args.max_train_steps - initial_step), smoothing=0, disable=False, desc="steps")
         
         def training_loop(break_at_steps, epoch):
             steps_done = 0
-            comfy_bar = ProgressBar(break_at_steps - self.global_step)
+            
             #accelerator.print(f"\nepoch {epoch+1}/{num_train_epochs}")
             progress_bar.set_description(f"Epoch {epoch + 1}/{num_train_epochs} - steps")
 
@@ -1235,7 +1236,7 @@ class NetworkTrainer:
                 if self.global_step >= break_at_steps:
                     break
                 steps_done += 1
-                comfy_bar.update(1)
+                self.comfy_pbar.update(1)
 
             if args.logging_dir is not None:
                 logs = {"loss/epoch": self.loss_recorder.moving_average}