Fix device issues

Author: tolgacangoz

examples/research_projects/anytext/auxiliary_latent_module.py

@@ -39,7 +39,7 @@ def __init__(self, glyph_channels=1, position_channels=1, model_channels=320, **
         super().__init__()
         self.font = ImageFont.truetype("Arial_Unicode.ttf", 60)
         self.use_fp16 = kwargs.get("use_fp16", False)
-        self.device = kwargs.get("device", "cpu")
+        self.device = kwargs.get("device", "cuda")
 
         self.glyph_block = nn.Sequential(
             nn.Conv2d(glyph_channels, 8, 3, padding=1),

examples/research_projects/anytext/embedding_manager.py

@@ -134,6 +134,7 @@ def __init__(
             self.position_encoder = EncodeNet(position_channels, token_dim)
         if emb_type == "ocr":
             self.proj = nn.Sequential(zero_module(nn.Linear(40 * 64, token_dim)), nn.LayerNorm(token_dim))
+            self.proj = self.proj.to(dtype=torch.float16 if kwargs.get("use_fp16", False) else torch.float32)
         if emb_type == "conv":
             self.glyph_encoder = EncodeNet(glyph_channels, token_dim)
 

examples/research_projects/anytext/frozen_clip_embedder_t3.py

@@ -20,7 +20,7 @@ def __init__(
     ):
         super().__init__()
         self.tokenizer = CLIPTokenizer.from_pretrained(version)
-        self.transformer = CLIPTextModel.from_pretrained(version).to(device)
+        self.transformer = CLIPTextModel.from_pretrained(version, torch_dtype=torch.float16).to(device)
         if use_vision:
             self.vit = CLIPVisionModelWithProjection.from_pretrained(version)
             self.processor = AutoProcessor.from_pretrained(version)

examples/research_projects/anytext/pipeline_anytext.py

@@ -1172,6 +1172,7 @@ def __call__(
                 num_images_per_prompt=num_images_per_prompt,
                 np_hint=np_hint,
             )
+            height, width = 512, 512
         # elif isinstance(controlnet, MultiControlNetModel):
         #     images = []
 

examples/research_projects/anytext/text_embedding_module.py

@@ -30,6 +30,7 @@ def __init__(self, use_fp16):
             "position_channels": 1,
             "add_pos": False,
             "placeholder_string": "*",
+            "use_fp16": self.use_fp16,
         }
         self.embedding_manager = EmbeddingManager(self.frozen_CLIP_embedder_t3, **self.embedding_manager_config)
         # TODO: Understand the reason of param.requires_grad = True
@@ -41,8 +42,10 @@ def __init__(self, use_fp16):
         args["rec_image_shape"] = "3, 48, 320"
         args["rec_batch_num"] = 6
         args["rec_char_dict_path"] = "ppocr_keys_v1.txt"
-        args["use_fp16"] = False
-        self.cn_recognizer = TextRecognizer(args, self.text_predictor)
+        args["use_fp16"] = True
+        self.cn_recognizer = TextRecognizer(
+            args, self.text_predictor.to(dtype=torch.float16 if use_fp16 else torch.float32)
+        )
         for param in self.text_predictor.parameters():
             param.requires_grad = False
         self.embedding_manager.recog = self.cn_recognizer
@@ -149,9 +152,9 @@ def forward(
                 glyphs = np.zeros((h * gly_scale, w * gly_scale, 1))
                 gly_line = np.zeros((80, 512, 1))
             pos = pre_pos[i]
-            text_info["glyphs"] += [self.arr2tensor(glyphs, len(prompt))]
-            text_info["gly_line"] += [self.arr2tensor(gly_line, len(prompt))]
-            text_info["positions"] += [self.arr2tensor(pos, len(prompt))]
+            text_info["glyphs"] += [self.arr2tensor(glyphs, num_images_per_prompt)]
+            text_info["gly_line"] += [self.arr2tensor(gly_line, num_images_per_prompt)]
+            text_info["positions"] += [self.arr2tensor(pos, num_images_per_prompt)]
 
         # hint = self.arr2tensor(np_hint, len(prompt))