fix(dependencies): protect peft import

This commit is contained in:
Steven Palma
2026-07-28 13:29:20 +02:00
parent 9b25b7fe0a
commit 74d177c8c8
7 changed files with 75 additions and 30 deletions
+8 -1
View File
@@ -44,12 +44,19 @@ from lerobot.utils.constants import (
POLICY_PREPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME,
) )
from lerobot.utils.feature_utils import dataset_to_policy_features from lerobot.utils.feature_utils import dataset_to_policy_features
from lerobot.utils.import_utils import _peft_available, require_package
from .evo1.configuration_evo1 import Evo1Config from .evo1.configuration_evo1 import Evo1Config
from .groot.configuration_groot import GrootConfig from .groot.configuration_groot import GrootConfig
from .pretrained import PreTrainedPolicy from .pretrained import PreTrainedPolicy
from .utils import validate_visual_features_consistency from .utils import validate_visual_features_consistency
if TYPE_CHECKING or _peft_available:
from peft import PeftConfig, PeftModel
else:
PeftConfig = None
PeftModel = None
def _reconnect_relative_absolute_steps( def _reconnect_relative_absolute_steps(
preprocessor: PolicyProcessorPipeline, postprocessor: PolicyProcessorPipeline preprocessor: PolicyProcessorPipeline, postprocessor: PolicyProcessorPipeline
@@ -334,7 +341,7 @@ def make_policy(
# Load a pretrained PEFT model on top of the policy. The pretrained path points to the folder/repo # Load a pretrained PEFT model on top of the policy. The pretrained path points to the folder/repo
# of the adapter and the adapter's config contains the path to the base policy. So we need the # of the adapter and the adapter's config contains the path to the base policy. So we need the
# adapter config first, then load the correct policy and then apply PEFT. # adapter config first, then load the correct policy and then apply PEFT.
from peft import PeftConfig, PeftModel require_package("peft", extra="peft")
logging.info("Loading policy's PEFT adapter.") logging.info("Loading policy's PEFT adapter.")
@@ -43,11 +43,22 @@ from torch.distributions import Beta
from lerobot.policies.pretrained import PreTrainedPolicy from lerobot.policies.pretrained import PreTrainedPolicy
from lerobot.utils.constants import ACTION from lerobot.utils.constants import ACTION
from lerobot.utils.import_utils import _scipy_available, _transformers_available, require_package from lerobot.utils.import_utils import (
_peft_available,
_scipy_available,
_transformers_available,
require_package,
)
from ..rtc.modeling_rtc import RTCProcessor from ..rtc.modeling_rtc import RTCProcessor
from .configuration_molmoact2 import MolmoAct2Config from .configuration_molmoact2 import MolmoAct2Config
if TYPE_CHECKING or _peft_available:
from peft import LoraConfig, get_peft_model
else:
LoraConfig = None
get_peft_model = None
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -1731,13 +1742,11 @@ class MolmoAct2Policy(PreTrainedPolicy):
def _build_inner_lora_config(self): def _build_inner_lora_config(self):
require_package("peft", extra="molmoact2") require_package("peft", extra="molmoact2")
from peft import LoraConfig
return LoraConfig(**self._get_inner_peft_targets()) return LoraConfig(**self._get_inner_peft_targets())
def _apply_lora_adapters(self) -> None: def _apply_lora_adapters(self) -> None:
require_package("peft", extra="molmoact2") require_package("peft", extra="molmoact2")
from peft import get_peft_model
peft_config = self._build_inner_lora_config() peft_config = self._build_inner_lora_config()
self._validate_peft_config(peft_config) self._validate_peft_config(peft_config)
+13 -5
View File
@@ -34,14 +34,22 @@ from lerobot.configs import PreTrainedConfig
from lerobot.configs.train import TrainPipelineConfig from lerobot.configs.train import TrainPipelineConfig
from lerobot.utils.device_utils import resolve_safetensors_device from lerobot.utils.device_utils import resolve_safetensors_device
from lerobot.utils.hub import HubMixin from lerobot.utils.hub import HubMixin
from lerobot.utils.import_utils import _peft_available, require_package
from .utils import log_model_loading_keys from .utils import log_model_loading_keys
T = TypeVar("T", bound="PreTrainedPolicy") if TYPE_CHECKING or _peft_available:
from peft import PEFT_TYPE_TO_CONFIG_MAPPING, PeftType, get_peft_model
else:
PEFT_TYPE_TO_CONFIG_MAPPING = None
PeftType = None
get_peft_model = None
if TYPE_CHECKING: if TYPE_CHECKING:
from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
T = TypeVar("T", bound="PreTrainedPolicy")
def _build_card_context( def _build_card_context(
cfg: TrainPipelineConfig | None, cfg: TrainPipelineConfig | None,
@@ -384,7 +392,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
peft_cli_overrides: Optional dict of CLI overrides (method_type, target_modules, r, etc.) peft_cli_overrides: Optional dict of CLI overrides (method_type, target_modules, r, etc.)
These are merged with policy defaults to build the final config. These are merged with policy defaults to build the final config.
""" """
from peft import get_peft_model require_package("peft", extra="peft")
# If user provided a complete config, use it directly (with overrides) # If user provided a complete config, use it directly (with overrides)
if peft_config is not None: if peft_config is not None:
@@ -455,7 +463,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
Returns: Returns:
Preprocessed dict with renamed keys and init_type mapped to method-specific key. Preprocessed dict with renamed keys and init_type mapped to method-specific key.
""" """
from peft import PeftType require_package("peft", extra="peft")
cli_overrides = cli_overrides.copy() cli_overrides = cli_overrides.copy()
@@ -480,7 +488,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
def _build_peft_config(self, cli_overrides: dict): def _build_peft_config(self, cli_overrides: dict):
"""Build a PEFT config from policy defaults and CLI overrides.""" """Build a PEFT config from policy defaults and CLI overrides."""
from peft import PEFT_TYPE_TO_CONFIG_MAPPING, PeftType require_package("peft", extra="peft")
# Determine PEFT method type (default to LORA) # Determine PEFT method type (default to LORA)
method_type_str = cli_overrides.get("method_type") or "lora" method_type_str = cli_overrides.get("method_type") or "lora"
@@ -507,7 +515,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
def _apply_peft_cli_overrides(self, peft_config, cli_overrides: dict): def _apply_peft_cli_overrides(self, peft_config, cli_overrides: dict):
"""Apply CLI overrides to an existing PEFT config.""" """Apply CLI overrides to an existing PEFT config."""
from peft import PEFT_TYPE_TO_CONFIG_MAPPING, PeftType require_package("peft", extra="peft")
# Get method type from existing config or CLI override # Get method type from existing config or CLI override
method_type_str = cli_overrides.get("method_type") method_type_str = cli_overrides.get("method_type")
+9 -1
View File
@@ -24,6 +24,7 @@ from __future__ import annotations
import logging import logging
from dataclasses import dataclass, field from dataclasses import dataclass, field
from threading import Event from threading import Event
from typing import TYPE_CHECKING
import torch import torch
@@ -47,6 +48,7 @@ from lerobot.processor.relative_action_processor import RelativeActionsProcessor
from lerobot.robots import make_robot_from_config from lerobot.robots import make_robot_from_config
from lerobot.teleoperators import Teleoperator, make_teleoperator_from_config from lerobot.teleoperators import Teleoperator, make_teleoperator_from_config
from lerobot.utils.feature_utils import combine_feature_dicts, hw_to_dataset_features from lerobot.utils.feature_utils import combine_feature_dicts, hw_to_dataset_features
from lerobot.utils.import_utils import _peft_available, require_package
from .configs import BaseStrategyConfig, DAggerStrategyConfig, RolloutConfig from .configs import BaseStrategyConfig, DAggerStrategyConfig, RolloutConfig
from .inference import ( from .inference import (
@@ -57,6 +59,12 @@ from .inference import (
) )
from .robot_wrapper import ThreadSafeRobot from .robot_wrapper import ThreadSafeRobot
if TYPE_CHECKING or _peft_available:
from peft import PeftConfig, PeftModel
else:
PeftConfig = None
PeftModel = None
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -171,7 +179,7 @@ def _load_pretrained_policy(policy_config: PreTrainedConfig) -> PreTrainedPolicy
revision=pretrained_revision, revision=pretrained_revision,
) )
from peft import PeftConfig, PeftModel require_package("peft", extra="peft")
peft_path = policy_config.pretrained_path peft_path = policy_config.pretrained_path
peft_config = PeftConfig.from_pretrained(peft_path, revision=pretrained_revision) peft_config = PeftConfig.from_pretrained(peft_path, revision=pretrained_revision)
+10 -7
View File
@@ -62,7 +62,7 @@ from dataclasses import asdict
from functools import partial from functools import partial
from pathlib import Path from pathlib import Path
from pprint import pformat from pprint import pformat
from typing import Any, TypedDict from typing import TYPE_CHECKING, Any, TypedDict
import einops import einops
import gymnasium as gym import gymnasium as gym
@@ -87,7 +87,7 @@ from lerobot.processor import PolicyProcessorPipeline
from lerobot.types import PolicyAction from lerobot.types import PolicyAction
from lerobot.utils.constants import ACTION, DONE, OBS_IMAGE, OBS_IMAGES, OBS_STR, REWARD from lerobot.utils.constants import ACTION, DONE, OBS_IMAGE, OBS_IMAGES, OBS_STR, REWARD
from lerobot.utils.device_utils import get_safe_torch_device from lerobot.utils.device_utils import get_safe_torch_device
from lerobot.utils.import_utils import register_third_party_plugins from lerobot.utils.import_utils import _peft_available, register_third_party_plugins, require_package
from lerobot.utils.io_utils import write_video from lerobot.utils.io_utils import write_video
from lerobot.utils.random_utils import set_seed from lerobot.utils.random_utils import set_seed
from lerobot.utils.utils import ( from lerobot.utils.utils import (
@@ -95,6 +95,11 @@ from lerobot.utils.utils import (
inside_slurm, inside_slurm,
) )
if TYPE_CHECKING or _peft_available:
from peft import PeftModel
else:
PeftModel = None
def _env_features_to_dataset_features(env_features: dict) -> dict: def _env_features_to_dataset_features(env_features: dict) -> dict:
"""Convert EnvConfig.features to the dict format expected by LeRobotDataset.create().""" """Convert EnvConfig.features to the dict format expected by LeRobotDataset.create()."""
@@ -444,13 +449,11 @@ def eval_policy(
exc = ValueError( exc = ValueError(
f"Policy of type 'PreTrainedPolicy' is expected, but type '{type(policy)}' was provided." f"Policy of type 'PreTrainedPolicy' is expected, but type '{type(policy)}' was provided."
) )
try: if not _peft_available:
from peft import PeftModel raise exc
require_package("peft", extra="peft")
if not isinstance(policy, PeftModel): if not isinstance(policy, PeftModel):
raise exc raise exc
except ImportError:
raise exc from None
start = time.time() start = time.time()
# Preserve the mode for direct callers. eval_policy_all scopes the mode # Preserve the mode for direct callers. eval_policy_all scopes the mode
+7 -4
View File
@@ -57,7 +57,7 @@ from lerobot.optim.factory import make_optimizer_and_scheduler
from lerobot.policies import PreTrainedPolicy, make_policy, make_pre_post_processors from lerobot.policies import PreTrainedPolicy, make_policy, make_pre_post_processors
from lerobot.rewards import make_reward_pre_post_processors from lerobot.rewards import make_reward_pre_post_processors
from lerobot.utils.collate import lerobot_collate_fn from lerobot.utils.collate import lerobot_collate_fn
from lerobot.utils.import_utils import register_third_party_plugins from lerobot.utils.import_utils import _peft_available, register_third_party_plugins, require_package
from lerobot.utils.logging_utils import AverageMeter, MetricsTracker from lerobot.utils.logging_utils import AverageMeter, MetricsTracker
from lerobot.utils.random_utils import set_seed from lerobot.utils.random_utils import set_seed
from lerobot.utils.utils import ( from lerobot.utils.utils import (
@@ -68,6 +68,11 @@ from lerobot.utils.utils import (
inside_slurm, inside_slurm,
) )
if TYPE_CHECKING or _peft_available:
from peft import PeftModel
else:
PeftModel = None
from .lerobot_eval import eval_policy_all from .lerobot_eval import eval_policy_all
@@ -207,8 +212,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
if cfg.job.is_remote: if cfg.job.is_remote:
return submit_to_hf(cfg) return submit_to_hf(cfg)
from lerobot.utils.import_utils import require_package
require_package("accelerate", extra="training") require_package("accelerate", extra="training")
from accelerate import Accelerator from accelerate import Accelerator
from accelerate.utils import DistributedDataParallelKwargs, DistributedType from accelerate.utils import DistributedDataParallelKwargs, DistributedType
@@ -312,7 +315,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
if cfg.peft is not None: if cfg.peft is not None:
if cfg.is_reward_model_training: if cfg.is_reward_model_training:
raise ValueError("PEFT is only supported for policy training. ") raise ValueError("PEFT is only supported for policy training. ")
from peft import PeftModel require_package("peft", extra="peft")
if isinstance(policy, PeftModel): if isinstance(policy, PeftModel):
logging.info("PEFT adapter already loaded from checkpoint, skipping wrap_with_peft.") logging.info("PEFT adapter already loaded from checkpoint, skipping wrap_with_peft.")
+14 -7
View File
@@ -185,18 +185,25 @@ def test_load_pretrained_peft_policy_keeps_adapter_and_base_revisions_separate(m
peft_config_from_pretrained = MagicMock(return_value=peft_config) peft_config_from_pretrained = MagicMock(return_value=peft_config)
adapted_policy = MagicMock() adapted_policy = MagicMock()
peft_model_from_pretrained = MagicMock(return_value=adapted_policy) peft_model_from_pretrained = MagicMock(return_value=adapted_policy)
monkeypatch.setitem( require_package = MagicMock()
sys.modules, monkeypatch.setattr(rollout_context, "require_package", require_package)
"peft", monkeypatch.setattr(
SimpleNamespace( rollout_context,
PeftConfig=SimpleNamespace(from_pretrained=peft_config_from_pretrained), "PeftConfig",
PeftModel=SimpleNamespace(from_pretrained=peft_model_from_pretrained), SimpleNamespace(from_pretrained=peft_config_from_pretrained),
), raising=False,
)
monkeypatch.setattr(
rollout_context,
"PeftModel",
SimpleNamespace(from_pretrained=peft_model_from_pretrained),
raising=False,
) )
policy = rollout_context._load_pretrained_policy(policy_config) policy = rollout_context._load_pretrained_policy(policy_config)
assert policy is adapted_policy assert policy is adapted_policy
require_package.assert_called_once_with("peft", extra="peft")
peft_config_from_pretrained.assert_called_once_with("user/adapter", revision="adapter-sha") peft_config_from_pretrained.assert_called_once_with("user/adapter", revision="adapter-sha")
policy_class.from_pretrained.assert_called_once_with( policy_class.from_pretrained.assert_called_once_with(
pretrained_name_or_path="user/base-policy", pretrained_name_or_path="user/base-policy",