Safetensor Rename

modeling_diva.py

@@ -117,7 +117,7 @@ class DiVAModel(PreTrainedModel):
     ):
         if os.path.isdir(pretrained_model_name_or_path):
             via_path = (
-                pretrained_model_name_or_path + "/model-00001-of-00004.safetensors"
+                pretrained_model_name_or_path + "/model.safetensors"
             )
             config_path = pretrained_model_name_or_path + "/config.json"
         else:
@@ -126,7 +126,7 @@ class DiVAModel(PreTrainedModel):
 
             hf_hub_download(
                 repo_id=pretrained_model_name_or_path,
-                filename="model-00001-of-00004.safetensors",
+                filename="model.safetensors",
                 token=kwargs.get("token", None),
                 local_dir=os.path.dirname(__file__),
             )
@@ -136,7 +136,7 @@ class DiVAModel(PreTrainedModel):
                 token=kwargs.get("token", None),
                 local_dir=os.path.dirname(__file__),
             )
-            via_path = os.path.dirname(__file__) + "/model-00001-of-00004.safetensors"
+            via_path = os.path.dirname(__file__) + "/model.safetensors"
             config_path = os.path.dirname(__file__) + "/config.json"
         with open(config_path, "r") as f:
             config_dict = json.loads(f.read())