Added vae_slicing option
Browse files- 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,
|