第二步SAM模型接入

2026-09-18 20:29:29 +08:00
parent 1a7b2f14a9
commit f095c28858
16 changed files with 788 additions and 26 deletions
@@ -0,0 +1,85 @@
+"""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