Update optimized.py

optimized.py

@@ -6,7 +6,7 @@ from diffusers import FluxControlNetModel, FluxControlNetPipeline, AutoencoderKL
 import gradio as gr
 huggingface_token = os.getenv("HUGGINFACE_TOKEN")
 
-good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=torch.bfloat16, token=huggingface_token).to("cuda")
+good_vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=torch.bfloat16, use_safetensors=True, token=huggingface_token).to("cuda")
 
 # Load pipeline
 controlnet = FluxControlNetModel.from_pretrained(
@@ -19,6 +19,7 @@ pipe = FluxControlNetPipeline.from_pretrained(
     torch_dtype=torch.bfloat16,
     device_map="auto",
     vae=good_vae,
+    use_safetensors=True, 
     token=huggingface_token
 )
 # Add to your pipeline initialization: