Update simple_app.py

simple_app.py

@@ -16,14 +16,15 @@ def infer(prompt, progress=gr.Progress(track_tqdm=True)):
     
     total_process_steps = 11
     irrelevant_steps = 4
-    relevant_steps = total_process_steps - irrelevant_steps  # 7 steps
+    # Only steps 5 through 11 (i.e. 7 steps) count.
+    relevant_steps = total_process_steps - irrelevant_steps
 
-    # Overall progress bar for the process steps.
+    # Create the overall progress bar for the steps.
     overall_bar = tqdm(total=relevant_steps, desc="Overall Process", position=1,
                        ncols=120, dynamic_ncols=False, leave=True)
     processed_steps = 0
 
-    # Regex for detecting video generation progress lines (e.g., "10%|...| 5/50")
+    # Regex for detecting video generation progress lines (e.g. "10%|...| 5/50")
     progress_pattern = re.compile(r"(\d+)%\|.*\| (\d+)/(\d+)")
     gen_progress_bar = None
 
@@ -34,40 +35,38 @@ def infer(prompt, progress=gr.Progress(track_tqdm=True)):
     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.
+        # Update sub-bar once per second for up to 20 seconds,
+        # unless 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():
+        # (Closing and overall-bar update are handled externally.)
+
+    def cancel_sub_bar():
         nonlocal current_sub_bar, current_sub_thread, current_cancel_event
         with sub_lock:
+            if current_cancel_event is not None:
+                current_cancel_event.set()
+            if current_sub_thread is not None:
+                current_sub_thread.join(timeout=1)
+                current_sub_thread = None
             if current_sub_bar is not None:
-                try:
-                    # 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
+                # Complete any remaining ticks.
+                remaining = current_sub_bar.total - current_sub_bar.n
+                if remaining > 0:
+                    current_sub_bar.update(remaining)
                 current_sub_bar.close()
-                overall_bar.update(1)
-                overall_bar.refresh()
                 current_sub_bar = 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
+            # Update overall progress by one step.
+            overall_bar.update(1)
+            overall_bar.refresh()
+            current_cancel_event = None
 
     command = [
-        "python", "-u", "-m", "generate",  # using -u for unbuffered output and module name without .py
+        "python", "-u", "-m", "generate",  # using unbuffered mode
         "--task", "t2v-1.3B",
         "--size", "832*480",
         "--ckpt_dir", "./Wan2.1-T2V-1.3B",
@@ -90,7 +89,7 @@ def infer(prompt, progress=gr.Progress(track_tqdm=True)):
         if not stripped_line:
             continue
 
-        # Check if this is a video generation progress line.
+        # Check for video generation progress lines.
         progress_match = progress_pattern.search(stripped_line)
         if progress_match:
             current = int(progress_match.group(2))
@@ -104,38 +103,37 @@ def infer(prompt, progress=gr.Progress(track_tqdm=True)):
 
         # Check for INFO lines.
         if "INFO:" in stripped_line:
+            # Extract the INFO message (the text after "INFO:")
             parts = stripped_line.split("INFO:", 1)
             msg = parts[1].strip() if len(parts) > 1 else ""
-            tqdm.write(stripped_line)
+            tqdm.write(stripped_line)  # Log the full line
 
             if processed_steps < irrelevant_steps:
                 processed_steps += 1
             else:
                 with sub_lock:
-                    # If a sub-bar is already active, cancel its update and close it.
+                    # If a sub-bar is active, cancel it immediately.
                     if current_sub_bar is not None:
-                        if current_cancel_event is not None:
-                            current_cancel_event.set()
-                        close_sub_bar()
-                    # Now create a new sub-bar for the current step.
+                        cancel_sub_bar()
+                    # Create a new sub-progress bar for this step (lasting up to 20 seconds).
                     current_sub_bar = tqdm(total=20, desc=msg, position=2,
                                            ncols=120, dynamic_ncols=False, leave=True)
                     current_cancel_event = threading.Event()
-                    current_sub_thread = threading.Thread(target=update_sub_bar,
-                                                          args=(current_sub_bar, current_cancel_event))
+                    current_sub_thread = threading.Thread(target=update_sub_bar, args=(current_sub_bar, current_cancel_event))
+                    current_sub_thread.daemon = True
                     current_sub_thread.start()
             continue
 
         else:
             tqdm.write(stripped_line)
 
+    # Process has ended; cancel any active sub-progress bar.
     process.wait()
-    # 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()
+            cancel_sub_bar()
 
     if gen_progress_bar:
         gen_progress_bar.close()