modeling.py

"""
Self-contained PyTorch model definition for the Multilingual OCR recogniser.

Architecture
------------
    Input image   (B, 3, H, W)
        |
    SVTRv2-LNConv-Two33 encoder        (3-stage hybrid Conv/Attention backbone)
        |
    ScriptMoE decoder                  (Transformer decoder with MoE FFN +
                                        sample-level script-aware routing)
        |
    Logits (B, T, vocab-2)

This module reproduces *only what is needed for inference*. Any training-only
branches (load-balancing aux loss, script classifier supervision, etc.) are
intentionally dropped to keep the file small and easy to read. The state dict
key names, however, are kept identical to the original implementation, so a
checkpoint trained with the full OpenOCR-MOE codebase loads here as-is.
"""

from __future__ import annotations

import math

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.init import kaiming_normal_, ones_, trunc_normal_, zeros_


# =============================================================================
# Common building blocks
# =============================================================================
class DropPath(nn.Module):
    """Stochastic depth — no-op at eval time, kept for state-dict compatibility."""

    def __init__(self, drop_prob: float = 0.0):
        super().__init__()
        self.drop_prob = drop_prob

    def forward(self, x):
        if self.drop_prob == 0.0 or not self.training:
            return x
        keep = 1 - self.drop_prob
        shape = (x.shape[0],) + (1,) * (x.ndim - 1)
        mask = x.new_empty(shape).bernoulli_(keep)
        if keep > 0.0:
            mask.div_(keep)
        return x * mask


class Identity(nn.Module):
    def forward(self, x):
        return x


class Mlp(nn.Module):
    def __init__(self, in_features, hidden_features=None, out_features=None,
                 act_layer=nn.GELU, drop=0.0):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x


# =============================================================================
# SVTRv2 encoder (LNConv variant with two 3×3 stem convs)
# =============================================================================
class ConvBNLayer(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1,
                 padding=0, bias=False, groups=1, act=nn.GELU):
        super().__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size,
                              stride=stride, padding=padding,
                              groups=groups, bias=bias)
        self.norm = nn.BatchNorm2d(out_channels)
        self.act = act()

    def forward(self, x):
        return self.act(self.norm(self.conv(x)))


class _Attention(nn.Module):
    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None,
                 attn_drop=0.0, proj_drop=0.0):
        super().__init__()
        self.num_heads = num_heads
        self.dim = dim
        self.head_dim = dim // num_heads
        self.scale = qk_scale or self.head_dim ** -0.5
        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

    def forward(self, x):
        B, N, _ = x.shape
        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)
        attn = (q @ k.transpose(-2, -1)) * self.scale
        attn = attn.softmax(dim=-1)
        attn = self.attn_drop(attn)
        x = (attn @ v).transpose(1, 2).reshape(B, N, self.dim)
        return self.proj_drop(self.proj(x))


class _Block(nn.Module):
    def __init__(self, dim, num_heads, mlp_ratio=4.0, qkv_bias=False,
                 qk_scale=None, drop=0.0, attn_drop=0.0, drop_path=0.0,
                 act_layer=nn.GELU, norm_layer=nn.LayerNorm, eps=1e-6):
        super().__init__()
        self.norm1 = norm_layer(dim, eps=eps)
        self.mixer = _Attention(dim, num_heads, qkv_bias, qk_scale,
                                attn_drop, drop)
        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else Identity()
        self.norm2 = norm_layer(dim, eps=eps)
        self.mlp = Mlp(dim, int(dim * mlp_ratio), act_layer=act_layer, drop=drop)

    def forward(self, x):
        x = self.norm1(x + self.drop_path(self.mixer(x)))
        x = self.norm2(x + self.drop_path(self.mlp(x)))
        return x


