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()
|
| 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 |
|