From fcf0ee92ba97ae0bd40afd49487c279ac35bdabc Mon Sep 17 00:00:00 2001 From: chflame163 Date: Wed, 11 Dec 2024 20:18:10 +0800 Subject: [PATCH] fix vitmatte model loader --- py/imagefunc.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/py/imagefunc.py b/py/imagefunc.py index e388f1c..7cb86ac 100644 --- a/py/imagefunc.py +++ b/py/imagefunc.py @@ -1526,12 +1526,13 @@ class VITMatteModel: self.processor = processor def load_VITMatte_model(model_name:str, local_files_only:bool=False) -> object: - if local_files_only: - model_name = Path(os.path.join(folder_paths.models_dir, "vitmatte")) - # model_name = Path(os.path.join(folder_paths.models_dir, "vitmatte")) + model_name = "vitmatte" + model_path = os.path.join(folder_paths.models_dir, model_name) + model_repo = "hustvl/vitmatte-small-composition-1k" + check_and_download_model(model_name, model_repo) from transformers import VitMatteImageProcessor, VitMatteForImageMatting - model = VitMatteForImageMatting.from_pretrained(model_name, local_files_only=local_files_only) - processor = VitMatteImageProcessor.from_pretrained(model_name, local_files_only=local_files_only) + model = VitMatteForImageMatting.from_pretrained(model_path, local_files_only=local_files_only) + processor = VitMatteImageProcessor.from_pretrained(model_path, local_files_only=local_files_only) vitmatte = VITMatteModel(model, processor) return vitmatte