Update pipeline.py

pipeline.py

@@ -625,35 +625,20 @@ class FluxWithCFGPipeline(DiffusionPipeline, FluxLoraLoaderMixin, FromSingleFile
                     if torch.backends.mps.is_available():
                         # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
                         latents = latents.to(latents_dtype)
-                        
-                if callback_on_step_end is not None:
-                    callback_kwargs = {}
-                    for k in callback_on_step_end_tensor_inputs:
-                        callback_kwargs[k] = locals()[k]
-                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
-
-                    latents = callback_outputs.pop("latents", latents)
-                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
 
+            # call the callback, if provided
                 if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
                     progress_bar.update()
-                    
-                if XLA_AVAILABLE:
-                    xm.mark_step()
-
-        if output_type == "latent":
-            image = latents
-            
-        else:
-            latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
-            latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
-            image = self.vae.decode(latents, return_dict=False)[0]
-            image = self.image_processor.postprocess(image, output_type=output_type)
 
-        # Offload all models
+        # Final image
+        return self._decode_latents_to_image(latents, height, width, output_type)
         self.maybe_free_model_hooks()
-
-        if not return_dict:
-            return (image,)
-
-        return FluxPipelineOutput(images=image)
+        torch.cuda.empty_cache()
+
+    def _decode_latents_to_image(self, latents, height, width, output_type, vae=None):
+        """Decodes the given latents into an image."""
+        vae = vae or self.vae
+        latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
+        latents = (latents / vae.config.scaling_factor) + vae.config.shift_factor
+        image = vae.decode(latents, return_dict=False)[0]
+        return self.image_processor.postprocess(image, output_type=output_type)[0]