Force MoveNet to run on CPU to avoid GPU JIT errors on HF Spaces

app.py

@@ -156,7 +156,7 @@ def predict_pose_cnn(img_path):
     finally:
         cleanup_memory()
 
-@app.route('/static/uploads/<path:filename>')
+@app.route('/static/uploads/<path:filename>', endpoint='serve_video')
 def serve_video(filename):
     response = send_from_directory(app.config['UPLOAD_FOLDER'], filename, as_attachment=False)
     # Ensure correct content type, especially for Safari/iOS if issues arise
@@ -185,18 +185,11 @@ def process_video_movenet(video_path):
         total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
         print(f"[DEBUG] Video properties - FPS: {fps}, Width: {width}, Height: {height}, Total Frames: {total_frames}")
         
-        # Initialize MoveNet model on GPU if available
-        print("[DEBUG] Initializing MoveNet model")
-        if gpus:
-            print("[DEBUG] Using GPU for MoveNet")
-            with tf.device('/GPU:0'):
-                movenet_model = hub.load("https://tfhub.dev/google/movenet/singlepose/lightning/4")
-                movenet = movenet_model.signatures['serving_default']
-        else:
-            print("[DEBUG] No GPU found, using CPU for MoveNet")
-            with tf.device('/CPU:0'):
-                movenet_model = hub.load("https://tfhub.dev/google/movenet/singlepose/lightning/4")
-                movenet = movenet_model.signatures['serving_default']
+        # Force MoveNet to CPU to avoid GPU JIT error
+        print("[DEBUG] Forcing CPU for MoveNet (due to GPU JIT error)")
+        with tf.device('/CPU:0'):
+            movenet_model = hub.load("https://tfhub.dev/google/movenet/singlepose/lightning/4")
+            movenet = movenet_model.signatures['serving_default']
         
         # Create output video writer
         output_filename = f'output_movenet_lightning.mp4'
@@ -210,14 +203,18 @@ def process_video_movenet(video_path):
         
         frame_count = 0
         processed_frames = 0
+        first_frame_shape = None
         
         while cap.isOpened():
             ret, frame = cap.read()
-            if not ret:
+            print(f"[DEBUG] Frame {frame_count+1}: ret={ret}, frame is None: {frame is None}")
+            if not ret or frame is None:
+                print(f"[DEBUG] Stopping at frame {frame_count+1}: ret={ret}, frame is None: {frame is None}")
                 break
+            if first_frame_shape is None:
+                first_frame_shape = frame.shape
+                print(f"[DEBUG] First frame shape: {first_frame_shape}")
             frame_count += 1
-            if frame_count % 10 != 0:  # Process every 10th frame
-                continue
             try:
                 # Ensure frame size matches VideoWriter
                 if frame.shape[1] != width or frame.shape[0] != height:
@@ -227,15 +224,10 @@ def process_video_movenet(video_path):
                 img = frame.copy()
                 img = tf.image.resize_with_pad(tf.expand_dims(img, axis=0), 192, 192)
                 img = tf.cast(img, dtype=tf.int32)
-                # Run inference on GPU if available
-                if gpus:
-                    with tf.device('/GPU:0'):
-                        results = movenet(img)
-                        keypoints = results['output_0'].numpy()
-                else:
-                    with tf.device('/CPU:0'):
-                        results = movenet(img)
-                        keypoints = results['output_0'].numpy()
+                # Always run inference on CPU
+                with tf.device('/CPU:0'):
+                    results = movenet(img)
+                    keypoints = results['output_0'].numpy()
                 # Process keypoints and draw on frame
                 y, x, c = frame.shape
                 shaped = np.squeeze(keypoints)
@@ -260,7 +252,9 @@ def process_video_movenet(video_path):
         print(f"[DEBUG] Output video file size: {file_size} bytes")
         if processed_frames == 0 or file_size < 1000:
             raise ValueError(f"Output video file is empty or too small: {output_path}")
-        return url_for('serve_video', filename=output_filename, _external=False)
+        video_url = url_for('serve_video', filename=output_filename, _external=False)
+        print(f"[DEBUG] Returning video URL: {video_url}")
+        return video_url
     except Exception as e:
         print(f"[ERROR] Error in process_video_movenet: {str(e)}")
         traceback.print_exc()
@@ -400,7 +394,9 @@ def process_video_mediapipe(video_path):
         if frame_count == 0:
             raise ValueError("No frames were processed from the video by MediaPipe")
         logger.info(f"MediaPipe video processing completed. Processed {frame_count} frames. Output: {output_path}")
-        return url_for('serve_video', filename=output_filename, _external=False)
+        video_url = url_for('serve_video', filename=output_filename, _external=False)
+        print(f"[DEBUG] Returning video URL: {video_url}")
+        return video_url
     except Exception as e:
         logger.error(f'Error in process_video_mediapipe: {e}')
         traceback.print_exc()