class _ConvBlock(nn.Module):
    def __init__(self, dim, num_heads, mlp_ratio=4.0, drop=0.0, drop_path=0.0,
                 act_layer=nn.GELU, norm_layer=nn.LayerNorm, eps=1e-6,
                 num_conv=2, kernel_size=3):
        super().__init__()
        self.norm1 = norm_layer(dim, eps=eps)
        self.mixer = nn.Sequential(*[
            nn.Conv2d(dim, dim, kernel_size, 1,
                      kernel_size // 2, groups=num_heads)
            for _ in range(num_conv)
        ])
        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else Identity()
        self.norm2 = norm_layer(dim, eps=eps)
        self.mlp = Mlp(dim, int(dim * mlp_ratio), act_layer=act_layer, drop=drop)

    def forward(self, x):
        C, H, W = x.shape[1:]
        x = x + self.drop_path(self.mixer(x))
        x = self.norm1(x.flatten(2).transpose(1, 2))
        x = self.norm2(x + self.drop_path(self.mlp(x)))
        x = x.transpose(1, 2).reshape(-1, C, H, W)
        return x


class _FlattenTranspose(nn.Module):
    def forward(self, x):
        return x.flatten(2).transpose(1, 2)


class _SubSample2D(nn.Module):
    def __init__(self, in_channels, out_channels, stride=(2, 1)):
        super().__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, 3,
                              stride=stride, padding=1)
        self.norm = nn.LayerNorm(out_channels)

    def forward(self, x, sz):
        x = self.conv(x)
        C, H, W = x.shape[1:]
        x = self.norm(x.flatten(2).transpose(1, 2))
        x = x.transpose(1, 2).reshape(-1, C, H, W)
        return x, [H, W]


class _SubSample1D(nn.Module):
    def __init__(self, in_channels, out_channels, stride=(2, 1)):
        super().__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, 3,
                              stride=stride, padding=1)
        self.norm = nn.LayerNorm(out_channels)

    def forward(self, x, sz):
        C = x.shape[-1]
        x = x.transpose(1, 2).reshape(-1, C, sz[0], sz[1])
        x = self.conv(x)
        C, H, W = x.shape[1:]
        x = self.norm(x.flatten(2).transpose(1, 2))
        return x, [H, W]


class _IdentitySize(nn.Module):
    def forward(self, x, sz):
        return x, sz


class _SVTRStage(nn.Module):
    def __init__(self, dim, out_dim, depth, mixer, kernel_sizes, sub_k,
                 num_heads, mlp_ratio, qkv_bias, qk_scale, drop_rate,
                 attn_drop_rate, drop_path, norm_layer, act, eps,
                 num_conv, downsample):
        super().__init__()
        self.dim = dim
        self.blocks = nn.Sequential()
        for i in range(depth):
            if mixer[i] == 'Conv':
                self.blocks.append(_ConvBlock(
                    dim=dim, kernel_size=kernel_sizes[i], num_heads=num_heads,
                    mlp_ratio=mlp_ratio, drop=drop_rate, act_layer=act,
                    drop_path=drop_path[i], norm_layer=norm_layer, eps=eps,
                    num_conv=num_conv[i]))
            else:
                if mixer[i] == 'FGlobal':
                    self.blocks.append(_FlattenTranspose())
                self.blocks.append(_Block(
                    dim=dim, num_heads=num_heads, mlp_ratio=mlp_ratio,
                    qkv_bias=qkv_bias, qk_scale=qk_scale, drop=drop_rate,
                    act_layer=act, attn_drop=attn_drop_rate,
                    drop_path=drop_path[i], norm_layer=norm_layer, eps=eps))

        if downsample:
            if mixer[-1] == 'Conv':
                self.downsample = _SubSample2D(dim, out_dim, stride=sub_k)
            else:
                self.downsample = _SubSample1D(dim, out_dim, stride=sub_k)
        else:
            self.downsample = _IdentitySize()

    def forward(self, x, sz):
        for blk in self.blocks:
            x = blk(x)
        return self.downsample(x, sz)


