44 lines
1.5 KiB
Python
44 lines
1.5 KiB
Python
"""Resource-limited CPU single-file decoder; no network or neighboring model lookup."""
|
|
|
|
import contextlib
|
|
import json
|
|
import os
|
|
import resource
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
# Apply limits before importing tensor/protobuf runtimes. Thread env also set by parent.
|
|
resource.setrlimit(resource.RLIMIT_CPU, (40, 40))
|
|
resource.setrlimit(resource.RLIMIT_AS, (8 * 1024**3, 8 * 1024**3))
|
|
resource.setrlimit(resource.RLIMIT_FSIZE, (16 * 1024**2, 16 * 1024**2))
|
|
resource.setrlimit(resource.RLIMIT_NOFILE, (64, 64))
|
|
resource.setrlimit(resource.RLIMIT_CORE, (0, 0))
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = ""
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
for root in (ROOT, ROOT.parent):
|
|
sys.path.insert(0, str(root))
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
payload = json.loads(sys.stdin.read(4096))
|
|
with contextlib.redirect_stdout(sys.stderr):
|
|
import torch
|
|
from pretrained_upload import import_upload
|
|
|
|
torch.set_num_threads(1)
|
|
manifest = import_upload(
|
|
payload["path"], payload["format"], payload["template"], payload["directory"]
|
|
)
|
|
print(json.dumps(manifest))
|
|
except Exception as error:
|
|
# Do not echo untrusted pickle/tensor names or absolute local paths.
|
|
from pretrained import PretrainedError
|
|
|
|
message = (
|
|
str(error)
|
|
if isinstance(error, PretrainedError)
|
|
else "文件解析失败；请提供受支持的自包含Go2 legacy47 actor"
|
|
)
|
|
print(json.dumps({"error": message[:500]}))
|
|
sys.exit(1)
|