diff --git a/py/imagefunc.py b/py/imagefunc.py
index a91bbc4..30d1fa8 100644
--- a/py/imagefunc.py
+++ b/py/imagefunc.py
@@ -1539,9 +1539,9 @@ 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"))
+    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"))
     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)