sec_env Environment: A separate virtual environment named sec_envwas created to isolate the project's dependencies.
Specific Package Versions: The following package versions were installed within sec_env. It's very important to note that, while we initially tried older versions, it appears the newer versions are working, in combination with load_file. If you created sec_env using your old environment's python, you could use
pip install torch==1.13.1 torchvision==0.14.1 torchaudio==0.13.1 numpy==1.23.5 pip install transformers safetensors pillow ``` to install the package we mentioned before. Or keep current version and use: ```bash pip install torch torchvision torchaudio numpy pip install transformers safetensors pillow ```Use code with caution.Bash
process_image_rmbg_local.sh Script:
PYTHON_PATH: Correctly points to the sec_env Python executable: /Users/mbp/Documents/sec_env/bin/python
Model Directory: LOCAL_MODEL_DIR is set to the correct path: /Users/mbp/Documents/models/rmbg/models--briaai--RMBG-1.4/snapshots/bfa28aaa79846d9292bb9222fbe7070123b468a4
Image Resizing: Uses your original, correct logic for extending the canvas to a 3:4 aspect ratio without cropping.
Python Script (Here Document):
Uses a "here document" (<< EOF) to create the /tmp/rmbg_script.py file.
Imports necessary libraries, including safetensors.torch.
Loads the model configuration using AutoConfig.from_pretrainedwith trust_remote_code=True and local_files_only=True.
Instantiates the model structure using AutoModelForImageSegmentation.from_config with trust_remote_code=True.
Loads the model weights using safetensors.torch.load_file(weights_path, device="cpu"). This is the crucial part that avoids the pickle error.
Loads the weights into the model using model.load_state_dict(state_dict).
Performs image preprocessing, inference, and postprocessing.
Saves the output image.
Model Files:
The model.safetensors file was downloaded directly from the Hugging Face Hub's snapshot folder.
The pytorch_model.bin file was deleted.
The config.json file is present and correct.
Explanation of the Change:
from safetensors.torch import load_file: Imports the load_filefunction. This function is specifically designed to load data from .safetensors files.
state_dict = load_file(weights_path, device="cpu"): This line uses load_file to load the model's weights (the state_dict) directly from the model.safetensors file. The device="cpu" argument ensures the weights are loaded onto the CPU.
model.load_state_dict(state_dict): This line, which remains unchanged, takes the loaded state_dict (a Python dictionary containing the model's parameters) and loads it into the model's structure.
0