Update app.py

app.py

@@ -13,6 +13,7 @@ import spaces
 import diffusers
 from diffusers import ZImagePipeline, DiffusionPipeline, AutoencoderTiny, AutoencoderKL
 from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+from diffusers.models import AutoencoderKL as DiffusersAutoencoderKL
 #from transformers import CLIPModel, CLIPProcessor, CLIPTextModel, CLIPTokenizer, CLIPConfig, T5EncoderModel, T5Tokenizer
 #from diffusers.models.transformers import FluxTransformer2DModel
 import copy
@@ -50,6 +51,8 @@ pipe = diffusers.ZImagePipeline.from_pretrained("dimitribarbot/Z-Image-Turbo-BF1
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
+pipe.vae = DiffusersAutoencoderKL.from_pretrained("kaiyuyue/FLUX.2-dev-vae", torch_dtype=torch.float16).to("cuda")
+
 #pipe.enable_model_cpu_offload()
 
 try: # A temp hack for some version diffusers lora loading problem