Update simple_app.py

simple_app.py

@@ -14,7 +14,7 @@ snapshot_download(
 
 def infer(prompt, progress=gr.Progress(track_tqdm=True)):
     
-    total_process_steps = 11
+   total_process_steps = 11
     irrelevant_steps = 4
     relevant_steps = total_process_steps - irrelevant_steps  # 7 steps
 
@@ -29,26 +29,45 @@ def infer(prompt, progress=gr.Progress(track_tqdm=True)):
 
     # Variables for managing the sub-progress bar for each step.
     current_sub_bar = None
-    current_timer = None
+    current_sub_thread = None
+    current_cancel_event = None
     sub_lock = threading.Lock()
 
+    def update_sub_bar(sub_bar, cancel_event):
+        # This function updates the sub_bar once per second up to 20 seconds,
+        # unless the cancel_event is set.
+        for i in range(20):
+            if cancel_event.is_set():
+                break
+            time.sleep(1)
+            sub_bar.update(1)
+            sub_bar.refresh()
+        # When done (or canceled), do nothing here;
+        # closing will be done in close_sub_bar() from the main thread.
+    
     def close_sub_bar():
-        nonlocal current_sub_bar, current_timer, overall_bar
+        nonlocal current_sub_bar, current_sub_thread, current_cancel_event
         with sub_lock:
             if current_sub_bar is not None:
                 try:
-                    # Ensure the sub-bar is complete.
-                    current_sub_bar.update(1 - current_sub_bar.n)
+                    # Complete any remaining ticks (if any)
+                    remaining = current_sub_bar.total - current_sub_bar.n
+                    if remaining > 0:
+                        current_sub_bar.update(remaining)
                 except Exception:
                     pass
                 current_sub_bar.close()
                 overall_bar.update(1)
                 overall_bar.refresh()
                 current_sub_bar = None
-                current_timer = None
-    
+            if current_sub_thread is not None:
+                current_sub_thread.join()
+                current_sub_thread = None
+            if current_cancel_event is not None:
+                current_cancel_event = None
+
     command = [
-        "python", "-u", "-m", "generate",  # using -u for unbuffered output and omitting .py extension
+        "python", "-u", "-m", "generate",  # using -u for unbuffered output and module name without .py
         "--task", "t2v-1.3B",
         "--size", "832*480",
         "--ckpt_dir", "./Wan2.1-T2V-1.3B",
@@ -58,7 +77,6 @@ def infer(prompt, progress=gr.Progress(track_tqdm=True)):
         "--save_file", "generated_video.mp4"
     ]
 
-    # Start the process with unbuffered output and combine stdout and stderr.
     process = subprocess.Popen(
         command,
         stdout=subprocess.PIPE,
@@ -88,34 +106,37 @@ def infer(prompt, progress=gr.Progress(track_tqdm=True)):
         if "INFO:" in stripped_line:
             parts = stripped_line.split("INFO:", 1)
             msg = parts[1].strip() if len(parts) > 1 else ""
-            tqdm.write(stripped_line)  # Print the log line
+            tqdm.write(stripped_line)
 
             if processed_steps < irrelevant_steps:
                 processed_steps += 1
             else:
                 with sub_lock:
-                    # If a sub-bar is active, cancel its timer and close it immediately.
+                    # If a sub-bar is already active, cancel its update and close it.
                     if current_sub_bar is not None:
-                        if current_timer is not None:
-                            current_timer.cancel()
+                        if current_cancel_event is not None:
+                            current_cancel_event.set()
                         close_sub_bar()
-                    # Create a new sub-bar for the current step.
-                    current_sub_bar = tqdm(total=1, desc=msg, position=2,
+                    # Now create a new sub-bar for the current step.
+                    current_sub_bar = tqdm(total=20, desc=msg, position=2,
                                            ncols=120, dynamic_ncols=False, leave=True)
-                    # Start a timer to automatically close this sub-bar after 20 seconds.
-                    current_timer = threading.Timer(20, close_sub_bar)
-                    current_timer.start()
+                    current_cancel_event = threading.Event()
+                    current_sub_thread = threading.Thread(target=update_sub_bar,
+                                                          args=(current_sub_bar, current_cancel_event))
+                    current_sub_thread.start()
             continue
 
         else:
             tqdm.write(stripped_line)
 
     process.wait()
-    # Clean up: if a sub-bar is still active, close it.
-    if current_timer is not None:
-        current_timer.cancel()
-    if current_sub_bar is not None:
-        close_sub_bar()
+    # After process ends, if a sub-bar is still active, cancel and close it.
+    with sub_lock:
+        if current_cancel_event is not None:
+            current_cancel_event.set()
+        if current_sub_bar is not None:
+            close_sub_bar()
+
     if gen_progress_bar:
         gen_progress_bar.close()
     overall_bar.close()