diff --git a/py/imagefunc.py b/py/imagefunc.py
index 18fa25a..cb276d2 100644
--- a/py/imagefunc.py
+++ b/py/imagefunc.py
@@ -98,6 +98,7 @@ def check_and_download_model(model_path, repo_id):
         print(f"Downloading {repo_id} model...")
         from huggingface_hub import snapshot_download
         snapshot_download(repo_id=repo_id, local_dir=model_path, ignore_patterns=["*.md", "*.txt", "onnx", ".git"])
+    return model_path
 
 '''Converter'''
 
@@ -1527,9 +1528,9 @@ class VITMatteModel:
 
 def load_VITMatte_model(model_name:str, local_files_only:bool=False) -> object:
     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)
+    model_path = check_and_download_model(model_name, model_repo)
+
     from transformers import VitMatteImageProcessor, VitMatteForImageMatting
     model = VitMatteForImageMatting.from_pretrained(model_path, local_files_only=local_files_only)
     processor = VitMatteImageProcessor.from_pretrained(model_path, local_files_only=local_files_only)
diff --git a/pyproject.toml b/pyproject.toml
index 52775a5..1c4c4ee 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,7 +1,7 @@
 [project]
 name = "comfyui_layerstyle"
 description = "A set of nodes for ComfyUI it generate image like Adobe Photoshop's Layer Style. the Drop Shadow is first completed node, and follow-up work is in progress."
-version = "2.0.4"
+version = "2.0.5"
 license = "MIT"
 dependencies = ["numpy", "pillow", "torch", "matplotlib", "Scipy", "scikit_image", "scikit_learn", "opencv-contrib-python", "pymatting", "timm", "colour-science",  "transformers",  "blend_modes",  "huggingface_hub", "loguru"]