Update app.py

app.py

@@ -6,7 +6,7 @@ import gradio as gr
 import numpy as np
 import spaces
 import torch
-from diffusers import FluxControlNetModel, FluxPipeline, AutoencoderTiny
+from diffusers import FluxControlNetModel, FluxPipeline, AutoencoderTiny, AutoencoderKL
 from transformers import T5Tokenizer
 from diffusers.pipelines import FluxControlNetPipeline
 from diffusers.utils import load_image
@@ -41,13 +41,13 @@ model_path = snapshot_download(
 
 tokenizer_2 = T5Tokenizer.from_pretrained("LPX55/FLUX.1-merged_uncensored", subfolder="tokenizer_2", token=huggingface_token)
 taef1 = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=torch.bfloat16).to(device)
-
+good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=torch.bfloat16, token=huggingface_token).to(device)
 # Load pipeline
 controlnet = FluxControlNetModel.from_pretrained(
     "jasperai/Flux.1-dev-Controlnet-Upscaler", torch_dtype=torch.bfloat16
 ).to(device)
 pipe = FluxControlNetPipeline.from_pretrained(
-    "LPX55/FLUX.1-merged_uncensored", controlnet=controlnet, torch_dtype=torch.bfloat16, vae=taef1, token=huggingface_token,
+    "LPX55/FLUX.1-merged_uncensored", controlnet=controlnet, torch_dtype=torch.bfloat16, vae=good_vae, token=huggingface_token,
 )
 # pipe.load_lora_weights(
 #     hf_hub_download("ByteDance/Hyper-SD", "Hyper-FLUX.1-dev-8steps-lora.safetensors"), adapter_name="hyper-sd"