Fix huggingface model loading problem

.gitignore

@@ -0,0 +1 @@
+/__pycache__

canvas.py

@@ -13,7 +13,7 @@ from pathlib import Path
 import numpy as np
 from matplotlib import pyplot as plt
 from sklearn.cluster import KMeans
-
+from huggingface_hub import hf_hub_download
 class BaseModelYamlJsonMixin:
     """
     BaseModel with helper methods for loading and saving to yaml/json format.
@@ -86,7 +86,8 @@ class KMeansActionTokenizer():
     def from_pretrained(cls, model_path: str | Path):
         model_path = Path(model_path)
         self = cls()
-        with open(model_path / "tokenizer.pkl", "rb") as file:
+        action_tokenizer_path = hf_hub_download(repo_id=str(model_path), filename="tokenizer.pkl")
+        with open(action_tokenizer_path, "rb") as file:
             self.kmeans = pickle.load(file)
         self.action_count = self.kmeans.n_clusters
         # assert self.action_count == 32
@@ -201,12 +202,16 @@ class Idefics2Pipeline():
 
     @classmethod
     def from_pretrained(cls, pretrained_model_name_or_path: str):
-        if not isinstance(pretrained_model_name_or_path, Path):
-            pretrained_model_name_or_path = Path(pretrained_model_name_or_path)
 
+        pipeline_config_path = hf_hub_download(repo_id=pretrained_model_name_or_path, filename="pipeline_config.json")
+        pipeline_config_path = Path(pipeline_config_path)
         config = Idefics2PipelineConfig.model_validate_json(
-            (pretrained_model_name_or_path / "pipeline_config.json").read_text()
+            (pipeline_config_path).read_text()
         )
+
+        if not isinstance(pretrained_model_name_or_path, Path):
+            pretrained_model_name_or_path = Path(pretrained_model_name_or_path)
+
         model = Idefics2ForConditionalGeneration.from_pretrained(pretrained_model_name_or_path)
         processor = Idefics2Processor.from_pretrained(pretrained_model_name_or_path)
         model.eval()

requirements.txt

@@ -5,4 +5,5 @@ numpy==2.1.3
 torch==2.4.0
 pydantic==2.9.2
 scikit-learn==1.5.2
-matplotlib==3.9.3
+matplotlib==3.9.3
+gradio==5.7.1