nodes/chflame163-ComfyUI_LayerStyle_Advance
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix vitmatte model loader

Browse Source
This commit is contained in:
chflame163
2024-12-11 20:18:10 +08:00
parent b62cd55250
commit fcf0ee92ba
1 changed files with 6 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
py/imagefunc.py
+6 -5
View File
@@ -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