Update app.py

app.py

@@ -60,9 +60,10 @@ class HFEmbedder(nn.Module):
         return outputs[self.output_key]
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
-t5 = HFEmbedder("DeepFloyd/t5-v1_1-xxl", max_length=512, torch_dtype=torch.bfloat16).to(device)
-clip = HFEmbedder("openai/clip-vit-large-patch14", max_length=77, torch_dtype=torch.bfloat16).to(device)
-ae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=torch.bfloat16).to(device)
+# Force using safetensors to avoid the torch.load security issue
+t5 = HFEmbedder("DeepFloyd/t5-v1_1-xxl", max_length=512, torch_dtype=torch.bfloat16, use_safetensors=True).to(device)
+clip = HFEmbedder("openai/clip-vit-large-patch14", max_length=77, torch_dtype=torch.bfloat16, use_safetensors=True).to(device)
+ae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae", torch_dtype=torch.bfloat16, use_safetensors=True).to(device)
 
 # ---------------- NF4 ----------------