Update app.py

app.py

@@ -96,6 +96,7 @@ def load_and_prepare_model(model_id):
 
     # Get the dtype based on the model_id
     dtype = model_dtypes.get(model_id, torch.bfloat16)  # Default to float32 if not found
+    vae = AutoencoderKL.from_pretrained("ford442/sdxl-vae-bf16", torch_dtype=torch.bfloat16).to('cuda')
 
     # Load the pipeline with the determined dtype
     pipe = StableDiffusionXLPipeline.from_pretrained(
@@ -103,6 +104,7 @@ def load_and_prepare_model(model_id):
         torch_dtype=torch.bfloat16,
         use_safetensors=True,
         add_watermarker=False,
+        vae=vae
     )
     pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
     pipe.to('cuda')