Ukawka
86 lines
4.6 KiB
Python
86 lines
4.6 KiB
Python
"""Local SAM image detector from a SAM 3.1 multiplex checkpoint.
|
|
|
|
No downloads, service, video tracker, or silently uninitialized parameters.
|
|
"""
|
|
import gc
|
|
import hashlib
|
|
import importlib.metadata
|
|
import os
|
|
from pathlib import Path
|
|
import sys
|
|
|
|
|
|
def check_environment(code_dir):
|
|
if Path(sys.prefix).name != 'ai-spine' or not (Path(sys.prefix) / 'conda-meta').is_dir():
|
|
raise RuntimeError('Run inside Conda ai-spine: conda run -n ai-spine ...')
|
|
code = Path(code_dir).resolve()
|
|
if not (code / 'sam3' / 'model_builder.py').is_file():
|
|
raise FileNotFoundError(f'SAM source checkout missing: {code}')
|
|
sys.path.insert(0, str(code))
|
|
os.environ['HF_HUB_OFFLINE'] = '1'
|
|
os.environ['TRANSFORMERS_OFFLINE'] = '1'
|
|
import torch
|
|
import sam3
|
|
from sam3.model_builder import build_sam3_image_model # noqa: F401
|
|
from sam3.model.sam3_image_processor import Sam3Processor # noqa: F401
|
|
if Path(sam3.__file__).resolve().parent != code / 'sam3':
|
|
raise RuntimeError('Imported SAM package does not match --sam-code-dir')
|
|
if not torch.cuda.is_available():
|
|
raise RuntimeError('CUDA GPU unavailable; this stage requires GPU inference')
|
|
source_hash = hashlib.sha256()
|
|
for file in sorted((code / 'sam3').rglob('*')):
|
|
if file.is_file() and (file.suffix == '.py' or file.name.endswith('.txt.gz')):
|
|
source_hash.update(file.relative_to(code).as_posix().encode())
|
|
source_hash.update(file.read_bytes())
|
|
versions = {}
|
|
for name in ['torch', 'torchvision', 'numpy', 'Pillow', 'opencv-python', 'timm', 'triton-windows']:
|
|
try:
|
|
versions[name] = importlib.metadata.version(name)
|
|
except importlib.metadata.PackageNotFoundError:
|
|
versions[name] = None
|
|
return {'python': sys.executable, 'sam_source': str(code), 'sam_source_sha256': source_hash.hexdigest(),
|
|
'versions': versions, 'cuda': torch.version.cuda, 'gpu': torch.cuda.get_device_name(0),
|
|
'gpu_total_bytes': torch.cuda.get_device_properties(0).total_memory,
|
|
'bf16_supported': torch.cuda.is_bf16_supported(), 'network_requested': False}
|
|
|
|
|
|
def load_detector(checkpoint, precision='bf16'):
|
|
import torch
|
|
from sam3.model_builder import build_sam3_image_model
|
|
model = build_sam3_image_model(device='cpu', load_from_HF=False,
|
|
enable_inst_interactivity=False, compile=False)
|
|
weights = torch.load(checkpoint, map_location='cpu', weights_only=True, mmap=True)
|
|
weights = weights.get('model', weights)
|
|
detector = {key.removeprefix('detector.'): value for key, value in weights.items() if key.startswith('detector.')}
|
|
if not detector:
|
|
raise RuntimeError('Checkpoint has no detector.* weights; refusing an uninitialized model')
|
|
# The image builder uses four scales then scalp=1 discards the last one.
|
|
# SAM 3.1 stores only the three used detection scales. Remove the unused
|
|
# fourth scale and scalp, preserving exactly the same three feature maps.
|
|
multiplex = any(k.startswith('backbone.vision_backbone.interactive_convs.') for k in detector)
|
|
if multiplex:
|
|
neck = model.backbone.vision_backbone
|
|
if len(neck.convs) != 4 or model.backbone.scalp != 1:
|
|
raise RuntimeError('Unsupported SAM image backbone layout; review adapter for this source version')
|
|
neck.convs = torch.nn.ModuleList(list(neck.convs)[:3])
|
|
neck.scale_factors = list(neck.scale_factors)[:3]
|
|
model.backbone.scalp = 0
|
|
ignored_prefixes = ('backbone.vision_backbone.interactive_convs.',
|
|
'backbone.vision_backbone.propagation_convs.') if multiplex else ()
|
|
selected = {k: v for k, v in detector.items() if not k.startswith(ignored_prefixes)}
|
|
# strict=True checks every active parameter and every tensor shape.
|
|
model.load_state_dict(selected, strict=True)
|
|
info = {'checkpoint_kind': 'sam3.1_multiplex_detector' if multiplex else 'sam3_detector',
|
|
'strict_load': True, 'loaded_tensors': len(selected),
|
|
'ignored_video_neck_tensors': len(detector) - len(selected),
|
|
'instance_interactivity': False, 'compile': False, 'precision': precision}
|
|
del selected, detector, weights
|
|
gc.collect()
|
|
if precision == 'bf16' and not torch.cuda.is_bf16_supported():
|
|
raise RuntimeError('BF16 unsupported on this GPU; use --sam-precision fp32')
|
|
# Keep FP32 parameters and complex RoPE buffers intact. The upstream
|
|
# decoder explicitly disables autocast for its FP32 FFN; casting the whole
|
|
# model to BF16 would both break it and discard complex imaginary values.
|
|
model = model.to(device='cuda').eval()
|
|
return model, info
|