Up

Author: tolgacangoz

examples/research_projects/anytext/auxiliary_latent_module.py

@@ -4,6 +4,7 @@
 import numpy as np
 import torch
 from PIL import ImageFont
+from safetensors.torch import load_file
 from torch import nn
 
 from diffusers.utils import logging
@@ -34,18 +35,12 @@ def retrieve_latents(
 
 
 class AuxiliaryLatentModule(nn.Module):
-    def __init__(self, dims=2, glyph_channels=1, position_channels=1, model_channels=320, **kwargs):
+    def __init__(self, glyph_channels=1, position_channels=1, model_channels=320, **kwargs):
         super().__init__()
-        self.font = ImageFont.truetype("/home/cosmos/Documents/gits/AnyText/font/Arial_Unicode.ttf", 60)
+        self.font = ImageFont.truetype("Arial_Unicode.ttf", 60)
         self.use_fp16 = kwargs.get("use_fp16", False)
         self.device = kwargs.get("device", "cpu")
-        self.model_channels = model_channels
-        time_embed_dim = model_channels * 4
-        self.time_embed = nn.Sequential(
-            nn.Linear(model_channels, time_embed_dim),
-            nn.SiLU(),
-            nn.Linear(time_embed_dim, time_embed_dim),
-        )
+
         self.glyph_block = nn.Sequential(
             nn.Conv2d(glyph_channels, 8, 3, padding=1),
             nn.SiLU(),
@@ -83,20 +78,21 @@ def __init__(self, dims=2, glyph_channels=1, position_channels=1, model_channels
             nn.Conv2d(32, 64, 3, padding=1, stride=2),
             nn.SiLU(),
         )
-        self.time_embed = self.time_embed.to(device="cuda", dtype=torch.float16)
+        self.glyph_block.load_state_dict(load_file("glyph_block.safetensors"))
+        self.position_block.load_state_dict(load_file("position_block.safetensors"))
         self.glyph_block = self.glyph_block.to(device="cuda", dtype=torch.float16)
         self.position_block = self.position_block.to(device="cuda", dtype=torch.float16)
 
         self.vae = kwargs.get("vae")
         self.vae.eval()
 
         self.fuse_block = zero_module(nn.Conv2d(256 + 64 + 4, model_channels, 3, padding=1))
+        self.fuse_block.load_state_dict(load_file("fuse_block.safetensors"))
         self.fuse_block = self.fuse_block.to(device="cuda", dtype=torch.float16)
 
     @torch.no_grad()
     def forward(
         self,
-        context,
         text_info,
         mode,
         draw_pos,

examples/research_projects/anytext/pipeline_anytext.py

@@ -1165,7 +1165,6 @@ def __call__(
             # )
             # height, width = image.shape[-2:]
             guided_hint = self.auxiliary_latent_module(
-                context=prompt_embeds[1],
                 text_info=text_info,
                 mode=mode,
                 draw_pos=draw_pos,

examples/research_projects/anytext/text_embedding_module.py

@@ -21,7 +21,7 @@ def __init__(self, use_fp16):
         self.use_fp16 = use_fp16
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         # TODO: Learn if the recommended font file is free to use
-        self.font = ImageFont.truetype("/home/cosmos/Documents/gits/AnyText/font/Arial_Unicode.ttf", 60)
+        self.font = ImageFont.truetype("Arial_Unicode.ttf", 60)
         self.frozen_CLIP_embedder_t3 = FrozenCLIPEmbedderT3(device=self.device)
         self.embedding_manager_config = {
             "valid": True,
@@ -35,12 +35,12 @@ def __init__(self, use_fp16):
         # TODO: Understand the reason of param.requires_grad = True
         for param in self.embedding_manager.embedding_parameters():
             param.requires_grad = True
-        rec_model_dir = "/home/cosmos/Documents/gits/AnyText/ocr_weights/ppv3_rec.pth"
+        rec_model_dir = "ppv3_rec.pth"
         self.text_predictor = create_predictor(rec_model_dir).eval()
         args = {}
         args["rec_image_shape"] = "3, 48, 320"
         args["rec_batch_num"] = 6
-        args["rec_char_dict_path"] = "/home/cosmos/Documents/gits/AnyText/ocr_weights/ppocr_keys_v1.txt"
+        args["rec_char_dict_path"] = "ppocr_keys_v1.txt"
         args["use_fp16"] = False
         self.cn_recognizer = TextRecognizer(args, self.text_predictor)
         for param in self.text_predictor.parameters():