AlanB commited on
Commit
e97d8f2
·
1 Parent(s): a551d9c

Added vae_slicing option

Browse files
Files changed (1) hide show
  1. pipeline.py +15 -0
pipeline.py CHANGED
@@ -542,6 +542,21 @@ class StableDiffusionLongPromptWeightingPipeline(DiffusionPipeline):
542
  # set slice_size = `None` to disable `attention slicing`
543
  self.enable_attention_slicing(None)
544
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
545
  def enable_sequential_cpu_offload(self):
546
  r"""
547
  Offloads all models to CPU using accelerate, significantly reducing memory usage. When called, unet,
 
542
  # set slice_size = `None` to disable `attention slicing`
543
  self.enable_attention_slicing(None)
544
 
545
+ def enable_vae_slicing(self):
546
+ r"""
547
+ Enable sliced VAE decoding.
548
+ When this option is enabled, the VAE will split the input tensor in slices to compute decoding in several
549
+ steps. This is useful to save some memory and allow larger batch sizes.
550
+ """
551
+ self.vae.enable_slicing()
552
+
553
+ def disable_vae_slicing(self):
554
+ r"""
555
+ Disable sliced VAE decoding. If `enable_vae_slicing` was previously invoked, this method will go back to
556
+ computing decoding in one step.
557
+ """
558
+ self.vae.disable_slicing()
559
+
560
  def enable_sequential_cpu_offload(self):
561
  r"""
562
  Offloads all models to CPU using accelerate, significantly reducing memory usage. When called, unet,