Update simple_app.py

simple_app.py

@@ -2,7 +2,6 @@ import gradio as gr
 import re
 import subprocess
 import time
-import select
 from tqdm import tqdm
 from huggingface_hub import snapshot_download
 import torch
@@ -59,30 +58,58 @@ def infer(prompt, progress=gr.Progress(track_tqdm=True)):
                                text=True,
                                bufsize=1)
     
-    stdout, stderr = process.communicate()
-
     # Print logs
-    print("🔹 Output Logs:\n", stdout)
-    if stderr:
-      print("🔺 Error Logs:\n", stderr)
-    print("Process started successfully.")
-    
-
+    print("🔹 Output Logs:")
+    stdout = process.stdout
+    stderr = process.stderr
     while True:
-        # Poll stdout with a 40ms timeout.
-        rlist, _, _ = select.select([process.stdout], [], [], 0.04)
-        if rlist:
-            line = process.stdout.readline()
-            if not line:
-                break
-            stripped_line = line.strip()
-            if not stripped_line:
+        line = stdout.readline()
+        if not line:
+            break
+        stripped_line = line.strip()
+        if not stripped_line:
+            continue
+
+        # Check for video generation progress (Level 3)
+        progress_match = progress_pattern.search(stripped_line)
+        
+        if progress_match:
+            if sub_bar is not None:
+                if sub_ticks < sub_tick_total:
+                    sub_bar.update(sub_tick_total - sub_ticks)
+                sub_bar.close()
+                overall_bar.update(1)
+                overall_bar.refresh()
+                sub_bar = None
+                sub_ticks = 0
+            video_phase = True
+            current = int(progress_match.group(2))
+            total = int(progress_match.group(3))
+            if video_progress_bar is None:
+                video_progress_bar = tqdm(total=total, desc="Video Generation", position=0,
+                                          ncols=120, dynamic_ncols=True, leave=True)
+            video_progress_bar.update(current - video_progress_bar.n)
+            video_progress_bar.refresh()
+            if video_progress_bar.n >= video_progress_bar.total:
+                video_phase = False
+                overall_bar.update(1)
+                overall_bar.refresh()
+                video_progress_bar.close()
+                video_progress_bar = None
+            continue
+
+        # Process INFO messages (Level 2 sub-step)
+        if "INFO:" in stripped_line:
+            parts = stripped_line.split("INFO:", 1)
+            msg = parts[1].strip() if len(parts) > 1 else ""
+            print(f"[INFO]: {msg}")  # Log the message
+
+            # For the first 4 INFO messages, simply count them.
+            if processed_steps < irrelevant_steps:
+                processed_steps += 1
                 continue
-
-            # Check for video generation progress (Level 3)
-            progress_match = progress_pattern.search(stripped_line)
-            
-            if progress_match:
+            else:
+                # A new relevant INFO message has arrived.
                 if sub_bar is not None:
                     if sub_ticks < sub_tick_total:
                         sub_bar.update(sub_tick_total - sub_ticks)
@@ -91,63 +118,21 @@ def infer(prompt, progress=gr.Progress(track_tqdm=True)):
                     overall_bar.refresh()
                     sub_bar = None
                     sub_ticks = 0
-                video_phase = True
-                current = int(progress_match.group(2))
-                total = int(progress_match.group(3))
-                if video_progress_bar is None:
-                    video_progress_bar = tqdm(total=total, desc="Video Generation", position=0,
-                                              ncols=120, dynamic_ncols=True, leave=True)
-                video_progress_bar.update(current - video_progress_bar.n)
-                video_progress_bar.refresh()
-                if video_progress_bar.n >= video_progress_bar.total:
-                    video_phase = False
-                    overall_bar.update(1)
-                    overall_bar.refresh()
-                    video_progress_bar.close()
-                    video_progress_bar = None
-                continue
-
-            # Process INFO messages (Level 2 sub-step)
-            if "INFO:" in stripped_line:
-                parts = stripped_line.split("INFO:", 1)
-                msg = parts[1].strip() if len(parts) > 1 else ""
-                print(f"[INFO]: {msg}")  # Log the message
-
-                # For the first 4 INFO messages, simply count them.
-                if processed_steps < irrelevant_steps:
-                    processed_steps += 1
-                    continue
-                else:
-                    # A new relevant INFO message has arrived.
-                    if sub_bar is not None:
-                        if sub_ticks < sub_tick_total:
-                            sub_bar.update(sub_tick_total - sub_ticks)
-                        sub_bar.close()
-                        overall_bar.update(1)
-                        overall_bar.refresh()
-                        sub_bar = None
-                        sub_ticks = 0
-                    # Start a new sub-step bar for the current INFO message.
-                    sub_bar = tqdm(total=sub_tick_total, desc=msg, position=2,
-                                   ncols=120, dynamic_ncols=False, leave=True)
-                    sub_ticks = 0
-                continue
-            else:
-                print(stripped_line)
+                # Start a new sub-step bar for the current INFO message.
+                sub_bar = tqdm(total=sub_tick_total, desc=msg, position=2,
+                               ncols=120, dynamic_ncols=False, leave=True)
+                sub_ticks = 0
+            continue
         else:
-            # No new data within 40ms.
-            if sub_bar is not None:
-                if sub_ticks < sub_tick_total:
-                    sub_bar.update(1)
-                    sub_ticks += 1
-                    sub_bar.refresh()
-        if process.poll() is not None:
-            break
+            print(stripped_line)
 
-    # Drain any remaining output.
+    # Drain any remaining output
     for line in process.stdout:
         print(line.strip())
+
     process.wait()
+
+    # Finalize progress bars
     if video_progress_bar is not None:
         video_progress_bar.close()
     if sub_bar is not None: