44from io import BytesIO
55
66import shutil
7+ from tempfile import TemporaryDirectory
78
89from diskcache import Cache
910
@@ -114,7 +115,9 @@ def push_wrapper(cat_path: str, description: str) -> None:
114115 # and then cache it locally
115116 base_name = os .path .basename (cat_path )
116117 model_hash = base_name .removesuffix (".zip" )
117- self [model_hash ] = cat_path
118+ with open (cat_path , "rb" ) as f :
119+ model_bytes = f .read ()
120+ self .insert_raw (model_hash , model_bytes )
118121
119122 den ._push_model_from_file = push_wrapper # type: ignore
120123 # wrap fetch_model to check cache first
@@ -128,8 +131,16 @@ def fetch_wrapper(model_info: ModelInfo) -> CATWrapper:
128131 model_path , model_info = model_info ))
129132 cat = orig_fetch (model_info )
130133 # cache it
131- model_path = f"{ model_hash } .zip"
132- self [model_hash ] = model_path
134+ with TemporaryDirectory () as tmpdir :
135+ model_path = cat .save_model_pack (
136+ tmpdir ,
137+ pack_name = model_hash ,
138+ add_hash_to_pack_name = False ,
139+ force_save_local = True
140+ )
141+ with open (model_path , "rb" ) as f :
142+ model_bytes = f .read ()
143+ self .insert_raw (model_hash , model_bytes )
133144 return cat
134145
135146 den .fetch_model = fetch_wrapper # type: ignore
0 commit comments