Update app.py

app.py

@@ -6,62 +6,57 @@ import datetime
 import sys
 
 def run_command(command):
-    """Run a shell command and print its output."""
-    print(f"Running command: {' '.join(command)}")
+    """Run a shell command and return its output and error status."""
+    print(f"Running command: {command}")
     try:
-        subprocess.check_call(command, shell=True)
+        result = subprocess.run(command, shell=True, check=True, capture_output=True, text=True)
+        return True, result.stdout
     except subprocess.CalledProcessError as e:
-        print(f"Error running command {command}: {e}")
-        sys.exit(1)
+        return False, f"Error running command: {e}\nOutput: {e.output}\nError: {e.stderr}"
 
 def check_for_mp4_in_outputs(given_folder):
-    # Define the path to the outputs folder
     outputs_folder = given_folder
-    
-    # Check if the outputs folder exists
     if not os.path.exists(outputs_folder):
         return None
-    
-    # Check if there is a .mp4 file in the outputs folder
     mp4_files = [f for f in os.listdir(outputs_folder) if f.endswith('.mp4')]
-    
-    # Return the path to the mp4 file if it exists
-    if mp4_files:
-        return os.path.join(outputs_folder, mp4_files[0])
-    else:
-        return None
-
+    return os.path.join(outputs_folder, mp4_files[0]) if mp4_files else None
 
 def infer(input_video, cropped_and_aligned):
+    try:
+        torch.cuda.empty_cache()
 
-    torch.cuda.empty_cache()
+        filepath = input_video
+        timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
+        output_folder_name = f"results_{timestamp}"
 
-    filepath = input_video
-    # Get the current timestamp
-    timestamp = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
+        if cropped_and_aligned:
+            command = f"{sys.executable} inference_keep.py -i={filepath} -o={output_folder_name} --has_aligned --save_video -s=1"
+        else:
+            command = f"{sys.executable} inference_keep.py -i={filepath} -o={output_folder_name} --draw_box --save_video -s=1 --bg_upsampler=realesrgan"
 
-    output_folder_name = f"results_{timestamp}"
+        success, output = run_command(command)
+        if not success:
+            return None, output  # Return None for the video and the error message
 
-    if cropped_and_aligned is True:
-        # Example: Run the inference script (replace with your actual command)
-        run_command(f"{sys.executable} inference_keep.py -i={filepath} -o={output_folder_name} --has_aligned --save_video -s=1")
-    else:
-        run_command(f"{sys.executable} inference_keep.py -i={filepath} -o={output_folder_name} --draw_box --save_video -s=1 --bg_upsampler=realesrgan")
-    
-    torch.cuda.empty_cache()
-    
-    # Call the function and print the result
-    this_infer_folder = os.path.splitext(os.path.basename(filepath))[0]
-    joined_path = os.path.join(output_folder_name, this_infer_folder)
-    mp4_file_path = check_for_mp4_in_outputs(joined_path)
-    print(mp4_file_path)
+        torch.cuda.empty_cache()
 
-    print(f"RESULT: {mp4_file_path}")
-    
-    return mp4_file_path
+        this_infer_folder = os.path.splitext(os.path.basename(filepath))[0]
+        joined_path = os.path.join(output_folder_name, this_infer_folder)
+        mp4_file_path = check_for_mp4_in_outputs(joined_path)
 
+        if mp4_file_path:
+            print(f"RESULT: {mp4_file_path}")
+            return mp4_file_path, "Processing completed successfully."
+        else:
+            return None, "Processing completed, but no output video was found."
 
+    except Exception as e:
+        return None, f"An unexpected error occurred: {str(e)}"
+
+# Gradio interface setup
 result_video = gr.Video()
+error_output = gr.Textbox(label="Status/Error")
+
 with gr.Blocks() as demo:
     with gr.Column():
         gr.Markdown("# KEEP")
@@ -94,19 +89,19 @@ with gr.Blocks() as demo:
                     ],
                     fn = infer,
                     inputs = [input_video, is_cropped_and_aligned],
-                    outputs = [result_video],
+                    outputs = [result_video, error_output],
                     run_on_click = False,
                     cache_examples = "lazy"
                 )
             
             with gr.Column():
                 result_video.render()
-                
+                error_output.render()
 
     submit_btn.click(
         fn = infer, 
         inputs = [input_video, is_cropped_and_aligned],
-        outputs = [result_video],
+        outputs = [result_video, error_output],
         show_api=False
     )