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)]