Update inferencer.py

inferencer.py

@@ -49,20 +49,20 @@ class UniPicV2Inferencer:
         else:
             transformer = SD3Transformer2DKontextModel.from_pretrained(
                 self.model_path, subfolder="transformer",
-                torch_dtype=torch.float16, device_map="auto", low_cpu_mem_usage=True
+                torch_dtype=torch.bfloat16, device_map="auto", low_cpu_mem_usage=True
             )
 
         # ===== 3. Load VAE =====
         vae = AutoencoderKL.from_pretrained(
             self.model_path, subfolder="vae",
-            torch_dtype=torch.float16, device_map="auto", low_cpu_mem_usage=True
+            torch_dtype=torch.bfloat16, device_map="auto", low_cpu_mem_usage=True
         ).to(self.device)
 
         # ===== 4. Load Qwen2.5-VL (LMM) =====
         try:
             self.lmm = Qwen2_5_VLForConditionalGeneration.from_pretrained(
                 self.qwen_vl_path,
-                torch_dtype=torch.float16,
+                torch_dtype=torch.bfloat16,
                 attn_implementation="flash_attention_2",
                 device_map="auto",
             ).to(self.device)
@@ -70,7 +70,7 @@ class UniPicV2Inferencer:
         except Exception:
             self.lmm = Qwen2_5_VLForConditionalGeneration.from_pretrained(
                 self.qwen_vl_path,
-                torch_dtype=torch.float16,
+                torch_dtype=torch.bfloat16,
                 attn_implementation="sdpa",
                 device_map="auto",
             ).to(self.device)
@@ -87,7 +87,7 @@ class UniPicV2Inferencer:
         # ===== 6. Load Conditioner =====
         self.conditioner = StableDiffusion3Conditioner.from_pretrained(
             self.model_path, subfolder="conditioner",
-            torch_dtype=torch.float16, low_cpu_mem_usage=True
+            torch_dtype=torch.bfloat16, low_cpu_mem_usage=True
         ).to(self.device)
 
         # ===== 7. Load Scheduler =====