example/model.py

r"""
```text
Input glyph image (1 x 64 x 64, grayscale)
        |
   Shared Encoder (ResNet-style, stage1~5)
        |
   Global Average Pooling
     /        \
content_proj   style_proj
    |              |
content code    style code
 (128-dim)      (512-dim, L2 normalized)
    |              |
 Hangul head     Font head
 cho/jung/jong    font logits + embedding
```

"""

from __future__ import annotations

import math
from dataclasses import dataclass

import torch
from torch import Tensor, nn
from torch.nn import functional as F


CHAR_SIZE = 64


def build_hangul_table() -> list[str]:
    """Generate the 2,350 KS X 1001 (wansseong/precomposed) Hangul syllables
    in code/print order.

    This is the same logic as the identically named function in
    scan-font-browser.py (see docs/scan-font-browser.md section 2.3). Since
    the two scripts are separate execution units, this small pure function is
    duplicated as-is to avoid unnecessary module coupling.
    """

    table = []
    for code in range(0xAC00, 0xD7A4):
        ch = chr(code)
        try:
            ch.encode("iso2022_kr")
        except UnicodeEncodeError:
            continue
        table.append(ch)
    return table


HANGUL_TABLE = build_hangul_table()

# Number of cho (initial), jung (medial), and jong (final, including "no
# final") consonant/vowel slots. These are the same constants used in the
# formula `code = 0xAC00 + (cho*NUM_JUNG + jung)*NUM_JONG + jong` from
# docs/model-design.md section 3.7.
NUM_CHO = 19
NUM_JUNG = 21
NUM_JONG = 28


def decompose_hangul_syllable(char: str) -> tuple[int, int, int]:
    """Decompose a single precomposed Hangul syllable into (cho, jung, jong)
    indices.

    This follows the arithmetic formula from docs/model-design.md section
    3.7 exactly. It's implemented and reused from this one place so that
    training label generation and inference decoding always use the same
    formula.
    """

    code = ord(char)
    if not (0xAC00 <= code <= 0xD7A3):
        raise ValueError(f"Not a modern Hangul syllable: {char!r}")

    index = code - 0xAC00
    cho, remainder = divmod(index, NUM_JUNG * NUM_JONG)
    jung, jong = divmod(remainder, NUM_JONG)
    return cho, jung, jong


CONTENT_DIM = 128
STYLE_DIM = 512
POOLED_DIM = 512  # stage5 channel count = vector dim right after GlobalAvgPool

# Number of GroupNorm groups.
_GROUP_NORM_GROUPS = 32


def _norm(num_channels: int) -> nn.Module:
    return nn.GroupNorm(min(_GROUP_NORM_GROUPS, num_channels), num_channels)


def _mlp(in_dim: int, hidden_dim: int, out_dim: int) -> nn.Sequential:
    return nn.Sequential(
        nn.Linear(in_dim, hidden_dim),
        nn.ReLU(inplace=True),
        nn.Linear(hidden_dim, out_dim),
    )


class ResidualBlock(nn.Module):
    """Standard ResNet residual block (conv-norm-act-conv-norm + shortcut).
    When `stride=2`, the first conv downsamples, and if channels/resolution
    change, the shortcut is also matched with a 1x1 conv + norm
    (model-design.md section 3.2)."""

    def __init__(self, in_channels: int, out_channels: int, stride: int = 1) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, 3, stride=stride,
                               padding=1, bias=False)
        self.norm1 = _norm(out_channels)
        self.conv2 = nn.Conv2d(out_channels, out_channels, 3, stride=1,
                               padding=1, bias=False)
        self.norm2 = _norm(out_channels)
        self.act = nn.ReLU(inplace=True)

        if stride != 1 or in_channels != out_channels:
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_channels, out_channels, 1,
                          stride=stride, bias=False),
                _norm(out_channels),
            )
        else:
            self.shortcut = nn.Identity()

        # Standard ResNet technique of zero-initializing the last norm in
        # each residual branch ("zero-init the last BN in each residual
        # branch") - this lets each block start close to an identity
        # function early in training, helping stability for very deep
        # networks.
        nn.init.zeros_(self.norm2.weight)

    def forward(self, x: Tensor) -> Tensor:
        identity = self.shortcut(x)
        out = self.act(self.norm1(self.conv1(x)))
        out = self.norm2(self.conv2(out))
        return self.act(out + identity)


