braindecode 1.3.0.dev175415232__py3-none-any.whl → 1.3.0.dev175955015__py3-none-any.whl

braindecode/eegneuralnet.py

@@ -189,6 +189,8 @@ class _EEGNeuralNet(NeuralNet, abc.ABC):
                 "Skipping setting signal-related parameters from data."
             )
             return
+        if classes is None:
+            classes = getattr(self, "classes", None)
         # get kwargs from signal:
         signal_kwargs = dict()
         # Using shape to work both with torch.tensor and numpy.array:

braindecode/models/attentionbasenet.py

@@ -381,6 +381,8 @@ class AttentionBaseNet(EEGModuleMixin, nn.Module):
         for k, pl, ps in zip(kernel_lengths, pool_lengths, pool_strides):
             out = math.floor(out + 2 * (k // 2) - k + 1)
             out = math.floor((out - pl) / ps + 1)
+            # Ensure output is at least 1 to avoid zero-sized tensors
+            out = max(1, out)
             seq_lengths.append(int(out))
         return seq_lengths
 

braindecode/models/base.py

@@ -5,15 +5,35 @@
 
 from __future__ import annotations
 
+import json
 import warnings
 from collections import OrderedDict
-from typing import Dict, Iterable, Optional
+from pathlib import Path
+from typing import Dict, Iterable, Optional, Type, Union
 
 import numpy as np
 import torch
 from docstring_inheritance import NumpyDocstringInheritanceInitMeta
+from mne.utils import _soft_import
 from torchinfo import ModelStatistics, summary
 
+from braindecode.version import __version__
+
+huggingface_hub = _soft_import(
+    "huggingface_hub", "Hugging Face Hub integration", strict=False
+)
+
+HAS_HF_HUB = huggingface_hub is not False
+
+
+class _BaseHubMixin:
+    pass
+
+
+# Define base class for hub mixin
+if HAS_HF_HUB:
+    _BaseHubMixin: Type = huggingface_hub.PyTorchModelHubMixin  # type: ignore
+
 
 def deprecated_args(obj, *old_new_args):
     out_args = []
@@ -32,10 +52,14 @@ def deprecated_args(obj, *old_new_args):
     return out_args
 
 
-class EEGModuleMixin(metaclass=NumpyDocstringInheritanceInitMeta):
+class EEGModuleMixin(_BaseHubMixin, metaclass=NumpyDocstringInheritanceInitMeta):
     """
     Mixin class for all EEG models in braindecode.
 
+    This class integrates with Hugging Face Hub when the ``huggingface_hub`` package
+    is installed, enabling models to be pushed to and loaded from the Hub using
+    :func:`push_to_hub()` and :func:`from_pretrained()` methods.
+
     Parameters
     ----------
     n_outputs : int
@@ -62,8 +86,87 @@ class EEGModuleMixin(metaclass=NumpyDocstringInheritanceInitMeta):
     -----
     If some input signal-related parameters are not specified,
     there will be an attempt to infer them from the other parameters.
+
+    Hugging Face Hub Integration
+    -----------------------------
+    When the optional ``huggingface_hub`` package is installed, all models
+    automatically gain the ability to be pushed to and loaded from the
+    Hugging Face Hub. Install with::
+
+        pip install braindecode[hug]
+
+    **Pushing a model to the Hub:**
+
+    .. code-block:: python
+
+        from braindecode.models import EEGNetv4
+
+        # Train your model
+        model = EEGNetv4(n_chans=22, n_outputs=4, n_times=1000)
+        # ... training code ...
+
+        # Push to the Hub
+        model.push_to_hub(
+            repo_id="username/my-eegnet-model", commit_message="Initial model upload"
+        )
+
+    **Loading a model from the Hub:**
+
+    .. code-block:: python
+
+        from braindecode.models import EEGNetv4
+
+        # Load pretrained model
+        model = EEGNetv4.from_pretrained("username/my-eegnet-model")
+
+    The integration automatically handles EEG-specific parameters (n_chans,
+    n_times, sfreq, chs_info, etc.) by saving them in a config file alongside
+    the model weights. This ensures that loaded models are correctly configured
+    for their original data specifications.
+
+    .. important::
+        Currently, only EEG-specific parameters (n_outputs, n_chans, n_times,
+        input_window_seconds, sfreq, chs_info) are saved to the Hub. Model-specific
+        parameters (e.g., dropout rates, activation functions, number of filters)
+        are not preserved and will use their default values when loading from the Hub.
+
+        To use non-default model parameters, specify them explicitly when calling
+        :func:`from_pretrained()`::
+
+            model = EEGNet.from_pretrained("user/model", dropout=0.3, activation='relu')
+
+        Full parameter serialization will be addressed in a future update.
     """
 
+    def __init_subclass__(cls, **kwargs):
+        if not HAS_HF_HUB:
+            super().__init_subclass__(**kwargs)
+            return
+
+        base_tags = ["braindecode", cls.__name__]
+        user_tags = kwargs.pop("tags", None)
+        tags = list(user_tags) if user_tags is not None else []
+        for tag in base_tags:
+            if tag not in tags:
+                tags.append(tag)
+
+        docs_url = kwargs.pop(
+            "docs_url",
+            f"https://braindecode.org/stable/generated/braindecode.models.{cls.__name__}.html",
+        )
+        repo_url = kwargs.pop("repo_url", "https://braindecode.org")
+        library_name = kwargs.pop("library_name", "braindecode")
+        license = kwargs.pop("license", "bsd-3-clause")
+        # TODO: model_card_template can be added in the future for custom model cards
+        super().__init_subclass__(
+            tags=tags,
+            docs_url=docs_url,
+            repo_url=repo_url,
+            library_name=library_name,
+            license=license,
+            **kwargs,
+        )
+
     def __init__(
         self,
         n_outputs: Optional[int] = None,  # type: ignore[assignment]
@@ -73,6 +176,16 @@ class EEGModuleMixin(metaclass=NumpyDocstringInheritanceInitMeta):
         input_window_seconds: Optional[float] = None,  # type: ignore[assignment]
         sfreq: Optional[float] = None,  # type: ignore[assignment]
     ):
+        # Deserialize chs_info if it comes as a list of dicts (from Hub)
+        if chs_info is not None and isinstance(chs_info, list):
+            if len(chs_info) > 0 and isinstance(chs_info[0], dict):
+                # Check if it needs deserialization (has 'loc' as list)
+                if "loc" in chs_info[0] and isinstance(chs_info[0]["loc"], list):
+                    chs_info = self._deserialize_chs_info(chs_info)
+                    warnings.warn(
+                        "Modifying chs_info argument using the _deserialize_chs_info() method"
+                    )
+
         if n_chans is not None and chs_info is not None and len(chs_info) != n_chans:
             raise ValueError(f"{n_chans=} different from {chs_info=} length")
         if (
@@ -294,3 +407,168 @@ class EEGModuleMixin(metaclass=NumpyDocstringInheritanceInitMeta):
 
     def __str__(self) -> str:
         return str(self.get_torchinfo_statistics())
+
+    @staticmethod
+    def _serialize_chs_info(chs_info):
+        """
+        Serialize MNE channel info to JSON-compatible format.
+
+        Parameters
+        ----------
+        chs_info : list of dict or None
+            Channel information from MNE Info object.
+
+        Returns
+        -------
+        list of dict or None
+            Serialized channel information that can be saved to JSON.
+        """
+        if chs_info is None:
+            return None
+
+        serialized = []
+        for ch in chs_info:
+            # Extract serializable fields from MNE channel info
+            ch_dict = {
+                "ch_name": ch.get("ch_name", ""),
+            }
+
+            # Handle kind field - can be either string or integer
+            kind_val = ch.get("kind")
+            if kind_val is not None:
+                ch_dict["kind"] = (
+                    kind_val if isinstance(kind_val, str) else int(kind_val)
+                )
+
+            # Add numeric fields with safe conversion
+            coil_type = ch.get("coil_type")
+            if coil_type is not None:
+                ch_dict["coil_type"] = int(coil_type)
+
+            unit = ch.get("unit")
+            if unit is not None:
+                ch_dict["unit"] = int(unit)
+
+            cal = ch.get("cal")
+            if cal is not None:
+                ch_dict["cal"] = float(cal)
+
+            range_val = ch.get("range")
+            if range_val is not None:
+                ch_dict["range"] = float(range_val)
+
+            # Serialize location array if present
+            if "loc" in ch and ch["loc"] is not None:
+                ch_dict["loc"] = (
+                    ch["loc"].tolist()
+                    if hasattr(ch["loc"], "tolist")
+                    else list(ch["loc"])
+                )
+            serialized.append(ch_dict)
+
+        return serialized
+
+    @staticmethod
+    def _deserialize_chs_info(chs_info_dict):
+        """
+        Deserialize channel info from JSON-compatible format to MNE-like structure.
+
+        Parameters
+        ----------
+        chs_info_dict : list of dict or None
+            Serialized channel information.
+
+        Returns
+        -------
+        list of dict or None
+            Deserialized channel information compatible with MNE.
+        """
+        if chs_info_dict is None:
+            return None
+
+        deserialized = []
+        for ch_dict in chs_info_dict:
+            ch = ch_dict.copy()
+            # Convert location back to numpy array if present
+            if "loc" in ch and ch["loc"] is not None:
+                ch["loc"] = np.array(ch["loc"])
+            deserialized.append(ch)
+
+        return deserialized
+
+    def _save_pretrained(self, save_directory):
+        """
+        Save model configuration and weights to the Hub.
+
+        This method is called by PyTorchModelHubMixin.push_to_hub() to save
+        model-specific configuration alongside the model weights.
+
+        Parameters
+        ----------
+        save_directory : str or Path
+            Directory where the configuration should be saved.
+        """
+        if not HAS_HF_HUB:
+            return
+
+        save_directory = Path(save_directory)
+
+        # Collect EEG-specific configuration
+        config = {
+            "n_outputs": self._n_outputs,
+            "n_chans": self._n_chans,
+            "n_times": self._n_times,
+            "input_window_seconds": self._input_window_seconds,
+            "sfreq": self._sfreq,
+            "chs_info": self._serialize_chs_info(self._chs_info),
+            "braindecode_version": __version__,
+        }
+
+        # Save to config.json
+        config_path = save_directory / "config.json"
+        with open(config_path, "w") as f:
+            json.dump(config, f, indent=2)
+
+        # Save model weights with standard Hub filename
+        weights_path = save_directory / "pytorch_model.bin"
+        torch.save(self.state_dict(), weights_path)
+
+        # Also save in safetensors format using parent's implementation
+        try:
+            super()._save_pretrained(save_directory)
+        except (ImportError, RuntimeError) as e:
+            # Fallback to pytorch_model.bin if safetensors saving fails
+            warnings.warn(
+                f"Could not save model in safetensors format: {e}. "
+                "Model weights saved in pytorch_model.bin instead.",
+                stacklevel=2,
+            )
+
+    if HAS_HF_HUB:
+
+        @classmethod
+        def _from_pretrained(
+            cls,
+            *,
+            model_id: str,
+            revision: Optional[str],
+            cache_dir: Optional[Union[str, Path]],
+            force_download: bool,
+            local_files_only: bool,
+            token: Union[str, bool, None],
+            map_location: str = "cpu",
+            strict: bool = False,
+            **model_kwargs,
+        ):
+            model_kwargs.pop("braindecode_version", None)
+            return super()._from_pretrained(  # type: ignore
+                model_id=model_id,
+                revision=revision,
+                cache_dir=cache_dir,
+                force_download=force_download,
+                local_files_only=local_files_only,
+                token=token,
+                map_location=map_location,
+                strict=strict,
+                **model_kwargs,
+            )

braindecode/models/labram.py

@@ -61,13 +61,24 @@ class Labram(EEGModuleMixin, nn.Module):
 
     .. versionadded:: 0.9
 
+
+    Examples on how to load pre-trained weights:
+    --------------------------------------------
+    >>> import torch
+    >>> from braindecode.models import Labram
+    >>> model = Labram(n_times=1600, n_chans=64, n_outputs=4)
+    >>> url = 'https://huggingface.co/braindecode/Labram-Braindecode/blob/main/braindecode_labram_base.pt'
+    >>> state = torch.hub.load_state_dict_from_url(url, progress=True)
+    >>> model.load_state_dict(state)
+
+
     Parameters
     ----------
     patch_size : int
         The size of the patch to be used in the patch embedding.
     emb_size : int
         The dimension of the embedding.
-    in_channels : int
+    in_conv_channels : int
         The number of convolutional input channels.
     out_channels : int
         The number of convolutional output channels.
@@ -79,8 +90,10 @@ class Labram(EEGModuleMixin, nn.Module):
         The expansion ratio of the mlp layer
     qkv_bias :  bool (default=False)
         If True, add a learnable bias to the query, key, and value tensors.
-    qk_norm : Pytorch Normalize layer (default=None)
-        If not None, apply LayerNorm to the query and key tensors
+    qk_norm : Pytorch Normalize layer (default=nn.LayerNorm)
+        If not None, apply LayerNorm to the query and key tensors.
+        Default is nn.LayerNorm for better weight transfer from original LaBraM.
+        Set to None to disable Q,K normalization.
     qk_scale : float (default=None)
         If not None, use this value as the scale factor. If None,
         use head_dim**-0.5, where head_dim = dim // num_heads.
@@ -92,9 +105,10 @@ class Labram(EEGModuleMixin, nn.Module):
         Dropout rate for the attention weights used on DropPath.
     norm_layer : Pytorch Normalize layer (default=nn.LayerNorm)
         The normalization layer to be used.
-    init_values : float (default=None)
+    init_values : float (default=0.1)
         If not None, use this value to initialize the gamma_1 and gamma_2
-        parameters.
+        parameters for residual scaling. Default is 0.1 for better weight
+        transfer from original LaBraM. Set to None to disable.
     use_abs_pos_emb : bool (default=True)
         If True, use absolute position embedding.
     use_mean_pooling : bool (default=True)
@@ -135,19 +149,19 @@ class Labram(EEGModuleMixin, nn.Module):
         input_window_seconds=None,
         patch_size=200,
         emb_size=200,
-        in_channels=1,
+        in_conv_channels=1,
         out_channels=8,
         n_layers=12,
         att_num_heads=10,
         mlp_ratio=4.0,
         qkv_bias=False,
-        qk_norm=None,
+        qk_norm=nn.LayerNorm,
         qk_scale=None,
         drop_prob=0.0,
         attn_drop_prob=0.0,
         drop_path_prob=0.0,
         norm_layer=nn.LayerNorm,
-        init_values=None,
+        init_values=0.1,
         use_abs_pos_emb=True,
         use_mean_pooling=True,
         init_scale=0.001,
@@ -183,15 +197,15 @@ class Labram(EEGModuleMixin, nn.Module):
             self.patch_size = patch_size
         self.n_path = self.n_times // self.patch_size
 
-        if neural_tokenizer and in_channels != 1:
+        if neural_tokenizer and in_conv_channels != 1:
             warn(
                 "The model is in Neural Tokenizer mode, but the variable "
-                + "`in_channels` is different from the default values."
-                + "`in_channels` is only needed for the Neural Decoder mode."
-                + "in_channels is not used in the Neural Tokenizer mode.",
+                + "`in_conv_channels` is different from the default values."
+                + "`in_conv_channels` is only needed for the Neural Decoder mode."
+                + "in_conv_channels is not used in the Neural Tokenizer mode.",
                 UserWarning,
             )
-            in_channels = 1
+            in_conv_channels = 1
             # If you can use the model in Neural Tokenizer mode,
         # temporal conv layer will be use over the patched dataset
         if neural_tokenizer:
@@ -228,7 +242,7 @@ class Labram(EEGModuleMixin, nn.Module):
                 _PatchEmbed(
                     n_times=self.n_times,
                     patch_size=patch_size,
-                    in_channels=in_channels,
+                    in_channels=in_conv_channels,
                     emb_dim=self.emb_size,
                 ),
             )
@@ -373,8 +387,7 @@ class Labram(EEGModuleMixin, nn.Module):
         Parameters
         ----------
         x : torch.Tensor
-            The input data with shape (batch, n_chans, n_patches, patch size),
-            if neural decoder or (batch, n_chans, n_times), if neural tokenizer.
+            The input data with shape (batch, n_chans, n_times).
         input_chans : int
             The number of input channels.
         return_patch_tokens : bool
@@ -387,37 +400,72 @@ class Labram(EEGModuleMixin, nn.Module):
         x : torch.Tensor
             The output of the model.
         """
+        batch_size = x.shape[0]
+
         if self.neural_tokenizer:
-            batch_size, nch, n_patch, temporal = self.patch_embed.segment_patch(x).shape
+            # For neural tokenizer: input is (batch, n_chans, n_times)
+            # patch_embed returns (batch, n_chans, emb_dim)
+            x = self.patch_embed(x)
+            # x shape: (batch, n_chans, emb_dim)
+            n_patch = self.n_chans
+            temporal = self.emb_size
         else:
-            batch_size, nch, n_patch = self.patch_embed(x).shape
-        x = self.patch_embed(x)
+            # For neural decoder: input is (batch, n_chans, n_times)
+            # patch_embed returns (batch, n_patchs, emb_dim)
+            x = self.patch_embed(x)
+            # x shape: (batch, n_patchs, emb_dim)
+            batch_size, n_patch, temporal = x.shape
+
         # add the [CLS] token to the embedded patch tokens
         cls_tokens = self.cls_token.expand(batch_size, -1, -1)
 
+        # Concatenate cls token with patch/channel embeddings
         x = torch.cat((cls_tokens, x), dim=1)
 
         # Positional Embedding
-        if input_chans is not None:
-            pos_embed_used = self.position_embedding[:, input_chans]
-        else:
-            pos_embed_used = self.position_embedding
-
         if self.position_embedding is not None:
-            pos_embed = self._adj_position_embedding(
-                pos_embed_used=pos_embed_used, batch_size=batch_size
-            )
+            if self.neural_tokenizer:
+                # In tokenizer mode, use channel-based position embedding
+                if input_chans is not None:
+                    pos_embed_used = self.position_embedding[:, input_chans]
+                else:
+                    pos_embed_used = self.position_embedding
+
+                pos_embed = self._adj_position_embedding(
+                    pos_embed_used=pos_embed_used, batch_size=batch_size
+                )
+            else:
+                # In decoder mode, we have different number of patches
+                # Adapt position embedding for n_patch patches
+                # Use the first n_patch+1 positions from position_embedding
+                n_pos = min(self.position_embedding.shape[1], n_patch + 1)
+                pos_embed_used = self.position_embedding[:, :n_pos, :]
+                pos_embed = pos_embed_used.expand(batch_size, -1, -1)
+
             x += pos_embed
 
         # The time embedding is added across the channels after the [CLS] token
         if self.neural_tokenizer:
             num_ch = self.n_chans
+            time_embed = self._adj_temporal_embedding(
+                num_ch=num_ch, batch_size=batch_size, dim_embed=temporal
+            )
+            x[:, 1:, :] += time_embed
         else:
-            num_ch = n_patch
-        time_embed = self._adj_temporal_embedding(
-            num_ch=num_ch, batch_size=batch_size, dim_embed=temporal
-        )
-        x[:, 1:, :] += time_embed
+            # In decoder mode, we have n_patch patches and don't need to expand
+            # Just broadcast the temporal embedding
+            if temporal is None:
+                temporal = self.emb_size
+
+            # Get temporal embeddings for n_patch patches
+            n_time_tokens = min(n_patch, self.temporal_embedding.shape[1] - 1)
+            time_embed = self.temporal_embedding[
+                :, 1 : n_time_tokens + 1, :
+            ]  # (1, n_patch, emb_dim)
+            time_embed = time_embed.expand(
+                batch_size, -1, -1
+            )  # (batch, n_patch, emb_dim)
+            x[:, 1:, :] += time_embed
 
         x = self.pos_drop(x)
 
@@ -428,10 +476,10 @@ class Labram(EEGModuleMixin, nn.Module):
         if self.fc_norm is not None:
             if return_all_tokens:
                 return self.fc_norm(x)
-            temporal = x[:, 1:, :]
+            tokens = x[:, 1:, :]
             if return_patch_tokens:
-                return self.fc_norm(temporal)
-            return self.fc_norm(temporal.mean(1))
+                return self.fc_norm(tokens)
+            return self.fc_norm(tokens.mean(1))
         else:
             if return_all_tokens:
                 return x
@@ -505,14 +553,16 @@ class Labram(EEGModuleMixin, nn.Module):
     def _adj_temporal_embedding(self, num_ch, batch_size, dim_embed=None):
         """
         Adjust the dimensions of the time embedding to match the
-        number of channels.
+        number of channels or patches.
 
         Parameters
         ----------
         num_ch : int
-            The number of channels or number of code books vectors.
+            The number of channels or number of patches.
         batch_size : int
             Batch size of the input data.
+        dim_embed : int
+            The embedding dimension (temporal feature dimension).
 
         Returns
         -------
@@ -523,17 +573,24 @@ class Labram(EEGModuleMixin, nn.Module):
         if dim_embed is None:
             cut_dimension = self.patch_size
         else:
-            cut_dimension = dim_embed
-        # first step will be match the time_embed to the number of channels
-        temporal_embedding = self.temporal_embedding[:, 1:cut_dimension, :]
+            cut_dimension = min(dim_embed, self.temporal_embedding.shape[1] - 1)
+
+        # Get the temporal embedding: (1, temporal_embedding_dim, emb_size)
+        # Slice to cut_dimension: (1, cut_dimension, emb_size)
+        temporal_embedding = self.temporal_embedding[:, 1 : cut_dimension + 1, :]
+
         # Add a new dimension to the time embedding
-        # e.g. (batch, 62, 200) -> (batch, 1, 62, 200)
+        # e.g. (1, 5, 200) -> (1, 1, 5, 200)
         temporal_embedding = temporal_embedding.unsqueeze(1)
-        # Expand the time embedding to match the number of channels
-        # or number of patches from
+
+        # Expand the time embedding to match the number of channels or patches
+        # (1, 1, cut_dimension, 200) -> (batch_size, num_ch, cut_dimension, 200)
         temporal_embedding = temporal_embedding.expand(batch_size, num_ch, -1, -1)
+
         # Flatten the intermediate dimensions
+        # (batch_size, num_ch, cut_dimension, 200) -> (batch_size, num_ch * cut_dimension, 200)
         temporal_embedding = temporal_embedding.flatten(1, 2)
+
         return temporal_embedding
 
     def _adj_position_embedding(self, pos_embed_used, batch_size):
@@ -679,25 +736,27 @@ class _SegmentPatch(nn.Module):
 
 
 class _PatchEmbed(nn.Module):
-    """EEG to Patch Embedding.
+    """EEG to Patch Embedding for Neural Decoder mode.
 
     This code is used when we want to apply the patch embedding
-    after the codebook layer.
+    after the codebook layer (Neural Decoder mode).
+
+    The input is expected to be in the format (Batch, n_channels, n_times),
+    but the original LaBraM expects pre-patched data (Batch, n_channels, n_patches, patch_size).
+    This class reshapes the input to the pre-patched format, then applies a 2D
+    convolution to project this pre-patched data to the embedding dimension,
+    and finally flattens across channels to produce a unified embedding.
 
     Parameters:
     -----------
     n_times: int (default=2000)
-        Number of temporal components of the input tensor.
+        Number of temporal components of the input tensor (used for dimension calculation).
     patch_size: int (default=200)
         Size of the patch, default is 1-seconds with 200Hz.
     in_channels: int (default=1)
-        Number of input channels for to be used in the convolution.
+        Number of input channels (from VQVAE codebook).
     emb_dim: int (default=200)
-        Number of out_channes to be used in the convolution, here,
-        we used the same as patch_size.
-    n_codebooks: int (default=62)
-        Number of patches to be used in the convolution, here,
-        we used the same as n_times // patch_size.
+        Number of output embedding dimension.
     """
 
     def __init__(
@@ -707,10 +766,13 @@ class _PatchEmbed(nn.Module):
         self.n_times = n_times
         self.patch_size = patch_size
         self.patch_shape = (1, self.n_times // self.patch_size)
-        n_patchs = n_codebooks * (self.n_times // self.patch_size)
-
-        self.n_patchs = n_patchs
+        self.n_patchs = self.n_times // self.patch_size
+        self.emb_dim = emb_dim
+        self.in_channels = in_channels
 
+        # 2D Conv to project the pre-patched data
+        # Input: (Batch, in_channels, n_patches, patch_size)
+        # After proj: (Batch, emb_dim, n_patches, 1)
         self.proj = nn.Conv2d(
             in_channels=in_channels,
             out_channels=emb_dim,
@@ -718,27 +780,64 @@ class _PatchEmbed(nn.Module):
             stride=(1, self.patch_size),
         )
 
-        self.merge_transpose = Rearrange(
-            "Batch ch patch spatch -> Batch patch spatch ch",
-        )
-
     def forward(self, x):
         """
-        Apply the convolution to the input tensor.
-        then merge the output tensor to the desired shape.
+        Apply the temporal projection to the input tensor after grouping channels.
 
-        Parameters:
-        -----------
-        x: torch.Tensor
-            Input tensor of shape (Batch, Channels, n_patchs, patch_size).
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input tensor of shape (Batch, n_channels, n_times) or
+            (Batch, n_channels, n_patches, patch_size).
 
-        Return:
+        Returns
         -------
-        x: torch.Tensor
-            Output tensor of shape (Batch, n_patchs, patch_size, channels).
+        torch.Tensor
+            Output tensor of shape (Batch, n_patchs, emb_dim).
         """
+        if x.ndim == 4:
+            batch_size, n_channels, n_patchs, patch_len = x.shape
+            if patch_len != self.patch_size:
+                raise ValueError(
+                    "When providing a 4D tensor, the last dimension "
+                    f"({patch_len}) must match patch_size ({self.patch_size})."
+                )
+            n_times = n_patchs * patch_len
+            x = x.reshape(batch_size, n_channels, n_times)
+        elif x.ndim == 3:
+            batch_size, n_channels, n_times = x.shape
+        else:
+            raise ValueError(
+                "Input must be either 3D (batch, channels, times) or "
+                "4D (batch, channels, n_patches, patch_size)."
+            )
+
+        if n_times % self.patch_size != 0:
+            raise ValueError(
+                f"n_times ({n_times}) must be divisible by patch_size ({self.patch_size})."
+            )
+        if n_channels % self.in_channels != 0:
+            raise ValueError(
+                "The input channel dimension "
+                f"({n_channels}) must be divisible by in_channels ({self.in_channels})."
+            )
+
+        group_size = n_channels // self.in_channels
+
+        # Reshape so Conv2d sees `in_channels` feature maps and uses the grouped
+        # EEG channels as the spatial height dimension.
+        # Shape after view: (Batch, in_channels, group_size, n_times)
+        x = x.view(batch_size, self.in_channels, group_size, n_times)
+
+        # Apply the temporal projection per group.
+        # Output shape: (Batch, emb_dim, group_size, n_patchs)
         x = self.proj(x)
-        x = self.merge_transpose(x)
+
+        # THIS IS braindecode's MODIFICATION:
+        # Average over the grouped channel dimension and permute to (Batch, n_patchs, emb_dim)
+        x = x.mean(dim=2)
+        x = x.transpose(1, 2).contiguous()
+
         return x
 
 

braindecode/models/signal_jepa.py

@@ -5,7 +5,8 @@ from __future__ import annotations
 
 import math
 from copy import deepcopy
-from typing import Any, Sequence
+from pathlib import Path
+from typing import Any, Optional, Sequence
 
 import torch
 from einops.layers.torch import Rearrange
@@ -319,25 +320,50 @@ class SignalJEPA_Contextual(_BaseSignalJEPA):
     @classmethod
     def from_pretrained(
         cls,
-        model: SignalJEPA,
-        n_outputs: int,
+        model: Optional[SignalJEPA | str | Path] = None,  # type: ignore
+        n_outputs: Optional[int] = None,  # type: ignore
         n_spat_filters: int = 4,
-        chs_info: list[dict[str, Any]] | None = None,
+        chs_info: Optional[list[dict[str, Any]]] = None,  # type: ignore
+        **kwargs,
     ):
-        """Instantiate a new model from a pre-trained :class:`SignalJEPA` model.
+        """Instantiate a new model from a pre-trained :class:`SignalJEPA` model or from Hub.
 
         Parameters
         ----------
-        model: SignalJEPA
-            Pre-trained model.
-        n_outputs: int
-            Number of classes for the new model.
+        model: SignalJEPA, str, Path, or None
+            Either a pre-trained :class:`SignalJEPA` model, a string/Path to a local directory
+            (for Hub-style loading), or None (for Hub loading via kwargs).
+        n_outputs: int or None
+            Number of classes for the new model. Required when loading from a SignalJEPA model,
+            optional when loading from Hub (will be read from config).
         n_spat_filters: int
             Number of spatial filters.
         chs_info: list of dict | None
             Information about each individual EEG channel. This should be filled with
             ``info["chs"]``. Refer to :class:`mne.Info` for more details.
+        **kwargs
+            Additional keyword arguments passed to the parent class for Hub loading.
         """
+        # Check if this is a Hub-style load (from a directory path)
+        if isinstance(model, (str, Path)) or (model is None and kwargs):
+            # This is a Hub load, delegate to parent class
+            if isinstance(model, (str, Path)):
+                # model is actually the repo_id or directory path
+                return super().from_pretrained(model, **kwargs)
+            else:
+                # model is None, treat as hub-style load
+                return super().from_pretrained(**kwargs)
+
+        # This is the original SignalJEPA transfer learning case
+        if not isinstance(model, SignalJEPA):
+            raise TypeError(
+                f"model must be a SignalJEPA instance, a path string, or Path object, got {type(model)}"
+            )
+        if n_outputs is None:
+            raise ValueError(
+                "n_outputs must be provided when loading from a SignalJEPA model"
+            )
+
         feature_encoder = model.feature_encoder
         pos_encoder = model.pos_encoder
         transformer = model.transformer
@@ -463,22 +489,47 @@ class SignalJEPA_PostLocal(_BaseSignalJEPA):
 
     @classmethod
     def from_pretrained(
-        cls, model: SignalJEPA, n_outputs: int, n_spat_filters: int = 4
+        cls,
+        model: SignalJEPA | str | Path = None,  # type: ignore
+        n_outputs: int = None,  # type: ignore
+        n_spat_filters: int = 4,
+        **kwargs,
     ):
-        """Instantiate a new model from a pre-trained :class:`SignalJEPA` model.
+        """Instantiate a new model from a pre-trained :class:`SignalJEPA` model or from Hub.
 
         Parameters
         ----------
-        model: SignalJEPA
-            Pre-trained model.
-        n_outputs: int
-            Number of classes for the new model.
+        model: SignalJEPA, str, Path, or None
+            Either a pre-trained :class:`SignalJEPA` model, a string/Path to a local directory
+            (for Hub-style loading), or None (for Hub loading via kwargs).
+        n_outputs: int or None
+            Number of classes for the new model. Required when loading from a SignalJEPA model,
+            optional when loading from Hub (will be read from config).
         n_spat_filters: int
             Number of spatial filters.
-        chs_info: list of dict | None
-            Information about each individual EEG channel. This should be filled with
-            ``info["chs"]``. Refer to :class:`mne.Info` for more details.
+        **kwargs
+            Additional keyword arguments passed to the parent class for Hub loading.
         """
+        # Check if this is a Hub-style load (from a directory path)
+        if isinstance(model, (str, Path)) or (model is None and kwargs):
+            # This is a Hub load, delegate to parent class
+            if isinstance(model, (str, Path)):
+                # model is actually the repo_id or directory path
+                return super().from_pretrained(model, **kwargs)
+            else:
+                # model is None, treat as hub-style load
+                return super().from_pretrained(**kwargs)
+
+        # This is the original SignalJEPA transfer learning case
+        if not isinstance(model, SignalJEPA):
+            raise TypeError(
+                f"model must be a SignalJEPA instance, a path string, or Path object, got {type(model)}"
+            )
+        if n_outputs is None:
+            raise ValueError(
+                "n_outputs must be provided when loading from a SignalJEPA model"
+            )
+
         feature_encoder = model.feature_encoder
         assert feature_encoder is not None
         new_model = cls(
@@ -597,22 +648,47 @@ class SignalJEPA_PreLocal(_BaseSignalJEPA):
 
     @classmethod
     def from_pretrained(
-        cls, model: SignalJEPA, n_outputs: int, n_spat_filters: int = 4
+        cls,
+        model: SignalJEPA | str | Path = None,  # type: ignore
+        n_outputs: int = None,  # type: ignore
+        n_spat_filters: int = 4,
+        **kwargs,
     ):
-        """Instantiate a new model from a pre-trained :class:`SignalJEPA` model.
+        """Instantiate a new model from a pre-trained :class:`SignalJEPA` model or from Hub.
 
         Parameters
         ----------
-        model: SignalJEPA
-            Pre-trained model.
-        n_outputs: int
-            Number of classes for the new model.
+        model: SignalJEPA, str, Path, or None
+            Either a pre-trained :class:`SignalJEPA` model, a string/Path to a local directory
+            (for Hub-style loading), or None (for Hub loading via kwargs).
+        n_outputs: int or None
+            Number of classes for the new model. Required when loading from a SignalJEPA model,
+            optional when loading from Hub (will be read from config).
         n_spat_filters: int
             Number of spatial filters.
-        chs_info: list of dict | None
-            Information about each individual EEG channel. This should be filled with
-            ``info["chs"]``. Refer to :class:`mne.Info` for more details.
+        **kwargs
+            Additional keyword arguments passed to the parent class for Hub loading.
         """
+        # Check if this is a Hub-style load (from a directory path)
+        if isinstance(model, (str, Path)) or (model is None and kwargs):
+            # This is a Hub load, delegate to parent class
+            if isinstance(model, (str, Path)):
+                # model is actually the repo_id or directory path
+                return super().from_pretrained(model, **kwargs)
+            else:
+                # model is None, treat as hub-style load
+                return super().from_pretrained(**kwargs)
+
+        # This is the original SignalJEPA transfer learning case
+        if not isinstance(model, SignalJEPA):
+            raise TypeError(
+                f"model must be a SignalJEPA instance, a path string, or Path object, got {type(model)}"
+            )
+        if n_outputs is None:
+            raise ValueError(
+                "n_outputs must be provided when loading from a SignalJEPA model"
+            )
+
         feature_encoder = model.feature_encoder
         assert feature_encoder is not None
         new_model = cls(

braindecode/version.py

@@ -1 +1 @@
-__version__ = "1.3.0.dev175415232"
+__version__ = "1.3.0.dev175955015"

{braindecode-1.3.0.dev175415232.dist-info → braindecode-1.3.0.dev175955015.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braindecode
-Version: 1.3.0.dev175415232
+Version: 1.3.0.dev175955015
 Summary: Deep learning software to decode EEG, ECG or MEG signals
 Author-email: Robin Tibor Schirrmeister <robintibor@gmail.com>
 Maintainer-email: Alexandre Gramfort <agramfort@meta.com>, Bruno Aristimunha Pinto <b.aristimunha@gmail.com>, Robin Tibor Schirrmeister <robintibor@gmail.com>
@@ -40,6 +40,8 @@ Requires-Dist: linear_attention_transformer
 Requires-Dist: docstring_inheritance
 Provides-Extra: moabb
 Requires-Dist: moabb>=1.2.0; extra == "moabb"
+Provides-Extra: hug
+Requires-Dist: huggingface_hub[torch]>=0.20.0; extra == "hug"
 Provides-Extra: tests
 Requires-Dist: pytest; extra == "tests"
 Requires-Dist: pytest-cov; extra == "tests"
@@ -65,7 +67,7 @@ Requires-Dist: pre-commit; extra == "docs"
 Requires-Dist: openneuro-py; extra == "docs"
 Requires-Dist: plotly; extra == "docs"
 Provides-Extra: all
-Requires-Dist: braindecode[docs,moabb,tests]; extra == "all"
+Requires-Dist: braindecode[docs,hug,moabb,tests]; extra == "all"
 Dynamic: license-file
 
 .. image:: https://badges.gitter.im/braindecodechat/community.svg

{braindecode-1.3.0.dev175415232.dist-info → braindecode-1.3.0.dev175955015.dist-info}/RECORD

@@ -1,9 +1,9 @@
 braindecode/__init__.py,sha256=Ac3LEEyIHWFY_fFh3eAY1GZUqXcUxVSJwOSUCwGEDvQ,182
 braindecode/classifier.py,sha256=k9vSCtfQbld0YVleDi5rrrmk6k_k5JYEPPBYcNxYjZ8,9807
-braindecode/eegneuralnet.py,sha256=dz8k_-2jV7WqkaX4bQG-dmr-vRT7ZtOwJqomXyC9PTw,15287
+braindecode/eegneuralnet.py,sha256=U6kRdT2u8A2Ca0axMTR8IAESBsvgjLMusAbYappKAOk,15368
 braindecode/regressor.py,sha256=VLfrpiXklwI4onkwue3QmzlBWcvspu0tlrLo9RT1Oiw,9375
 braindecode/util.py,sha256=J-tBcDJNlMTIFW2mfOy6Ko0nsgdP4obRoEVDeg2rFH0,12686
-braindecode/version.py,sha256=j0FY0WrH2HJ4RqRpEcu8_b9lb6Ys0OgPvlC3GXwoqvI,35
+braindecode/version.py,sha256=_vqrBwFXZ8aTOQzNbG2-qCrZTT7wOQJENwjT0xpmojI,35
 braindecode/augmentation/__init__.py,sha256=LG7ONqCufYAF9NZt8POIp10lYXb8iSueYkF-CWGK2Ls,1001
 braindecode/augmentation/base.py,sha256=gg7wYsVfa9jfqBddtE03B5ZrPHFFmPl2sa3LOrRnGfo,7325
 braindecode/augmentation/functional.py,sha256=lPhGpZcVtgfQ3oV6p6IQLBCWM_Psa60TwxH3Wj1WyOQ,41133
@@ -29,9 +29,9 @@ braindecode/functional/functions.py,sha256=CoEweM6YLhigx0tNmmz6yAc8iQ078sTFY2GeC
 braindecode/functional/initialization.py,sha256=BUSC7y2TMsfShpMYBVwm3xg3ODFqWp-STH7yD4sn8zk,1388
 braindecode/models/__init__.py,sha256=vB0ZFhucH1cRQPoAAAcc3S-hVTnAy674Eu0FjjjKJp0,2543
 braindecode/models/atcnet.py,sha256=H2IWMscm3IM4PH8DA_iLkUaeMXgA120DmVld4jBFOCM,32242
-braindecode/models/attentionbasenet.py,sha256=AK78VvwrZXyJY20zadzDUHl17C-5zcWCd5xPRN7Lr4o,30385
+braindecode/models/attentionbasenet.py,sha256=_bml0Ofy7yB12X19a026EYkcLuzZIab0v3sQTqZ5HGQ,30485
 braindecode/models/attn_sleep.py,sha256=m6sdFfD4en2hHf_TpotLPC1hVweJcYZvjgf12bV5FZg,17822
-braindecode/models/base.py,sha256=9icrWNZBGbh_VLyB9m8g_K1QyK7s3mh8X-hJ29gEbWs,10802
+braindecode/models/base.py,sha256=KjsHVQDdUCAJB4nS-a6eze-H7ayvU4565tsFUcDVxVQ,20212
 braindecode/models/biot.py,sha256=d2P1i_8k98SU3FkN_dKPXcCoFVmyQIIrBbI1-F3g-8E,17509
 braindecode/models/contrawr.py,sha256=eeR_ik4gNZ3rJLM6Mw9gJ2gTMkZ8CU8C4rN_GQMQTAE,10044
 braindecode/models/ctnet.py,sha256=ce5F31q2weBKvg7PL80iDm7za9fhGaCFvNfHoJW_dtg,17315
@@ -51,12 +51,12 @@ braindecode/models/fblightconvnet.py,sha256=d5MwhawhkjilAMo0ckaYMxJhdGMEuorWgHX-
 braindecode/models/fbmsnet.py,sha256=9bZn2_n1dTrI1Qh3Sz9zMZnH_a-Yq-13UHYSmF6r_UE,11659
 braindecode/models/hybrid.py,sha256=hA8jwD3_3LL71BxUjRM1dkhqlHU9E9hjuDokh-jBq-4,4024
 braindecode/models/ifnet.py,sha256=Y2bwfko3SDjD74AzgUEzgMhKJFGCCw_Q_Noh5VONEjQ,15137
-braindecode/models/labram.py,sha256=vcrpwiu4F-djtIPscFbtP2Y0jTosyR_cXnOMQQRGPLw,41798
+braindecode/models/labram.py,sha256=dnZpHbuB60pKZWZHNQaM01eNajGG0tkZB2iutT882PM,46563
 braindecode/models/msvtnet.py,sha256=hxeCLkHS6w2w89YlLfEPCyQ4XQQpt45bEYPiQJ9SFzY,12642
 braindecode/models/patchedtransformer.py,sha256=9TY9l2X4EoCuE9IoOObjubKFRdmsN5lbrVQLnmr66VY,23444
 braindecode/models/sccnet.py,sha256=C7vdwIR5cI6wJCl5f8TnGQG6qinq21y4HG6l-D5AwbY,11971
 braindecode/models/shallow_fbcsp.py,sha256=7U07DJBrm2JHV8v5ja-xuE5-IH5tfmryhJtrfO1n4jk,7531
-braindecode/models/signal_jepa.py,sha256=UeSkeAM3Qmx8bbAqHCj5nP-PtZM00_5SGA8ibo9mptc,37079
+braindecode/models/signal_jepa.py,sha256=ObP8-AauGZHG9tXxRGnvEnlSiwZ1YssbARuXpUl7swk,41013
 braindecode/models/sinc_shallow.py,sha256=Ilv8K1XhMGiRTBtQdq7L595i6cEFYOBe0_UDv-LqL7s,11907
 braindecode/models/sleep_stager_blanco_2020.py,sha256=vXulnDYutEFLM0UPXyAI0YIj5QImUMVEmYZb78j34H8,6034
 braindecode/models/sleep_stager_chambon_2018.py,sha256=8w8IR2PsfG0jSc3o0YVopgHpOvCHNIuMi7-QRJOYEW4,5245
@@ -95,9 +95,9 @@ braindecode/training/scoring.py,sha256=WRkwqbitA3m_dzRnGp2ZIZPge5Nhx9gAEQhIHzeH4
 braindecode/visualization/__init__.py,sha256=4EER_xHqZIDzEvmgUEm7K1bgNKpyZAIClR9ZCkMuY4M,240
 braindecode/visualization/confusion_matrices.py,sha256=qIWMLEHow5CJ7PhGggD8mnD55Le6xhma9HSzt4R33fc,9509
 braindecode/visualization/gradients.py,sha256=KZo-GA0uwiwty2_94j2IjmCR2SKcfPb1Bi3sQq7vpTk,2170
-braindecode-1.3.0.dev175415232.dist-info/licenses/LICENSE.txt,sha256=7rg7k6hyj8m9whQ7dpKbqnCssoOEx_Mbtqb4uSOjljE,1525
-braindecode-1.3.0.dev175415232.dist-info/licenses/NOTICE.txt,sha256=sOxuTbalPxTM8H6VqtvGbXCt_BoOF7JevEYG_knqbm4,620
-braindecode-1.3.0.dev175415232.dist-info/METADATA,sha256=qepe7bWcaDI2Wg5yj-uymfg3oqu-vIw7Iy8lz9rtZqA,7129
-braindecode-1.3.0.dev175415232.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-braindecode-1.3.0.dev175415232.dist-info/top_level.txt,sha256=pHsWQmSy0uhIez62-HA9j0iaXKvSbUL39ifFRkFnChA,12
-braindecode-1.3.0.dev175415232.dist-info/RECORD,,
+braindecode-1.3.0.dev175955015.dist-info/licenses/LICENSE.txt,sha256=7rg7k6hyj8m9whQ7dpKbqnCssoOEx_Mbtqb4uSOjljE,1525
+braindecode-1.3.0.dev175955015.dist-info/licenses/NOTICE.txt,sha256=sOxuTbalPxTM8H6VqtvGbXCt_BoOF7JevEYG_knqbm4,620
+braindecode-1.3.0.dev175955015.dist-info/METADATA,sha256=E4357DmrBx2DdYRbANQcKTrbPLlxlkAPfVRur72RBlU,7215
+braindecode-1.3.0.dev175955015.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+braindecode-1.3.0.dev175955015.dist-info/top_level.txt,sha256=pHsWQmSy0uhIez62-HA9j0iaXKvSbUL39ifFRkFnChA,12
+braindecode-1.3.0.dev175955015.dist-info/RECORD,,