suvadityamuk commited on
Commit
706d9f4
1 Parent(s): a5edd99

Update pipeline.py

Browse files
Files changed (1) hide show
  1. pipeline.py +8 -2
pipeline.py CHANGED
@@ -119,8 +119,14 @@ class StableDiffusionComparisonPipeline(DiffusionPipeline):
119
  if slice_size == "auto":
120
  # half the attention head size is usually a good trade-off between
121
  # speed and memory
122
- slice_size = self.unet.config.attention_head_dim // 2
123
- self.unet.set_attention_slice(slice_size)
 
 
 
 
 
 
124
 
125
  def disable_attention_slicing(self):
126
  r"""
 
119
  if slice_size == "auto":
120
  # half the attention head size is usually a good trade-off between
121
  # speed and memory
122
+ slice_size_1 = self.pipe1.unet.config.attention_head_dim // 2
123
+ slice_size_2 = self.pipe2.unet.config.attention_head_dim // 2
124
+ slice_size_3 = self.pipe3.unet.config.attention_head_dim // 2
125
+ slice_size_4 = self.pipe4.unet.config.attention_head_dim // 2
126
+ self.pipe1.unet.set_attention_slice(slice_size_1)
127
+ self.pipe2.unet.set_attention_slice(slice_size_2)
128
+ self.pipe3.unet.set_attention_slice(slice_size_3)
129
+ self.pipe4.unet.set_attention_slice(slice_size_4)
130
 
131
  def disable_attention_slicing(self):
132
  r"""