File size: 2,600 Bytes
8e2a9eb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 | """
FontID — SigLIP2 NaFlex backbone + four linear-probe-style classification heads.
Self-contained: rebuilds the vision tower from the vision_config embedded in
config.json (no dependency on the base repo at load time) and loads
model.safetensors. Predicts font / language / color / style.
from modeling_fontid import FontIDModel
model = FontIDModel.from_pretrained("path/to/dir") # or the HF repo dir
"""
import os, json
import torch
import torch.nn as nn
from safetensors.torch import load_file
from transformers.models.siglip2.modeling_siglip2 import (
Siglip2VisionTransformer, Siglip2VisionConfig,
)
HEADS = ["font", "lang", "color", "style"]
class FontIDModel(nn.Module):
"""Pooled feature from the SigLIP2 attention-pooling head (d=768) -> shared
Dropout(0.2) -> per head: Linear(768->512) -> LayerNorm -> GELU ->
Dropout(0.1) -> Linear(512->n)."""
def __init__(self, vision_config: dict, num_classes: dict, heads=HEADS):
super().__init__()
self.heads = list(heads)
cfg = Siglip2VisionConfig(**vision_config)
# Rebuilt from a dict -> _attn_implementation is unset; pin it so the
# attention dispatch resolves (matches the trained eager forward).
cfg._attn_implementation = "eager"
self.vision = Siglip2VisionTransformer(cfg)
d = cfg.hidden_size # 768
self.dropout = nn.Dropout(0.2)
def make_head(n):
return nn.Sequential(
nn.Linear(d, 512),
nn.LayerNorm(512),
nn.GELU(),
nn.Dropout(0.1),
nn.Linear(512, n),
)
self.head_modules = nn.ModuleDict({h: make_head(num_classes[h]) for h in self.heads})
@torch.no_grad()
def forward(self, pixel_values, pixel_attention_mask, spatial_shapes):
vis = self.vision(
pixel_values=pixel_values,
attention_mask=pixel_attention_mask, # HF NaFlex arg name
spatial_shapes=spatial_shapes,
)
feat = self.dropout(vis.pooler_output) # d=768
return {h: self.head_modules[h](feat) for h in self.heads}
@classmethod
def from_pretrained(cls, path, device="cpu"):
with open(os.path.join(path, "config.json")) as f:
cfg = json.load(f)
model = cls(cfg["vision_config"], cfg["num_classes"], cfg.get("heads", HEADS))
sd = load_file(os.path.join(path, "model.safetensors"))
missing, unexpected = model.load_state_dict(sd, strict=True)
model.eval().to(device)
return model
|