diff --git a/hook_kohya_ss_run.py b/hook_kohya_ss_run.py
index c8f3339..d0488e6 100644
--- a/hook_kohya_ss_run.py
+++ b/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)
 
diff --git a/mz_train_tools_core.py b/mz_train_tools_core.py
index 6db82ec..89bd55d 100644
--- a/mz_train_tools_core.py
+++ b/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:
diff --git a/mz_train_tools_utils.py b/mz_train_tools_utils.py
index ce938de..49ac5eb 100644
--- a/mz_train_tools_utils.py
+++ b/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: