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)

""" 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

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:

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