class _POPatchEmbed(nn.Module):
    def __init__(self, in_channels=3, embed_dim=128, flatten=False, bias=False):
        super().__init__()
        self.patch_embed = nn.Sequential(
            ConvBNLayer(in_channels, embed_dim // 2,
                        kernel_size=3, stride=2, padding=1, bias=bias),
            ConvBNLayer(embed_dim // 2, embed_dim,
                        kernel_size=3, stride=2, padding=1, bias=bias),
        )
        if flatten:
            self.patch_embed.append(_FlattenTranspose())

    def forward(self, x):
        sz = x.shape[2:]
        x = self.patch_embed(x)
        return x, [sz[0] // 4, sz[1] // 4]


class SVTRv2LNConvTwo33(nn.Module):
    """SVTRv2-LNConv backbone with two 3×3 stem convs.

    Default config below matches the trained checkpoint, derived from
    ``configs/rec/scriptmoe/svtrv2_scriptmoe_mlt_all.yml``.
    """

    def __init__(
        self,
        in_channels: int = 3,
        out_channels: int = 256,
        depths=(6, 6, 6),
        dims=(128, 256, 384),
        mixer=(
            ('Conv',) * 6,
            ('Conv', 'Conv', 'FGlobal', 'Global', 'Global', 'Global'),
            ('Global',) * 6,
        ),
        sub_k=((1, 1), (2, 1), (-1, -1)),
        num_heads=(4, 8, 12),
        mlp_ratio: float = 4.0,
        qkv_bias: bool = True,
        drop_rate: float = 0.0,
        attn_drop_rate: float = 0.0,
        drop_path_rate: float = 0.1,
        eps: float = 1e-6,
    ):
        super().__init__()
        num_stages = len(depths)
        self.num_features = dims[-1]

        self.pope = _POPatchEmbed(
            in_channels=in_channels,
            embed_dim=dims[0],
            flatten=mixer[0][0] != 'Conv',
            bias=False,
        )

        dpr = np.linspace(0, drop_path_rate, sum(depths))
        kernel_sizes = [[3] * d for d in depths]
        num_convs = [[2] * d for d in depths]

        self.stages = nn.ModuleList()
        for i in range(num_stages):
            stage = _SVTRStage(
                dim=dims[i],
                out_dim=dims[i + 1] if i < num_stages - 1 else 0,
                depth=depths[i],
                mixer=list(mixer[i]),
                kernel_sizes=kernel_sizes[i],
                sub_k=sub_k[i],
                num_heads=num_heads[i],
                mlp_ratio=mlp_ratio,
                qkv_bias=qkv_bias,
                qk_scale=None,
                drop_rate=drop_rate,
                attn_drop_rate=attn_drop_rate,
                drop_path=dpr[sum(depths[:i]):sum(depths[:i + 1])],
                norm_layer=nn.LayerNorm,
                act=nn.GELU,
                eps=eps,
                num_conv=num_convs[i],
                downsample=(i != num_stages - 1),
            )
            self.stages.append(stage)

        self.out_channels = self.num_features
        self.apply(self._init_weights)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, mean=0, std=0.02)
            if m.bias is not None:
                zeros_(m.bias)
        elif isinstance(m, nn.LayerNorm):
            zeros_(m.bias)
            ones_(m.weight)
        elif isinstance(m, nn.Conv2d):
            kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')

    def forward(self, x):
        x, sz = self.pope(x)
        for stage in self.stages:
            x, sz = stage(x, sz)
        return x


# =============================================================================
# ScriptMoE decoder (inference-only)
# =============================================================================
class _ExpertFFN(nn.Module):
    def __init__(self, d_model, dim_feedforward, dropout=0.1, act_layer=nn.ReLU):
        super().__init__()
        self.fc1 = nn.Linear(d_model, dim_feedforward)
        self.act = act_layer()
        self.fc2 = nn.Linear(dim_feedforward, d_model)
        self.drop = nn.Dropout(dropout)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x


