fix

t2i_inference.py

@@ -18,6 +18,7 @@ class T2IConfig:
     decoder_path: str = 'vq_ds16_t2i.pt'
     
     device: str = "cuda:0"
+    dtype: torch.dtype = torch.bfloat16
     # generation parameters
     scale: int = 0  # choose from [0, 1, 2]
     seq_len: int = 729  # choose from [729, 169, 81]
@@ -33,7 +34,7 @@ class TextToImageInference:
         self._load_models()
         
     def _load_models(self):
-        self.model = Qwen2ForCausalLM.from_pretrained(self.config.model_path).to(self.device)
+        self.model = Qwen2ForCausalLM.from_pretrained(self.config.model_path, torch_dtype=self.config.dtype).to(self.device)
         self.tokenizer = AutoTokenizer.from_pretrained(self.config.model_path)
         
         # Initialize visual tokenizer
@@ -44,7 +45,7 @@ class TextToImageInference:
             encoder_args={'input_type': 'rec'},
             decoder_args={},
         )
-        self.visual_tokenizer = MMAutoEncoder(**config).eval().to(self.device)
+        self.visual_tokenizer = MMAutoEncoder(**config).eval().to(dtype=self.config.dtype, device=self.device)
         self.visual_tokenizer.ar_model.cls_token_num = self.config.seq_len
         self.visual_tokenizer.encoder.pool_scale = self.config.scale + 1