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,
|