| 14 | |
| 15 | |
| 16 | def main(): |
| 17 | try: |
| 18 | from modelscope import snapshot_download |
| 19 | except ImportError: |
| 20 | print("Install modelscope first: pip install modelscope") |
| 21 | sys.exit(1) |
| 22 | |
| 23 | print(f"Downloading {RMBG_MODEL_ID} from ModelScope...") |
| 24 | cache_dir = snapshot_download(RMBG_MODEL_ID) |
| 25 | # Model may be at repo root or under onnx/ |
| 26 | candidates = [ |
| 27 | os.path.join(cache_dir, "model.onnx"), |
| 28 | os.path.join(cache_dir, "onnx", "model.onnx"), |
| 29 | ] |
| 30 | src = None |
| 31 | for p in candidates: |
| 32 | if os.path.isfile(p): |
| 33 | src = p |
| 34 | break |
| 35 | if not src: |
| 36 | # Search for any .onnx file |
| 37 | for root, _, files in os.walk(cache_dir): |
| 38 | for f in files: |
| 39 | if f.endswith(".onnx"): |
| 40 | src = os.path.join(root, f) |
| 41 | break |
| 42 | if src: |
| 43 | break |
| 44 | if not src or not os.path.isfile(src): |
| 45 | print(f"model.onnx not found under {cache_dir}; download from ModelScope and place at {TARGET_PATH}") |
| 46 | sys.exit(1) |
| 47 | |
| 48 | os.makedirs(MODEL_DIR, exist_ok=True) |
| 49 | shutil.copy2(src, TARGET_PATH) |
| 50 | print(f"Saved: {TARGET_PATH}") |
| 51 | print("For RMBG inference also install: pip install onnxruntime # or onnxruntime-gpu") |
| 52 | |
| 53 | |
| 54 | if __name__ == "__main__": |