admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-24 18:26:11 +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 5c10d0011b
commit e8ea7b9fa0
1 changed files with 2 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lerobot/policies/vla_jepa/action_head.py
+2 -1
View File
@@ -9,7 +9,7 @@ import torch.nn.functional as F # noqa: N812
from torch import nn from torch import nn
from torch.distributions import Beta 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: if TYPE_CHECKING or _diffusers_available:
from diffusers import ConfigMixin, ModelMixin from diffusers import ConfigMixin, ModelMixin
@@ -70,6 +70,7 @@ class ActionEncoder(nn.Module):
class TimestepEncoder(nn.Module): class TimestepEncoder(nn.Module):
def __init__(self, embedding_dim: int): def __init__(self, embedding_dim: int):
super().__init__() 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.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) self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)