Update app.py

app.py

@@ -52,7 +52,7 @@ pipe = diffusers.ZImagePipeline.from_pretrained("dimitribarbot/Z-Image-Turbo-BF1
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 #pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=torch.float16).to("cuda")
-pipe.vae = AutoencoderKL.from_pretrained("AlekseyCalvin/AnimeVAE_by_Anzhc_for_Flux_ZiT", torch_dtype=torch.bfloat16).to("cuda")
+pipe.vae = AutoencoderKL.from_pretrained("AlekseyCalvin/AnimeVAE_by_Anzhc_for_Flux_ZiT", torch_dtype=torch.bfloat16, device_map="cuda")
 
 
 #pipe.vae = AutoencoderKL.from_pretrained("REPA-E/e2e-flux-vae", torch_dtype=torch.bfloat16).to("cuda")