Upload folder using huggingface_hub

config.json

@@ -0,0 +1,89 @@
+{
+  "_comment": "EgaraNet: DINOv3 ViT-L backbone + StyleNet (Transposed Transformer Block) composite model config.",
+
+  "model_type": "egara_net",
+  "architectures": ["EgaraNetModel"],
+  "torch_dtype": "float32",
+  "transformers_version": "4.56.0.dev0",
+
+  "auto_map": {
+    "AutoConfig": "configuration_egara_net.EgaraNetConfig",
+    "AutoModel": "modeling_egara_net.EgaraNetModel"
+  },
+
+  "_section_backbone": "--- DINOv3 ViT Backbone (nested sub-config) ---",
+  "backbone_config": {
+    "model_type": "dinov3_vit",
+    "architectures": ["DINOv3ViTModel"],
+
+    "hidden_size": 1024,
+    "num_hidden_layers": 24,
+    "num_attention_heads": 16,
+    "intermediate_size": 4096,
+    "hidden_act": "gelu",
+
+    "image_size": 224,
+    "patch_size": 16,
+    "num_channels": 3,
+    "num_register_tokens": 4,
+
+    "attention_dropout": 0.0,
+    "drop_path_rate": 0.0,
+    "layer_norm_eps": 1e-05,
+    "layerscale_value": 1.0,
+
+    "pos_embed_rescale": 2.0,
+    "pos_embed_jitter": null,
+    "pos_embed_shift": null,
+    "rope_theta": 100.0,
+
+    "key_bias": false,
+    "query_bias": true,
+    "value_bias": true,
+    "proj_bias": true,
+    "mlp_bias": true,
+    "use_gated_mlp": false,
+
+    "initializer_range": 0.02,
+    "torch_dtype": "float32",
+    "transformers_version": "4.56.0.dev0"
+  },
+
+  "_section_stylenet": "--- StyleNet: Transposed Transformer Block (TTB) head ---",
+
+  "ttb_input_dim": null,
+  "_ttb_input_dim_note": "null = auto-inferred from backbone_config.hidden_size at model init",
+
+  "ttb_hidden_dim": 1024,
+  "_ttb_hidden_dim_note": "Internal channel dimension of TTB layers. Matches backbone hidden_size for ViT-L.",
+
+  "ttb_output_dim": 1024,
+  "_ttb_output_dim_note": "Final L2-normalised style vector dimension.",
+
+  "ttb_num_layers": 3,
+  "_ttb_num_layers_note": "Number of stacked TransposedTransformerBlock layers.",
+
+  "ttb_num_heads": 16,
+  "_ttb_num_heads_note": "Number of attention heads in TTB. Must divide ttb_hidden_dim evenly. (1024 / 16 = 64)",
+
+  "_section_ttb_internals": "--- TTB internals (derived / documented for reference) ---",
+  "ttb_rms_norm_eps": 1e-05,
+  "ttb_swiglu_multiple": 64,
+  "_ttb_swiglu_note": "SwiGLU hidden = round_up(floor(hidden_dim * 8/3), multiple). E.g. 768->2048, 1024->2752.",
+
+  "_section_attnpool": "--- Attention Pooling ---",
+  "attn_pool_num_heads": 8,
+  "_attn_pool_note": "nn.MultiheadAttention heads used in AttentionPooling. Must divide ttb_hidden_dim.",
+
+  "_section_head": "--- Projection Head (hidden -> output) ---",
+  "head_act": "silu",
+  "_head_note": "Linear(hidden) -> SiLU -> Linear(output). Output is L2-normalised.",
+
+  "_section_preprocessing": "--- Default inference preprocessing ---",
+  "image_size": 512,
+  "keep_aspect_ratio": true,
+  "_keep_aspect_ratio_note": "true = MaxResizeMod16(image_size); false = square resize.",
+  "image_mean": [0.485, 0.456, 0.406],
+  "image_std":  [0.229, 0.224, 0.225],
+  "_image_stats_note": "ImageNet stats; must match backbone preprocessor_config.json."
+}

configuration_egara_net.py

