ductai199x
commited on
Commit
•
c2e7f35
1
Parent(s):
0a89f59
minor
Browse files- convert_sam_hq_to_hf.py +1 -7
convert_sam_hq_to_hf.py
CHANGED
@@ -27,11 +27,7 @@ import re
|
|
27 |
import torch
|
28 |
from safetensors.torch import save_model
|
29 |
from huggingface_hub import hf_hub_download
|
30 |
-
from transformers import
|
31 |
-
SamImageProcessor,
|
32 |
-
SamProcessor,
|
33 |
-
SamVisionConfig,
|
34 |
-
)
|
35 |
from sam_hq_vit_huge.modeling_sam_hq import SamHQModel
|
36 |
from sam_hq_vit_huge.configuration_sam_hq import SamHQConfig
|
37 |
|
@@ -132,7 +128,6 @@ def convert_sam_checkpoint(model_name, checkpoint_path, output_dir):
|
|
132 |
|
133 |
state_dict = torch.load(checkpoint_path, map_location="cpu")
|
134 |
state_dict = replace_keys(state_dict)
|
135 |
-
# print(state_dict.keys())
|
136 |
|
137 |
hf_model = SamHQModel(config)
|
138 |
hf_model.eval()
|
@@ -167,6 +162,5 @@ if __name__ == "__main__":
|
|
167 |
checkpoint_path = args.checkpoint_path
|
168 |
else:
|
169 |
checkpoint_path = hf_hub_download("lkeab/hq-sam", f"{args.model_name}.pth")
|
170 |
-
print(checkpoint_path)
|
171 |
|
172 |
convert_sam_checkpoint(args.model_name, checkpoint_path, args.output_dir)
|
|
|
27 |
import torch
|
28 |
from safetensors.torch import save_model
|
29 |
from huggingface_hub import hf_hub_download
|
30 |
+
from transformers import SamVisionConfig
|
|
|
|
|
|
|
|
|
31 |
from sam_hq_vit_huge.modeling_sam_hq import SamHQModel
|
32 |
from sam_hq_vit_huge.configuration_sam_hq import SamHQConfig
|
33 |
|
|
|
128 |
|
129 |
state_dict = torch.load(checkpoint_path, map_location="cpu")
|
130 |
state_dict = replace_keys(state_dict)
|
|
|
131 |
|
132 |
hf_model = SamHQModel(config)
|
133 |
hf_model.eval()
|
|
|
162 |
checkpoint_path = args.checkpoint_path
|
163 |
else:
|
164 |
checkpoint_path = hf_hub_download("lkeab/hq-sam", f"{args.model_name}.pth")
|
|
|
165 |
|
166 |
convert_sam_checkpoint(args.model_name, checkpoint_path, args.output_dir)
|