Update app.py

app.py

@@ -13,11 +13,15 @@ resolutions = ["1024 1024","1280 768","1344 768","768 1344","768 1280" ]
 
 # Load pipeline 
 
+vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
 unet = UNet2DConditionModel.from_pretrained("briaai/BRIA-2.2-FAST", torch_dtype=torch.float16)
-pipe = DiffusionPipeline.from_pretrained("briaai/BRIA-2.2", unet=unet, torch_dtype=torch.float16)
-del unet
+pipe = DiffusionPipeline.from_pretrained("briaai/BRIA-2.2", torch_dtype=torch.float16, unet=unet, vae=vae)
 pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
-pipe = pipe.to("cuda")
+pipe.to('cuda')
+del unet
+del vae
+
+
 pipe.force_zeros_for_empty_prompt = False
 
 print("Optimizing BRIA 2.2 FAST - this could take a while")