refactor: Update to() method in FrozenCLIPEmbedderT3 and TextEmbeddingModule

Author: tolgacangoz

examples/research_projects/anytext/frozen_clip_embedder_t3.py

@@ -207,3 +207,8 @@ def split_chunks(self, input_ids, chunk_size=75):
             remaining_group_pad = torch.cat((id_start, remaining_group, padding, id_end), dim=1)
             tokens_list.append(remaining_group_pad)
         return tokens_list
+
+    def to(self, *args, **kwargs):
+        self.transformer = self.transformer.to(*args, **kwargs)
+        self.device = self.transformer.device
+        return self

examples/research_projects/anytext/text_embedding_module.py

@@ -297,9 +297,9 @@ def insert_spaces(self, string, nSpace):
             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)
+    def to(self, *args, **kwargs):
+        self.frozen_CLIP_embedder_t3 = self.frozen_CLIP_embedder_t3.to(*args, **kwargs)
+        self.embedding_manager = self.embedding_manager.to(*args, **kwargs)
+        self.text_predictor = self.text_predictor.to(*args, **kwargs)
+        self.device = self.frozen_CLIP_embedder_t3.device
         return self