Hey, Computer, Make Me a Font
import string
from dataclasses import dataclass

import torch


@dataclass
class FontogenConfig:
    d_model: int
    nhead: int
    device: str
    num_layers: int
    max_text_tokens: int
    max_glyph_tokens: int
    max_font_tokens: int
    glyphs: str
    glyph_res: int
    learning_rate: float
    batch_size: int


def fontogen_config() -> FontogenConfig:
    d_model = 512
    nhead = 8
    num_layers = 16
    glyphs = string.ascii_uppercase + string.digits + '.,!?'
    max_text_tokens = 16
    max_glyph_tokens = 320
    # ensure a power of two to use block sparse attention from Triton.
    max_font_tokens = 8192 - max_text_tokens
    glyph_res = 150

    learning_rate = 0.0003
    batch_size = 2  # proper output

    device = "cuda" if torch.cuda.is_available() else "mps" if torch.has_mps else "cpu"

    return FontogenConfig(
        d_model=d_model,
        nhead=nhead,
        device=device,
        num_layers=num_layers,
        max_text_tokens=max_text_tokens,
        max_glyph_tokens=max_glyph_tokens,
        max_font_tokens=max_font_tokens,
        glyphs=glyphs,
        glyph_res=glyph_res,
        learning_rate=learning_rate,
        batch_size=batch_size,
    )

config.py
SerCeMan/fontogen
from dataclasses import dataclass
from typing import List

import torch
from torch.utils.data import Dataset


@dataclass
class FontDatasetEntry:
    font_tokens: torch.Tensor
    text_tokens: torch.Tensor
    text_embeddings: torch.Tensor


class FontDataset(Dataset):
    def __init__(self, ds_file, device):
        self.data: List[FontDatasetEntry] = torch.load(ds_file)
        self.device = device

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        row = self.data[idx]
        return {
            'font_tokens': row.font_tokens,
            'text_tokens': row.text_tokens,
            'text_embeddings': row.text_embeddings,
        }

dataset/loader.py
SerCeMan/fontogen
from typing import Dict, Tuple

import torch

from fonts import Font, GlyphPaths


