File size: 4,210 Bytes
6f059aa | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 | #!/usr/bin/env python3
"""
Convert CLIP and SAM3 checkpoints to safetensors format.
Run from repo root: python convert_to_safetensors.py
"""
import argparse
from pathlib import Path
import torch
from safetensors.torch import save_file
def convert_clip(source_dir: Path, output_dir: Path = None) -> Path:
"""Convert CLIP pytorch_model.bin to model.safetensors."""
output_dir = output_dir or source_dir
bin_path = source_dir / "pytorch_model.bin"
out_path = output_dir / "model.safetensors"
if out_path.exists():
print(f"CLIP: {out_path} already exists, skip")
return out_path
if bin_path.exists():
print(f"CLIP: Loading from {bin_path}...")
state_dict = torch.load(bin_path, map_location="cpu", weights_only=True)
else:
print(f"CLIP: Loading from HuggingFace...")
from transformers import CLIPModel
model = CLIPModel.from_pretrained(str(source_dir) if source_dir.exists() else "openai/clip-vit-base-patch16")
state_dict = model.state_dict()
state_dict = {k: v.float() if v.dtype in (torch.float16, torch.bfloat16) else v for k, v in state_dict.items()}
save_file(state_dict, str(out_path))
print(f"CLIP: Saved to {out_path}")
if bin_path.exists():
bin_path.unlink()
print(f"CLIP: Removed {bin_path}")
return out_path
def _extract_sam3_state_dict(ckpt: dict) -> dict:
"""Extract SAM3 image model state dict from checkpoint (same logic as sam3._load_checkpoint)."""
if "model" in ckpt and isinstance(ckpt["model"], dict):
ckpt = ckpt["model"]
sam3_image_ckpt = {
k.replace("detector.", ""): v for k, v in ckpt.items() if "detector" in k
}
return sam3_image_ckpt
def convert_sam3(pt_path: Path, output_path: Path = None) -> Path:
"""Convert SAM3 sam3.pt to model.safetensors (image model weights only)."""
output_path = output_path or pt_path.parent / "model.safetensors"
print(f"SAM3: Loading from {pt_path}...")
ckpt = torch.load(pt_path, map_location="cpu", weights_only=True)
state_dict = _extract_sam3_state_dict(ckpt)
state_dict = {k: v.float() if v.dtype in (torch.float16, torch.bfloat16) else v for k, v in state_dict.items()}
save_file(state_dict, str(output_path))
print(f"SAM3: Saved to {output_path}")
return output_path
def copy_sam3_safetensors(source: Path, dest_dir: Path) -> Path:
"""Copy HF model.safetensors (detector_model.* keys) to SegEarth-OV. Pipeline maps keys on load."""
dest = dest_dir / "model.safetensors"
if source.exists():
import shutil
shutil.copy2(source, dest)
print(f"SAM3: Copied {source} -> {dest}")
return dest
return None
def main():
parser = argparse.ArgumentParser(description="Convert CLIP and SAM3 checkpoints to safetensors")
parser.add_argument("--clip", action="store_true", help="Convert CLIP only")
parser.add_argument("--sam3", action="store_true", help="Convert SAM3 only")
parser.add_argument("--all", action="store_true", help="Convert all (default when no --clip/--sam3)")
args = parser.parse_args()
repo = Path(__file__).parent
do_both = not args.clip and not args.sam3
if args.clip or do_both:
clip_dir = repo / "OV" / "weights" / "backbone" / "clip-vit-base-patch16"
if clip_dir.exists():
convert_clip(clip_dir)
else:
print(f"CLIP: {clip_dir} not found, skip")
if args.sam3 or do_both:
sam3_dir = repo / "OV-3" / "weights" / "backbone" / "sam3"
hf_safetensors = Path("/data/projects/models/hf_models/facebook/sam3/model.safetensors")
sam3_pt = sam3_dir / "sam3.pt"
if hf_safetensors.exists():
copy_sam3_safetensors(hf_safetensors, sam3_dir)
if sam3_pt.exists():
sam3_pt.unlink()
print(f"SAM3: Removed {sam3_pt}")
elif sam3_pt.exists():
st_path = convert_sam3(sam3_pt)
sam3_pt.unlink()
print(f"SAM3: Removed {sam3_pt}")
else:
print(f"SAM3: Neither {hf_safetensors} nor {sam3_pt} found, skip")
print("Done.")
if __name__ == "__main__":
main()
|