Could you help fix the deserialization vulnerability caused by a risky pre-trained models used in this repo?

[Rockstar292]

Hi, @JunyuanDeng I'd like to report that a potentially risky pretrained model is being used in this project, which may pose deserialization threats. Please check the following code example:

• DMCalib/infer.py

    parser.add_argument(
        "--pretrained_model_path",
        type=str,
        default="juneyoung9/DM-Calib",
        help="pretrained model path from hugging face or local dir",
    )
    
     checkpoint_path = args.pretrained_model_path
        
                vae = AutoencoderKL.from_pretrained(stable_diffusion_repo_path, subfolder="vae")
        unet = UNet2DConditionModel.from_pretrained(checkpoint_path, subfolder="depth")
        vae.decoder = torch.load(os.path.join(checkpoint_path, "depth", "vae_decoder.pth"))

Issue Description

As shown above, in the DMCalib/infer.py file, the model "juneyoung9/DM-Calib" and file ``vae_decoder.pthis downloaded and loaded by thetorch.load` method.

This model has been flagged as risky on the HuggingFace platform. Specifically, its vae_decoder.pth file is marked as malicious and may trigger deserialization threats. Once model is load, the vulnerability could be activated.

I speculate that this is because the executable code in the model file contains suspicious modules. In the latest versions of PyTorch, loading such files will likely be flagged as a dangerous operation and may even be blocked entirely.

Related Risk Reports:juneyoung9/DM-Calib_model

Suggested Repair Methods

1. Replace these models with safer alternatives, such as PooryaPiroozfar/Flair-Persian-NER
2. Convert the model to safer safetensors format and re-upload
3. Remove the suspicious modules in the executable code of the suspicious files and re-upload them
4. Try using the torch.load method from PyTorch 2.6 or later to load the model weights. If it loads successfully, then there should be no issues.

As a popular machine learning projects, every potential risk could be propagated and amplified. Could you please address the above issues?

Thanks for your help~

Best regards,
Rockstar