class FontCodec:
    def __init__(self, glyps: str, glyph_res: int, max_font_tokens: int, max_glyph_tokens: int) -> None:
        self.glyphs = glyps
        self.glyph_res = glyph_res
        self.max_glyph_tokens = max_glyph_tokens
        self.max_font_tokens = max_font_tokens
        self.first_system_token = glyph_res ** 2
        self.glyph_vocab = range(self.first_system_token, self.first_system_token + len(glyps))
        system_tokens = list(glyps) + [
            "moveTo",
            "lineTo",
            "qCurveTo",
            "curveTo",
            "closePath",
            "<EOS>",
            "<PAD>",
        ]
        self.mapping: Dict[str, int] = {
            token: self.first_system_token + i for i, token in enumerate(system_tokens)
        }
        self.vocab_size = glyph_res ** 2 + len(self.mapping)
        self.reverse_mapping = {value: key for key, value in self.mapping.items()}
        self.pad_token = self.mapping["<PAD>"]
        self.eos_token = self.mapping["<EOS>"]

    def encode_point(self, point: Tuple[int, int]) -> int:
        x, y = point
        assert 0 <= x < self.glyph_res, f"x={x} is out of bounds"
        assert 0 <= y < self.glyph_res, f"y={y} is out of bounds"
        return int(x + self.glyph_res * y)

    def encode_font(self, font: Font, limit: int = -1, pad_glyphs: bool = False) -> Tuple[torch.Tensor, torch.Tensor]:
        # fill with padding tokens.
        tokens_size = self.max_font_tokens if limit == -1 else limit
        tokens = torch.fill(torch.zeros(tokens_size, dtype=torch.long), self.pad_token)
        idx = 0

        def next_idx():
            nonlocal idx
            idx += 1
            return idx - 1

        last_glyph_letter = list(font.glyph_paths.keys())[-1]

        for (glyph_letter, glyph_paths) in font.glyph_paths.items():
            glyph_paths: GlyphPaths
            glyph_start_idx = idx
            tokens[next_idx()] = self.mapping[glyph_letter]
            for command, data in glyph_paths:
                if command == "moveTo":
                    assert len(data) == 1, f"moveTo should have 1 point, got {len(data)}"
                    tokens[next_idx()] = self.mapping[command]
                    tokens[next_idx()] = self.encode_point(data[0])
                elif command == "lineTo":
                    assert len(data) == 1, f"lineTo should have 1 point, got {len(data)}"
                    tokens[next_idx()] = self.mapping[command]
                    tokens[next_idx()] = self.encode_point(data[0])
                elif command == "qCurveTo":
                    control_points = data
                    tokens[next_idx()] = self.mapping[command]
                    for point in control_points:
                        tokens[next_idx()] = self.encode_point(point)
                elif command == "curveTo":
                    assert len(data) == 3, f"curveTo should have 3 points, got {len(data)}"
                    control_points = data
                    tokens[next_idx()] = self.mapping[command]
                    tokens[next_idx()] = self.encode_point(control_points[0])
                    tokens[next_idx()] = self.encode_point(control_points[1])
                    tokens[next_idx()] = self.encode_point(control_points[2])
                elif command == "closePath":
                    assert len(data) == 0, f"closePath should have 0 points, got {len(data)}"
                    tokens[next_idx()] = self.mapping[command]
                else:
                    raise Exception(f'unknown command {command}')
                if idx - glyph_start_idx > self.max_glyph_tokens:
                    raise Exception(f'too many tokens for glyph {glyph_letter}')
            # Add padding for the current glyph, do not pad the last glyph so that we have enough space for EOS.
            while pad_glyphs and glyph_letter != last_glyph_letter and idx - glyph_start_idx < self.max_glyph_tokens:
                tokens[next_idx()] = self.pad_token
        tokens[next_idx()] = self.mapping["<EOS>"]
        mask = tokens.not_equal(self.pad_token)
        return tokens, mask

    def decode_point(self, token: int) -> Tuple[int, int]:
        x = int(token % self.glyph_res)
        y = int(token // self.glyph_res)
        assert 0 <= x < self.glyph_res, f"x={x} is out of bounds, token={token}"
        assert 0 <= y < self.glyph_res, f"y={y} is out of bounds, token={token}"
        return x, y

    def decode_font(self, tokens: torch.Tensor) -> Font:
        glyph_paths: Dict[str, GlyphPaths] = {}
        current_glyph_path: GlyphPaths = []

        idx = 0

        def next_idx():
            nonlocal idx
            idx += 1
            return idx - 1

        token = tokens[next_idx()].item()
        # the first token should always be a letter
        glyph_letter = self.reverse_mapping[token]
        assert glyph_letter in self.glyphs, f"first token should be a glyph, got {glyph_letter}"

        while idx < len(tokens):
            command = self.reverse_mapping[tokens[next_idx()].item()]
            if command == "<EOS>":
                glyph_paths[glyph_letter] = current_glyph_path
                return Font(glyph_paths)
            elif command == "<PAD>":
                # skip PAD tokens
                continue
            elif command in self.glyphs:
                print(f'decoding glyph {glyph_letter}')
                glyph_paths[glyph_letter] = current_glyph_path
                glyph_letter = command
                current_glyph_path = []
            elif command == "closePath":
                current_glyph_path.append((command, tuple()))
            elif command == "curveTo":
                data = tuple([self.decode_point(tokens[next_idx()].item()),
                              self.decode_point(tokens[next_idx()].item()),
                              self.decode_point(tokens[next_idx()].item())])
                current_glyph_path.append((command, data))
            elif command == "moveTo":
                data = tuple([self.decode_point(tokens[next_idx()].item())])
                current_glyph_path.append((command, data))
            elif command == "lineTo":
                data = tuple([self.decode_point(tokens[next_idx()].item())])
                current_glyph_path.append((command, data))
            elif command == "qCurveTo":
                control_points = []
                while True:
                    if tokens[idx].item() >= self.first_system_token:
                        break
                    control_point = self.decode_point(tokens[next_idx()].item())
                    control_points.append(control_point)
                current_glyph_path.append((command, tuple(control_points)))
            else:
                raise Exception(f'unknown command {command}')

        return Font(glyph_paths)

font_codec.py
SerCeMan/fontogen
import string
import sys
from dataclasses import dataclass
from typing import List, Tuple, Dict

import matplotlib.pyplot as plt
import numpy as np
from fontTools.cu2qu import curve_to_quadratic
from fontTools.fontBuilder import FontBuilder
from fontTools.pens.recordingPen import RecordingPen, DecomposingRecordingPen
from fontTools.pens.ttGlyphPen import TTGlyphPen
from fontTools.ttLib import TTFont
from fontTools.ufoLib.glifLib import Glyph

GlyphPaths = List[Tuple[str, Tuple[Tuple[int, int]]]]


@dataclass
class Font:
    glyph_paths: Dict[str, GlyphPaths]


class Fonts:
    def __init__(self, glyps: str, glyph_res: int) -> None:
        self.glyphs = glyps
        self.glyph_res = glyph_res
        self.scaler = FontScaler(glyph_res)

    def extract_path_data_using_pen(self, font: TTFont, char: str) -> GlyphPaths | None:
        """Extract glyph path data using the pen API."""
        cmap = font.getBestCmap()
        glyph_set = font.getGlyphSet()
        char_idx = ord(char)
        if char_idx not in cmap:
            return None
        glyph = glyph_set[cmap[char_idx]]

        pen = DecomposingRecordingPen(glyph_set)
        glyph.draw(pen)

        # sometimes the glyphs include useless moveTo commands right before a closePath
        # running through a drawing loop optimises them away.
        ttf_pen = TTGlyphPen()
        pen.replay(ttf_pen)
        optimised_pen = RecordingPen()
        ttf_pen.glyph().draw(optimised_pen, font.getGlyphSet())

        return optimised_pen.value

    def load_font(self, path: str, allow_missing: bool = False) -> Font:
        ttfont = TTFont(path)
        glyph_paths = {}
        for glyph in self.glyphs:
            font_data = self.extract_path_data_using_pen(ttfont, glyph)
            if font_data is None:
                if glyph in string.ascii_uppercase or not allow_missing:
                    # always fail if the glyph is an letter to skip wild fonts
                    raise Exception(f'character is missing in the font: {glyph}')
                else:
                    continue
            assert font_data is not None, f'font_data must not be None for glyph {glyph}'
            glyph_paths[glyph] = font_data
        return Font(glyph_paths)

    def record_glyph_paths(self, glyph_paths: Dict[str, GlyphPaths]) -> Dict[str, Glyph]:
        glyphs = {".notdef": TTGlyphPen(None).glyph()}
        for glyph_name, path_data in glyph_paths.items():
            pen = TTGlyphPen(None)
            curr_point = None
            for command, data in path_data:
                if command == "moveTo":
                    pen.moveTo(*data)
                    curr_point = data[-1]
                elif command == "lineTo":
                    pen.lineTo(*data)
                    curr_point = data[-1]
                elif command == "qCurveTo":
                    pen.qCurveTo(*data)
                    curr_point = data[-1]
                elif command == "curveTo":
                    pen.qCurveTo(*curve_to_quadratic(tuple([curr_point] + list(data)), 0.5))
                    curr_point = data[-1]
                elif command == "closePath":
                    pen.closePath()
                else:
                    raise Exception(f'Unknown command: {command}')
            glyphs[glyph_name] = pen.glyph()
        return glyphs

    def save_as_ttf(self, path: str, font: Font):
        fb = FontBuilder(self.glyph_res, isTTF=True)
        glyphs = self.record_glyph_paths(font.glyph_paths)

        cmap = {}
        glyph_order = []
        for glyph in glyphs.keys():
            if glyph != ".notdef":
                cmap[ord(glyph)] = glyph
            glyph_order.append(glyph)

        fb.setupGlyphOrder(glyph_order)
        fb.setupCharacterMap(cmap)
        fb.setupGlyf(glyphs)

        metrics = {}
        glyph_table = fb.font["glyf"]
        for glyph in glyphs.keys():
            width_boundary = self.glyph_res / 10
            glyph_width = width_boundary + (glyph_table[glyph].xMax - glyph_table[glyph].xMin)
            metrics[glyph] = (glyph_width, glyph_table[glyph].xMin)
        fb.setupHorizontalMetrics(metrics)
        # name/horiz are required by Chrome
        fb.setupNameTable(dict(
            familyName=dict(en="Fontogen"),
            styleName=dict(en="Classy"),
        ))
        fb.setupHorizontalHeader(ascent=self.glyph_res, descent=0)
        fb.setupOS2()
        fb.setupPost()
        fb.setupDummyDSIG()

        fb.save(path)


class FontScaler:
    def __init__(self, glyph_res: int) -> None:
        self.glyph_res = glyph_res

    def find_boundaries(self, font) -> Tuple[int, int]:
        min_y_min = sys.maxsize
        max_dim = 0
        for (glyph, pen_data) in font.glyph_paths.items():
            all_coords = [coord for command, data in pen_data for coord in data]
            x_coords, y_coords = zip(*all_coords)
            x_min, x_max = min(x_coords), max(x_coords)
            y_min, y_max = min(y_coords), max(y_coords)
            width = x_max - x_min
            height = y_max - y_min
            if min_y_min > y_min:
                min_y_min = y_min
            if max_dim < max(width, height):
                max_dim = max(width, height)
        assert max_dim > 0, "max_dim should be greater than 0"
        return max_dim, min_y_min

    def scale_and_translate_path_data(self, pen_data: GlyphPaths, max_dim: int, min_y_min: int) -> GlyphPaths:
        """
        Scale the path data to fit within the target range, round to integers,
        and then translate it to make all coordinates non-negative.
        """
        target_range = self.glyph_res
        all_coords = [coord for command, data in pen_data for coord in data]
        x_coords, y_coords = zip(*all_coords)
        # apply the vertical offset from the glyph
        x_min, x_max = min(x_coords), max(x_coords)
        y_min, y_max = min(y_coords), max(y_coords)
        y_min = min(y_min, min_y_min)
        scale_factor = target_range / max_dim
        translated_and_scaled_pen_data = []
        for command, data in pen_data:
            scaled_data = tuple([(min(round((x - x_min) * scale_factor), target_range - 1),
                                  min(round((y - y_min) * scale_factor), target_range - 1)) for x, y in data])
            translated_and_scaled_pen_data.append((command, scaled_data))
        return translated_and_scaled_pen_data

    def normalize_font(self, font: Font) -> Font:
        scaled_glyphs = {}
        max_dim, min_y_min = self.find_boundaries(font)
        for (glyph, glyph_path) in font.glyph_paths.items():
            scaled_glyphs[glyph] = self.scale_and_translate_path_data(glyph_path, max_dim, min_y_min)
        return Font(scaled_glyphs)


class FontPlotter:
    def __init__(self) -> None:
        super().__init__()

    @staticmethod
    def plot_quadratic_bezier_curve(ax, control_points):
        """Plot a quadratic Bézier curve given its control points."""
        p0, p1, p2 = control_points
        t = np.linspace(0, 1, 100)
        curve = np.outer((1 - t) ** 2, p0) + np.outer(2 * (1 - t) * t, p1) + np.outer(t ** 2, p2)
        ax.plot(curve[:, 0], curve[:, 1], 'r-')

    @staticmethod
    def plot_cubic_bezier_curve(ax, control_points):
        """Plot a cubic Bézier curve given its control points."""
        p0, p1, p2, p3 = control_points
        t = np.linspace(0, 1, 100)
        curve = (
                np.outer((1 - t) ** 3, p0) +
                np.outer(3 * (1 - t) ** 2 * t, p1) +
                np.outer(3 * (1 - t) * t ** 2, p2) +
                np.outer(t ** 3, p3)
        )
        ax.plot(curve[:, 0], curve[:, 1], 'g-')

    @staticmethod
    def plot_glyph_from_pen_data(pen_data: GlyphPaths, title: str) -> None:
        """Plot glyph from pen data."""
        ax = plt.gca()
        start_point: Tuple[int, int] = None
        current_point: Tuple[int, int] = None

        for command, data in pen_data:
            if command == "moveTo":
                start_point = data[0]
                current_point = start_point
            elif command == "lineTo":
                end_point = data[0]
                ax.plot([current_point[0], end_point[0]], [current_point[1], end_point[1]], 'b-')
                current_point = end_point
            elif command == "qCurveTo":
                control_points = data
                num_points = len(control_points)
                # Process quadratic Bézier curves with pairs of control points
                for i in range(0, num_points - 1, 2):
                    if i + 2 < num_points:
                        # If there are more control points ahead, compute the implied on-curve point
                        implied_point = (round((control_points[i][0] + control_points[i + 1][0]) / 2),
                                         round((control_points[i][1] + control_points[i + 1][1]) / 2))
                        FontPlotter.plot_quadratic_bezier_curve(ax, (current_point, control_points[i], implied_point))
                        current_point = implied_point
                    else:
                        # If these are the last two control points, use the last point as on-curve
                        FontPlotter.plot_quadratic_bezier_curve(ax, (
                            current_point, control_points[i], control_points[i + 1]))
                        current_point = control_points[i + 1]
            elif command == "curveTo":
                FontPlotter.plot_cubic_bezier_curve(ax, (current_point, data[0], data[1], data[2]))
                current_point = data[2]
            elif command == "closePath":
                ax.plot([current_point[0], start_point[0]], [current_point[1], start_point[1]], 'b-')
                current_point = start_point
            else:
                raise Exception(f'unknown command {command}')

        ax.set_title(title)
        ax.set_aspect('equal', 'box')
        ax.axis('off')  # Hide axis

    def plt_font(self, font: Font):
        # Plot original and scaled glyphs side by side for comparison
        for glyph in font.glyph_paths.keys():
            plt.figure(figsize=(6, 6))
            self.plot_glyph_from_pen_data(font.glyph_paths[glyph], f"glyph={glyph}")
            plt.show()

fonts.py
SerCeMan/fontogen
import os
from dataclasses import dataclass, asdict
from typing import List, Any, Dict, Union, Tuple

import lightning.pytorch as pl
import torch
import torch.nn.functional as F
from torch import nn
from torch.utils.data import DataLoader, random_split
from transformers import BertTokenizer, get_linear_schedule_with_warmup
from xformers.components.attention import BlockSparseAttention
from xformers.components.attention.sparsity_config import BigBirdSparsityConfig
from xformers.components.positional_embedding import SinePositionalEmbedding
from xformers.factory import xFormerEncoderConfig, xFormer

from config import FontogenConfig
from dataset.loader import FontDataset
from font_codec import FontCodec
from model.text_embedding import BERT_MODEL


class FontEmbeddings(nn.Module):
    def __init__(self, d_model: int, font_codec: FontCodec, glyph_res: int, max_seq_len: int, device: str):
        super(FontEmbeddings, self).__init__()
        self.d_model = d_model
        self.coord_size = glyph_res
        self.max_seq_len = max_seq_len
        vocab_size = font_codec.vocab_size
        self.embedding = nn.Embedding(vocab_size, d_model, padding_idx=font_codec.pad_token)
        # Two extra learnable matrices
        self.coord_pad = 0
        self.coord_embedding_x = nn.Embedding(glyph_res + 1, d_model, padding_idx=self.coord_pad)
        self.coord_embedding_y = nn.Embedding(glyph_res + 1, d_model, padding_idx=self.coord_pad)
        self.first_system_token = font_codec.first_system_token
        self.positional_encoding = SinePositionalEmbedding(d_model, seq_len=max_seq_len)
        self.device = device

    def font_to_vec(self, tokens: torch.Tensor) -> [torch.Tensor, torch.Tensor]:
        token_inds = tokens < self.first_system_token

        points = torch.zeros_like(tokens, device=self.device)
        points[token_inds] = tokens[token_inds]

        x_coords = (points % self.coord_size) + 1
        y_coords = (points // self.coord_size) + 1

        return x_coords, y_coords

    def forward(self, font_tokens: torch.Tensor):
        # x: (batch_size, seq_len)
        x_coords, y_coords = self.font_to_vec(font_tokens)

        # Embedding for the token
        token_embedding = self.embedding(font_tokens)  # (batch_size, seq_len, d_model)

        # Embedding for the coordinates
        coord_embedding_x = self.coord_embedding_x(x_coords)  # (batch_size, seq_len, d_model)
        coord_embedding_y = self.coord_embedding_y(y_coords)  # (batch_size, seq_len, d_model)

        # Final embedding
        embedding = token_embedding + coord_embedding_x + coord_embedding_y

        embedding = self.positional_encoding(embedding)

        return embedding


class FontogenDataModule(pl.LightningDataModule):
    def __init__(self, config: FontogenConfig, dataset_path: str, val_perc=0.04) -> None:
        super().__init__()
        self.config = config
        self.dataset = FontDataset(dataset_path, config.device)
        dataset_size = len(self.dataset)

        val_size = int(dataset_size * val_perc)
        train_size = dataset_size - val_size

        # split the dataset, ensure consistency
        g = torch.Generator().manual_seed(42)
        self.train_dataset, self.val_dataset = random_split(self.dataset, [train_size, val_size], generator=g)

    def train_dataloader(self) -> DataLoader:
        return DataLoader(self.train_dataset,
                          pin_memory=True,
                          num_workers=int(os.cpu_count() / 3),
                          # num_workers=0,
                          batch_size=self.config.batch_size)

    def val_dataloader(self) -> DataLoader:
        return DataLoader(self.val_dataset,
                          num_workers=int(os.cpu_count() / 3),
                          # num_workers=0,
                          batch_size=self.config.batch_size)


@dataclass
class StepLoss:
    loss: torch.Tensor


class FontogenModule(pl.LightningModule):
    def __init__(self, config: FontogenConfig) -> None:
        super().__init__()
        self.save_hyperparameters(asdict(config))
        self.config = config
        self.font_codec = FontCodec(config.glyphs, config.glyph_res, config.max_font_tokens, config.max_glyph_tokens)
        self.model = Fontogen(config, self.font_codec)
        self.bert_tokenizer_pad_token_id = 0

    def forward(self,
                text_embeddings: torch.Tensor,
                font_tokens: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        text_out, font_out = self.model(text_embeddings, font_tokens)
        return text_out, font_out

    def make_a_step(self, batch) -> StepLoss:
        font_tokens = batch['font_tokens']
        text_embeddings = batch['text_embeddings']

        text_out, font_out = self.forward(text_embeddings, font_tokens)

        B, L, D = font_out.shape
        font_loss = F.cross_entropy(font_out.reshape(B * L, D), font_tokens.view(-1),
                                    ignore_index=self.font_codec.pad_token)

        return StepLoss(font_loss)

    def training_step(self, batch: Dict[str, List[str]], batch_idx: int) -> Dict[str, torch.Tensor]:
        step_loss: StepLoss = self.make_a_step(batch)
        self.log_dict({
            "loss": step_loss.loss,
        }, batch_size=self.config.batch_size, on_step=True, prog_bar=True, logger=True)
        return {"loss": step_loss.loss}

    def validation_step(self, batch: Dict[str, List[str]], batch_idx: int) -> Dict[str, torch.Tensor]:
        step_loss: StepLoss = self.make_a_step(batch)
        self.log_dict({
            "val_loss": step_loss.loss,
        }, batch_size=self.config.batch_size, on_step=True, prog_bar=True, logger=True)
        return {"val_loss": step_loss.loss}

    def configure_optimizers(self) -> Union[torch.optim.Optimizer, Tuple[List[torch.optim.Optimizer], List[Any]]]:
        print(f'creating an optimizer with LR: {self.config.learning_rate}')
        optimizer = torch.optim.AdamW(self.model.parameters(), lr=self.config.learning_rate)
        total_steps = self.trainer.estimated_stepping_batches
        scheduler = {
            'scheduler': get_linear_schedule_with_warmup(
                optimizer,
                num_warmup_steps=1000,
                num_training_steps=total_steps,
            ),
            'interval': 'step',
        }
        return [optimizer], [scheduler]


class Fontogen(nn.Module):
    def __init__(self, config: FontogenConfig, font_codec: FontCodec):
        super(Fontogen, self).__init__()
        self.font_embedding = FontEmbeddings(
            config.d_model,
            font_codec,
            config.glyph_res,
            config.max_font_tokens,
            config.device
        )
        self.tokenizer: BertTokenizer = BertTokenizer.from_pretrained(BERT_MODEL)
        self.text_vocab_size = self.tokenizer.vocab_size
        self.transformer = FontogenTransformer(
            config.d_model,
            config.nhead,
            config.num_layers,
            self.text_vocab_size,
            font_codec.vocab_size,
            config.max_text_tokens,
            config.max_font_tokens,
            config.max_glyph_tokens,
            config.glyph_res,
            config.device,
        )
        self.text_padding_token = self.tokenizer.pad_token_id

    def font_embeddings(self, font_tokens: torch.Tensor) -> [torch.Tensor, torch.Tensor]:
        return self.font_embedding(font_tokens)

    def forward(self, text_embeddings, font_tokens):
        font_embeddings = self.font_embeddings(font_tokens)
        return self.transformer(text_embeddings, font_embeddings)


class FontogenTransformer(nn.Module):
    def __init__(self, d_model: int, nhead: int, num_layers: int, text_vocab_size: int, font_vocab_size: int,
                 max_text_tokens: int, max_font_tokens: int, max_glyph_tokens: int, glyph_res: int, device: str):
        super(FontogenTransformer, self).__init__()
        self.d_model = d_model
        self.seq_len = max_text_tokens + max_font_tokens
        self.max_text_len = max_text_tokens
        # let's try to match the text size
        block_size = 16
        assert block_size == max_text_tokens
        x_config = xFormerEncoderConfig(
            num_layers=num_layers,
            dim_model=d_model,
            feedforward_config={
                "name": "MLP",
                "dropout": 0.1,
                "activation": "gelu",
                "hidden_layer_multiplier": 4,
            },
            multi_head_config={
                "num_heads": nhead,
                "attention": BlockSparseAttention(
                    num_heads=nhead,
                    layout=BigBirdSparsityConfig(
                        num_heads=nhead,
                        block_size=block_size,
                        num_global_blocks=1,
                        num_random_blocks=2,
                        num_sliding_window_blocks=int(3 * max_glyph_tokens / block_size)
                    ).make_layout(seq_len=self.seq_len),
                    block_size=block_size,
                    dropout=0.1,
                    causal=True,
                )
            },
        )
        self.transformer = xFormer(x_config)
        self.text_vocab_size = text_vocab_size
        self.svg_vocab_size = font_vocab_size
        self.to_logits = nn.Sequential(
            nn.LayerNorm(d_model),
            # Note, we use nn.Linear here instead of FusedLinear as FusedLinear makes the model not converge.
            # bias is redundant after layer norm
            nn.Linear(d_model, font_vocab_size, bias=False),
        )
        self.device = device

    def forward(self, text_embeddings, font_embeddings):
        sos_value = 0.42  # replace with your special value
        batch_size, _, _ = text_embeddings.shape
        sos_token = torch.full((batch_size, 1, self.d_model), sos_value, device=self.device)

        tgt = torch.cat([
            sos_token,
            text_embeddings,
            font_embeddings[:, :-1, :]
        ], dim=1)

        output = self.transformer(
            src=tgt,
            # in practice, the mask doesn't work, see DALLE paper :(
            # > In preliminary experiments on Conceptual Captions (Sharma et al., 2018), we found that this
            # > resulted in higher validation loss, but better performance on out-of-distribution captions.
            # encoder_input_mask=tgt_key_padding_mask,
        )

        logits = self.to_logits(output)
        text_probabilities = logits[:, :self.max_text_len, :]
        svg_probabilities = logits[:, self.max_text_len:, :]
        return text_probabilities, svg_probabilities

model/model.py
SerCeMan/fontogen
from typing import List

import torch
from transformers import BertTokenizer, BertModel


class TextEmbedder:
    def __init__(self, max_text_token_len: int, device: str):
        self.max_token_len = max_text_token_len
        self.tokenizer = BertTokenizer.from_pretrained(BERT_MODEL)
        self.bert = BertModel.from_pretrained(BERT_MODEL)
        self.device = device

    def tokenize_batch(self, text_batch: List[str]) -> [torch.Tensor, torch.Tensor]:
        encoded_input = self.tokenizer.batch_encode_plus(
            text_batch,
            add_special_tokens=True,
            max_length=self.max_token_len,
            padding='max_length',
            return_attention_mask=True,
            return_tensors='pt',
            truncation=True
        ).to(self.device)
        input_ids = encoded_input['input_ids']
        attn_mask_ids = encoded_input['attention_mask']
        return input_ids, attn_mask_ids

    def embed_tokens(self, text_tokens: torch.Tensor):
        # Predict hidden states features for each layer
        text_tokens = text_tokens.to(self.bert.device)
        with torch.no_grad():
            batch_embeddings = self.bert.embeddings(text_tokens)
        # prepare for transformers:
        # -> (sequence_length, batch_size, embedding_dimension)
        return batch_embeddings

    def embed_text(self, text_batch: List[str]):
        input_ids, attn_mask_ids = self.tokenize_batch(text_batch)
        # ignore attention mask as advised by DALLE paper
        return self.embed_tokens(input_ids), attn_mask_ids


BERT_MODEL = 'google/bert_uncased_L-12_H-512_A-8'

model/text_embedding.py
SerCeMan/fontogen
import os

import torch

from config import fontogen_config
from dataset.loader import FontDatasetEntry
from font_codec import FontCodec
from fonts import FontScaler, Fonts
from model.text_embedding import TextEmbedder


def main():
    device = 'cpu'
    config = fontogen_config()
    glyphs = config.glyphs
    glyph_res = config.glyph_res
    max_text_length = config.max_text_tokens
    max_glyph_tokens = config.max_glyph_tokens
    max_font_tokens = config.max_font_tokens
    dataset_path = 'example/processed_dataset/fonts.ds'

    scaler = FontScaler(glyph_res)
    fonts = Fonts(glyphs, glyph_res)
    codec = FontCodec(glyphs, glyph_res, max_font_tokens, max_glyph_tokens)
    text_embedder = TextEmbedder(max_text_length, device)

    dataset = []
    for dirpath, dirnames, filenames in os.walk('example/dataset'):
        for filename in filenames:
            if filename.endswith('.ttf'):
                filepath = os.path.join(dirpath, filename)
                font = fonts.load_font(filepath)
                font = scaler.normalize_font(font)
                with open(filepath.replace('.ttf', '.txt'), 'r') as file:
                    text = file.read()
                text_tokens, _ = text_embedder.tokenize_batch([text])
                text_embeddings = text_embedder.embed_tokens(text_tokens)
                font_tokens, _ = codec.encode_font(font)
                dataset.append(FontDatasetEntry(
                    font_tokens,
                    text_tokens[0],
                    text_embeddings[0],
                ))
    torch.save(dataset, dataset_path)

if __name__ == '__main__':
    main()
process_dataset.py
SerCeMan/fontogen
import os
import random
import traceback

import lightning.pytorch as pl
import torch
import torch.nn.functional as F

from config import fontogen_config, FontogenConfig
from fonts import Fonts, Font
from model.model import FontogenModule
from model.text_embedding import TextEmbedder


class SamplingCallback(pl.Callback):
    def __init__(self, config: FontogenConfig, sample_every_epoch: int, out_folder: str):
        super().__init__()
        self.config = config
        self.device = config.device
        self.out_folder = out_folder
        self.sample_every_epoch = sample_every_epoch

    def on_train_epoch_start(self, trainer: pl.Trainer, pl_module: pl.LightningModule):
        if (trainer.current_epoch + 1) % self.sample_every_epoch == 0:
            model: FontogenModule = pl_module
            model.eval()
            try:
                sampler = FontogenSampler(model, self.config, self.out_folder)
                sampler.sample('bold sans', trainer.current_epoch)
            except Exception as e:
                print(f'failed to sample', e)
            model.train()


# Source: https://github.com/samxuxiang/SkexGen/blob/c38f30e8ac40aabfa2a71d6842cc585faa9b9862/model/code.py#L12
def top_k_top_p_filtering(logits: torch.Tensor, top_k=0, top_p=0.0, filter_value=-float('Inf')):
    """ Filter a distribution of logits using top-k and/or nucleus (top-p) filtering
        Args:
            logits: logits distribution shape (vocabulary size)
            top_k >0: keep only top k tokens with the highest probability (top-k filtering).
            top_p >0.0: keep the top tokens with cumulative probability >= top_p (nucleus filtering).
                Nucleus filtering is described in Holtzman et al. (http://arxiv.org/abs/1904.09751)
    """
    assert logits.dim() == 1  # batch size 1 for now - could be updated for more but the code would be less clear
    top_k = min(top_k, logits.size(-1))  # Safety check
    if top_k > 0:
        # Remove all tokens with a probability less than the last token of the top-k
        indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
        logits[indices_to_remove] = filter_value

    if top_p > 0.0:
        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
        cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)

        # Remove tokens with cumulative probability above the threshold
        sorted_indices_to_remove = cumulative_probs > top_p
        # Shift the indices to the right to keep also the first token above the threshold
        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
        sorted_indices_to_remove[..., 0] = 0

        indices_to_remove = sorted_indices[sorted_indices_to_remove]
        logits[indices_to_remove] = filter_value
    return logits


class FontogenSampler:
    def __init__(self, model: FontogenModule, config: FontogenConfig, out_folder: str, glyphs: str = None):
        self.model = model
        if glyphs is None:
            glyphs = config.glyphs
        self.glyphs = glyphs
        self.glyph_res = config.glyph_res
        # import string
        # self.glyphs = 'ABCDEFGHIJK'
        self.font_codec = model.font_codec
        self.max_glyph_tokens = config.max_glyph_tokens
        self.max_font_tokens = config.max_font_tokens
        self.out_folder = out_folder
        self.device = config.device
        # do not consume GPU memory for the text module
        self.text_embedder = TextEmbedder(config.max_text_tokens, self.device)

    def sample_next_token(self, logit: torch.Tensor, strategy: str = 'greedy', temperature: float = 1.0):
        max_index: torch.Tensor
        logit /= temperature
        max_index: torch.Tensor
        if strategy == 'multinomial':
            probabilities = torch.softmax(logit, dim=0)
            max_index = torch.multinomial(probabilities, num_samples=1)
        elif strategy == 'greedy':
            probabilities = torch.softmax(logit, dim=0)
            max_index = torch.argmax(probabilities, dim=0)
        elif strategy == 'topknuc':
            # Top-K and/or Nucleus Filtering
            top_k = 10
            top_p = 0.9
            logit = top_k_top_p_filtering(logit, top_k=top_k, top_p=top_p)

            probabilities = torch.softmax(logit, dim=0)
            max_index = torch.multinomial(probabilities, num_samples=1)
        else:
            raise ValueError(f'Unknown type {strategy}')
        return max_index.item()

    def sample_font(self, text_embeddings: torch.Tensor, temperature: int = None, strategy: str = None) -> Font:
        if temperature is None:
            temperature = 0.6
        if strategy is None:
            strategy = 'multinomial'
        font_tokens, _ = self.font_codec.encode_font(Font({'A': []}))
        font_tokens = font_tokens.to(self.device)
        font_tokens = torch.stack([font_tokens])
        max_length = self.max_font_tokens
        curr_token = 'A'
        tokens_per_glyph = 0
        token_idx = 1
        token_attempt = 1
        i = 1
        while i < max_length:
            if tokens_per_glyph > self.max_glyph_tokens * 2:
                if token_attempt > 5:
                    raise Exception(f'too many tokens for glyph {curr_token}')
                else:
                    print(f'too many tokens for glyph {curr_token}, trying again')
                    token_attempt += 1
                    tokens_per_glyph = 0
                    i = token_idx
                    font_tokens[0][token_idx] = self.font_codec.eos_token

            _, font_out = self.model(text_embeddings, font_tokens)

            last_predicted = font_out[0][i]  # B, S, Prob

            next_token = self.sample_next_token(last_predicted, strategy=strategy, temperature=temperature)
            if next_token in self.font_codec.glyph_vocab or next_token == self.font_codec.eos_token:
                if self.glyphs.index(curr_token) + 1 == len(self.glyphs):
                    #
                    next_token = self.font_codec.eos_token
                else:
                    curr_token = self.glyphs[self.glyphs.index(curr_token) + 1]
                    next_token = self.font_codec.mapping[str(curr_token)]
                    print(f'Starting building {self.font_codec.reverse_mapping[next_token]}')
                    tokens_per_glyph = 0
                    token_idx = i + 1
                    token_attempt = 1

            if next_token == self.font_codec.eos_token:
                break

            # put EOS at the end
            eos = font_tokens[0][i].item()
            assert eos == self.font_codec.eos_token, f'unexpected eos at position {i}'
            font_tokens[0][i] = next_token
            if i + 1 < max_length:
                font_tokens[0][i + 1] = eos
            if i == max_length - 1:
                raise ValueError('Max length reached')
            tokens_per_glyph += 1
            i += 1

        return self.font_codec.decode_font(font_tokens[0])

    def sample(self, text: str, step: int = -1, temperature: int = None, strategy: str = None) -> str:
        print(f'sampling {text}')
        out_folder = self.out_folder
        with torch.no_grad():
            text_tokens, _ = self.text_embedder.tokenize_batch([text])

            text_tokens = text_tokens.to(self.device)

            text_embeddings = self.text_embedder.embed_tokens(text_tokens)
            text_embeddings = text_embeddings.to(self.device)

            out_font: Font = self.sample_font(text_embeddings, temperature, strategy)
            step_str = "" if step < 0 else f"{step}_"
            out_path = f'{out_folder}/{step_str}{text.replace(" ", "_")}_{random.randint(0, 100000)}.ttf'
            self.try_save_font_tokens(out_path, out_font)
            return out_path

    def try_save_font_tokens(self, font_path: str, font: Font):
        print(f'saving font as {font_path}')
        try:
            fonts = Fonts(self.glyphs, self.glyph_res)
            os.makedirs(os.path.dirname(font_path), exist_ok=True)
            fonts.save_as_ttf(font_path, font)
        except Exception as e:
            print(f'failed to save font', e)


def create_sampler(
        out_folder: str,
        glyphs: str = None,
        checkpoint_path: str = 'models/fontogen_78958.ckpt') -> FontogenSampler:
    config = fontogen_config()
    device = "cuda" if torch.cuda.is_available() else "mps" if torch.has_mps else "cpu"
    model = FontogenModule.load_from_checkpoint(
        checkpoint_path,
        map_location=device,
        config=config,
    )
    model.eval()
    return FontogenSampler(model, config, out_folder, glyphs)


if __name__ == '__main__':
    texts = [
        'bold sans',
        'canvalove sans',
        'basic, serif, new times',
        'Beauty, Script, Calligraphy',
        'techno, sci-fi, extrabold',
        'handwritten, beauty, script',
        'soviet, block, russian',
        'serif, newspaper, news',
        'tech blog, perfect',
        'fancy, horror, scary, halloween',
        'horror scary',
        'anime, Fancy, Horror',
        'bold, cyrrilic',
        'basic,serif,harry, potter',
        'taco, bell',
        'celtic, magic',
        'elvish, runes',
        'fancy, kids',
        'Script,Handwritten,Painter',
        'medieval, gothic',
        'newspaper, typewriter',
        'comic, superman',
        'sci-fi, futuristic',
        '_',
        '',
    ]

    sampler = create_sampler('training/samples')
    for (ind, text) in enumerate(texts):
        for i in range(2):
            try:
                sampler.sample(text, -1, strategy='multinomial')
            except Exception as e:
                traceback.print_exc()
                print(f'Failed to generate font for {text}', e)

sampler.py
SerCeMan/fontogen
import glob
import os
from typing import List

import gradio as gr
from fastapi import FastAPI
from fastapi.staticfiles import StaticFiles

from config import fontogen_config
from sampler import create_sampler

# Import your sampler here. For instance:
# from your_module import sampler

FONT_DIR = "./training/samples/"
app = FastAPI()
sampler = None
config = fontogen_config()

CHECKPOINT_DIR = "."
available_checkpoints = sorted(
    [os.path.relpath(f, CHECKPOINT_DIR) for f in glob.glob(f"{CHECKPOINT_DIR}/**/*.ckpt", recursive=True)])
curr_checkpoint_path = None


def generate_font_styles(fonts):
    styles = []
    for font in fonts:
        name = font.split('.')[0]
        style_name = "FontoGen-" + name
        style = f"""
        @font-face {{
            font-family: '{style_name}';
            src: url(/fonts/{font});
        }}
        .{style_name.replace(",", "_").replace(";", "_")} {{
            font-family: '{style_name}', serif;
            font-size: 32px;
            width: 100%;
        }}
        """
        styles.append(style)
    return "\n".join(styles)


def generate_font_paragraphs(glyphs: str, fonts: List[str]):
    paragraphs = []
    for font in fonts:
        name = font.split('.')[0]
        style_name = "FontoGen-" + name
        paragraph = f"""
        <p>
            {name}
        </p>
        <p class="{style_name.replace(",", "_").replace(";", "_")}">
            {glyphs}
        </p>
        """
        paragraphs.append(paragraph)
    return "\n".join(paragraphs)


def font_generator(font_txt, temperature: int = None, glyphs: str = None, strategy: str = None,
                   checkpoint_path: str = None):
    # If a description is provided, generate a new font using the sampler
    if font_txt != '$all':
        # Call your sampler here.
        global sampler
        global curr_checkpoint_path
        if sampler is None or sampler.glyphs != glyphs or curr_checkpoint_path != checkpoint_path:
            sampler = create_sampler(FONT_DIR, glyphs=glyphs, checkpoint_path=checkpoint_path)
            curr_checkpoint_path = checkpoint_path
        font_path = sampler.sample(font_txt, -1, temperature=temperature, strategy=strategy)
        fonts = [os.path.basename(font_path)]
    else:
        fonts = sorted([f for f in os.listdir(FONT_DIR) if f.endswith('.ttf')],
                       key=lambda x: os.path.getctime(os.path.join(FONT_DIR, x)),
                       reverse=True)
    glyphs = config.glyphs

    # Return the styles and paragraphs (updated if a new font was generated)
    return f"""
    <style>
    {generate_font_styles(fonts)}
    </style>
    {generate_font_paragraphs(glyphs, fonts)}
    """


with gr.Blocks() as demo:
    title = gr.Markdown(
        f"<h1 style='text-align: center; margin-bottom: 1rem'>FontoGen</h1>"
    )
    model_dropdown = gr.Dropdown(label='Select Model Checkpoint', choices=available_checkpoints)
    custom_glyphs_input = gr.Textbox(label="Glyphs Set", value=config.glyphs)
    name = gr.Textbox(label="Font Style")
    with gr.Row():
        temperature_slider = gr.Slider(minimum=0, maximum=3.0, value=0.6, label='Temperature', interactive=True)
        sampling_strategy_dropdown = gr.Dropdown(
            label='Sampling Strategy',
            value='multinomial',
            interactive=True,
            choices=[
                'multinomial',
                'greedy',
                'topknuc',
            ])
    gen_button = gr.Button("Generate", variant='primary')
    refresh_button = gr.Button("Refresh Fonts", variant='secondary')  # New refresh button
    result = gr.HTML()
    all_fonts = gr.HTML(value=font_generator('$all'))


    @gen_button.click(inputs=[name, model_dropdown], outputs=[result])  # Update inputs here
    def generate(font_text, selected_model):
        checkpoint_path = os.path.join(CHECKPOINT_DIR, selected_model)
        gen_font_result = font_generator(
            font_text,
            temperature=temperature_slider.value,
            glyphs=custom_glyphs_input.value,
            checkpoint_path=checkpoint_path
        )
        return gen_font_result


    @refresh_button.click(outputs=[all_fonts])
    def refresh():
        return font_generator('$all')

app.mount("/fonts/", StaticFiles(directory='./samples/'), name="fonts")
app.mount("/", demo.app)

if __name__ == '__main__':
    import uvicorn

    uvicorn.run(app, host='0.0.0.0', port=8888)

server.py
SerCeMan/fontogen
import string
import unittest

from font_codec import FontCodec
from fonts import Fonts, FontScaler


class TestFontCodec(unittest.TestCase):
    def test_font_coding_decoding_ttf(self):
        glyph_res = 150
        glyphs = string.ascii_uppercase + string.digits + '.,;:!?\'"/&+@()-_'
        scaler = FontScaler(glyph_res)
        fonts = Fonts(glyphs, glyph_res)
        font = fonts.load_font('./OpenSans-Medium.ttf')
        font = scaler.normalize_font(font)

        codec = FontCodec(glyphs, glyph_res, 200 * len(glyphs), 200)
        tokens, mask = codec.encode_font(font)
        decoded_font = codec.decode_font(tokens)
        print(font)
        print(decoded_font)
        self.assertEqual(font, decoded_font)

    def test_font_coding_decoding_otf(self):
        glyph_res = 150
        glyphs = string.ascii_uppercase + string.digits + '.,;:!?\'"/&+@()-_'
        scaler = FontScaler(glyph_res)
        fonts = Fonts(glyphs, glyph_res)
        font = fonts.load_font('./Roboto-Medium.ttf')
        font = scaler.normalize_font(font)

        codec = FontCodec(glyphs, glyph_res, 200 * len(glyphs), 200)
        tokens, mask = codec.encode_font(font)
        decoded_font = codec.decode_font(tokens)
        print(font)
        print(decoded_font)
        self.assertEqual(font, decoded_font)


if __name__ == '__main__':
    unittest.main()

tests/test_codec.py
SerCeMan/fontogen
import string
import unittest

from fonts import Fonts, FontScaler, FontPlotter


class TestFontLoader(unittest.TestCase):
    def test_font_saving_ttf(self):
        fonts = [
            './Roboto-Medium.ttf',
            './OpenSans-Medium.ttf',
            # TODO: learn how to deal with Dune Rise and the Nones
            # './Dune_Rise.ttf',
        ]
        for f in fonts:
            glyph_res = 160
            glyphs = string.ascii_uppercase + string.digits + '.,;:!?\'"/&+@()-_'
            fonts = Fonts(glyphs, glyph_res)
            font = fonts.load_font(f)

            scaler = FontScaler(glyph_res)
            font = scaler.normalize_font(font)

            fonts.save_as_ttf('./converted.ttf', font)
            font2 = fonts.load_font('./converted.ttf')
            if f.endswith('.otf'):
                # With OTF fonts, cubic bézier curves have to be converted to quadratic bézier curves
                # so comparison isn't possible
                continue
            print(font.glyph_paths)
            print(font2.glyph_paths)
            self.assertEqual(font.glyph_paths, font2.glyph_paths)

    def test_font_plotting_ttf(self):
        glyph_res = 160
        fonts = Fonts('ABCD', glyph_res)
        font = fonts.load_font('./OpenSans-Medium.ttf')

        scaler = FontScaler(glyph_res)
        scaled_font = scaler.normalize_font(font)

        font_plotter = FontPlotter()
        font_plotter.plt_font(font)
        font_plotter.plt_font(scaled_font)

    def test_font_plotting_otf(self):
        glyph_res = 160
        fonts = Fonts('.,;:!?\'"/&+@()-_', glyph_res)
        font = fonts.load_font('./Roboto-Medium.ttf')

        scaler = FontScaler(glyph_res)
        scaled_font = scaler.normalize_font(font)

        font_plotter = FontPlotter()
        font_plotter.plt_font(font)
        font_plotter.plt_font(scaled_font)


if __name__ == '__main__':
    unittest.main()

tests/test_fonts.py
SerCeMan/fontogen
import argparse
import string

import lightning.pytorch as pl
import torch
from lightning.pytorch.callbacks import ModelCheckpoint, LearningRateMonitor
from lightning.pytorch.loggers import WandbLogger

from config import fontogen_config, FontogenConfig
from model.model import FontogenModule, FontogenDataModule
from sampler import SamplingCallback


# from sampler import SamplingCallback


# Install the magic combination of nightly deps for increased compatibility and maximum performance.
#
# pip install -U --index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/Triton-Nightly/pypi/simple/ triton-nightly==2.1.0.dev20230801015042 --no-deps
#

def main():
    parser = argparse.ArgumentParser(description="train")
    parser.add_argument("--dataset_path", type=str)
    args = parser.parse_args()

    config = fontogen_config()

    device = "cuda" if torch.cuda.is_available() else "mps" if torch.has_mps else "cpu"
    precision = "bf16-mixed" if device == "cuda" else 32
    use_wandb = False
    checkpoint_path = None
    dataset_path = args.dataset_path

    torch.set_float32_matmul_precision("medium")

    model = FontogenModule(config) if checkpoint_path is None else FontogenModule.load_from_checkpoint(
        checkpoint_path,
        config=config,
    )

    data_module = FontogenDataModule(config, dataset_path)

    checkpoint_callback = ModelCheckpoint(
        save_last=True,
        save_top_k=10,
        monitor='val_loss',
    )

    lr_monitor = LearningRateMonitor(logging_interval='step')
    num_epochs = 200

    trainer = pl.Trainer(
        max_epochs=num_epochs,
        accelerator=device,
        accumulate_grad_batches=16,
        gradient_clip_val=0.5,
        precision=precision,
        # enable_checkpointing=False,
        callbacks=[
            checkpoint_callback,
            lr_monitor,
            SamplingCallback(config, 10, 'training/samples'),
        ],
        logger=WandbLogger(project='fontogen3', dir='training') if use_wandb else None,
    )
    trainer.fit(
        model,
        data_module,
        # ckpt_path=checkpoint_path,
    )


if __name__ == "__main__":
    main()

train.py
SerCeMan/fontogen