Use torch.autocast for automatic mixed precision inference

src/flux/xflux_pipeline.py

@@ -365,44 +365,46 @@ class XFluxPipeline:
             if neg_image_proj is not None:
                 neg_image_proj = neg_image_proj.to(inference_dtype)
 
-            if self.controlnet_loaded:
-                x = denoise_controlnet(
-                    self.model,
-                    **inp_cond,
-                    controlnet=self.controlnet,
-                    timesteps=timesteps,
-                    guidance=guidance,
-                    controlnet_cond=controlnet_image,
-                    timestep_to_start_cfg=timestep_to_start_cfg,
-                    neg_txt=neg_inp_cond['txt'],
-                    neg_txt_ids=neg_inp_cond['txt_ids'],
-                    neg_vec=neg_inp_cond['vec'],
-                    true_gs=true_gs,
-                    controlnet_gs=control_weight,
-                    image_proj=image_proj,
-                    neg_image_proj=neg_image_proj,
-                    ip_scale=ip_scale,
-                    neg_ip_scale=neg_ip_scale,
-                )
-            else:
-                x = denoise(
-                    self.model,
-                    **inp_cond,
-                    timesteps=timesteps,
-                    guidance=guidance,
-                    cond_latent=cond_latent,
-                    cond_txt_latent=cond_txt_latent,
-                    timestep_to_start_cfg=timestep_to_start_cfg,
-                    neg_txt=neg_inp_cond['txt'],
-                    neg_txt_ids=neg_inp_cond['txt_ids'],
-                    neg_vec=neg_inp_cond['vec'],
-                    true_gs=true_gs,
-                    image_proj=image_proj,
-                    neg_image_proj=neg_image_proj,
-                    ip_scale=ip_scale,
-                    neg_ip_scale=neg_ip_scale,
-                    is_generation=is_generation,
-                )
+            # Use autocast for automatic mixed precision - handles fp16/fp32 fallback
+            with torch.autocast(device_type='cuda', dtype=torch.float16):
+                if self.controlnet_loaded:
+                    x = denoise_controlnet(
+                        self.model,
+                        **inp_cond,
+                        controlnet=self.controlnet,
+                        timesteps=timesteps,
+                        guidance=guidance,
+                        controlnet_cond=controlnet_image,
+                        timestep_to_start_cfg=timestep_to_start_cfg,
+                        neg_txt=neg_inp_cond['txt'],
+                        neg_txt_ids=neg_inp_cond['txt_ids'],
+                        neg_vec=neg_inp_cond['vec'],
+                        true_gs=true_gs,
+                        controlnet_gs=control_weight,
+                        image_proj=image_proj,
+                        neg_image_proj=neg_image_proj,
+                        ip_scale=ip_scale,
+                        neg_ip_scale=neg_ip_scale,
+                    )
+                else:
+                    x = denoise(
+                        self.model,
+                        **inp_cond,
+                        timesteps=timesteps,
+                        guidance=guidance,
+                        cond_latent=cond_latent,
+                        cond_txt_latent=cond_txt_latent,
+                        timestep_to_start_cfg=timestep_to_start_cfg,
+                        neg_txt=neg_inp_cond['txt'],
+                        neg_txt_ids=neg_inp_cond['txt_ids'],
+                        neg_vec=neg_inp_cond['vec'],
+                        true_gs=true_gs,
+                        image_proj=image_proj,
+                        neg_image_proj=neg_image_proj,
+                        ip_scale=ip_scale,
+                        neg_ip_scale=neg_ip_scale,
+                        is_generation=is_generation,
+                    )
 
             if self.offload:
                 self.offload_model_to_cpu(self.model)