@@ -0,0 +1,307 @@
+"""
+configuration_egara_net.py
+─────────────────────────
+Hugging Face PretrainedConfig for the EgaraNet composite model.
+
+Model structure
+───────────────
+  EgaraNetModel
+  ├backbone  : DINOv3 ViT (frozen or fine-tuned)
+  │               config stored under `backbone_config`
+  └style_net : StyleNet
+      ├input_proj        : Identity or Linear (input_dim → hidden_dim)
+      ├ttb_layers [×N]   : TransposedTransformerBlock
+      │   ├RMSNorm + TTB attention (cross-covariance / transposed)
+      │   └RMSNorm + SwiGLU FFN
+      ├attn_pool         : AttentionPooling (learned query → MHA)
+      └head              : Linear → SiLU → Linear  (L2-normalised output)
+
+Usage
+─────
+    from configuration_egara_net import EgaraNetConfig
+
+    # Build from scratch
+    cfg = EgaraNetConfig(
+        ttb_hidden_dim=1024,
+        ttb_output_dim=1024,
+        ttb_num_layers=3,
+        ttb_num_heads=16,
+    )
+    cfg.save_pretrained("./my_egara_net/")
+
+    # Load from directory
+    cfg = EgaraNetConfig.from_pretrained("./my_egara_net/")
+
+    # Load backbone config separately
+    backbone_cfg = cfg.backbone_config   # PretrainedConfig instance
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Any
+
+from transformers import AutoConfig, PretrainedConfig
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Helpers
+# ──────────────────────────────────────────────────────────────────────────────
+
+def _swiglu_hidden_dim(dim: int, multiple: int = 64) -> int:
+    """Compute SwiGLU hidden dimension: round_up(floor(dim * 8/3), multiple)."""
+    raw = int(dim * 8 / 3)
+    return ((raw + multiple - 1) // multiple) * multiple
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Config
+# ──────────────────────────────────────────────────────────────────────────────
+
+class EgaraNetConfig(PretrainedConfig):
+    r"""
+    Configuration for the EgaraNet style extractor.
+
+    Args
+    ────
+    backbone_config (dict | PretrainedConfig | None):
+        Config for the DINOv3 ViT backbone.  A dict is automatically
+        converted to the appropriate ``PretrainedConfig`` subclass via
+        ``AutoConfig``.  If ``None``, a default ``dinov3_vit`` config is
+        loaded from ``backbone_model_id``.
+
+    backbone_model_id (str):
+        HF hub ID used when ``backbone_config`` is ``None``, e.g.
+        ``"facebook/dinov3-vitl16-pretrain-lvd1689m"``.  Ignored when
+        ``backbone_config`` is provided explicitly.
+
+    ttb_input_dim (int | None):
+        Input channel dimension fed into StyleNet.  ``None`` (default)
+        means it is inferred automatically from
+        ``backbone_config.hidden_size`` at model initialisation.
+
+    ttb_hidden_dim (int):
+        Internal channel width of every TTB layer and the AttentionPooling
+        query.  Should match ``backbone_config.hidden_size`` (1024 for ViT-L).
+
+    ttb_output_dim (int):
+        Dimension of the final L2-normalised style vector.
+
+    ttb_num_layers (int):
+        Number of stacked ``TransposedTransformerBlock`` layers.
+
+    ttb_num_heads (int):
+        Number of attention heads inside each TTB.
+        Must satisfy ``ttb_hidden_dim % ttb_num_heads == 0``.
+
+    ttb_rms_norm_eps (float):
+        ε used in all ``RMSNorm`` layers.
+
+    ttb_swiglu_multiple (int):
+        SwiGLU hidden dim is rounded up to the nearest multiple of this
+        value.  ``hidden = round_up(floor(ttb_hidden_dim * 8/3), multiple)``.
+
+    attn_pool_num_heads (int):
+        Number of heads for ``nn.MultiheadAttention`` inside
+        ``AttentionPooling``.  Must divide ``ttb_hidden_dim``.
+
+    head_act (str):
+        Activation between the two linear layers in the projection head.
+        Currently only ``"silu"`` is supported.
+
+    image_size (int):
+        Default long-edge size used at inference time.
+
+    keep_aspect_ratio (bool):
+        ``True``  → ``MaxResizeMod16(image_size)`` (preserves aspect ratio,
+                     snaps to multiples of 16).
+        ``False`` → square ``Resize((image_size, image_size))``.
+
+    image_mean / image_std (list[float]):
+        Per-channel normalisation statistics.  Must match the backbone's
+        ``preprocessor_config.json``.  Defaults to ImageNet statistics.
+    """
+
+    model_type = "egara_net"
+    is_composition = True  # tells HF this is a composite config
+
+    def __init__(
+        self,
+        # Backbone
+        backbone_config: dict | PretrainedConfig | None = None,
+        backbone_model_id: str = "facebook/dinov3-vitl16-pretrain-lvd1689m",
+        # StyleNet / TTB
+        ttb_input_dim: int | None = None,
+        ttb_hidden_dim: int = 1024,
+        ttb_output_dim: int = 1024,
+        ttb_num_layers: int = 3,
+        ttb_num_heads: int = 16,
+        ttb_rms_norm_eps: float = 1e-5,
+        ttb_swiglu_multiple: int = 64,
+        # Attention Pooling
+        attn_pool_num_heads: int = 8,
+        # Projection Head
+        head_act: str = "silu",
+        # Default Inference Preprocessing
+        image_size: int = 512,
+        keep_aspect_ratio: bool = True,
+        image_mean: list[float] | None = None,
+        image_std: list[float] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        # Validate constraints
+        if ttb_hidden_dim % ttb_num_heads != 0:
+            raise ValueError(
+                f"ttb_hidden_dim ({ttb_hidden_dim}) must be divisible by "
+                f"ttb_num_heads ({ttb_num_heads})."
+            )
+        if ttb_hidden_dim % attn_pool_num_heads != 0:
+            raise ValueError(
+                f"ttb_hidden_dim ({ttb_hidden_dim}) must be divisible by "
+                f"attn_pool_num_heads ({attn_pool_num_heads})."
+            )
+        if ttb_num_layers < 1:
+            raise ValueError(f"ttb_num_layers must be >= 1, got {ttb_num_layers}.")
+
+        # Backbone config
+        if isinstance(backbone_config, dict):
+            # AutoConfig dispatches on "model_type" key
+            backbone_config = AutoConfig.for_model(**backbone_config)
+        self.backbone_config: PretrainedConfig | None = backbone_config
+        self.backbone_model_id: str = backbone_model_id
+
+        # StyleNet params
+        self.ttb_input_dim: int | None = ttb_input_dim
+        self.ttb_hidden_dim: int = ttb_hidden_dim
+        self.ttb_output_dim: int = ttb_output_dim
+        self.ttb_num_layers: int = ttb_num_layers
+        self.ttb_num_heads: int = ttb_num_heads
+        self.ttb_rms_norm_eps: float = ttb_rms_norm_eps
+        self.ttb_swiglu_multiple: int = ttb_swiglu_multiple
+
+        # Attention Pooling
+        self.attn_pool_num_heads: int = attn_pool_num_heads
+
+        # Projection Head
+        self.head_act: str = head_act
+
+        # Preprocessing defaults
+        self.image_size: int = image_size
+        self.keep_aspect_ratio: bool = keep_aspect_ratio
+        self.image_mean: list[float] = image_mean or [0.485, 0.456, 0.406]
+        self.image_std: list[float] = image_std  or [0.229, 0.224, 0.225]
+
+        super().__init__(**kwargs)
+
+    # Derived / read-only properties
+
+    @property
+    def backbone_hidden_size(self) -> int:
+        """Backbone output channel width (inferred from backbone_config)."""
+        if self.backbone_config is not None:
+            return self.backbone_config.hidden_size
+        raise AttributeError(
+            "backbone_config is not set; cannot infer backbone_hidden_size."
+        )
+
+    @property
+    def effective_ttb_input_dim(self) -> int:
+        """Actual input_dim used by StyleNet input_proj."""
+        return self.ttb_input_dim if self.ttb_input_dim is not None \
+               else self.backbone_hidden_size
+
+    @property
+    def ttb_swiglu_hidden_dim(self) -> int:
+        """Derived SwiGLU intermediate dimension (not stored in JSON)."""
+        return _swiglu_hidden_dim(self.ttb_hidden_dim, self.ttb_swiglu_multiple)
+
+    @property
+    def ttb_head_dim(self) -> int:
+        """Per-head dimension inside TTB attention."""
+        return self.ttb_hidden_dim // self.ttb_num_heads
+
+    # Serialisation
+
+    def to_dict(self) -> dict[str, Any]:
+        output = super().to_dict()
+        # Serialise nested backbone config as a plain dict
+        if self.backbone_config is not None:
+            output["backbone_config"] = self.backbone_config.to_dict()
+        return output
+
+    @classmethod
+    def from_dict(cls, config_dict: dict[str, Any], **kwargs) -> "EgaraNetConfig":
+        # AutoConfig.for_model is invoked inside __init__ for backbone_config
+        return super().from_dict(config_dict, **kwargs)
+
+    # Pretty repr
+
+    def __repr__(self) -> str:
+        backbone_name = (
+            getattr(self.backbone_config, "model_type", "?")
+            if self.backbone_config else self.backbone_model_id
+        )
+        return (
+            f"EgaraNetConfig(\n"
+            f"  backbone          : {backbone_name}"
+            f" (hidden={getattr(self.backbone_config, 'hidden_size', '?')})\n"
+            f"  ttb_input_dim     : {self.effective_ttb_input_dim}  "
+            f"({'auto' if self.ttb_input_dim is None else 'explicit'})\n"
+            f"  ttb_hidden_dim    : {self.ttb_hidden_dim}\n"
+            f"  ttb_output_dim    : {self.ttb_output_dim}\n"
+            f"  ttb_num_layers    : {self.ttb_num_layers}\n"
+            f"  ttb_num_heads     : {self.ttb_num_heads}  "
+            f"(head_dim={self.ttb_head_dim})\n"
+            f"  ttb_swiglu_hidden : {self.ttb_swiglu_hidden_dim}\n"
+            f"  attn_pool_heads   : {self.attn_pool_num_heads}\n"
+            f"  head_act          : {self.head_act}\n"
+            f"  image_size        : {self.image_size}  "
+            f"keep_ratio={self.keep_aspect_ratio}\n"
+            f")"
+        )
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Auto-registration  (call once at import time, safe to repeat)
+# ──────────────────────────────────────────────────────────────────────────────
+try:
+    AutoConfig.register("egara_net", EgaraNetConfig)
+except ValueError:
+    pass  # already registered (e.g. module reimported in notebook)
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Quick self-test
+# ──────────────────────────────────────────────────────────────────────────────
+if __name__ == "__main__":
+    import json, pathlib
+
+    here = pathlib.Path(__file__).parent
+
+    # Build from the JSON files supplied with the model
+    backbone_dict = json.loads((here / "config.json").read_text())
+    backbone_cfg  = backbone_dict.get("backbone_config", backbone_dict)
+
+    cfg = EgaraNetConfig(
+        backbone_config=backbone_cfg,
+        ttb_hidden_dim=backbone_cfg.get("hidden_size", 1024),
+        ttb_output_dim=backbone_cfg.get("hidden_size", 1024),
+        ttb_num_layers=3,
+        ttb_num_heads=16,
+        attn_pool_num_heads=8,
+    )
+    print(cfg)
+
+    # Round-trip JSON serialisation
+    tmp = pathlib.Path("/tmp/egara_net_test")
+    cfg.save_pretrained(tmp)
+    cfg2 = EgaraNetConfig.from_pretrained(tmp)
+    assert cfg.ttb_hidden_dim == cfg2.ttb_hidden_dim
+    assert cfg.ttb_output_dim == cfg2.ttb_output_dim
+    assert cfg.ttb_num_layers == cfg2.ttb_num_layers
+    print("Round-trip OK ✓")
+
+    # Derived properties
+    print(f"SwiGLU hidden dim : {cfg.ttb_swiglu_hidden_dim}")
+    print(f"TTB head dim      : {cfg.ttb_head_dim}")
+    print(f"Effective input   : {cfg.effective_ttb_input_dim}")

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b9b0d29d8f3bca26aa2429d3a14684e90b4298aafb8f039dc078d16518f8c073
+size 1389584856

modeling_egara_net.py

@@ -0,0 +1,259 @@
+"""
+modeling_egara_net.py
+────────────────────
+Hugging Face PreTrainedModel implementation for EgaraNet.
+
+EgaraNet is a composite model:
+    backbone  : DINOv3 ViT  (frozen feature extractor)
+    style_net : StyleNet    (Transposed Transformer Block head)
+
+Usage:
+    from modeling_egara_net import EgaraNetModel
+    model = EgaraNetModel.from_pretrained("path/to/model_dir", trust_remote_code=True)
+    # or after uploading to Hub:
+    model = EgaraNetModel.from_pretrained("user/egara-net", trust_remote_code=True)
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import AutoModel, PreTrainedModel
+from transformers.modeling_outputs import ModelOutput
+
+from .configuration_egara_net import EgaraNetConfig
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Output dataclass
+# ──────────────────────────────────────────────────────────────────────────────
+
+@dataclass
+class EgaraNetOutput(ModelOutput):
+    """
+    Output of EgaraNetModel.
+
+    Attributes:
+        style_embedding: L2-normalised style vector [B, output_dim]
+        backbone_features: raw backbone output [B, N, hidden_size] (optional)
+    """
+    style_embedding: torch.FloatTensor = None
+    backbone_features: Optional[torch.FloatTensor] = None
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Building blocks
+# ──────────────────────────────────────────────────────────────────────────────
+
+class RMSNorm(nn.Module):
+    """Root Mean Square Layer Normalization."""
+
+    def __init__(self, dim: int, eps: float = 1e-5):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        orig_dtype = x.dtype
+        x_fp32 = x.to(torch.float32)
+        rms = torch.sqrt(x_fp32.pow(2).mean(-1, keepdim=True) + self.eps)
+        out = (x_fp32 / rms) * self.weight.to(torch.float32)
+        return out.to(orig_dtype)
+
+
+class SwiGLU(nn.Module):
+    """SwiGLU Feed-Forward Network."""
+
+    def __init__(self, dim: int, multiple: int = 64):
+        super().__init__()
+        hidden_dim = int(dim * 8 / 3)
+        hidden_dim = ((hidden_dim + multiple - 1) // multiple) * multiple
+        self.w_gate = nn.Linear(dim, hidden_dim, bias=False)
+        self.w_up = nn.Linear(dim, hidden_dim, bias=False)
+        self.w_down = nn.Linear(hidden_dim, dim, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x))
+
+
+class TransposedTransformerBlock(nn.Module):
+    """
+    TTB — Transposed Transformer Block.
+    Cross-covariance attention in (HeadDim × HeadDim) space.
+    """
+
+    def __init__(self, dim: int, num_heads: int, eps: float = 1e-5,
+                 swiglu_multiple: int = 64):
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.scale = self.head_dim ** -0.5
+
+        self.norm = RMSNorm(dim, eps=eps)
+        self.qkv = nn.Linear(dim, dim * 3, bias=False)
+        self.q_norm = RMSNorm(self.head_dim, eps=eps)
+        self.k_norm = RMSNorm(self.head_dim, eps=eps)
+        self.proj = nn.Linear(dim, dim)
+
+        self.norm_ffn = RMSNorm(dim, eps=eps)
+        self.ffn = SwiGLU(dim, multiple=swiglu_multiple)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, N, C = x.shape
+        shortcut = x
+        x = self.norm(x)
+
+        qkv = (self.qkv(x)
+               .reshape(B, N, 3, self.num_heads, self.head_dim)
+               .permute(2, 0, 3, 1, 4))
+        q, k, v = qkv.unbind(0)
+
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+
+        # Transposed attention: (HeadDim, N) @ (N, HeadDim) → (HeadDim, HeadDim)
+        q = q.transpose(-2, -1)
+        attn = (q @ k) * self.scale
+        attn = attn.softmax(dim=-1)
+
+        v = v.transpose(-2, -1)
+        x = attn @ v  # (HeadDim, HeadDim) @ (HeadDim, N) → (HeadDim, N)
+
+        x = x.transpose(-2, -1).reshape(B, N, C)
+        x = self.proj(x)
+        x = x + shortcut
+
+        x = x + self.ffn(self.norm_ffn(x))
+        return x
+
+
+class AttentionPooling(nn.Module):
+    """Learned-query attention pooling: [B, N, C] → [B, C]."""
+
+    def __init__(self, dim: int, num_heads: int = 8):
+        super().__init__()
+        self.pool_query = nn.Parameter(torch.randn(1, 1, dim))
+        self.pool_attn = nn.MultiheadAttention(
+            embed_dim=dim, num_heads=num_heads, batch_first=True
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B = x.size(0)
+        query = self.pool_query.expand(B, -1, -1)
+        out, _ = self.pool_attn(query, x, x)
+        return out.squeeze(1)
+
+
+class StyleNet(nn.Module):
+    """
+    StyleNet head: TTB layers → AttentionPooling → Projection → L2 norm.
+    """
+
+    def __init__(self, config: EgaraNetConfig):
+        super().__init__()
+        input_dim = config.effective_ttb_input_dim
+        hidden_dim = config.ttb_hidden_dim
+
+        self.input_proj = (
+            nn.Identity()
+            if input_dim == hidden_dim
+            else nn.Linear(input_dim, hidden_dim)
+        )
+
+        self.tab_layers = nn.ModuleList([
+            TransposedTransformerBlock(
+                dim=hidden_dim,
+                num_heads=config.ttb_num_heads,
+                eps=config.ttb_rms_norm_eps,
+                swiglu_multiple=config.ttb_swiglu_multiple,
+            )
+            for _ in range(config.ttb_num_layers)
+        ])
+
+        self.attn_pool = AttentionPooling(
+            dim=hidden_dim, num_heads=config.attn_pool_num_heads
+        )
+
+        self.head = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_dim),
+            nn.SiLU(),
+            nn.Linear(hidden_dim, config.ttb_output_dim),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.input_proj(x)
+        for layer in self.tab_layers:
+            x = layer(x)
+        x = self.attn_pool(x)
+        x = self.head(x)
+        return F.normalize(x, p=2, dim=-1)
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+# EgaraNet Model
+# ──────────────────────────────────────────────────────────────────────────────
+
+class EgaraNetModel(PreTrainedModel):
+    """
+    EgaraNet: DINOv3 ViT backbone + StyleNet head.
+
+    Accepts ``pixel_values`` [B, C, H, W] and returns an
+    ``EgaraNetOutput`` containing the L2-normalised style embedding.
+    """
+
+    config_class = EgaraNetConfig
+    supports_gradient_checkpointing = False
+
+    def __init__(self, config: EgaraNetConfig):
+        super().__init__(config)
+
+        # Backbone
+        if config.backbone_config is not None:
+            self.backbone = AutoModel.from_config(config.backbone_config)
+        else:
+            self.backbone = AutoModel.from_pretrained(config.backbone_model_id)
+
+        # StyleNet head
+        self.style_net = StyleNet(config)
+
+        # Post-init (weight init for any uninitialised params)
+        self.post_init()
+
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        output_backbone_features: bool = False,
+    ) -> EgaraNetOutput:
+        """
+        Args:
+            pixel_values: [B, C, H, W] — preprocessed image tensor.
+            output_backbone_features: if True, also return raw backbone features.
+
+        Returns:
+            EgaraNetOutput with ``style_embedding`` [B, output_dim].
+        """
+        backbone_out = self.backbone(pixel_values=pixel_values)
+        features = backbone_out.last_hidden_state  # [B, N, hidden_size]
+
+        style_emb = self.style_net(features)  # [B, output_dim]
+
+        return EgaraNetOutput(
+            style_embedding=style_emb,
+            backbone_features=features if output_backbone_features else None,
+        )
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Auto-registration
+# ──────────────────────────────────────────────────────────────────────────────
+
+try:
+    from transformers import AutoModel as _AM
+    _AM.register(EgaraNetConfig, EgaraNetModel)
+except (ValueError, ImportError):
+    pass  # already registered or not available