Update app.py

app.py

@@ -5,6 +5,7 @@ import spaces
 import torch
 import time
 from diffusers import DiffusionPipeline, AutoencoderTiny
+from diffusers.models.attention_processor import AttnProcessor2_0
 from custom_pipeline import FluxWithCFGPipeline
 
 torch.backends.cuda.matmul.allow_tf32 = True
@@ -23,7 +24,7 @@ pipe = FluxWithCFGPipeline.from_pretrained(
 )
 pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype)
 pipe.to("cuda")
-pipe.enable_xformers_memory_efficient_attention()
+pipe.unet.set_attn_processor(AttnProcessor2_0())
 
 torch.cuda.empty_cache()