NLag commited on
Commit
4c34134
·
1 Parent(s): 691a8e2

refactor(mediapipe): simplify GPU delegate preference logic and remove unused CUDA visibility check

Browse files
src/pozify/steps/pose_backends/mediapipe.py CHANGED
@@ -117,7 +117,7 @@ class _MediaPipeTasksPoseAdapter:
117
  return gpu_delegate
118
  if configured_delegate not in {"", "auto"}:
119
  return self._cpu_delegate()
120
- if zero_gpu_enabled() or _cuda_device_visible():
121
  return gpu_delegate
122
  return self._cpu_delegate()
123
 
@@ -195,22 +195,6 @@ def _create_legacy_pose(mediapipe_module: Any) -> Any:
195
  raise _native_library_error(exc) from exc
196
 
197
 
198
- def _cuda_device_visible() -> bool:
199
- for env_name in ("CUDA_VISIBLE_DEVICES", "NVIDIA_VISIBLE_DEVICES"):
200
- value = os.getenv(env_name)
201
- if value and value.strip().lower() not in {"", "-1", "none", "void", "no"}:
202
- return True
203
-
204
- try:
205
- import torch
206
- except Exception:
207
- return False
208
-
209
- cuda = getattr(torch, "cuda", None)
210
- is_available = getattr(cuda, "is_available", None)
211
- return bool(callable(is_available) and is_available())
212
-
213
-
214
  def _native_library_error(exc: OSError) -> PoseBackendUnavailableError:
215
  return PoseBackendUnavailableError(
216
  "MediaPipe could not load native system libraries. On Hugging Face Spaces, "
 
117
  return gpu_delegate
118
  if configured_delegate not in {"", "auto"}:
119
  return self._cpu_delegate()
120
+ if zero_gpu_enabled():
121
  return gpu_delegate
122
  return self._cpu_delegate()
123
 
 
195
  raise _native_library_error(exc) from exc
196
 
197
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
198
  def _native_library_error(exc: OSError) -> PoseBackendUnavailableError:
199
  return PoseBackendUnavailableError(
200
  "MediaPipe could not load native system libraries. On Hugging Face Spaces, "
tests/test_pose_steps.py CHANGED
@@ -325,12 +325,6 @@ class MediaPipeTasksDelegateTests(unittest.TestCase):
325
  with patch.dict(os.environ, {"SPACES_ZERO_GPU": "1"}, clear=True):
326
  self.assertEqual(adapter._preferred_delegate(), "gpu")
327
 
328
- def test_mediapipe_tasks_prefers_gpu_when_cuda_device_is_visible(self) -> None:
329
- adapter = self._adapter(self._fake_mediapipe())
330
-
331
- with patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "0"}, clear=True):
332
- self.assertEqual(adapter._preferred_delegate(), "gpu")
333
-
334
  def test_mediapipe_delegate_env_can_force_gpu(self) -> None:
335
  adapter = self._adapter(self._fake_mediapipe())
336
 
 
325
  with patch.dict(os.environ, {"SPACES_ZERO_GPU": "1"}, clear=True):
326
  self.assertEqual(adapter._preferred_delegate(), "gpu")
327
 
 
 
 
 
 
 
328
  def test_mediapipe_delegate_env_can_force_gpu(self) -> None:
329
  adapter = self._adapter(self._fake_mediapipe())
330