refactor: Simplify code for loading models and handling data types

Author: tolgacangoz

examples/research_projects/anytext/auxiliary_latent_module.py

@@ -29,10 +29,10 @@ def retrieve_latents(
 
 class AuxiliaryLatentModule(nn.Module):
     def __init__(
-        self, glyph_channels=1, position_channels=1, model_channels=320, vae=None, device="cpu", use_fp16=False
+        self, font_path, 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.font = ImageFont.truetype(font_path, 60)
         self.use_fp16 = use_fp16
         self.device = device
 
@@ -79,12 +79,12 @@ def __init__(
         self.fuse_block = nn.Conv2d(256 + 64 + 4, model_channels, 3, padding=1)
 
         self.glyph_block.load_state_dict(
-            load_file("AuxiliaryLatentModule/glyph_block.safetensors", device=self.device)
+            load_file("glyph_block.safetensors", device=str(self.device))
         )
         self.position_block.load_state_dict(
-            load_file("AuxiliaryLatentModule/position_block.safetensors", device=self.device)
+            load_file("position_block.safetensors", device=str(self.device))
         )
-        self.fuse_block.load_state_dict(load_file("AuxiliaryLatentModule/fuse_block.safetensors", device=self.device))
+        self.fuse_block.load_state_dict(load_file("fuse_block.safetensors", device=str(self.device)))
 
         if use_fp16:
             self.glyph_block = self.glyph_block.to(dtype=torch.float16)

examples/research_projects/anytext/embedding_manager.py

@@ -46,7 +46,7 @@ def __init__(
         self.token_dim = token_dim
 
         self.proj = nn.Linear(40 * 64, token_dim)
-        self.proj.load_state_dict(load_file("EmbeddingManager/embedding_manager.safetensors", device=self.device))
+        self.proj.load_state_dict(load_file("proj.safetensors", device=str(embedder.device)))
         if use_fp16:
             self.proj = self.proj.to(dtype=torch.float16)
 
@@ -65,7 +65,7 @@ def encode_text(self, text_info):
 
         if len(gline_list) > 0:
             recog_emb = self.get_recog_emb(gline_list)
-            enc_glyph = self.proj(recog_emb.reshape(recog_emb.shape[0], -1))
+            enc_glyph = self.proj(recog_emb.reshape(recog_emb.shape[0], -1).to(self.proj.weight.device))
 
         self.text_embs_all = []
         n_idx = 0

examples/research_projects/anytext/ocr_recog/RecModel.py

@@ -34,6 +34,8 @@ def load_3rd_state_dict(self, _3rd_name, _state):
         self.head.load_3rd_state_dict(_3rd_name, _state)
 
     def forward(self, x):
+        import torch
+        x = x.to(torch.float32)
         x = self.backbone(x)
         x = self.neck(x)
         x = self.head(x)

examples/research_projects/anytext/pipeline_anytext.py

@@ -208,6 +208,7 @@ class AnyTextPipeline(
 
     def __init__(
         self,
+        font_path: str,
         vae: AutoencoderKL,
         text_encoder: CLIPTextModel,
         tokenizer: CLIPTokenizer,
@@ -218,7 +219,6 @@ def __init__(
         feature_extractor: CLIPImageProcessor,
         image_encoder: CLIPVisionModelWithProjection = None,
         requires_safety_checker: bool = True,
-        font_path: str = "font/Arial_Unicode.ttf",
     ):
         super().__init__()
         self.text_embedding_module = TextEmbeddingModule(
@@ -257,13 +257,15 @@ def __init__(
             safety_checker=safety_checker,
             feature_extractor=feature_extractor,
             image_encoder=image_encoder,
+            # text_embedding_module=text_embedding_module,
+            # auxiliary_latent_module=auxiliary_latent_module,
         )
         self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
         self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor, do_convert_rgb=True)
         self.control_image_processor = VaeImageProcessor(
             vae_scale_factor=self.vae_scale_factor, do_convert_rgb=True, do_normalize=False
         )
-        self.register_to_config(requires_safety_checker=requires_safety_checker)
+        self.register_to_config(requires_safety_checker=requires_safety_checker, font_path=font_path)
 
     def modify_prompt(self, prompt):
         prompt = prompt.replace("“", '"')

examples/research_projects/anytext/recognizer.py

@@ -98,9 +98,7 @@ def create_predictor(model_dir=None, model_lang="ch", device="cpu", use_fp16=Fal
 
     rec_model = RecModel(rec_config)
     if model_file_path is not None:
-        rec_model.load_state_dict(load_file(model_file_path, device=device)).to(
-            dtype=torch.float16 if use_fp16 else torch.float32
-        )
+        rec_model.load_state_dict(torch.load(model_file_path, map_location=device))
     return rec_model
 
 

examples/research_projects/anytext/text_embedding_module.py

@@ -24,10 +24,9 @@ def __init__(self, font_path, use_fp16=False, device="cpu"):
         self.font = ImageFont.truetype(font_path, 60)
         self.frozen_CLIP_embedder_t3 = FrozenCLIPEmbedderT3(device=self.device, use_fp16=self.use_fp16)
         self.embedding_manager = EmbeddingManager(self.frozen_CLIP_embedder_t3, use_fp16=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 = "OCR/ppv3_rec.safetensors"
+        # for param in self.embedding_manager.embedding_parameters():
+        #     param.requires_grad = True
+        rec_model_dir = "OCR/ppv3_rec.pth"
         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
@@ -36,8 +35,7 @@ def __init__(self, font_path, use_fp16=False, device="cpu"):
         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)
-        self.embedding_manager.recog = self.cn_recognizer
+        self.embedding_manager.recog = TextRecognizer(args, self.text_predictor)
 
     @torch.no_grad()
     def forward(
@@ -290,3 +288,18 @@ def draw_glyph2(self, font, text, polygon, vertAng=10, scale=1, width=512, heigh
         img.paste(rotated_layer, (x_offset, y_offset), rotated_layer)
         img = np.expand_dims(np.array(img.convert("1")), axis=2).astype(np.float64)
         return img
+
+    def insert_spaces(self, string, nSpace):
+        if nSpace == 0:
+            return string
+        new_string = ""
+        for char in string:
+            new_string += char + " " * nSpace
+        return new_string[:-nSpace]
+
+    def to(self, device):
+        self.device = device
+        self.frozen_CLIP_embedder_t3.to(device)
+        self.embedding_manager.to(device)
+        self.text_predictor.to(device)
+        return self