class UpsampleResidualBlock(nn.Module):
    """Residual block for the decoder. Uses nearest-neighbor upsample + conv
    instead of a transposed conv to avoid checkerboard artifacts
    (model-design.md section 3.5 - a standard practice, though not
    explicitly specified in the document)."""

    def __init__(self, in_channels: int, out_channels: int) -> None:
        super().__init__()
        self.upsample = nn.Upsample(scale_factor=2, mode="nearest")
        self.conv1 = nn.Conv2d(in_channels, out_channels,
                               3, padding=1, bias=False)
        self.norm1 = _norm(out_channels)
        self.conv2 = nn.Conv2d(out_channels, out_channels,
                               3, padding=1, bias=False)
        self.norm2 = _norm(out_channels)
        self.act = nn.ReLU(inplace=True)
        self.shortcut = (
            nn.Conv2d(in_channels, out_channels, 1, bias=False)
            if in_channels != out_channels else nn.Identity()
        )
        nn.init.zeros_(self.norm2.weight)

    def forward(self, x: Tensor) -> Tensor:
        x = self.upsample(x)
        identity = self.shortcut(x)
        out = self.act(self.norm1(self.conv1(x)))
        out = self.norm2(self.conv2(out))
        return self.act(out + identity)


class Encoder(nn.Module):
    """Implements the stage table from docs/model-design.md section 3.2
    directly.

    | stage  | resolution | channels | blocks                   |
    | ------ | ---------- | -------- | ------------------------ |
    | stem   | 64x64      | 32       | conv3x3                  |
    | stage1 | 64x64      | 32       | residual x2              |
    | stage2 | 32x32      | 64       | residual x2 (downsample) |
    | stage3 | 16x16      | 128      | residual x2 (downsample) |
    | stage4 | 8x8        | 256      | residual x3 (downsample) |
    | stage5 | 4x4        | 512      | residual x3 (downsample) |
    """

    def __init__(self) -> None:
        super().__init__()
        self.stem = nn.Sequential(
            nn.Conv2d(1, 32, 3, padding=1, bias=False),
            _norm(32),
            nn.ReLU(inplace=True),
        )
        self.stage1 = nn.Sequential(
            ResidualBlock(32, 32), ResidualBlock(32, 32))
        self.stage2 = nn.Sequential(ResidualBlock(
            32, 64, stride=2), ResidualBlock(64, 64))
        self.stage3 = nn.Sequential(ResidualBlock(
            64, 128, stride=2), ResidualBlock(128, 128))
        self.stage4 = nn.Sequential(
            ResidualBlock(128, 256, stride=2), ResidualBlock(
                256, 256), ResidualBlock(256, 256),
        )
        self.stage5 = nn.Sequential(
            ResidualBlock(256, 512, stride=2), ResidualBlock(
                512, 512), ResidualBlock(512, 512),
        )

    def forward(self, x: Tensor) -> Tensor:
        x = self.stem(x)
        x = self.stage1(x)
        x = self.stage2(x)
        x = self.stage3(x)
        x = self.stage4(x)
        x = self.stage5(x)
        return x  # (B, 512, 4, 4)


