fix: use FP32 dtype on CPU (no bfloat16 autocast with INT8 model)

infer/inference.py

@@ -442,10 +442,11 @@ class ImageGenerator:
                 self.ae = self.ae.cpu()
                 cudagc()
 
+        _dtype = torch.float32 if self.device.type == "cpu" else torch.bfloat16
         if args is not None and hasattr(args, 'single_denoise') and not args.single_denoise:
-            x = torch.randn(num_samples,16,height // 8,width // 8,device=self.device,dtype=torch.bfloat16,generator=torch.Generator(device=self.device).manual_seed(seed),)
+            x = torch.randn(num_samples,16,height // 8,width // 8,device=self.device,dtype=_dtype,generator=torch.Generator(device=self.device).manual_seed(seed),)
         else:
-            x= torch.zeros(num_samples,16,height // 8,width // 8,device=self.device,dtype=torch.bfloat16,)
+            x= torch.zeros(num_samples,16,height // 8,width // 8,device=self.device,dtype=_dtype,)
 
 
         timesteps = sampling.get_schedule(num_steps, x.shape[-1] * x.shape[-2] // 4, shift=True)
@@ -470,7 +471,7 @@ class ImageGenerator:
             ref_image_raw=ref_images_raw,
             empty_llm=empty_llm)
 
-        with torch.autocast(device_type=self.device.type,dtype=torch.bfloat16):
+        with torch.autocast(device_type=self.device.type, dtype=torch.float32) if self.device.type == "cpu" else torch.autocast(device_type=self.device.type, dtype=torch.bfloat16):
             # Lpred,Rpred = self.double_denoise(**inputs,cfg_guidance=cfg_guidance,timesteps=timesteps,height=height,width=width)#图像中包括ref image
             Lpred,Rpred = self.denoise(**inputs,cfg_guidance=cfg_guidance,timesteps=timesteps,show_progress=show_progress,timesteps_truncate=1.0,)#图像中包括ref image
             Lpred=self.unpack(Lpred.float(),height,width)