Autodisable xformers if using Torch2 in readme example

README.md

@@ -75,32 +75,32 @@ And we can now run the model locally.
 
 By default `diffusers` makes use of [model cpu offloading](https://huggingface.co/docs/diffusers/optimization/fp16#model-offloading-for-fast-inference-and-memory-savings) to run the whole IF pipeline with as little as 14 GB of VRAM.
 
-If you are using `torch>=2.0.0`, make sure to **delete all** `enable_xformers_memory_efficient_attention()`
-functions.
-
 ```py
 from diffusers import DiffusionPipeline
 from diffusers.utils import pt_to_pil
 import torch
 
 # stage 1
 stage_1 = DiffusionPipeline.from_pretrained("DeepFloyd/IF-I-XL-v1.0", variant="fp16", torch_dtype=torch.float16)
-stage_1.enable_xformers_memory_efficient_attention()  # remove line if torch.__version__ >= 2.0.0
 stage_1.enable_model_cpu_offload()
 
 # stage 2
 stage_2 = DiffusionPipeline.from_pretrained(
     "DeepFloyd/IF-II-L-v1.0", text_encoder=None, variant="fp16", torch_dtype=torch.float16
 )
-stage_2.enable_xformers_memory_efficient_attention()  # remove line if torch.__version__ >= 2.0.0
 stage_2.enable_model_cpu_offload()
 
 # stage 3
 safety_modules = {"feature_extractor": stage_1.feature_extractor, "safety_checker": stage_1.safety_checker, "watermarker": stage_1.watermarker}
 stage_3 = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-x4-upscaler", **safety_modules, torch_dtype=torch.float16)
-stage_3.enable_xformers_memory_efficient_attention()  # remove line if torch.__version__ >= 2.0.0
 stage_3.enable_model_cpu_offload()
 
+# xformers memory efficient attention shouldn't be used with PyTorch2
+if not torch.__version__.startswith('2'):
+    stage_1.enable_xformers_memory_efficient_attention()
+    stage_2.enable_xformers_memory_efficient_attention()
+    stage_3.enable_xformers_memory_efficient_attention()
+
 prompt = 'a photo of a kangaroo wearing an orange hoodie and blue sunglasses standing in front of the eiffel tower holding a sign that says "very deep learning"'
 
 # text embeds