class ProjectionHeads(nn.Module):
    """Extracts a content code (128-dim) and a style code (512-dim, L2
    normalized) from `pooled` (512-dim) (model-design.md section 3.3). The
    hidden dimension isn't specified in the document, so it's a choice made
    in this file - content is compressed halfway down since it's on its way
    to the 68-way (19+21+28) combination space, while style keeps its final
    dimension (512) unchanged."""

    def __init__(self, pooled_dim: int = POOLED_DIM,
                 content_dim: int = CONTENT_DIM, style_dim: int = STYLE_DIM) -> None:
        super().__init__()
        self.content_proj = _mlp(
            pooled_dim, (pooled_dim + content_dim) // 2, content_dim)
        self.style_proj = _mlp(pooled_dim, style_dim, style_dim)

    def forward(self, pooled: Tensor) -> tuple[Tensor, Tensor]:
        content = self.content_proj(pooled)
        style = F.normalize(self.style_proj(pooled), dim=-1)
        return content, style


class HangulHead(nn.Module):
    """Produces cho/jung/jong logits from content (128) via three
    independent MLPs (128->64->N) (model-design.md section 3.6 - the hidden
    dimension of 64 is the value specified in the document)."""

    def __init__(self, content_dim: int = CONTENT_DIM, hidden_dim: int = 64) -> None:
        super().__init__()
        self.cho = _mlp(content_dim, hidden_dim, NUM_CHO)
        self.jung = _mlp(content_dim, hidden_dim, NUM_JUNG)
        self.jong = _mlp(content_dim, hidden_dim, NUM_JONG)

    def forward(self, content: Tensor) -> tuple[Tensor, Tensor, Tensor]:
        return self.cho(content), self.jung(content), self.jong(content)


class FontHead(nn.Module):
    """Passes the normalized style (512) through one more nonlinear MLP
    before predicting font logits. This is option 1 from section 4.1 of
    docs/model-design-enhancement-strategy.md: it lets a deeper head
    separate the 3,000+ fine-grained font classes instead of relying on a
    single linear boundary, and it also lets the final classifier's input
    feature scale vary freely again, which eases logit dynamic-range
    saturation."""

    def __init__(
        self,
        num_font_classes: int,
        style_dim: int = STYLE_DIM,
        hidden_dim: int = 1024,
        dropout: float = 0.1,
    ) -> None:
        super().__init__()
        self.mlp = nn.Sequential(
            nn.Linear(style_dim, hidden_dim),
            nn.LayerNorm(hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, style_dim),
            nn.LayerNorm(style_dim),
            nn.GELU(),
        )
        self.classifier = nn.Linear(style_dim, num_font_classes)

    def forward(self, style: Tensor) -> Tensor:
        return self.classifier(self.mlp(style))


class Decoder(nn.Module):
    """concat(content, style) -> 4x4 -> 4 upsample stages -> 64x64
    (model-design.md section 3.5). There's no cross skip, so it never
    receives any encoder features at all - this is exactly why mode A/B
    (section 4.2) can share the identical decoder. The output uses sigmoid
    so it lands in the same [0,1] range as the input produced by
    `FontGlyphDataset` (sigmoid was chosen from the document's "sigmoid or
    clipped grayscale" options)."""

    def __init__(self, content_dim: int = CONTENT_DIM, style_dim: int = STYLE_DIM) -> None:
        super().__init__()
        in_dim = content_dim + style_dim
        self.fc = _mlp(in_dim, in_dim, 512 * 4 * 4)
        self.up1 = UpsampleResidualBlock(512, 256)  # 4 -> 8
        self.up2 = UpsampleResidualBlock(256, 128)  # 8 -> 16
        self.up3 = UpsampleResidualBlock(128, 64)   # 16 -> 32
        self.up4 = UpsampleResidualBlock(64, 32)    # 32 -> 64
        self.out_conv = nn.Conv2d(32, 1, kernel_size=3, padding=1)

    def forward(self, content: Tensor, style: Tensor) -> Tensor:
        z = torch.cat([content, style], dim=-1)
        x = self.fc(z).view(-1, 512, 4, 4)
        x = self.up1(x)
        x = self.up2(x)
        x = self.up3(x)
        x = self.up4(x)
        return torch.sigmoid(self.out_conv(x))


@dataclass
class FontModelOutput:
    content: Tensor
    style: Tensor
    cho_logits: Tensor
    jung_logits: Tensor
    jong_logits: Tensor
    font_logits: Tensor
    reconstruction: Tensor | None = None


class FontRecognitionModel(nn.Module):
    """The top-level model tying together all of docs/model-design.md
    section 3.

    `num_font_classes` isn't hardcoded here - callers must pass it
    explicitly. Since annotation work is still ongoing, the actual number of
    fonts (currently around 3,480) changes over time, so this is meant to be
    taken directly from `FontGlyphDataset.num_font_classes`
    (font_classifier/dataset_loader.py).
    """

    def __init__(self, num_font_classes: int) -> None:
        super().__init__()
        self.encoder = Encoder()
        self.pool = nn.AdaptiveAvgPool2d(1)
        self.projection = ProjectionHeads()
        self.hangul_head = HangulHead()
        self.font_head = FontHead(num_font_classes)
        self.decoder = Decoder()
        self._init_weights()

    def _init_weights(self) -> None:
        for module in self.modules():
            if isinstance(module, nn.Conv2d):
                nn.init.kaiming_normal_(
                    module.weight, mode="fan_out", nonlinearity="relu")
            elif isinstance(module, nn.Linear):
                nn.init.kaiming_uniform_(module.weight, a=math.sqrt(5))
                if module.bias is not None:
                    nn.init.zeros_(module.bias)
        # ResidualBlock/UpsampleResidualBlock's __init__ has already zeroed
        # out its own final norm (zero-init residual), so we don't overwrite
        # it again here.

    def encode(self, x: Tensor) -> FontModelOutput:
        """Extracts content/style codes and cho/jung/jong/font logits from a
        single glyph image (or a batch). When content and style need to be
        pulled separately from different source images, as in mode B
        (model-design.md section 4.2), call this method twice and then mix
        the results with `decode()`."""

        if x.shape[-2:] != (CHAR_SIZE, CHAR_SIZE):
            raise ValueError(
                f"expected {CHAR_SIZE}x{CHAR_SIZE} input, got {tuple(x.shape[-2:])}"
            )

        features = self.encoder(x)
        pooled = self.pool(features).flatten(1)
        content, style = self.projection(pooled)
        cho_logits, jung_logits, jong_logits = self.hangul_head(content)
        font_logits = self.font_head(style)
        return FontModelOutput(
            content=content, style=style,
            cho_logits=cho_logits, jung_logits=jung_logits, jong_logits=jong_logits,
            font_logits=font_logits,
        )

    def decode(self, content: Tensor, style: Tensor) -> Tensor:
        """Reconstructs a 64x64 image from content/style codes. `content`
        and `style` may come from the same image (mode A) or from different
        images (mode B) - since the decoder has no cross skip
        (model-design.md section 3.4), the two cases don't need to be
        distinguished."""

        return self.decoder(content, style)

    def forward(self, x: Tensor) -> FontModelOutput:
        """Convenience method that chains `encode(x)`'s content/style
        straight into `decode` (corresponds to mode A / evaluation) - if the
        training script is at a stage that doesn't use reconstruction loss
        (model-design.md section 4.5 Phase 1), you can call `encode()` alone
        and skip the decoder computation entirely."""

        output = self.encode(x)
        output.reconstruction = self.decode(output.content, output.style)
        return output


# --------------------------------------------------------------------------
# Cho/jung/jong composition and decoding
# --------------------------------------------------------------------------

def compose_hangul_syllable(cho: int, jung: int, jong: int) -> str:
    """Composes a (cho, jung, jong) index triple into a single precomposed
    Hangul syllable. This is the inverse of
    `font_dataset.decompose_hangul_syllable` and uses the same arithmetic
    formula - the formula is kept in exactly these two places
    (font_dataset.py handles decomposition, this function handles
    composition) so that training label generation and inference decoding
    always agree."""

    code = 0xAC00 + (cho * NUM_JUNG + jung) * NUM_JONG + jong
    return chr(code)


# Pre-decompose every character in HANGUL_TABLE (2,350 chars) into
# (cho, jung, jong) indices - since HANGUL_TABLE is fixed at module load
# time, there's no need to recompute this every time `decode_restricted` is
# called (e.g. every validation batch).
_TABLE_CHO, _TABLE_JUNG, _TABLE_JONG = zip(
    *(decompose_hangul_syllable(ch) for ch in HANGUL_TABLE)
)


def decode_restricted(cho_logits: Tensor, jung_logits: Tensor, jong_logits: Tensor) -> list[str]:
    """Restricted decoding: picks the character over the KS X 1001 2,350
    characters (`HANGUL_TABLE`) that maximizes
    `log P(cho) + log P(jung) + log P(jong)`. Used as an evaluation metric
    since it can be compared directly against training labels. Input is
    batched logits (`(B, N)`); the return value is a list of characters of
    the same batch size."""

    log_p_cho = F.log_softmax(cho_logits, dim=-1)
    log_p_jung = F.log_softmax(jung_logits, dim=-1)
    log_p_jong = F.log_softmax(jong_logits, dim=-1)

    device = cho_logits.device
    cho_idx = torch.as_tensor(_TABLE_CHO, device=device)
    jung_idx = torch.as_tensor(_TABLE_JUNG, device=device)
    jong_idx = torch.as_tensor(_TABLE_JONG, device=device)

    # (B, 2350) = joint log-probability for each character in the table
    scores = (
        log_p_cho[:, cho_idx] + log_p_jung[:,
                                           jung_idx] + log_p_jong[:, jong_idx]
    )
    best = scores.argmax(dim=-1).tolist()
    return [HANGUL_TABLE[i] for i in best]


def decode_open(cho_logits: Tensor, jung_logits: Tensor, jong_logits: Tensor) -> list[str]:
    """Open decoding: takes the argmax of cho/jung/jong independently and
    plugs them straight into the composition formula to get one of all
    11,172 syllables. Characters outside the 2,350-character table can also
    be recovered as long as the combination is valid."""

    cho = cho_logits.argmax(dim=-1).tolist()
    jung = jung_logits.argmax(dim=-1).tolist()
    jong = jong_logits.argmax(dim=-1).tolist()
    return [compose_hangul_syllable(c, j, g) for c, j, g in zip(cho, jung, jong)]