|
6 | 6 | # /// script |
7 | 7 | # requires-python = ">=3.11" |
8 | 8 | # dependencies = [ |
9 | | -# "coreai-core==1.0.0b2", |
10 | | -# "coreai-torch==0.4.1", |
| 9 | +# "coreai-models", |
| 10 | +# "transformers>=5.5.4,<5.10.1", |
11 | 11 | # "tokenizers<0.23.0rc", |
| 12 | +# "huggingface-hub>=1.5.0,<2.0", |
12 | 13 | # "torchvision", |
13 | | -# "transformers>=5.5.4,<5.10.1", |
14 | 14 | # ] |
15 | 15 | # |
16 | 16 | # [tool.uv] |
17 | 17 | # index-url = "https://pypi.org/simple" |
18 | 18 | # prerelease = "allow" |
19 | 19 | # index-strategy = "unsafe-best-match" |
| 20 | +# # Force these past the workspace's `coreai-models` pins: |
| 21 | +# # - transformers: workspace is <5.0, SAM3 needs Sam3Model from >=5.5.4 |
| 22 | +# # - huggingface-hub: workspace is <1.0, transformers 5.x requires >=1.5.0,<2.0 |
| 23 | +# override-dependencies = [ |
| 24 | +# "transformers>=5.5.4,<5.10.1", |
| 25 | +# "huggingface-hub>=1.5.0,<2.0", |
| 26 | +# ] |
| 27 | +# |
| 28 | +# [tool.uv.sources] |
| 29 | +# # Resolve `coreai-models` against the workspace checkout instead of PyPI so |
| 30 | +# # this script picks up local edits to `coreai_models.segmentation.*` and |
| 31 | +# # `coreai_models.models.ios.sam3.*`. |
| 32 | +# coreai-models = { path = "../../python", editable = true } |
20 | 33 | # /// |
21 | | -import argparse |
22 | | -import json |
23 | | -import shutil |
24 | | -import time |
25 | | -from pathlib import Path |
26 | | - |
27 | | -import torch |
28 | | -import transformers |
29 | | -from coreai.runtime import AIModelAssetMetadata |
30 | | -from coreai_torch import TorchConverter, get_decomp_table |
31 | | - |
32 | | - |
33 | | -def reference_inputs(model_name: str, dtype: torch.dtype) -> dict[str, torch.Tensor]: |
34 | | - processor = transformers.Sam3Processor.from_pretrained(model_name) |
35 | | - text_inputs = processor.tokenizer(["dummy"], return_tensors="pt") |
36 | | - return { |
37 | | - "pixel_values": torch.randn(1, 3, 1008, 1008).to(dtype), |
38 | | - "input_ids": text_inputs["input_ids"].to(torch.int32), |
39 | | - } |
40 | | - |
41 | | - |
42 | | -class Sam3Module(torch.nn.Module): |
43 | | - def __init__(self, model_id: str = "facebook/sam3"): |
44 | | - super().__init__() |
45 | | - self._model = transformers.Sam3Model.from_pretrained(model_id) |
46 | | - |
47 | | - def forward(self, pixel_values, input_ids): |
48 | | - outputs = self._model(pixel_values=pixel_values, input_ids=input_ids) |
49 | | - return ( |
50 | | - outputs.pred_masks, |
51 | | - outputs.pred_boxes, |
52 | | - outputs.pred_logits, |
53 | | - outputs.presence_logits, |
54 | | - outputs.semantic_seg, |
55 | | - ) |
56 | | - |
57 | | - |
58 | | -def _default_output_dir() -> str: |
59 | | - return str(Path(__file__).resolve().parents[2] / "exports") |
60 | | - |
61 | | - |
62 | | -def _variant_name(model_name: str, dtype: torch.dtype) -> str: |
63 | | - safe_name = Path(model_name).name |
64 | | - dtype_name = str(dtype).split(".")[-1] |
65 | | - return f"{safe_name}_{dtype_name}" |
66 | | - |
67 | | - |
68 | | -def _bundle_paths( |
69 | | - output_dir: str, model_name: str, dtype: torch.dtype |
70 | | -) -> tuple[Path, Path]: |
71 | | - """Return (bundle_dir, model_path) where the .aimodel sits inside the bundle dir.""" |
72 | | - variant = _variant_name(model_name, dtype) |
73 | | - bundle_dir = Path(output_dir) / variant |
74 | | - return bundle_dir, bundle_dir / f"{variant}.aimodel" |
75 | | - |
76 | | - |
77 | | -def _build_aimodel_metadata() -> AIModelAssetMetadata: |
78 | | - # Source: https://github.com/facebookresearch/sam3 |
79 | | - metadata = AIModelAssetMetadata() |
80 | | - metadata.author = "N. Carion et al." |
81 | | - metadata.license = "SAM License" |
82 | | - metadata.model_description = "SAM 3 is a unified foundation model for promptable segmentation in images and videos. It can detect, segment, and track objects using text or visual prompts such as points, boxes, and masks. This variant is explicitly for image segmentation. Source: https://github.com/facebookresearch/sam3" |
83 | | - metadata.creation_date = int(time.time()) |
84 | | - return metadata |
85 | | - |
86 | | - |
87 | | -def _save_asset(coreai_program, model_path: Path, overwrite: bool) -> None: |
88 | | - bundle_dir = model_path.parent |
89 | | - if bundle_dir.exists(): |
90 | | - if not overwrite: |
91 | | - raise FileExistsError( |
92 | | - f"{bundle_dir} already exists. Pass --overwrite to replace it." |
93 | | - ) |
94 | | - shutil.rmtree(bundle_dir) |
95 | | - bundle_dir.mkdir(parents=True, exist_ok=True) |
96 | | - coreai_program.save_asset(model_path, _build_aimodel_metadata()) |
97 | | - |
98 | | - |
99 | | -def _write_tokenizer(dest: Path, model_id: str) -> None: |
100 | | - print(f"[INFO] Saving tokenizer from {model_id} to {dest}...") |
101 | | - tokenizer = transformers.AutoTokenizer.from_pretrained(model_id) |
102 | | - tokenizer.save_pretrained(str(dest)) |
103 | | - |
104 | | - |
105 | | -def _write_bundle_metadata(bundle_dir: Path, variant: str) -> None: |
106 | | - metadata = { |
107 | | - "metadata_version": "0.2", |
108 | | - "kind": "segmenter", |
109 | | - "name": variant, |
110 | | - "assets": {"main": f"{variant}.aimodel"}, |
111 | | - } |
112 | | - metadata_path = bundle_dir / "metadata.json" |
113 | | - with open(metadata_path, "w") as f: |
114 | | - json.dump(metadata, f, indent=2) |
115 | | - print(f"[INFO] Wrote metadata to {metadata_path}.") |
116 | | - |
117 | | - |
118 | | -def create_sam3( |
119 | | - output_dir: str, |
120 | | - model_name: str, |
121 | | - dtype: torch.dtype, |
122 | | - overwrite: bool, |
123 | | -): |
124 | | - print("[INFO] Sourcing model...") |
125 | | - model = Sam3Module(model_id=model_name) |
126 | | - model.eval() |
127 | | - model.to(dtype) |
128 | | - print("[INFO] Model sourced. Running torch export with decompositions...") |
129 | | - |
130 | | - example_inputs = reference_inputs(model_name, dtype) |
131 | | - |
132 | | - with torch.autocast(device_type="cpu", dtype=dtype): |
133 | | - exported = torch.export.export( |
134 | | - model, |
135 | | - args=(), |
136 | | - kwargs=example_inputs, |
137 | | - ) |
138 | | - exported = exported.run_decompositions(get_decomp_table()) |
139 | | - print("[INFO] Model exported. Converting to Core AI...") |
140 | 34 |
|
141 | | - converter = TorchConverter().add_exported_program( |
142 | | - exported_program=exported, |
143 | | - input_names=["pixel_values", "input_ids"], |
144 | | - output_names=[ |
145 | | - "pred_masks", |
146 | | - "pred_boxes", |
147 | | - "pred_logits", |
148 | | - "presence_logits", |
149 | | - "semantic_seg", |
150 | | - ], |
151 | | - ) |
152 | | - coreai_program = converter.to_coreai() |
153 | | - print("[INFO] Model converted.") |
154 | | - coreai_program.optimize() |
155 | | - print("[INFO] Model optimized.") |
| 35 | +"""SAM3 export entry point — runs as a uv inline-script. |
156 | 36 |
|
157 | | - bundle_dir, model_path = _bundle_paths(output_dir, model_name, dtype) |
158 | | - _save_asset(coreai_program, model_path, overwrite) |
159 | | - print(f"[INFO] Successfully created and saved Core AI model to {model_path}.") |
| 37 | +The export pipeline lives in ``coreai_models.segmentation.pipeline``; this |
| 38 | +script is a thin wrapper whose only job is to give uv the right |
| 39 | +``transformers`` resolution (SAM3 needs >= 5.5.4, the workspace pins < 5.0). |
| 40 | +The PEP 723 metadata at the top of this file is what makes that possible — |
| 41 | +``uv run models/sam3/export.py`` resolves an isolated env per-script. |
160 | 42 |
|
161 | | - _write_tokenizer(bundle_dir / "tokenizer", model_name) |
162 | | - _write_bundle_metadata(bundle_dir, _variant_name(model_name, dtype)) |
| 43 | +Usage: |
| 44 | + uv run models/sam3/export.py [--image-size N] [--n-bits N] [--output-dir PATH] ... |
| 45 | +""" |
163 | 46 |
|
| 47 | +import sys |
164 | 48 |
|
165 | | -def main(): |
166 | | - parser = argparse.ArgumentParser( |
167 | | - description="Create and save a Core AI AIProgram for SAM3." |
168 | | - ) |
169 | | - parser.add_argument( |
170 | | - "--model", |
171 | | - choices=["facebook/sam3"], |
172 | | - default="facebook/sam3", |
173 | | - help="Model variant to convert.", |
174 | | - ) |
175 | | - parser.add_argument( |
176 | | - "--output-dir", |
177 | | - default=None, |
178 | | - help="Output directory for the .aimodel bundle (default: <repo-root>/exports/)", |
179 | | - ) |
180 | | - parser.add_argument( |
181 | | - "--dtype", |
182 | | - choices=["float16", "float32"], |
183 | | - default="float32", |
184 | | - help="Torch dtype to use for the model.", |
185 | | - ) |
186 | | - parser.add_argument( |
187 | | - "--overwrite", |
188 | | - action="store_true", |
189 | | - help="Overwrite an existing .aimodel asset at the output path.", |
190 | | - ) |
191 | | - args = parser.parse_args() |
192 | 49 |
|
193 | | - dtype = { |
194 | | - "float16": torch.float16, |
195 | | - "float32": torch.float32, |
196 | | - }[args.dtype] |
| 50 | +def main() -> None: |
| 51 | + # Lazy import so the inline-script header parses cleanly even if the |
| 52 | + # workspace package isn't on sys.path yet (uv handles that before main()). |
| 53 | + from coreai_models.segmentation.export import main as segmentation_main |
197 | 54 |
|
198 | | - output_dir = args.output_dir or _default_output_dir() |
199 | | - create_sam3(output_dir, args.model, dtype, args.overwrite) |
| 55 | + # The shared CLI takes a `model` positional. This script only handles SAM3, |
| 56 | + # so inject "sam3" and forward any user-supplied flags untouched. |
| 57 | + if len(sys.argv) == 1 or sys.argv[1].startswith("-"): |
| 58 | + sys.argv.insert(1, "sam3") |
| 59 | + segmentation_main() |
200 | 60 |
|
201 | 61 |
|
202 | 62 | if __name__ == "__main__": |
|
0 commit comments