ydy0615/lerobot-clone
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-05-30 10:21:24 +00:00
Code Issues Packages Projects Releases Wiki Activity

adding guard for diffusers

Browse Source
This commit is contained in:
Maximellerbach
2026-05-21 15:41:26 +02:00
parent c01a00a972
commit dd13eda002
1 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
src/lerobot/policies/vla_jepa/action_head.py
View File

@@ -9,7 +9,7 @@ import torch.nn.functional as F # noqa: N812
from torch import nn
from torch.distributions import Beta
from lerobot.utils.import_utils import _diffusers_available
from lerobot.utils.import_utils import _diffusers_available, require_package
if TYPE_CHECKING or _diffusers_available:
from diffusers import ConfigMixin, ModelMixin
@@ -70,6 +70,7 @@ class ActionEncoder(nn.Module):
class TimestepEncoder(nn.Module):
def __init__(self, embedding_dim: int):
super().__init__()
require_package("diffusers", extra="vla_jepa")
self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=1)
self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)