diff --git a/src/lerobot/utils/import_utils.py b/src/lerobot/utils/import_utils.py index b9a9e6825..7d461fb88 100644 --- a/src/lerobot/utils/import_utils.py +++ b/src/lerobot/utils/import_utils.py @@ -62,6 +62,7 @@ def is_package_available(pkg_name: str, return_version: bool = False) -> tuple[b _transformers_available = is_package_available("transformers") +_timm_available = is_package_available("timm") _peft_available = is_package_available("peft") diff --git a/tests/policies/xvla/test_xvla_original_vs_lerobot.py b/tests/policies/xvla/test_xvla_original_vs_lerobot.py index ac37af146..a893c9988 100644 --- a/tests/policies/xvla/test_xvla_original_vs_lerobot.py +++ b/tests/policies/xvla/test_xvla_original_vs_lerobot.py @@ -26,9 +26,18 @@ import numpy as np import pytest import torch -from lerobot.policies.xvla.configuration_xvla import XVLAConfig -from lerobot.policies.xvla.modeling_xvla import XVLAPolicy -from lerobot.policies.xvla.processor_xvla import make_xvla_pre_post_processors +# Conditional import for type checking and lazy loading +from lerobot.utils.import_utils import _timm_available, _transformers_available + +if _timm_available and _transformers_available: + from lerobot.policies.xvla.configuration_xvla import XVLAConfig + from lerobot.policies.xvla.modeling_xvla import XVLAPolicy + from lerobot.policies.xvla.processor_xvla import make_xvla_pre_post_processors +else: + XVLAConfig = None + XVLAPolicy = None + make_xvla_pre_post_processors = None + from lerobot.processor import PolicyAction, PolicyProcessorPipeline # noqa: E402 from lerobot.utils.constants import OBS_IMAGES, OBS_STATE # noqa: E402 from tests.utils import require_package # noqa: E402