Up

Author: tolgacangoz

examples/research_projects/anytext/auxiliary_latent_module.py

@@ -35,11 +35,13 @@ def retrieve_latents(
 
 
 class AuxiliaryLatentModule(nn.Module):
-    def __init__(self, glyph_channels=1, position_channels=1, model_channels=320, **kwargs):
+    def __init__(
+        self, glyph_channels=1, position_channels=1, model_channels=320, vae=None, device="cpu", use_fp16=False
+    ):
         super().__init__()
         self.font = ImageFont.truetype("font/Arial_Unicode.ttf", 60)
-        self.use_fp16 = kwargs.get("use_fp16", False)
-        self.device = kwargs.get("device", "cpu")
+        self.use_fp16 = use_fp16
+        self.device = device
 
         self.glyph_block = nn.Sequential(
             nn.Conv2d(glyph_channels, 8, 3, padding=1),
@@ -79,8 +81,7 @@ def __init__(self, glyph_channels=1, position_channels=1, model_channels=320, **
             nn.SiLU(),
         )
 
-        self.vae = kwargs.get("vae")
-        self.vae.eval()
+        self.vae = vae.eval()
 
         self.fuse_block = zero_module(nn.Conv2d(256 + 64 + 4, model_channels, 3, padding=1))
 

examples/research_projects/anytext/text_embedding_module.py

@@ -28,15 +28,15 @@ def __init__(self, font_path, use_fp16=False, device="cpu"):
         for param in self.embedding_manager.embedding_parameters():
             param.requires_grad = True
         rec_model_dir = "OCR/ppv3_rec.safetensors"
-        self.text_predictor = create_predictor(rec_model_dir, device=self.device, use_fp16=self.use_fp16)
+        self.text_predictor = create_predictor(rec_model_dir, device=self.device, use_fp16=self.use_fp16).eval()
+        for param in self.text_predictor.parameters():
+            param.requires_grad = False
         args = {}
         args["rec_image_shape"] = "3, 48, 320"
         args["rec_batch_num"] = 6
         args["rec_char_dict_path"] = "OCR/ppocr_keys_v1.txt"
         args["use_fp16"] = self.use_fp16
-        self.cn_recognizer = TextRecognizer(args, self.text_predictor, device=self.device, use_fp16=self.use_fp16)
-        for param in self.text_predictor.parameters():
-            param.requires_grad = False
+        self.cn_recognizer = TextRecognizer(args, self.text_predictor)
         self.embedding_manager.recog = self.cn_recognizer
 
     @torch.no_grad()