class _ScriptAwareRouter(nn.Module):
    """Predicts which experts to activate per sample.

    For inference we always run in 'sample' routing mode: a single set of
    experts is selected for the whole image, based on the mean-pooled encoder
    feature.
    """

    def __init__(self, d_model, num_experts, top_k=2):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        self.router = nn.Linear(d_model, num_experts, bias=False)

    def forward(self, x, memory):
        # memory: [B, N, D] -> routing input [B, 1, D]
        routing_input = memory.mean(dim=1, keepdim=True)
        router_logits = self.router(routing_input)
        router_probs = F.softmax(router_logits, dim=-1)

        top_k_weights, top_k_indices = torch.topk(router_probs,
                                                  self.top_k, dim=-1)
        top_k_weights = top_k_weights / (
            top_k_weights.sum(dim=-1, keepdim=True) + 1e-9)
        return top_k_indices, top_k_weights


class _MoEFFNLayer(nn.Module):
    def __init__(self, d_model, dim_feedforward, num_experts=4, top_k=2,
                 dropout=0.1, shared_expert_ratio=0.5):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        self.d_model = d_model

        self.router = _ScriptAwareRouter(d_model, num_experts, top_k)
        self.experts = nn.ModuleList([
            _ExpertFFN(d_model, dim_feedforward, dropout)
            for _ in range(num_experts)
        ])
        shared_dim = int(dim_feedforward * shared_expert_ratio)
        self.shared_expert = _ExpertFFN(d_model, shared_dim, dropout)
        self.shared_gate = nn.Linear(d_model, 1, bias=False)

    def forward(self, x, memory):
        B, T, D = x.shape

        top_k_indices, top_k_weights = self.router(x, memory)
        # Sample-level routing -> broadcast to every token in the sequence
        top_k_indices = top_k_indices.expand(-1, T, -1)
        top_k_weights = top_k_weights.expand(-1, T, -1)

        x_flat = x.reshape(B * T, D)
        idx_flat = top_k_indices.reshape(B * T, self.top_k)
        w_flat = top_k_weights.reshape(B * T, self.top_k)

        routed = torch.zeros_like(x_flat)
        for k in range(self.top_k):
            ek_idx = idx_flat[:, k]
            ek_w = w_flat[:, k].unsqueeze(-1)
            for e in range(self.num_experts):
                mask = (ek_idx == e)
                if mask.any():
                    routed[mask] += self.experts[e](x_flat[mask]) * ek_w[mask]
        routed = routed.reshape(B, T, D)

        shared = self.shared_expert(x)
        gate = torch.sigmoid(self.shared_gate(x))
        return gate * shared + (1 - gate) * routed


class _MultiheadAttention(nn.Module):
    def __init__(self, embed_dim, num_heads, dropout=0.0, self_attn=False):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        assert self.head_dim * num_heads == embed_dim
        self.scale = self.head_dim ** -0.5
        self.self_attn = self_attn
        if self_attn:
            self.qkv = nn.Linear(embed_dim, embed_dim * 3)
        else:
            self.q = nn.Linear(embed_dim, embed_dim)
            self.kv = nn.Linear(embed_dim, embed_dim * 2)
        self.attn_drop = nn.Dropout(dropout)
        self.out_proj = nn.Linear(embed_dim, embed_dim)

    def forward(self, query, key=None, attn_mask=None):
        B, qN = query.shape[:2]
        if self.self_attn:
            qkv = self.qkv(query).reshape(B, qN, 3, self.num_heads,
                                          self.head_dim).permute(2, 0, 3, 1, 4)
            q, k, v = qkv.unbind(0)
        else:
            kN = key.shape[1]
            q = self.q(query).reshape(B, qN, self.num_heads,
                                      self.head_dim).transpose(1, 2)
            kv = self.kv(key).reshape(B, kN, 2, self.num_heads,
                                      self.head_dim).permute(2, 0, 3, 1, 4)
            k, v = kv.unbind(0)

        attn = (q @ k.transpose(-2, -1)) * self.scale
        if attn_mask is not None:
            attn = attn + attn_mask
        attn = F.softmax(attn, dim=-1)
        attn = self.attn_drop(attn)
        x = (attn @ v).transpose(1, 2).reshape(B, qN, self.embed_dim)
        return self.out_proj(x)


