guard molmoact2 transformers imports

This commit is contained in:
hq-fang
2026-05-19 21:58:44 +00:00
parent 140cf2a420
commit e13c6a6110
@@ -22,7 +22,7 @@ import types
from collections import defaultdict, deque from collections import defaultdict, deque
from contextlib import nullcontext, suppress from contextlib import nullcontext, suppress
from pathlib import Path from pathlib import Path
from typing import Any from typing import TYPE_CHECKING, Any
import numpy as np import numpy as np
import torch import torch
@@ -34,11 +34,17 @@ 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 require_package from lerobot.utils.import_utils import _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 _transformers_available:
from transformers import AutoModelForImageTextToText, AutoProcessor
else:
AutoModelForImageTextToText = None
AutoProcessor = None
_MODEL_INPUT_KEYS = { _MODEL_INPUT_KEYS = {
"input_ids", "input_ids",
"pixel_values", "pixel_values",
@@ -714,7 +720,6 @@ class MolmoAct2Policy(PreTrainedPolicy):
def _load_hf_model(self) -> None: def _load_hf_model(self) -> None:
require_package("transformers", extra="molmoact2") require_package("transformers", extra="molmoact2")
from transformers import AutoModelForImageTextToText
checkpoint_location = _resolve_checkpoint_location( checkpoint_location = _resolve_checkpoint_location(
self.config.checkpoint_path, self.config.checkpoint_path,
@@ -1040,7 +1045,6 @@ class MolmoAct2Policy(PreTrainedPolicy):
def _load_discrete_action_tokenizer(self) -> Any: def _load_discrete_action_tokenizer(self) -> Any:
if self.action_tokenizer is None: if self.action_tokenizer is None:
require_package("transformers", extra="molmoact2") require_package("transformers", extra="molmoact2")
from transformers import AutoProcessor
self.action_tokenizer = AutoProcessor.from_pretrained( self.action_tokenizer = AutoProcessor.from_pretrained(
self.config.discrete_action_tokenizer, self.config.discrete_action_tokenizer,