Update app.py

app.py

@@ -4,6 +4,7 @@ import gradio as gr
 import AnimeGANv3_src
 import numpy as np
 import logging
+import onnxruntime as ort  # Added this import
 
 # Set up logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
@@ -13,6 +14,10 @@ class AnimeGANv3:
         os.makedirs('output', exist_ok=True)
         os.makedirs('frames', exist_ok=True)
         logging.info(f"Available ONNX Runtime providers: {ort.get_available_providers()}")
+        if 'CUDAExecutionProvider' in ort.get_available_providers():
+            logging.info("Running on GPU with CUDA")
+        else:
+            logging.info("Running on CPU")
 
     def process_frame(self, frame, style_code, det_face):
         frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
@@ -45,12 +50,12 @@ class AnimeGANv3:
             logging.info(f"Extracted {frame_count} frames at {fps} FPS to process")
 
             # Process in batches
-            batch_size = 50  # Adjust based on testing (e.g., 50 frames per batch)
+            batch_size = 50  # Adjust based on testing
             save_path = "output/out.mp4"
             fourcc = cv2.VideoWriter_fourcc(*'mp4v')
             width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
             height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-            out = None  # Video writer initialized later
+            out = None
 
             frame_idx = 0
             while cap.isOpened():
@@ -65,7 +70,6 @@ class AnimeGANv3:
                 if not batch_frames:
                     break
 
-                # Process batch
                 for idx, frame in enumerate(batch_frames):
                     stylized_frame = self.process_frame(frame, style_code, det_face)
                     if out is None:  # Initialize writer on first frame