class _MoETransformerBlock(nn.Module):
    def __init__(self, d_model, nhead, dim_feedforward,
                 attention_dropout_rate=0.0, residual_dropout_rate=0.1,
                 num_experts=4, top_k=2, shared_expert_ratio=0.5,
                 epsilon=1e-5):
        super().__init__()
        self.self_attn = _MultiheadAttention(d_model, nhead,
                                             attention_dropout_rate,
                                             self_attn=True)
        self.norm1 = nn.LayerNorm(d_model, eps=epsilon)
        self.dropout1 = nn.Dropout(residual_dropout_rate)

        self.cross_attn = _MultiheadAttention(d_model, nhead,
                                              attention_dropout_rate,
                                              self_attn=False)
        self.norm2 = nn.LayerNorm(d_model, eps=epsilon)
        self.dropout2 = nn.Dropout(residual_dropout_rate)

        self.moe_ffn = _MoEFFNLayer(
            d_model=d_model, dim_feedforward=dim_feedforward,
            num_experts=num_experts, top_k=top_k,
            dropout=residual_dropout_rate,
            shared_expert_ratio=shared_expert_ratio,
        )
        self.norm3 = nn.LayerNorm(d_model, eps=epsilon)
        self.dropout3 = nn.Dropout(residual_dropout_rate)

    def forward(self, tgt, memory, self_mask=None):
        tgt = self.norm1(tgt + self.dropout1(
            self.self_attn(tgt, attn_mask=self_mask)))
        tgt = self.norm2(tgt + self.dropout2(
            self.cross_attn(tgt, key=memory)))
        ffn_out = self.moe_ffn(tgt, memory=memory)
        tgt = self.norm3(tgt + self.dropout3(ffn_out))
        return tgt


class _PositionalEncoding(nn.Module):
    def __init__(self, dropout, dim, max_len=5000):
        super().__init__()
        self.dropout = nn.Dropout(p=dropout)
        pe = torch.zeros([max_len, dim])
        position = torch.arange(0, max_len, dtype=torch.float32).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, dim, 2).float()
                             * (-math.log(10000.0) / dim))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        self.register_buffer('pe', pe.unsqueeze(0))

    def forward(self, x):
        x = x + self.pe[:, :x.shape[1], :]
        return self.dropout(x)


class _Embeddings(nn.Module):
    def __init__(self, d_model, vocab, padding_idx=0, scale_embedding=True):
        super().__init__()
        self.embedding = nn.Embedding(vocab, d_model, padding_idx=padding_idx)
        self.embedding.weight.data.normal_(mean=0.0, std=d_model ** -0.5)
        self.d_model = d_model
        self.scale_embedding = scale_embedding

    def forward(self, x):
        x = self.embedding(x)
        if self.scale_embedding:
            x = x * math.sqrt(self.d_model)
        return x


