File size: 5,143 Bytes
843d12b | 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 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 | """
Download SAM2 ONNX models from HuggingFace.
Usage:
python models/sam2/download_models.py [variant]
Variants:
tiny - SAM2 Hiera-Tiny (smallest, fastest) ~155 MB
small - SAM2 Hiera-Small (good balance) ~183 MB
base_plus - SAM2 Hiera-Base+ (higher accuracy) ~360 MB
large - SAM2 Hiera-Large (most accurate) ~910 MB
Default: tiny (recommended for real-time LivePortrait use)
Models are downloaded from:
https://huggingface.co/vietanhdev/segment-anything-2-onnx-models
"""
import os
import sys
import urllib.request
import zipfile
HUGGINGFACE_BASE_URL = "https://huggingface.co/vietanhdev/segment-anything-2-onnx-models/resolve/main"
VARIANTS = {
"tiny": "sam2_hiera_tiny",
"small": "sam2_hiera_small",
"base_plus": "sam2_hiera_base_plus",
"large": "sam2_hiera_large",
}
def download_variant(variant_key: str = "tiny", output_dir: str = None):
"""Download and extract a SAM2 ONNX model variant."""
if output_dir is None:
output_dir = os.path.dirname(os.path.abspath(__file__))
if variant_key not in VARIANTS:
print(f"Unknown variant: {variant_key}")
print(f"Available variants: {', '.join(VARIANTS.keys())}")
return False
variant_name = VARIANTS[variant_key]
zip_filename = f"{variant_name}.zip"
zip_url = f"{HUGGINGFACE_BASE_URL}/{zip_filename}"
zip_path = os.path.join(output_dir, zip_filename)
encoder_path = os.path.join(output_dir, f"{variant_name}.encoder.onnx")
decoder_path = os.path.join(output_dir, f"{variant_name}.decoder.onnx")
# Check if already downloaded
if os.path.exists(encoder_path) and os.path.exists(decoder_path):
print(f"[SAM2] Models already exist for variant '{variant_key}':")
print(f" Encoder: {encoder_path}")
print(f" Decoder: {decoder_path}")
return True
print(f"[SAM2] Downloading {variant_key} variant from HuggingFace...")
print(f" URL: {zip_url}")
print(f" Destination: {output_dir}")
try:
# Download zip
def progress_hook(count, block_size, total_size):
percent = int(count * block_size * 100 / total_size) if total_size > 0 else 0
percent = min(percent, 100)
sys.stdout.write(f"\r Downloading: {percent}%")
sys.stdout.flush()
urllib.request.urlretrieve(zip_url, zip_path, reporthook=progress_hook)
print() # newline after progress
# Extract
print(f" Extracting...")
with zipfile.ZipFile(zip_path, 'r') as z:
z.extractall(output_dir)
# The zip extracts to a subdirectory; move files up if needed
extracted_dir = os.path.join(output_dir, variant_name)
if os.path.isdir(extracted_dir):
for fname in os.listdir(extracted_dir):
src = os.path.join(extracted_dir, fname)
dst = os.path.join(output_dir, f"{variant_name}.{fname}")
if not os.path.exists(dst):
os.rename(src, dst)
# Clean up extracted directory
try:
os.rmdir(extracted_dir)
except OSError:
pass
# Clean up zip
if os.path.exists(zip_path):
os.remove(zip_path)
# Verify
if os.path.exists(encoder_path) and os.path.exists(decoder_path):
print(f"[SAM2] Successfully downloaded {variant_key} variant!")
print(f" Encoder: {encoder_path}")
print(f" Decoder: {decoder_path}")
return True
else:
# Try alternate naming (files might be named without the variant prefix)
alt_encoder = os.path.join(output_dir, "encoder.onnx")
alt_decoder = os.path.join(output_dir, "decoder.onnx")
if os.path.exists(alt_encoder) and not os.path.exists(encoder_path):
os.rename(alt_encoder, encoder_path)
if os.path.exists(alt_decoder) and not os.path.exists(decoder_path):
os.rename(alt_decoder, decoder_path)
if os.path.exists(encoder_path) and os.path.exists(decoder_path):
print(f"[SAM2] Successfully downloaded {variant_key} variant!")
return True
else:
print(f"[SAM2] Warning: Expected files not found after extraction.")
print(f" Looking for: {encoder_path}")
print(f" Looking for: {decoder_path}")
print(f" Contents of {output_dir}:")
for f in os.listdir(output_dir):
print(f" {f}")
return False
except Exception as e:
print(f"\n[SAM2] Download failed: {e}")
# Clean up partial download
if os.path.exists(zip_path):
os.remove(zip_path)
return False
if __name__ == "__main__":
variant = sys.argv[1] if len(sys.argv) > 1 else "tiny"
success = download_variant(variant)
if not success:
sys.exit(1)
|