From 733f9768b587d0f4195796bdb9b66b336056145d Mon Sep 17 00:00:00 2001 From: hq-fang <71356829+hq-fang@users.noreply.github.com> Date: Thu, 21 May 2026 20:54:13 +0000 Subject: [PATCH] lazy import molmoact2 scipy --- .../molmoact2/hf_model/action_tokenizer.py | 8 ++++++-- .../policies/molmoact2/modeling_molmoact2.py | 17 ++++++++++++----- .../policies/molmoact2/processor_molmoact2.py | 12 ++++++++---- 3 files changed, 26 insertions(+), 11 deletions(-) diff --git a/src/lerobot/policies/molmoact2/hf_model/action_tokenizer.py b/src/lerobot/policies/molmoact2/hf_model/action_tokenizer.py index 77b8fe9a5..fad5c1fc1 100644 --- a/src/lerobot/policies/molmoact2/hf_model/action_tokenizer.py +++ b/src/lerobot/policies/molmoact2/hf_model/action_tokenizer.py @@ -22,8 +22,6 @@ from pathlib import Path from typing import ClassVar import numpy as np -from scipy.fft import dct -from scipy.fft import idct from tokenizers import ByteLevelBPETokenizer from tokenizers.trainers import BpeTrainer from huggingface_hub import snapshot_download @@ -86,6 +84,8 @@ class UniversalActionProcessor(ProcessorMixin): self.bpe_tokenizer = self.tokenizer def __call__(self, action_chunk: np.array) -> np.array: + from scipy.fft import dct + assert action_chunk.ndim <= 3, "Only 3 dimensions supported: [batch, timesteps, action_dim]" if action_chunk.ndim == 2: action_chunk = action_chunk[None, ...] @@ -109,6 +109,8 @@ class UniversalActionProcessor(ProcessorMixin): time_horizon: int | None = None, action_dim: int | None = None, ) -> np.array: + from scipy.fft import idct + self.time_horizon = time_horizon or self.time_horizon or self.called_time_horizon self.action_dim = action_dim or self.action_dim or self.called_action_dim @@ -150,6 +152,8 @@ class UniversalActionProcessor(ProcessorMixin): time_horizon: int | None = None, action_dim: int | None = None, ) -> "UniversalActionProcessor": + from scipy.fft import dct + # Run DCT over all inputs dct_tokens = [dct(a, axis=0, norm="ortho").flatten() for a in action_data] diff --git a/src/lerobot/policies/molmoact2/modeling_molmoact2.py b/src/lerobot/policies/molmoact2/modeling_molmoact2.py index 271e268b3..288837268 100644 --- a/src/lerobot/policies/molmoact2/modeling_molmoact2.py +++ b/src/lerobot/policies/molmoact2/modeling_molmoact2.py @@ -29,24 +29,30 @@ import torch.nn.functional as F # noqa: N812 from safetensors.torch import load_file as load_safetensors_file from torch import Tensor from torch.distributions import Beta -from transformers.utils import SAFE_WEIGHTS_INDEX_NAME, SAFE_WEIGHTS_NAME from lerobot.policies.pretrained import PreTrainedPolicy from lerobot.utils.constants import ACTION -from lerobot.utils.import_utils import _transformers_available, require_package +from lerobot.utils.import_utils import _scipy_available, _transformers_available, require_package from ..rtc.modeling_rtc import RTCProcessor from .configuration_molmoact2 import MolmoAct2Config, _hf_token, _resolve_checkpoint_location if TYPE_CHECKING or _transformers_available: - from .hf_model.action_tokenizer import UniversalActionProcessor + from transformers.utils import SAFE_WEIGHTS_INDEX_NAME, SAFE_WEIGHTS_NAME + from .hf_model.configuration_molmoact2 import MolmoAct2Config as HFMolmoAct2Config from .hf_model.modeling_molmoact2 import MolmoAct2ForConditionalGeneration else: - UniversalActionProcessor = None + SAFE_WEIGHTS_INDEX_NAME = "model.safetensors.index.json" + SAFE_WEIGHTS_NAME = "model.safetensors" HFMolmoAct2Config = None MolmoAct2ForConditionalGeneration = None +if TYPE_CHECKING or (_transformers_available and _scipy_available): + from .hf_model.action_tokenizer import UniversalActionProcessor +else: + UniversalActionProcessor = None + _MODEL_INPUT_KEYS = { "input_ids", "pixel_values", @@ -483,9 +489,10 @@ class MolmoAct2Policy(PreTrainedPolicy): def _load_discrete_action_tokenizer(self) -> Any: if self.action_tokenizer is None: require_package("transformers", extra="molmoact2") + require_package("scipy", extra="molmoact2") if UniversalActionProcessor is None: - raise RuntimeError("transformers is required to load MolmoAct2 action tokenizer.") + raise RuntimeError("transformers and scipy are required to load MolmoAct2 action tokenizer.") self.action_tokenizer = UniversalActionProcessor.from_pretrained_local( self.config.discrete_action_tokenizer, ) diff --git a/src/lerobot/policies/molmoact2/processor_molmoact2.py b/src/lerobot/policies/molmoact2/processor_molmoact2.py index ee9bd91fd..b393dd440 100644 --- a/src/lerobot/policies/molmoact2/processor_molmoact2.py +++ b/src/lerobot/policies/molmoact2/processor_molmoact2.py @@ -52,24 +52,27 @@ from lerobot.utils.constants import ( POLICY_POSTPROCESSOR_DEFAULT_NAME, POLICY_PREPROCESSOR_DEFAULT_NAME, ) -from lerobot.utils.import_utils import _transformers_available, require_package +from lerobot.utils.import_utils import _scipy_available, _transformers_available, require_package from .configuration_molmoact2 import MolmoAct2Config, infer_molmoact2_max_sequence_length if TYPE_CHECKING or _transformers_available: from transformers import Qwen2Tokenizer - from .hf_model.action_tokenizer import UniversalActionProcessor from .hf_model.image_processing_molmoact2 import MolmoAct2ImageProcessor from .hf_model.processing_molmoact2 import MolmoAct2Processor from .hf_model.video_processing_molmoact2 import MolmoAct2VideoProcessor else: Qwen2Tokenizer = None - UniversalActionProcessor = None MolmoAct2ImageProcessor = None MolmoAct2Processor = None MolmoAct2VideoProcessor = None +if TYPE_CHECKING or (_transformers_available and _scipy_available): + from .hf_model.action_tokenizer import UniversalActionProcessor +else: + UniversalActionProcessor = None + ACTION_OUTPUT_TOKEN = "" # nosec B105 ACTION_START_TOKEN = "" # nosec B105 ACTION_END_TOKEN = "" # nosec B105 @@ -586,8 +589,9 @@ class MolmoAct2PackInputsProcessorStep(ProcessorStep): self.processor = _load_local_molmoact2_processor(checkpoint_location) self.action_processor = None if self.action_mode in {"discrete", "both"}: + require_package("scipy", extra="molmoact2") if UniversalActionProcessor is None: - raise RuntimeError("transformers is required to load MolmoAct2 action tokenizer.") + raise RuntimeError("transformers and scipy are required to load MolmoAct2 action tokenizer.") self.action_processor = UniversalActionProcessor.from_pretrained_local( self.discrete_action_tokenizer, )