KingNish commited on
Commit
38f7479
1 Parent(s): bb27283

Update custom_pipeline.py

Browse files
Files changed (1) hide show
  1. custom_pipeline.py +1 -1
custom_pipeline.py CHANGED
@@ -155,7 +155,7 @@ class FLUXPipelineWithIntermediateOutputs(FluxPipeline):
155
  torch.cuda.empty_cache()
156
 
157
  # Final image
158
- yield self._decode_latents_to_image(latents, height, width, output_type)
159
  self.maybe_free_model_hooks()
160
  torch.cuda.empty_cache()
161
 
 
155
  torch.cuda.empty_cache()
156
 
157
  # Final image
158
+ return self._decode_latents_to_image(latents, height, width, output_type)
159
  self.maybe_free_model_hooks()
160
  torch.cuda.empty_cache()
161