From 5af36d21823add99985a695348f04b1d9f5dd4ce Mon Sep 17 00:00:00 2001
From: Vladimir Mandic <mandic00@live.com>
Date: Sun, 30 Jun 2024 14:09:48 -0400
Subject: [PATCH] track video processing total job count

---
 modules/control/run.py  | 1 +
 modules/progress.py     | 2 +-
 modules/shared_state.py | 3 +++
 3 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/modules/control/run.py b/modules/control/run.py
index 1c043725c..ba736a4a8 100644
--- a/modules/control/run.py
+++ b/modules/control/run.py
@@ -388,6 +388,7 @@ def control_run(units: List[unit.Unit] = [], inputs: List[Image.Image] = [], ini
                     codec = util.decode_fourcc(video.get(cv2.CAP_PROP_FOURCC))
                     status, frame = video.read()
                     if status:
+                        shared.state.frame_count = 1 + frames // (video_skip_frames + 1)
                         frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                     shared.log.debug(f'Control: input video: path={inputs} frames={frames} fps={fps} size={w}x{h} codec={codec}')
                 except Exception as e:
diff --git a/modules/progress.py b/modules/progress.py
index f5573374b..aeef195b4 100644
--- a/modules/progress.py
+++ b/modules/progress.py
@@ -62,7 +62,7 @@ def progressapi(req: ProgressRequest):
     paused = shared.state.paused
     if not active:
         return InternalProgressResponse(job=shared.state.job, active=active, queued=queued, paused=paused, completed=completed, id_live_preview=-1, textinfo="Queued..." if queued else "Waiting...")
-    shared.state.job_count = max(shared.state.job_count, shared.state.job_no)
+    shared.state.job_count = max(shared.state.frame_count, shared.state.job_count, shared.state.job_no)
     batch_x = max(shared.state.job_no, 0)
     batch_y = max(shared.state.job_count, 1)
     step_x = max(shared.state.sampling_step, 0)
diff --git a/modules/shared_state.py b/modules/shared_state.py
index 79ee20f19..82f8d21dc 100644
--- a/modules/shared_state.py
+++ b/modules/shared_state.py
@@ -12,6 +12,7 @@ class State:
     job = ""
     job_no = 0
     job_count = 0
+    frame_count = 0
     total_jobs = 0
     job_timestamp = '0'
     sampling_step = 0
@@ -71,6 +72,7 @@ class State:
         self.interrupted = False
         self.job = title
         self.job_count = -1
+        self.frame_count = -1
         self.job_no = 0
         self.job_timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
         self.paused = False
@@ -93,6 +95,7 @@ class State:
         self.job = ""
         self.job_count = 0
         self.job_no = 0
+        self.frame_count = 0
         self.paused = False
         self.interrupted = False
         self.skipped = False