'update'

hook_kohya_ss_run.py

@@ -112,6 +112,7 @@ def sample_images(self, *args, **kwargs):
     LOG({
         "type": "sample_images",
         "global_step": global_step,
+        "total_steps": cmd_args.max_train_steps,
         # "latent": noise_pred_latent_path,
     })
 
@@ -133,7 +134,7 @@ def run_lora_sd1_5():
     train_args = config2args(train_network.setup_parser(), train_config)
 
     LOG({
-        "type": "start_train", 
+        "type": "start_train",
     })
     trainer.train(train_args)
 
@@ -153,9 +154,9 @@ def run_lora_sdxl():
     trainer = sdxl_train_network.SdxlNetworkTrainer()
     train_config = json.loads(train_config_json)
     train_args = config2args(sdxl_train_network.setup_parser(), train_config)
-    
+
     LOG({
-        "type": "start_train", 
+        "type": "start_train",
     })
     trainer.train(train_args)
 

mz_train_tools_core.py

@@ -365,6 +365,7 @@ def run_hook_kohya_ss_run_file(kohya_ss_tool_dir, train_config, trainer_func, ot
         __file__), "hook_kohya_ss_run.py",)
     train_config_str = json.dumps(train_config)
     max_train_steps = train_config.get("max_train_steps")
+    max_train_epochs = train_config.get("max_train_epochs")
     is_running = True
 
     taesd_type = "sd1_5"
@@ -372,9 +373,12 @@ def run_hook_kohya_ss_run_file(kohya_ss_tool_dir, train_config, trainer_func, ot
         taesd_type = "sd1_5"
     if trainer_func.find("sdxl") != -1:
         taesd_type = "sdxl"
-    pb = Utils.progress_bar(train_config.get("max_train_steps"), taesd_type)
 
-    previewer = pb.get_previewer()
+
+    
+ 
+    pb = Utils.progress_bar(train_config.get("max_train_steps"), taesd_type)
+ 
     import traceback
 
     import comfy.model_management
@@ -395,9 +399,12 @@ def run_hook_kohya_ss_run_file(kohya_ss_tool_dir, train_config, trainer_func, ot
                 xy_img = get_sample_images(train_config)
 
                 max_side = max(xy_img.width, xy_img.height)
+                # print(f"global_step: {global_step}, max_train_steps: {max_train_steps}")
 
+
+                total_steps = resp.get("total_steps")
                 pb.update(
-                    int(global_step), int(max_train_steps), ("JPEG", xy_img, max_side))
+                    int(global_step), int(total_steps), ("JPEG", xy_img, max_side))
             else:
                 print(f"LOG: {log}")
         except Exception as e:

mz_train_tools_utils.py

@@ -717,6 +717,7 @@ class Utils:
                     pil_img_info = pil_img
                 try:
                     # print("pil_img_info:", type(pil_img), pil_img_info)
+                    # print("step:", step, "total_steps:", total_steps)
                     self.pbar.update_absolute(
                         step, total_steps, pil_img_info)
                 except Exception as e: