1inkusFace commited on
Commit
faea40f
·
verified ·
1 Parent(s): 9f656d9

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -0
app.py CHANGED
@@ -114,6 +114,9 @@ def load_and_prepare_model():
114
  pipe.to(device=device, dtype=torch.bfloat16)
115
  pipe.vae = vaeXL.to(device) #.to('cpu') #.to(torch.bfloat16)
116
 
 
 
 
117
  pipe.unet.set_attn_processor(AttnProcessor2_0())
118
  pipe.vae.set_default_attn_processor()
119
  return pipe
 
114
  pipe.to(device=device, dtype=torch.bfloat16)
115
  pipe.vae = vaeXL.to(device) #.to('cpu') #.to(torch.bfloat16)
116
 
117
+ # Enable attention slicing
118
+ pipe.enable_attention_slicing()
119
+
120
  pipe.unet.set_attn_processor(AttnProcessor2_0())
121
  pipe.vae.set_default_attn_processor()
122
  return pipe