Update app.py

app.py

@@ -5,7 +5,7 @@ from PIL import Image
 import spaces
 import torch
 from huggingface_hub import hf_hub_download, HfApi
-from diffusers import FluxPriorReduxPipeline, FluxPipeline
+from diffusers import FluxPriorReduxPipeline, FluxPipeline, AutoencoderKL
 from diffusers.utils import load_image
 import os
 api = HfApi(
@@ -13,11 +13,17 @@ api = HfApi(
 )
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
-
+good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", 
+                                         subfolder="vae", 
+                                         torch_dtype=torch.bfloat16,
+                                         use_safetensors=True, 
+                                         token=HF_TOKEN
+                                        ).to("cuda") 
 pipe = FluxPipeline.from_pretrained(
     "LPX55/FLUX.1-merged_uncensored", 
     torch_dtype=torch.bfloat16,
     token=os.getenv('HF_TOKEN'),
+    vae=good_vae
 ).to("cuda")
 pipe.load_lora_weights(hf_hub_download("ByteDance/Hyper-SD", "Hyper-FLUX.1-dev-8steps-lora.safetensors"), lora_scale=0.125)
 pipe.fuse_lora(lora_scale=0.125)