Update app.py

app.py

@@ -7,6 +7,7 @@ from PIL import Image
 from os import path
 from torchvision import transforms
 from dataclasses import dataclass
+from io import BytesIO
 import math
 from typing import Callable
 import spaces
@@ -51,10 +52,12 @@ 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 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=torch.float16).to("cuda")
 
-#pipe.vae = DiffusersAutoencoderKL.from_pretrained("kaiyuyue/FLUX.2-dev-vae", torch_dtype=torch.float16, scaling_factor = 0.3611, shift_factor = 0.1159).to("cuda")
+pipe.vae = AutoencoderKL.from_pretrained("REPA-E/e2e-flux-vae", torch_dtype=torch.float16).to("cuda")
 
+#pipe.vae = DiffusersAutoencoderKL.from_pretrained("kaiyuyue/FLUX.2-dev-vae", torch_dtype=torch.float16, scaling_factor = 0.3611, shift_factor = 0.1159).to("cuda")
+## Alas, the model would need to be retrained to work with the Flux2 vae, with its doubled channel count of 32.
 #pipe.enable_model_cpu_offload()
 
 try: # A temp hack for some version diffusers lora loading problem