class ScriptMoEDecoder(nn.Module):
    """Inference-only ScriptMoE decoder (greedy autoregressive)."""

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        nhead: int = 12,
        num_decoder_layers: int = 2,
        num_experts: int = 4,
        top_k: int = 2,
        shared_expert_ratio: float = 0.5,
        max_len: int = 25,
        attention_dropout_rate: float = 0.0,
        residual_dropout_rate: float = 0.1,
        scale_embedding: bool = True,
    ):
        super().__init__()
        self.out_channels = out_channels
        self.ignore_index = out_channels - 1
        self.bos = out_channels - 2
        self.eos = 0
        self.max_len = max_len

        d_model = in_channels
        dim_feedforward = d_model * 4

        self.embedding = _Embeddings(d_model, vocab=out_channels,
                                     padding_idx=0,
                                     scale_embedding=scale_embedding)
        self.positional_encoding = _PositionalEncoding(
            dropout=residual_dropout_rate, dim=d_model)

        # Optional encoder layers in the decoder block (kept None per yml).
        self.encoder = None

        self.decoder = nn.ModuleList([
            _MoETransformerBlock(
                d_model=d_model, nhead=nhead,
                dim_feedforward=dim_feedforward,
                attention_dropout_rate=attention_dropout_rate,
                residual_dropout_rate=residual_dropout_rate,
                num_experts=num_experts, top_k=top_k,
                shared_expert_ratio=shared_expert_ratio,
            ) for _ in range(num_decoder_layers)
        ])

        # Auxiliary script classifier head (kept for state-dict compatibility,
        # not used at inference).
        self.script_classifier = nn.Sequential(
            nn.Linear(d_model, d_model // 2),
            nn.ReLU(),
            nn.Dropout(0.1),
            nn.Linear(d_model // 2, num_experts),
        )

        self.tgt_word_prj = nn.Linear(d_model, out_channels - 2, bias=False)
        w0 = np.random.normal(0.0, d_model ** -0.5,
                              (d_model, out_channels - 2)).astype(np.float32)
        self.tgt_word_prj.weight.data = torch.from_numpy(w0.transpose())

    @staticmethod
    def _causal_mask(sz, device):
        mask = torch.triu(
            torch.full((sz, sz), float('-inf'), dtype=torch.float32),
            diagonal=1,
        )
        return mask.unsqueeze(0).unsqueeze(0).to(device)

    @torch.no_grad()
    def forward(self, src):
        """Greedy autoregressive decoding.

        Args:
            src: encoder output, [B, N, D]
        Returns:
            logits over the vocabulary, [B, T, out_channels-2]
        """
        bs = src.shape[0]
        memory = src

        dec_seq = torch.full((bs, self.max_len + 1), self.ignore_index,
                             dtype=torch.int64, device=src.device)
        dec_seq[:, 0] = self.bos

        logits = []
        for step in range(self.max_len):
            embed = self.embedding(dec_seq[:, :step + 1])
            embed = self.positional_encoding(embed)
            mask = self._causal_mask(embed.shape[1], src.device)

            tgt = embed
            for layer in self.decoder:
                tgt = layer(tgt, memory, self_mask=mask)

            word_prob = F.softmax(self.tgt_word_prj(tgt[:, -1:, :]), dim=-1)
            logits.append(word_prob)
            if step < self.max_len - 1:
                dec_seq[:, step + 1] = word_prob.squeeze(1).argmax(-1)
                if (dec_seq == self.eos).any(dim=-1).all():
                    break

        return torch.cat(logits, dim=1)


# =============================================================================
# End-to-end recogniser (encoder + decoder) — wraps both above
# =============================================================================
class ScriptMoERecModel(nn.Module):
    """SVTRv2 backbone + ScriptMoE decoder.

    The state-dict layout is identical to the original
    ``BaseRecognizer(encoder=SVTRv2LNConvTwo33, decoder=ScriptMoEDecoder)``
    so checkpoints can be loaded with ``strict=False``.
    """

    def __init__(self, num_classes: int, max_len: int = 25):
        super().__init__()
        self.encoder = SVTRv2LNConvTwo33(in_channels=3, out_channels=256)
        self.decoder = ScriptMoEDecoder(
            in_channels=self.encoder.out_channels,
            out_channels=num_classes,
            nhead=12,
            num_decoder_layers=2,
            num_experts=4,
            top_k=2,
            shared_expert_ratio=0.5,
            max_len=max_len,
        )

    def forward(self, x):
        feat = self.encoder(x)
        return self.decoder(feat)