Update app.py

app.py

@@ -6,10 +6,12 @@ import tempfile
 import huggingface_hub
 import subprocess
 import threading
+from tqdm import tqdm
 
-def stream_output(pipe):
+def stream_output(pipe, progress_bar):
     for line in iter(pipe.readline, ''):
         print(line, end='')
+        progress_bar.update(1)
 
 HF_TKN = os.environ.get("GATED_HF_TOKEN")
 huggingface_hub.login(token=HF_TKN)
@@ -77,7 +79,7 @@ def check_for_mp4_in_outputs():
     else:
         return None
 
-def infer(ref_video_in, ref_image_in, progress=gr.Progress()):
+def infer(ref_video_in, ref_image_in, progress=gr.Progress(track_tqdm=True)):
     # check if 'outputs' dir exists and empty it if necessary
     check_outputs_folder('./outputs')
     
@@ -132,9 +134,13 @@ def infer(ref_video_in, ref_image_in, progress=gr.Progress()):
         command = ['python', 'inference.py', '--inference_config', file_path]
         process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, bufsize=1)
     
+        # Create a tqdm progress bar
+        progress_bar = tqdm(total=100, desc="Inference Progress", unit="line", ncols=100)
+    
         # Create threads to handle stdout and stderr
-        stdout_thread = threading.Thread(target=stream_output, args=(process.stdout,))
-        stderr_thread = threading.Thread(target=stream_output, args=(process.stderr,))
+        stdout_thread = threading.Thread(target=stream_output, args=(process.stdout, progress_bar))
+        stderr_thread = threading.Thread(target=stream_output, args=(process.stderr, progress_bar))
+
     
         # Start the threads
         stdout_thread.start()
@@ -144,6 +150,9 @@ def infer(ref_video_in, ref_image_in, progress=gr.Progress()):
         process.wait()
         stdout_thread.join()
         stderr_thread.join()
+
+        # Close the progress bar
+        progress_bar.close()
     
         print("Inference script finished with return code:", process.returncode)