Unverified Commit 803e817e authored by Zoltan's avatar Zoltan Committed by GitHub
Browse files

Add vae slicing and tiling to flux pipeline (#9122)



add vae slicing and tiling to flux pipeline
Co-authored-by: default avatarSayak Paul <spsayakpaul@gmail.com>
parent 67f5cce2
...@@ -454,6 +454,35 @@ class FluxPipeline(DiffusionPipeline, FluxLoraLoaderMixin): ...@@ -454,6 +454,35 @@ class FluxPipeline(DiffusionPipeline, FluxLoraLoaderMixin):
return latents return latents
def enable_vae_slicing(self):
r"""
Enable sliced VAE decoding. When this option is enabled, the VAE will split the input tensor in slices to
compute decoding in several steps. This is useful to save some memory and allow larger batch sizes.
"""
self.vae.enable_slicing()
def disable_vae_slicing(self):
r"""
Disable sliced VAE decoding. If `enable_vae_slicing` was previously enabled, this method will go back to
computing decoding in one step.
"""
self.vae.disable_slicing()
def enable_vae_tiling(self):
r"""
Enable tiled VAE decoding. When this option is enabled, the VAE will split the input tensor into tiles to
compute decoding and encoding in several steps. This is useful for saving a large amount of memory and to allow
processing larger images.
"""
self.vae.enable_tiling()
def disable_vae_tiling(self):
r"""
Disable tiled VAE decoding. If `enable_vae_tiling` was previously enabled, this method will go back to
computing decoding in one step.
"""
self.vae.disable_tiling()
def prepare_latents( def prepare_latents(
self, self,
batch_size, batch_size,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment