From 842ac46b7a8043f8b613a264cebe8eb88ded8f7d Mon Sep 17 00:00:00 2001
From: Patrick von Platen <patrick.v.platen@gmail.com>
Date: Mon, 5 Sep 2022 17:12:10 +0000
Subject: [PATCH 1/7] up

---
 src/diffusers/models/attention.py | 33 ++++++++++++++++++++-----------
 1 file changed, 22 insertions(+), 11 deletions(-)

diff --git a/src/diffusers/models/attention.py b/src/diffusers/models/attention.py
index 8d52ee9bde92..e4336ef3df1b 100644
--- a/src/diffusers/models/attention.py
+++ b/src/diffusers/models/attention.py
@@ -63,18 +63,19 @@ def forward(self, hidden_states):
 
         # get scores
         scale = 1 / math.sqrt(math.sqrt(self.channels / self.num_heads))
+
         attention_scores = torch.matmul(query_states * scale, key_states.transpose(-1, -2) * scale)
         attention_probs = torch.softmax(attention_scores.float(), dim=-1).type(attention_scores.dtype)
 
         # compute attention output
-        context_states = torch.matmul(attention_probs, value_states)
+        hidden_states = torch.matmul(attention_probs, value_states)
 
-        context_states = context_states.permute(0, 2, 1, 3).contiguous()
-        new_context_states_shape = context_states.size()[:-2] + (self.channels,)
-        context_states = context_states.view(new_context_states_shape)
+        hidden_states = hidden_states.permute(0, 2, 1, 3).contiguous()
+        new_hidden_states_shape = hidden_states.size()[:-2] + (self.channels,)
+        hidden_states = hidden_states.view(new_hidden_states_shape)
 
         # compute next hidden_states
-        hidden_states = self.proj_attn(context_states)
+        hidden_states = self.proj_attn(hidden_states)
         hidden_states = hidden_states.transpose(-1, -2).reshape(batch, channel, height, width)
 
         # res connect and rescale
@@ -151,6 +152,7 @@ def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.
 
         self.scale = dim_head**-0.5
         self.heads = heads
+        self.chunk_size = None
 
         self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
         self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
@@ -186,20 +188,29 @@ def forward(self, x, context=None, mask=None):
         k = self.reshape_heads_to_batch_dim(k)
         v = self.reshape_heads_to_batch_dim(v)
 
-        sim = torch.einsum("b i d, b j d -> b i j", q, k) * self.scale
-
         if mask is not None:
+            # TODO(PVP) - this will fail when we supply a mask
+            # which we don't do yet though
             mask = mask.reshape(batch_size, -1)
             max_neg_value = -torch.finfo(sim.dtype).max
             mask = mask[:, None, :].repeat(h, 1, 1)
             sim.masked_fill_(~mask, max_neg_value)
 
+        batch_size_attention = q.shape[0]
+        chunk_size = self.chunk_size if self.chunk_size is not None else q.shape[0]
+        chunk_size = 1
+
         # attention, what we cannot get enough of
-        attn = sim.softmax(dim=-1)
+        hidden_states = torch.zeros((batch_size_attention, sequence_length, dim // self.heads), device=q.device, dtype=q.dtype)
+        for i in range(hidden_states.shape[0] // chunk_size):
+            attn_slice = torch.einsum("b i d, b j d -> b i j", q[i:i+chunk_size], k[i:i+chunk_size]) * self.scale
+            attn_slice = attn_slice.softmax(dim=-1)
+            attn_slice = torch.einsum("b i j, b j d -> b i d", attn_slice, v[i:i+chunk_size])
+
+            hidden_states[i:i+chunk_size] = attn_slice
 
-        out = torch.einsum("b i j, b j d -> b i d", attn, v)
-        out = self.reshape_batch_dim_to_heads(out)
-        return self.to_out(out)
+        hidden_states = self.reshape_batch_dim_to_heads(hidden_states)
+        return self.to_out(hidden_states)
 
 
 class FeedForward(nn.Module):

From c5599fab58a4694efaf11f299ad0a48c1ff83cf9 Mon Sep 17 00:00:00 2001
From: Patrick von Platen <patrick.v.platen@gmail.com>
Date: Tue, 6 Sep 2022 08:39:00 +0000
Subject: [PATCH 2/7] add tests

---
 src/diffusers/models/attention.py             | 51 +++++++++-----
 src/diffusers/models/unet_2d_condition.py     | 22 ++++++
 src/diffusers/models/unet_blocks.py           | 48 +++++++++++++
 .../pipeline_stable_diffusion.py              |  8 +++
 tests/test_pipelines.py                       | 70 +++++++++++++++++++
 5 files changed, 180 insertions(+), 19 deletions(-)

diff --git a/src/diffusers/models/attention.py b/src/diffusers/models/attention.py
index e4336ef3df1b..134dbbb27f30 100644
--- a/src/diffusers/models/attention.py
+++ b/src/diffusers/models/attention.py
@@ -108,6 +108,10 @@ def __init__(self, in_channels, n_heads, d_head, depth=1, dropout=0.0, context_d
 
         self.proj_out = nn.Conv2d(inner_dim, in_channels, kernel_size=1, stride=1, padding=0)
 
+    def _set_attention_chunk_size(self, chunk_size):
+        for block in self.transformer_blocks:
+            block._set_attention_chunk_size(chunk_size)
+
     def forward(self, x, context=None):
         # note: if no context is given, cross-attention defaults to self-attention
         b, c, h, w = x.shape
@@ -137,6 +141,10 @@ def __init__(self, dim, n_heads, d_head, dropout=0.0, context_dim=None, gated_ff
         self.norm3 = nn.LayerNorm(dim)
         self.checkpoint = checkpoint
 
+    def _set_attention_chunk_size(self, chunk_size):
+        self.attn1._chunk_size = chunk_size
+        self.attn2._chunk_size = chunk_size
+
     def forward(self, x, context=None):
         x = self.attn1(self.norm1(x)) + x
         x = self.attn2(self.norm2(x), context=context) + x
@@ -152,7 +160,10 @@ def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.
 
         self.scale = dim_head**-0.5
         self.heads = heads
-        self.chunk_size = None
+        # for chunk_size > 0 the attention score computation
+        # is split across the batch axis to save memory
+        # You can set chunk_size with `set_attention_chunk_size`
+        self._chunk_size = None
 
         self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
         self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
@@ -177,8 +188,6 @@ def reshape_batch_dim_to_heads(self, tensor):
     def forward(self, x, context=None, mask=None):
         batch_size, sequence_length, dim = x.shape
 
-        h = self.heads
-
         q = self.to_q(x)
         context = context if context is not None else x
         k = self.to_k(context)
@@ -188,29 +197,33 @@ def forward(self, x, context=None, mask=None):
         k = self.reshape_heads_to_batch_dim(k)
         v = self.reshape_heads_to_batch_dim(v)
 
-        if mask is not None:
-            # TODO(PVP) - this will fail when we supply a mask
-            # which we don't do yet though
-            mask = mask.reshape(batch_size, -1)
-            max_neg_value = -torch.finfo(sim.dtype).max
-            mask = mask[:, None, :].repeat(h, 1, 1)
-            sim.masked_fill_(~mask, max_neg_value)
-
-        batch_size_attention = q.shape[0]
-        chunk_size = self.chunk_size if self.chunk_size is not None else q.shape[0]
-        chunk_size = 1
+        # TODO(PVP) - mask is currently never used. Remember to re-implement when used
 
         # attention, what we cannot get enough of
-        hidden_states = torch.zeros((batch_size_attention, sequence_length, dim // self.heads), device=q.device, dtype=q.dtype)
+        hidden_states = self._attention(q, k, v, sequence_length, dim)
+
+        return self.to_out(hidden_states)
+
+    def _attention(self, query, key, value, sequence_length, dim):
+        batch_size_attention = query.shape[0]
+        hidden_states = torch.zeros(
+            (batch_size_attention, sequence_length, dim // self.heads), device=query.device, dtype=query.dtype
+        )
+        chunk_size = self._chunk_size if self._chunk_size is not None else hidden_states.shape[0]
         for i in range(hidden_states.shape[0] // chunk_size):
-            attn_slice = torch.einsum("b i d, b j d -> b i j", q[i:i+chunk_size], k[i:i+chunk_size]) * self.scale
+            start_idx = i * chunk_size
+            end_idx = (i + 1) * chunk_size
+            attn_slice = (
+                torch.einsum("b i d, b j d -> b i j", query[start_idx:end_idx], key[start_idx:end_idx]) * self.scale
+            )
             attn_slice = attn_slice.softmax(dim=-1)
-            attn_slice = torch.einsum("b i j, b j d -> b i d", attn_slice, v[i:i+chunk_size])
+            attn_slice = torch.einsum("b i j, b j d -> b i d", attn_slice, value[start_idx:end_idx])
 
-            hidden_states[i:i+chunk_size] = attn_slice
+            hidden_states[start_idx:end_idx] = attn_slice
 
+        # reshape hidden_states
         hidden_states = self.reshape_batch_dim_to_heads(hidden_states)
-        return self.to_out(hidden_states)
+        return hidden_states
 
 
 class FeedForward(nn.Module):
diff --git a/src/diffusers/models/unet_2d_condition.py b/src/diffusers/models/unet_2d_condition.py
index e9aeb5eb0649..704cdbd2fba5 100644
--- a/src/diffusers/models/unet_2d_condition.py
+++ b/src/diffusers/models/unet_2d_condition.py
@@ -133,6 +133,28 @@ def __init__(
         self.conv_act = nn.SiLU()
         self.conv_out = nn.Conv2d(block_out_channels[0], out_channels, 3, padding=1)
 
+    def set_attention_chunk_size(self, chunk_size):
+        if chunk_size is not None and self.config.attention_head_dim % chunk_size != 0:
+            raise ValueError(
+                f"Make sure chunk_size {chunk_size} is a divisor of "
+                f"the number of heads used in cross_attention {self.config.attention_head_dim}"
+            )
+        if chunk_size is not None and chunk_size > self.config.attention_head_dim:
+            raise ValueError(
+                f"Chunk_size {chunk_size} has to be smaller or equal to "
+                f"the number of heads used in cross_attention {self.config.attention_head_dim}"
+            )
+
+        for block in self.down_blocks:
+            if hasattr(block, "attentions") and block.attentions is not None:
+                block.set_attention_chunk_size(chunk_size)
+
+        self.mid_block.set_attention_chunk_size(chunk_size)
+
+        for block in self.up_blocks:
+            if hasattr(block, "attentions") and block.attentions is not None:
+                block.set_attention_chunk_size(chunk_size)
+
     def forward(
         self,
         sample: torch.FloatTensor,
diff --git a/src/diffusers/models/unet_blocks.py b/src/diffusers/models/unet_blocks.py
index a0729213dad1..193d5de562f4 100644
--- a/src/diffusers/models/unet_blocks.py
+++ b/src/diffusers/models/unet_blocks.py
@@ -295,6 +295,7 @@ def __init__(
         super().__init__()
 
         self.attention_type = attention_type
+        self.attn_num_head_channels = attn_num_head_channels
         resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)
 
         # there is always at least one resnet
@@ -342,6 +343,21 @@ def __init__(
         self.attentions = nn.ModuleList(attentions)
         self.resnets = nn.ModuleList(resnets)
 
+    def set_attention_chunk_size(self, chunk_size):
+        if chunk_size is not None and self.attn_num_head_channels % chunk_size != 0:
+            raise ValueError(
+                f"Make sure chunk_size {chunk_size} is a divisor of "
+                f"the number of heads used in cross_attention {self.attn_num_head_channels}"
+            )
+        if chunk_size is not None and chunk_size > self.attn_num_head_channels:
+            raise ValueError(
+                f"Chunk_size {chunk_size} has to be smaller or equal to "
+                f"the number of heads used in cross_attention {self.attn_num_head_channels}"
+            )
+
+        for attn in self.attentions:
+            attn._set_attention_chunk_size(chunk_size)
+
     def forward(self, hidden_states, temb=None, encoder_hidden_states=None):
         hidden_states = self.resnets[0](hidden_states, temb)
         for attn, resnet in zip(self.attentions, self.resnets[1:]):
@@ -457,6 +473,7 @@ def __init__(
         attentions = []
 
         self.attention_type = attention_type
+        self.attn_num_head_channels = attn_num_head_channels
 
         for i in range(num_layers):
             in_channels = in_channels if i == 0 else out_channels
@@ -497,6 +514,21 @@ def __init__(
         else:
             self.downsamplers = None
 
+    def set_attention_chunk_size(self, chunk_size):
+        if chunk_size is not None and self.attn_num_head_channels % chunk_size != 0:
+            raise ValueError(
+                f"Make sure chunk_size {chunk_size} is a divisor of "
+                f"the number of heads used in cross_attention {self.attn_num_head_channels}"
+            )
+        if chunk_size is not None and chunk_size > self.attn_num_head_channels:
+            raise ValueError(
+                f"Chunk_size {chunk_size} has to be smaller or equal to "
+                f"the number of heads used in cross_attention {self.attn_num_head_channels}"
+            )
+
+        for attn in self.attentions:
+            attn._set_attention_chunk_size(chunk_size)
+
     def forward(self, hidden_states, temb=None, encoder_hidden_states=None):
         output_states = ()
 
@@ -989,6 +1021,7 @@ def __init__(
         attentions = []
 
         self.attention_type = attention_type
+        self.attn_num_head_channels = attn_num_head_channels
 
         for i in range(num_layers):
             res_skip_channels = in_channels if (i == num_layers - 1) else out_channels
@@ -1025,6 +1058,21 @@ def __init__(
         else:
             self.upsamplers = None
 
+    def set_attention_chunk_size(self, chunk_size):
+        if chunk_size is not None and self.attn_num_head_channels % chunk_size != 0:
+            raise ValueError(
+                f"Make sure chunk_size {chunk_size} is a divisor of "
+                f"the number of heads used in cross_attention {self.attn_num_head_channels}"
+            )
+        if chunk_size is not None and chunk_size > self.attn_num_head_channels:
+            raise ValueError(
+                f"Chunk_size {chunk_size} has to be smaller or equal to "
+                f"the number of heads used in cross_attention {self.attn_num_head_channels}"
+            )
+
+        for attn in self.attentions:
+            attn._set_attention_chunk_size(chunk_size)
+
     def forward(self, hidden_states, res_hidden_states_tuple, temb=None, encoder_hidden_states=None):
         for resnet, attn in zip(self.resnets, self.attentions):
 
diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py
index 36ef4a195ff3..3aa13a4b9150 100644
--- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py
+++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py
@@ -36,6 +36,14 @@ def __init__(
             feature_extractor=feature_extractor,
         )
 
+    def set_attention_chunk_size(self, chunk_size: Optional[Union[str, int]] = "auto"):
+        # set chunk_size = `None` to disable `set_attention_chunk_size`
+        if chunk_size == "auto":
+            # half the attention head size is usually a good trade-off between
+            # speed and memory
+            chunk_size = self.unet.config.attention_head_dim // 2
+        self.unet.set_attention_chunk_size(chunk_size)
+
     @torch.no_grad()
     def __call__(
         self,
diff --git a/tests/test_pipelines.py b/tests/test_pipelines.py
index fbd0faf02c11..121b17874b0a 100644
--- a/tests/test_pipelines.py
+++ b/tests/test_pipelines.py
@@ -410,6 +410,38 @@ def test_stable_diffusion_k_lms(self):
         assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2
         assert np.abs(image_from_tuple_slice.flatten() - expected_slice).max() < 1e-2
 
+    def test_stable_diffusion_attention_chunk(self):
+        device = "cpu"  # ensure determinism for the device-dependent torch.Generator
+        unet = self.dummy_cond_unet
+        scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear")
+        vae = self.dummy_vae
+        bert = self.dummy_text_encoder
+        tokenizer = CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
+
+        # make sure here that pndm scheduler skips prk
+        sd_pipe = StableDiffusionPipeline(
+            unet=unet,
+            scheduler=scheduler,
+            vae=vae,
+            text_encoder=bert,
+            tokenizer=tokenizer,
+            safety_checker=self.dummy_safety_checker,
+            feature_extractor=self.dummy_extractor,
+        )
+        sd_pipe = sd_pipe.to(device)
+        sd_pipe.set_progress_bar_config(disable=None)
+
+        prompt = "A painting of a squirrel eating a burger"
+        generator = torch.Generator(device=device).manual_seed(0)
+        output_1 = sd_pipe([prompt], generator=generator, guidance_scale=6.0, num_inference_steps=2, output_type="np")
+
+        # make sure chunking the attention yields the same result
+        sd_pipe.set_attention_chunk_size(chunk_size=1)
+        generator = torch.Generator(device=device).manual_seed(0)
+        output_2 = sd_pipe([prompt], generator=generator, guidance_scale=6.0, num_inference_steps=2, output_type="np")
+
+        assert np.abs(output_2.images.flatten() - output_1.images.flatten()).max() < 1e-4
+
     def test_score_sde_ve_pipeline(self):
         unet = self.dummy_uncond_unet
         scheduler = ScoreSdeVeScheduler(tensor_format="pt")
@@ -1045,6 +1077,44 @@ def test_lms_stable_diffusion_pipeline(self):
         expected_slice = np.array([0.9077, 0.9254, 0.9181, 0.9227, 0.9213, 0.9367, 0.9399, 0.9406, 0.9024])
         assert np.abs(image_slice.flatten() - expected_slice).max() < 1e-2
 
+    @slow
+    @unittest.skipIf(torch_device == "cpu", "Stable diffusion is supposed to run on GPU")
+    def test_stable_diffusion_memory_chunking(self):
+        torch.cuda.reset_peak_memory_stats()
+        model_id = "CompVis/stable-diffusion-v1-4"
+        pipe = StableDiffusionPipeline.from_pretrained(
+            model_id, revision="fp16", torch_dtype=torch.float16, use_auth_token=True
+        ).to(torch_device)
+        pipe.set_progress_bar_config(disable=None)
+
+        prompt = "a photograph of an astronaut riding a horse"
+
+        # make attention efficient
+        pipe.set_attention_chunk_size()
+        generator = torch.Generator(device=torch_device).manual_seed(0)
+        with torch.autocast(torch_device):
+            image_chunked = pipe(
+                [prompt], generator=generator, guidance_scale=7.5, num_inference_steps=10, output_type="numpy"
+            )["sample"]
+
+        mem_bytes = torch.cuda.max_memory_allocated()
+        torch.cuda.reset_peak_memory_stats()
+        # make sure that less than 3.75 GB is allocated
+        assert mem_bytes < 3.75 * 10**9
+
+        # disable chunking
+        pipe.set_attention_chunk_size(chunk_size=None)
+        generator = torch.Generator(device=torch_device).manual_seed(0)
+        with torch.autocast(torch_device):
+            image = pipe(
+                [prompt], generator=generator, guidance_scale=7.5, num_inference_steps=10, output_type="numpy"
+            )["sample"]
+
+        # make sure that more than 3.75 GB is allocated
+        mem_bytes = torch.cuda.max_memory_allocated()
+        assert mem_bytes > 3.75 * 10**9
+        assert np.abs(image_chunked.flatten() - image.flatten()).max() < 1e-3
+
     @slow
     @unittest.skipIf(torch_device == "cpu", "Stable diffusion is supposed to run on GPU")
     def test_stable_diffusion_img2img_pipeline(self):

From 16b9ef643b82cb3d09f1e903004bfaa00079c99a Mon Sep 17 00:00:00 2001
From: Patrick von Platen <patrick.v.platen@gmail.com>
Date: Tue, 6 Sep 2022 08:40:44 +0000
Subject: [PATCH 3/7] correct

---
 tests/test_pipelines.py | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/tests/test_pipelines.py b/tests/test_pipelines.py
index 121b17874b0a..bfeb41794dd1 100644
--- a/tests/test_pipelines.py
+++ b/tests/test_pipelines.py
@@ -1093,9 +1093,10 @@ def test_stable_diffusion_memory_chunking(self):
         pipe.set_attention_chunk_size()
         generator = torch.Generator(device=torch_device).manual_seed(0)
         with torch.autocast(torch_device):
-            image_chunked = pipe(
+            output_chunked = pipe(
                 [prompt], generator=generator, guidance_scale=7.5, num_inference_steps=10, output_type="numpy"
-            )["sample"]
+            )
+            image_chunked = output_chunked.images
 
         mem_bytes = torch.cuda.max_memory_allocated()
         torch.cuda.reset_peak_memory_stats()
@@ -1106,9 +1107,10 @@ def test_stable_diffusion_memory_chunking(self):
         pipe.set_attention_chunk_size(chunk_size=None)
         generator = torch.Generator(device=torch_device).manual_seed(0)
         with torch.autocast(torch_device):
-            image = pipe(
+            output = pipe(
                 [prompt], generator=generator, guidance_scale=7.5, num_inference_steps=10, output_type="numpy"
-            )["sample"]
+            )
+            image = output.images
 
         # make sure that more than 3.75 GB is allocated
         mem_bytes = torch.cuda.max_memory_allocated()

From 4e6e05c26212b92f4d548145ec5dfa94c1646696 Mon Sep 17 00:00:00 2001
From: Patrick von Platen <patrick.v.platen@gmail.com>
Date: Tue, 6 Sep 2022 08:42:39 +0000
Subject: [PATCH 4/7] up

---
 README.md                                            |  5 ++++-
 src/diffusers/models/attention.py                    |  8 ++++----
 src/diffusers/models/unet_2d_condition.py            |  8 ++++----
 src/diffusers/models/unet_blocks.py                  | 12 ++++++------
 .../stable_diffusion/pipeline_stable_diffusion.py    |  6 +++---
 tests/test_pipelines.py                              |  6 +++---
 6 files changed, 24 insertions(+), 21 deletions(-)

diff --git a/README.md b/README.md
index 6fc017959591..d2a7f34fc6da 100644
--- a/README.md
+++ b/README.md
@@ -104,7 +104,9 @@ with autocast("cuda"):
     image = pipe(prompt).images[0]  
 ```
 
-If you are limited by GPU memory, you might want to consider using the model in `fp16`.
+If you are limited by GPU memory, you might want to consider using the model in `fp16` as 
+well as chunking the attention computation.
+The following snippet should result in less than 4GB VRAM.
 
 ```python
 pipe = StableDiffusionPipeline.from_pretrained(
@@ -116,6 +118,7 @@ pipe = StableDiffusionPipeline.from_pretrained(
 pipe = pipe.to("cuda")
 
 prompt = "a photo of an astronaut riding a horse on mars"
+pipe.set_attention_chunk()
 with autocast("cuda"):
     image = pipe(prompt).images[0]  
 ```
diff --git a/src/diffusers/models/attention.py b/src/diffusers/models/attention.py
index 134dbbb27f30..e7331a86eb52 100644
--- a/src/diffusers/models/attention.py
+++ b/src/diffusers/models/attention.py
@@ -108,9 +108,9 @@ def __init__(self, in_channels, n_heads, d_head, depth=1, dropout=0.0, context_d
 
         self.proj_out = nn.Conv2d(inner_dim, in_channels, kernel_size=1, stride=1, padding=0)
 
-    def _set_attention_chunk_size(self, chunk_size):
+    def _set_attention_chunk(self, chunk_size):
         for block in self.transformer_blocks:
-            block._set_attention_chunk_size(chunk_size)
+            block._set_attention_chunk(chunk_size)
 
     def forward(self, x, context=None):
         # note: if no context is given, cross-attention defaults to self-attention
@@ -141,7 +141,7 @@ def __init__(self, dim, n_heads, d_head, dropout=0.0, context_dim=None, gated_ff
         self.norm3 = nn.LayerNorm(dim)
         self.checkpoint = checkpoint
 
-    def _set_attention_chunk_size(self, chunk_size):
+    def _set_attention_chunk(self, chunk_size):
         self.attn1._chunk_size = chunk_size
         self.attn2._chunk_size = chunk_size
 
@@ -162,7 +162,7 @@ def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.
         self.heads = heads
         # for chunk_size > 0 the attention score computation
         # is split across the batch axis to save memory
-        # You can set chunk_size with `set_attention_chunk_size`
+        # You can set chunk_size with `set_attention_chunk`
         self._chunk_size = None
 
         self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
diff --git a/src/diffusers/models/unet_2d_condition.py b/src/diffusers/models/unet_2d_condition.py
index 704cdbd2fba5..41170db488c7 100644
--- a/src/diffusers/models/unet_2d_condition.py
+++ b/src/diffusers/models/unet_2d_condition.py
@@ -133,7 +133,7 @@ def __init__(
         self.conv_act = nn.SiLU()
         self.conv_out = nn.Conv2d(block_out_channels[0], out_channels, 3, padding=1)
 
-    def set_attention_chunk_size(self, chunk_size):
+    def set_attention_chunk(self, chunk_size):
         if chunk_size is not None and self.config.attention_head_dim % chunk_size != 0:
             raise ValueError(
                 f"Make sure chunk_size {chunk_size} is a divisor of "
@@ -147,13 +147,13 @@ def set_attention_chunk_size(self, chunk_size):
 
         for block in self.down_blocks:
             if hasattr(block, "attentions") and block.attentions is not None:
-                block.set_attention_chunk_size(chunk_size)
+                block.set_attention_chunk(chunk_size)
 
-        self.mid_block.set_attention_chunk_size(chunk_size)
+        self.mid_block.set_attention_chunk(chunk_size)
 
         for block in self.up_blocks:
             if hasattr(block, "attentions") and block.attentions is not None:
-                block.set_attention_chunk_size(chunk_size)
+                block.set_attention_chunk(chunk_size)
 
     def forward(
         self,
diff --git a/src/diffusers/models/unet_blocks.py b/src/diffusers/models/unet_blocks.py
index 193d5de562f4..4f33d2822c53 100644
--- a/src/diffusers/models/unet_blocks.py
+++ b/src/diffusers/models/unet_blocks.py
@@ -343,7 +343,7 @@ def __init__(
         self.attentions = nn.ModuleList(attentions)
         self.resnets = nn.ModuleList(resnets)
 
-    def set_attention_chunk_size(self, chunk_size):
+    def set_attention_chunk(self, chunk_size):
         if chunk_size is not None and self.attn_num_head_channels % chunk_size != 0:
             raise ValueError(
                 f"Make sure chunk_size {chunk_size} is a divisor of "
@@ -356,7 +356,7 @@ def set_attention_chunk_size(self, chunk_size):
             )
 
         for attn in self.attentions:
-            attn._set_attention_chunk_size(chunk_size)
+            attn._set_attention_chunk(chunk_size)
 
     def forward(self, hidden_states, temb=None, encoder_hidden_states=None):
         hidden_states = self.resnets[0](hidden_states, temb)
@@ -514,7 +514,7 @@ def __init__(
         else:
             self.downsamplers = None
 
-    def set_attention_chunk_size(self, chunk_size):
+    def set_attention_chunk(self, chunk_size):
         if chunk_size is not None and self.attn_num_head_channels % chunk_size != 0:
             raise ValueError(
                 f"Make sure chunk_size {chunk_size} is a divisor of "
@@ -527,7 +527,7 @@ def set_attention_chunk_size(self, chunk_size):
             )
 
         for attn in self.attentions:
-            attn._set_attention_chunk_size(chunk_size)
+            attn._set_attention_chunk(chunk_size)
 
     def forward(self, hidden_states, temb=None, encoder_hidden_states=None):
         output_states = ()
@@ -1058,7 +1058,7 @@ def __init__(
         else:
             self.upsamplers = None
 
-    def set_attention_chunk_size(self, chunk_size):
+    def set_attention_chunk(self, chunk_size):
         if chunk_size is not None and self.attn_num_head_channels % chunk_size != 0:
             raise ValueError(
                 f"Make sure chunk_size {chunk_size} is a divisor of "
@@ -1071,7 +1071,7 @@ def set_attention_chunk_size(self, chunk_size):
             )
 
         for attn in self.attentions:
-            attn._set_attention_chunk_size(chunk_size)
+            attn._set_attention_chunk(chunk_size)
 
     def forward(self, hidden_states, res_hidden_states_tuple, temb=None, encoder_hidden_states=None):
         for resnet, attn in zip(self.resnets, self.attentions):
diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py
index 3aa13a4b9150..28f518d93551 100644
--- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py
+++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py
@@ -36,13 +36,13 @@ def __init__(
             feature_extractor=feature_extractor,
         )
 
-    def set_attention_chunk_size(self, chunk_size: Optional[Union[str, int]] = "auto"):
-        # set chunk_size = `None` to disable `set_attention_chunk_size`
+    def set_attention_chunk(self, chunk_size: Optional[Union[str, int]] = "auto"):
+        # set chunk_size = `None` to disable `set_attention_chunk`
         if chunk_size == "auto":
             # half the attention head size is usually a good trade-off between
             # speed and memory
             chunk_size = self.unet.config.attention_head_dim // 2
-        self.unet.set_attention_chunk_size(chunk_size)
+        self.unet.set_attention_chunk(chunk_size)
 
     @torch.no_grad()
     def __call__(
diff --git a/tests/test_pipelines.py b/tests/test_pipelines.py
index bfeb41794dd1..f9a0062c5ef6 100644
--- a/tests/test_pipelines.py
+++ b/tests/test_pipelines.py
@@ -436,7 +436,7 @@ def test_stable_diffusion_attention_chunk(self):
         output_1 = sd_pipe([prompt], generator=generator, guidance_scale=6.0, num_inference_steps=2, output_type="np")
 
         # make sure chunking the attention yields the same result
-        sd_pipe.set_attention_chunk_size(chunk_size=1)
+        sd_pipe.set_attention_chunk(chunk_size=1)
         generator = torch.Generator(device=device).manual_seed(0)
         output_2 = sd_pipe([prompt], generator=generator, guidance_scale=6.0, num_inference_steps=2, output_type="np")
 
@@ -1090,7 +1090,7 @@ def test_stable_diffusion_memory_chunking(self):
         prompt = "a photograph of an astronaut riding a horse"
 
         # make attention efficient
-        pipe.set_attention_chunk_size()
+        pipe.set_attention_chunk()
         generator = torch.Generator(device=torch_device).manual_seed(0)
         with torch.autocast(torch_device):
             output_chunked = pipe(
@@ -1104,7 +1104,7 @@ def test_stable_diffusion_memory_chunking(self):
         assert mem_bytes < 3.75 * 10**9
 
         # disable chunking
-        pipe.set_attention_chunk_size(chunk_size=None)
+        pipe.set_attention_chunk(chunk_size=None)
         generator = torch.Generator(device=torch_device).manual_seed(0)
         with torch.autocast(torch_device):
             output = pipe(

From b75fbba4fb5e2c5e7872886f830c7d4ee7455a2e Mon Sep 17 00:00:00 2001
From: Patrick von Platen <patrick.v.platen@gmail.com>
Date: Tue, 6 Sep 2022 08:49:56 +0000
Subject: [PATCH 5/7] finish

---
 .../pipeline_stable_diffusion_img2img.py                 | 9 +++++++++
 .../pipeline_stable_diffusion_inpaint.py                 | 9 +++++++++
 2 files changed, 18 insertions(+)

diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py
index 84b427610bde..5f54a5d30754 100644
--- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py
+++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py
@@ -47,6 +47,15 @@ def __init__(
             feature_extractor=feature_extractor,
         )
 
+    def set_attention_chunk(self, chunk_size: Optional[Union[str, int]] = "auto"):
+        # set chunk_size = `None` to disable `set_attention_chunk`
+        if chunk_size == "auto":
+            # half the attention head size is usually a good trade-off between
+            # speed and memory
+            chunk_size = self.unet.config.attention_head_dim // 2
+
+        self.unet.set_attention_chunk(chunk_size)
+
     @torch.no_grad()
     def __call__(
         self,
diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py
index ca054adb5bd1..ba35fcef6c62 100644
--- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py
+++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py
@@ -61,6 +61,15 @@ def __init__(
             feature_extractor=feature_extractor,
         )
 
+    def set_attention_chunk(self, chunk_size: Optional[Union[str, int]] = "auto"):
+        # set chunk_size = `None` to disable `set_attention_chunk`
+        if chunk_size == "auto":
+            # half the attention head size is usually a good trade-off between
+            # speed and memory
+            chunk_size = self.unet.config.attention_head_dim // 2
+
+        self.unet.set_attention_chunk(chunk_size)
+
     @torch.no_grad()
     def __call__(
         self,

From 22b9e3c22b7ed5c0c1113f6480b304c16ace6785 Mon Sep 17 00:00:00 2001
From: Patrick von Platen <patrick.v.platen@gmail.com>
Date: Tue, 6 Sep 2022 15:46:26 +0000
Subject: [PATCH 6/7] better naming

---
 README.md                                     |  2 +-
 src/diffusers/models/attention.py             | 24 ++++++-------
 src/diffusers/models/unet_2d_condition.py     | 16 ++++-----
 src/diffusers/models/unet_blocks.py           | 36 +++++++++----------
 .../pipeline_stable_diffusion.py              | 13 ++++---
 .../pipeline_stable_diffusion_img2img.py      | 12 ++++---
 .../pipeline_stable_diffusion_inpaint.py      | 12 ++++---
 tests/test_pipelines.py                       |  7 ++--
 8 files changed, 64 insertions(+), 58 deletions(-)

diff --git a/README.md b/README.md
index d2a7f34fc6da..ca6a8f8ea8bf 100644
--- a/README.md
+++ b/README.md
@@ -118,7 +118,7 @@ pipe = StableDiffusionPipeline.from_pretrained(
 pipe = pipe.to("cuda")
 
 prompt = "a photo of an astronaut riding a horse on mars"
-pipe.set_attention_chunk()
+pipe.set_attention_slice()
 with autocast("cuda"):
     image = pipe(prompt).images[0]  
 ```
diff --git a/src/diffusers/models/attention.py b/src/diffusers/models/attention.py
index e7331a86eb52..a69d9014bdf6 100644
--- a/src/diffusers/models/attention.py
+++ b/src/diffusers/models/attention.py
@@ -108,9 +108,9 @@ def __init__(self, in_channels, n_heads, d_head, depth=1, dropout=0.0, context_d
 
         self.proj_out = nn.Conv2d(inner_dim, in_channels, kernel_size=1, stride=1, padding=0)
 
-    def _set_attention_chunk(self, chunk_size):
+    def _set_attention_slice(self, slice_size):
         for block in self.transformer_blocks:
-            block._set_attention_chunk(chunk_size)
+            block._set_attention_slice(slice_size)
 
     def forward(self, x, context=None):
         # note: if no context is given, cross-attention defaults to self-attention
@@ -141,9 +141,9 @@ def __init__(self, dim, n_heads, d_head, dropout=0.0, context_dim=None, gated_ff
         self.norm3 = nn.LayerNorm(dim)
         self.checkpoint = checkpoint
 
-    def _set_attention_chunk(self, chunk_size):
-        self.attn1._chunk_size = chunk_size
-        self.attn2._chunk_size = chunk_size
+    def _set_attention_slice(self, slice_size):
+        self.attn1._slice_size = slice_size
+        self.attn2._slice_size = slice_size
 
     def forward(self, x, context=None):
         x = self.attn1(self.norm1(x)) + x
@@ -160,10 +160,10 @@ def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.
 
         self.scale = dim_head**-0.5
         self.heads = heads
-        # for chunk_size > 0 the attention score computation
+        # for slice_size > 0 the attention score computation
         # is split across the batch axis to save memory
-        # You can set chunk_size with `set_attention_chunk`
-        self._chunk_size = None
+        # You can set slice_size with `set_attention_slice`
+        self._slice_size = None
 
         self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
         self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
@@ -209,10 +209,10 @@ def _attention(self, query, key, value, sequence_length, dim):
         hidden_states = torch.zeros(
             (batch_size_attention, sequence_length, dim // self.heads), device=query.device, dtype=query.dtype
         )
-        chunk_size = self._chunk_size if self._chunk_size is not None else hidden_states.shape[0]
-        for i in range(hidden_states.shape[0] // chunk_size):
-            start_idx = i * chunk_size
-            end_idx = (i + 1) * chunk_size
+        slice_size = self._slice_size if self._slice_size is not None else hidden_states.shape[0]
+        for i in range(hidden_states.shape[0] // slice_size):
+            start_idx = i * slice_size
+            end_idx = (i + 1) * slice_size
             attn_slice = (
                 torch.einsum("b i d, b j d -> b i j", query[start_idx:end_idx], key[start_idx:end_idx]) * self.scale
             )
diff --git a/src/diffusers/models/unet_2d_condition.py b/src/diffusers/models/unet_2d_condition.py
index c9f93fe36f8c..774b350283a7 100644
--- a/src/diffusers/models/unet_2d_condition.py
+++ b/src/diffusers/models/unet_2d_condition.py
@@ -133,27 +133,27 @@ def __init__(
         self.conv_act = nn.SiLU()
         self.conv_out = nn.Conv2d(block_out_channels[0], out_channels, 3, padding=1)
 
-    def set_attention_chunk(self, chunk_size):
-        if chunk_size is not None and self.config.attention_head_dim % chunk_size != 0:
+    def set_attention_slice(self, slice_size):
+        if slice_size is not None and self.config.attention_head_dim % slice_size != 0:
             raise ValueError(
-                f"Make sure chunk_size {chunk_size} is a divisor of "
+                f"Make sure slice_size {slice_size} is a divisor of "
                 f"the number of heads used in cross_attention {self.config.attention_head_dim}"
             )
-        if chunk_size is not None and chunk_size > self.config.attention_head_dim:
+        if slice_size is not None and slice_size > self.config.attention_head_dim:
             raise ValueError(
-                f"Chunk_size {chunk_size} has to be smaller or equal to "
+                f"Chunk_size {slice_size} has to be smaller or equal to "
                 f"the number of heads used in cross_attention {self.config.attention_head_dim}"
             )
 
         for block in self.down_blocks:
             if hasattr(block, "attentions") and block.attentions is not None:
-                block.set_attention_chunk(chunk_size)
+                block.set_attention_slice(slice_size)
 
-        self.mid_block.set_attention_chunk(chunk_size)
+        self.mid_block.set_attention_slice(slice_size)
 
         for block in self.up_blocks:
             if hasattr(block, "attentions") and block.attentions is not None:
-                block.set_attention_chunk(chunk_size)
+                block.set_attention_slice(slice_size)
 
     def forward(
         self,
diff --git a/src/diffusers/models/unet_blocks.py b/src/diffusers/models/unet_blocks.py
index 4f33d2822c53..9e062165357c 100644
--- a/src/diffusers/models/unet_blocks.py
+++ b/src/diffusers/models/unet_blocks.py
@@ -343,20 +343,20 @@ def __init__(
         self.attentions = nn.ModuleList(attentions)
         self.resnets = nn.ModuleList(resnets)
 
-    def set_attention_chunk(self, chunk_size):
-        if chunk_size is not None and self.attn_num_head_channels % chunk_size != 0:
+    def set_attention_slice(self, slice_size):
+        if slice_size is not None and self.attn_num_head_channels % slice_size != 0:
             raise ValueError(
-                f"Make sure chunk_size {chunk_size} is a divisor of "
+                f"Make sure slice_size {slice_size} is a divisor of "
                 f"the number of heads used in cross_attention {self.attn_num_head_channels}"
             )
-        if chunk_size is not None and chunk_size > self.attn_num_head_channels:
+        if slice_size is not None and slice_size > self.attn_num_head_channels:
             raise ValueError(
-                f"Chunk_size {chunk_size} has to be smaller or equal to "
+                f"Chunk_size {slice_size} has to be smaller or equal to "
                 f"the number of heads used in cross_attention {self.attn_num_head_channels}"
             )
 
         for attn in self.attentions:
-            attn._set_attention_chunk(chunk_size)
+            attn._set_attention_slice(slice_size)
 
     def forward(self, hidden_states, temb=None, encoder_hidden_states=None):
         hidden_states = self.resnets[0](hidden_states, temb)
@@ -514,20 +514,20 @@ def __init__(
         else:
             self.downsamplers = None
 
-    def set_attention_chunk(self, chunk_size):
-        if chunk_size is not None and self.attn_num_head_channels % chunk_size != 0:
+    def set_attention_slice(self, slice_size):
+        if slice_size is not None and self.attn_num_head_channels % slice_size != 0:
             raise ValueError(
-                f"Make sure chunk_size {chunk_size} is a divisor of "
+                f"Make sure slice_size {slice_size} is a divisor of "
                 f"the number of heads used in cross_attention {self.attn_num_head_channels}"
             )
-        if chunk_size is not None and chunk_size > self.attn_num_head_channels:
+        if slice_size is not None and slice_size > self.attn_num_head_channels:
             raise ValueError(
-                f"Chunk_size {chunk_size} has to be smaller or equal to "
+                f"Chunk_size {slice_size} has to be smaller or equal to "
                 f"the number of heads used in cross_attention {self.attn_num_head_channels}"
             )
 
         for attn in self.attentions:
-            attn._set_attention_chunk(chunk_size)
+            attn._set_attention_slice(slice_size)
 
     def forward(self, hidden_states, temb=None, encoder_hidden_states=None):
         output_states = ()
@@ -1058,20 +1058,20 @@ def __init__(
         else:
             self.upsamplers = None
 
-    def set_attention_chunk(self, chunk_size):
-        if chunk_size is not None and self.attn_num_head_channels % chunk_size != 0:
+    def set_attention_slice(self, slice_size):
+        if slice_size is not None and self.attn_num_head_channels % slice_size != 0:
             raise ValueError(
-                f"Make sure chunk_size {chunk_size} is a divisor of "
+                f"Make sure slice_size {slice_size} is a divisor of "
                 f"the number of heads used in cross_attention {self.attn_num_head_channels}"
             )
-        if chunk_size is not None and chunk_size > self.attn_num_head_channels:
+        if slice_size is not None and slice_size > self.attn_num_head_channels:
             raise ValueError(
-                f"Chunk_size {chunk_size} has to be smaller or equal to "
+                f"Chunk_size {slice_size} has to be smaller or equal to "
                 f"the number of heads used in cross_attention {self.attn_num_head_channels}"
             )
 
         for attn in self.attentions:
-            attn._set_attention_chunk(chunk_size)
+            attn._set_attention_slice(slice_size)
 
     def forward(self, hidden_states, res_hidden_states_tuple, temb=None, encoder_hidden_states=None):
         for resnet, attn in zip(self.resnets, self.attentions):
diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py
index 28f518d93551..f189055eb91a 100644
--- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py
+++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py
@@ -36,13 +36,16 @@ def __init__(
             feature_extractor=feature_extractor,
         )
 
-    def set_attention_chunk(self, chunk_size: Optional[Union[str, int]] = "auto"):
-        # set chunk_size = `None` to disable `set_attention_chunk`
-        if chunk_size == "auto":
+    def enable_attention_slicing(self, slice_size: Optional[Union[str, int]] = "auto"):
+        if slice_size == "auto":
             # half the attention head size is usually a good trade-off between
             # speed and memory
-            chunk_size = self.unet.config.attention_head_dim // 2
-        self.unet.set_attention_chunk(chunk_size)
+            slice_size = self.unet.config.attention_head_dim // 2
+        self.unet.set_attention_slice(slice_size)
+
+    def disable_attention_slicing(self):
+        # set slice_size = `None` to disable `set_attention_slice`
+        self.enable_attention_slice(None)
 
     @torch.no_grad()
     def __call__(
diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py
index 5f54a5d30754..414b7994892b 100644
--- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py
+++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_img2img.py
@@ -47,14 +47,16 @@ def __init__(
             feature_extractor=feature_extractor,
         )
 
-    def set_attention_chunk(self, chunk_size: Optional[Union[str, int]] = "auto"):
-        # set chunk_size = `None` to disable `set_attention_chunk`
-        if chunk_size == "auto":
+    def enable_attention_slicing(self, slice_size: Optional[Union[str, int]] = "auto"):
+        if slice_size == "auto":
             # half the attention head size is usually a good trade-off between
             # speed and memory
-            chunk_size = self.unet.config.attention_head_dim // 2
+            slice_size = self.unet.config.attention_head_dim // 2
+        self.unet.set_attention_slice(slice_size)
 
-        self.unet.set_attention_chunk(chunk_size)
+    def disable_attention_slicing(self):
+        # set slice_size = `None` to disable `set_attention_slice`
+        self.enable_attention_slice(None)
 
     @torch.no_grad()
     def __call__(
diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py
index ba35fcef6c62..5bf9c16a814d 100644
--- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py
+++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion_inpaint.py
@@ -61,14 +61,16 @@ def __init__(
             feature_extractor=feature_extractor,
         )
 
-    def set_attention_chunk(self, chunk_size: Optional[Union[str, int]] = "auto"):
-        # set chunk_size = `None` to disable `set_attention_chunk`
-        if chunk_size == "auto":
+    def enable_attention_slicing(self, slice_size: Optional[Union[str, int]] = "auto"):
+        if slice_size == "auto":
             # half the attention head size is usually a good trade-off between
             # speed and memory
-            chunk_size = self.unet.config.attention_head_dim // 2
+            slice_size = self.unet.config.attention_head_dim // 2
+        self.unet.set_attention_slice(slice_size)
 
-        self.unet.set_attention_chunk(chunk_size)
+    def disable_attention_slicing(self):
+        # set slice_size = `None` to disable `set_attention_slice`
+        self.enable_attention_slice(None)
 
     @torch.no_grad()
     def __call__(
diff --git a/tests/test_pipelines.py b/tests/test_pipelines.py
index 876bccff98d4..a05d57a73d41 100644
--- a/tests/test_pipelines.py
+++ b/tests/test_pipelines.py
@@ -153,7 +153,6 @@ def dummy_text_encoder(self):
         torch.manual_seed(0)
         config = CLIPTextConfig(
             bos_token_id=0,
-            chunk_size_feed_forward=0,
             eos_token_id=2,
             hidden_size=32,
             intermediate_size=37,
@@ -436,7 +435,7 @@ def test_stable_diffusion_attention_chunk(self):
         output_1 = sd_pipe([prompt], generator=generator, guidance_scale=6.0, num_inference_steps=2, output_type="np")
 
         # make sure chunking the attention yields the same result
-        sd_pipe.set_attention_chunk(chunk_size=1)
+        sd_pipe.enable_attention_slicing(slice_size=1)
         generator = torch.Generator(device=device).manual_seed(0)
         output_2 = sd_pipe([prompt], generator=generator, guidance_scale=6.0, num_inference_steps=2, output_type="np")
 
@@ -1090,7 +1089,7 @@ def test_stable_diffusion_memory_chunking(self):
         prompt = "a photograph of an astronaut riding a horse"
 
         # make attention efficient
-        pipe.set_attention_chunk()
+        pipe.enable_attention_slicing()
         generator = torch.Generator(device=torch_device).manual_seed(0)
         with torch.autocast(torch_device):
             output_chunked = pipe(
@@ -1104,7 +1103,7 @@ def test_stable_diffusion_memory_chunking(self):
         assert mem_bytes < 3.75 * 10**9
 
         # disable chunking
-        pipe.set_attention_chunk(chunk_size=None)
+        pipe.disable_attention_slicing()
         generator = torch.Generator(device=torch_device).manual_seed(0)
         with torch.autocast(torch_device):
             output = pipe(

From dcf9b0fe176403f5b36249a4a1e2fcb8f419efe4 Mon Sep 17 00:00:00 2001
From: Patrick von Platen <patrick.v.platen@gmail.com>
Date: Tue, 6 Sep 2022 18:06:31 +0200
Subject: [PATCH 7/7] Update README.md

Co-authored-by: Pedro Cuenca <pedro@huggingface.co>
---
 README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/README.md b/README.md
index ca6a8f8ea8bf..434d0cee2b3b 100644
--- a/README.md
+++ b/README.md
@@ -118,7 +118,7 @@ pipe = StableDiffusionPipeline.from_pretrained(
 pipe = pipe.to("cuda")
 
 prompt = "a photo of an astronaut riding a horse on mars"
-pipe.set_attention_slice()
+pipe.enable_attention_slicing()
 with autocast("cuda"):
     image = pipe(prompt).images[0]  
 ```